||| Primitive operations for fixed size, signed integers.
module Primitives.Int

import Primitives.Common
import Core.Context
import CairoCode.CairoCode
import Data.SortedSet
import Data.SortedMap
import CommonDef
import CodeGen.CodeGenHelper

---------------------------------------------------------------------------------------------------
-- UINT ADD
---------------------------------------------------------------------------------------------------
add_intX_name : Nat -> Name
add_intX_name nBits = makeBuiltinName "add_int\{show nBits}"

add_intX_code : Name -> Nat -> String
add_intX_code name nBits = """
    func \{cairoName name}(range_check_ptr, a : felt, b : felt) -> (res : felt, range_check_ptr):
        alloc_locals
        local res
        local overflow
        %{
            raw_res = (ids.a + ids.b) % PRIME
            if (raw_res >= \{show halfShift}) & (raw_res < \{show shift}):
                (ids.overflow, ids.res) = (1, (raw_res - \{show shift}) % PRIME)
            elif (raw_res < PRIME - \{show halfShift}) & (PRIME - \{show shift} <= raw_res):
                (ids.overflow, ids.res) = (-1 % PRIME, (raw_res + \{show shift}) % PRIME)
            else:
                (ids.overflow, ids.res) = (0, raw_res)
        %}

        assert overflow * overflow * overflow = overflow  # overflow is -1, 0 or 1
        assert res = a + b - overflow * \{show shift}

        #inlined int_check
        assert [range_check_ptr] = res + \{show halfShift}			#Ensure: -HALF_SHIFT <= res
        assert [range_check_ptr+1] = \{show (halfShift - 1)} - res	#Ensure: res < HALF_SHIFT
        let range_check_ptr = range_check_ptr + 2

        return (res, range_check_ptr)
    end

"""
  where shift : Integer
        shift = pow2 nBits
        halfShift : Integer
        halfShift = pow2 (nBits `minus` 1)

add_intX : Nat -> (Name, Lazy CairoDef)
add_intX nBits = (name, def)
  where name : Name
        name = add_intX_name nBits
        def : CairoDef
        def = ForeignDef (MkForeignInfo True Nothing [range_check_builtin] empty (add_intX_code name nBits)) 2 1

-- Implements
export
[add_int] PrimFn where
    eval [_, ConstValue (I 0)] = Just $ ArgValue 0
    eval [ConstValue (I 0), _] = Just $ ArgValue 1
    eval [ConstValue (I a), ConstValue (I b)] = Just $ NewValue $ ConstValue $ I (a+b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [add_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (add_intX_name 64) args]

export
[add_int8] PrimFn where
    eval [_, ConstValue (I8 0)] = Just $ ArgValue 0
    eval [ConstValue (I8 0), _] = Just $ ArgValue 1
    eval [ConstValue (I8 a), ConstValue (I8 b)] = Just $ NewValue $ ConstValue $ I8 (a+b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [add_intX 8]
    generateCode res args implicits = Instructions [CALL [res] implicits (add_intX_name 8) args]

export
[add_int16] PrimFn where
    eval [_, ConstValue (I16 0)] = Just $ ArgValue 0
    eval [ConstValue (I16 0), _] = Just $ ArgValue 1
    eval [ConstValue (I16 a), ConstValue (I16 b)] = Just $ NewValue $ ConstValue $ I16 (a+b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [add_intX 16]
    generateCode res args implicits = Instructions [CALL [res] implicits (add_intX_name 16) args]

export
[add_int32] PrimFn where
    eval [_, ConstValue (I32 0)] = Just $ ArgValue 0
    eval [ConstValue (I32 0), _] = Just $ ArgValue 1
    eval [ConstValue (I32 a), ConstValue (I32 b)] = Just $ NewValue $ ConstValue $ I32 (a+b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [add_intX 32]
    generateCode res args implicits = Instructions [CALL [res] implicits (add_intX_name 32) args]

export
[add_int64] PrimFn where
    eval [_, ConstValue (I64 0)] = Just $ ArgValue 0
    eval [ConstValue (I64 0), _] = Just $ ArgValue 1
    eval [ConstValue (I64 a), ConstValue (I64 b)] = Just $ NewValue $ ConstValue $ I64 (a+b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [add_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (add_intX_name 64) args]

---------------------------------------------------------------------------------------------------
-- UINT SUB
---------------------------------------------------------------------------------------------------
sub_intX_name : Nat -> Name
sub_intX_name nBits = makeBuiltinName "sub_int\{show nBits}"

sub_intX_code : Name -> Nat -> String
sub_intX_code name nBits = """
    # Subtraction.
    func \{cairoName name}(range_check_ptr, a : felt, b : felt) -> (res : felt, range_check_ptr):
        alloc_locals
        local res : felt
        local overflow : felt
        %{
            a_sub_b = (ids.a - ids.b) % PRIME
            if (a_sub_b >= \{show halfShift}) & (a_sub_b < \{show shift}):
                (ids.overflow, ids.res) = (1, (a_sub_b - \{show shift}) % PRIME)
            elif (a_sub_b < PRIME - \{show halfShift}) & (PRIME - \{show shift} <= a_sub_b):
                (ids.overflow, ids.res) = (-1 % PRIME, (a_sub_b + \{show shift}) % PRIME)
            else:
                (ids.overflow, ids.res) = (0, a_sub_b)
        %}

        assert overflow * overflow * overflow = overflow  # overflow is -1, 0 or 1
        assert res = a - b - overflow * \{show shift}

        #inlined int_check
        assert [range_check_ptr] = res + \{show halfShift}			#Ensure: -HALF_SHIFT <= res
        assert [range_check_ptr+1] = \{show (halfShift - 1)} - res	#Ensure: res < HALF_SHIFT
        let range_check_ptr = range_check_ptr + 2

        return (res,range_check_ptr)
    end

"""
  where shift : Integer
        shift = pow2 nBits
        halfShift : Integer
        halfShift = pow2 (nBits `minus` 1)

sub_intX : Nat -> (Name, Lazy CairoDef)
sub_intX nBits = (name, def)
  where name : Name
        name = sub_intX_name nBits
        def : CairoDef
        def = ForeignDef (MkForeignInfo True Nothing [range_check_builtin] empty (sub_intX_code name nBits)) 2 1

export
[sub_int] PrimFn where
    eval [_, ConstValue (I 0)] = Just $ ArgValue 0
    eval [ConstValue (I a), ConstValue (I b)] = Just $ NewValue $ ConstValue $ I (a-b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [sub_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (sub_intX_name 64) args]

export
[sub_int8] PrimFn where
    eval [_, ConstValue (I8 0)] = Just $ ArgValue 0
    eval [ConstValue (I8 a), ConstValue (I8 b)] = Just $ NewValue $ ConstValue $ I8 (a-b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [sub_intX 8]
    generateCode res args implicits = Instructions [CALL [res] implicits (sub_intX_name 8) args]

export
[sub_int16] PrimFn where
    eval [_, ConstValue (I16 0)] = Just $ ArgValue 0
    eval [ConstValue (I16 a), ConstValue (I16 b)] = Just $ NewValue $ ConstValue $ I16 (a-b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [sub_intX 16]
    generateCode res args implicits = Instructions [CALL [res] implicits (sub_intX_name 16) args]

export
[sub_int32] PrimFn where
    eval [_, ConstValue (I32 0)] = Just $ ArgValue 0
    eval [ConstValue (I32 a), ConstValue (I32 b)] = Just $ NewValue $ ConstValue $ I32 (a-b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [sub_intX 32]
    generateCode res args implicits = Instructions [CALL [res] implicits (sub_intX_name 32) args]

export
[sub_int64] PrimFn where
    eval [_, ConstValue (I64 0)] = Just $ ArgValue 0
    eval [ConstValue (I64 a), ConstValue (I64 b)] = Just $ NewValue $ ConstValue $ I64 (a-b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [sub_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (sub_intX_name 64) args]

---------------------------------------------------------------------------------------------------
-- UINT MUL
---------------------------------------------------------------------------------------------------
mul_intX_name : Nat -> Name
mul_intX_name nBits = makeBuiltinName "mul_int\{show nBits}"

mul_intX_code : Name -> Nat -> String
mul_intX_code name nBits = """
    # Multiplies two int.
    # Returns the result as int.
    func \{cairoName name}(range_check_ptr, a : felt, b : felt) -> (res : felt, range_check_ptr):
        alloc_locals
        local res
        local offset
        %{
            from starkware.cairo.common.math_utils import as_int
            rawMult = ((as_int(ids.a, PRIME) * as_int(ids.b, PRIME)) + \{show halfShift})
            (rawOffset, resRaw) = divmod(rawMult, \{show shift})
            ids.offset = rawOffset % PRIME
            ids.res = (resRaw - \{show halfShift}) % PRIME
        %}

        # Check offset can not trigger overflow / underflow
        #  Note: Max valid offset absolute = HALF_SHIFT*HALF_SHIFT/SHIFT = (SHIFT*SHIFT)/4*SHIFT = SHIFT/4
        #  For simplicity we weaken to -HALF_SHIFT <= offset <= HALF_SHIFT then no overflow (for 64bit): -2**128 <= +/-(2**63 * 2**63) +/- (2**63 * 2**64) <= +2**128 & overflow is at ~+/-PRIME/2 (with PRIME ~= 2**192)
        assert [range_check_ptr] = offset + \{show halfShift}		#Ensure: -HALF_SHIFT <= offset
        assert [range_check_ptr+1] = \{show halfShift} - offset	    #Ensure: offset <= HALF_SHIFT

        assert res = a * b - offset * \{show shift}

        # inlined int_check(res)
        assert [range_check_ptr+2] = res + \{show halfShift}		#Ensure: -HALF_SHIFT <= res
        assert [range_check_ptr+3] = \{show (halfShift - 1)} - res	#Ensure: res < HALF_SHIFT
        let range_check_ptr = range_check_ptr + 4

        return (res, range_check_ptr)
    end

"""
  where shift : Integer
        shift = pow2 nBits
        halfShift : Integer
        halfShift = pow2 (nBits `minus` 1)

mul_intX : Nat -> (Name, Lazy CairoDef)
mul_intX nBits = (name, def)
  where name : Name
        name = mul_intX_name nBits
        def : CairoDef
        def = ForeignDef (MkForeignInfo True Nothing [range_check_builtin] empty (mul_intX_code name nBits)) 2 1

export
[mul_int] PrimFn where
    eval [_, ConstValue (I 1)] = Just $ ArgValue 0
    eval [ConstValue (I 1), _] = Just $ ArgValue 1
    eval [_, ConstValue (I 0)] = Just $ NewValue $ ConstValue $ I 0
    eval [ConstValue (I 0), _] = Just $ NewValue $ ConstValue $ I 0
    eval [ConstValue (I a), ConstValue (I b)] = Just $ NewValue $ ConstValue $ I (a*b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [mul_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (mul_intX_name 64) args]

export
[mul_int8] PrimFn where
    eval [_, ConstValue (I8 1)] = Just $ ArgValue 0
    eval [ConstValue (I8 1), _] = Just $ ArgValue 1
    eval [_, ConstValue (I8 0)] = Just $ NewValue $ ConstValue $ I8 0
    eval [ConstValue (I8 0), _] = Just $ NewValue $ ConstValue $ I8 0
    eval [ConstValue (I8 a), ConstValue (I8 b)] = Just $ NewValue $ ConstValue $ I8 (a*b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [mul_intX 8]
    generateCode res args implicits = Instructions [CALL [res] implicits (mul_intX_name 8) args]

export
[mul_int16] PrimFn where
    eval [_, ConstValue (I16 1)] = Just $ ArgValue 0
    eval [ConstValue (I16 1), _] = Just $ ArgValue 1
    eval [_, ConstValue (I16 0)] = Just $ NewValue $ ConstValue $ I16 0
    eval [ConstValue (I16 0), _] = Just $ NewValue $ ConstValue $ I16 0
    eval [ConstValue (I16 a), ConstValue (I16 b)] = Just $ NewValue $ ConstValue $ I16 (a*b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [mul_intX 16]
    generateCode res args implicits = Instructions [CALL [res] implicits (mul_intX_name 16) args]

export
[mul_int32] PrimFn where
    eval [_, ConstValue (I32 1)] = Just $ ArgValue 0
    eval [ConstValue (I32 1), _] = Just $ ArgValue 1
    eval [_, ConstValue (I32 0)] = Just $ NewValue $ ConstValue $ I32 0
    eval [ConstValue (I32 0), _] = Just $ NewValue $ ConstValue $ I32 0
    eval [ConstValue (I32 a), ConstValue (I32 b)] = Just $ NewValue $ ConstValue $ I32 (a*b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [mul_intX 32]
    generateCode res args implicits = Instructions [CALL [res] implicits (mul_intX_name 32) args]

export
[mul_int64] PrimFn where
    eval [_, ConstValue (I64 1)] = Just $ ArgValue 0
    eval [ConstValue (I64 1), _] = Just $ ArgValue 1
    eval [_, ConstValue (I64 0)] = Just $ NewValue $ ConstValue $ I64 0
    eval [ConstValue (I64 0), _] = Just $ NewValue $ ConstValue $ I64 0
    eval [ConstValue (I64 a), ConstValue (I64 b)] = Just $ NewValue $ ConstValue $ I64 (a*b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [mul_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (mul_intX_name 64) args]

---------------------------------------------------------------------------------------------------
-- UINT DIV
---------------------------------------------------------------------------------------------------
div_intX_name : Nat -> Name
div_intX_name nBits = makeBuiltinName "div_int\{show nBits}"

div_intX_code : Name -> Nat -> String
div_intX_code name nBits = """
    # Division between int.
    # Returns the quotient.
    # Conforms to Idris specifications: division by 0 yields error.
    func \{cairoName name}(range_check_ptr, a : felt, b : felt) -> (quotient : felt, range_check_ptr):
        alloc_locals
        local res
        local rem
        local sign
        local signB
        %{
            from starkware.cairo.common.math_utils import as_int
            signB = 1 if as_int(ids.b, PRIME) >= 0 else -1
            sign = signB if as_int(ids.a, PRIME) >= 0 else -signB
            ids.signB = signB % PRIME
            ids.sign = sign % PRIME
            #Note fails if b = 0, which is ok as this is expected (Hint can not find a solution)
            # Question: is it ok to fail in hint or is a check with assert 0 =  1 needed
            (res, rem) = divmod(abs(as_int(ids.a,PRIME)), abs(as_int(ids.b,PRIME)))
            ids.res = res*sign % PRIME
            ids.rem = rem*sign % PRIME
        %}

        #check that the reverse is correct (this proofs correctnes if rem is valid)
        assert res * b + signB*rem = a

        #check the sign helpers are -1 or 1
        assert sign*sign = 1
        assert signB*signB = 1

        # check that rem & res have same sign & that the sign is correct
        assert [range_check_ptr] = sign*rem 				#rem is zero or has same sign as sign
        assert [range_check_ptr+1] = sign*res 				#res is zero or has same sign as sign

        #check that rem is between b and 0 (exclusive b, inclusive 0)
        # Note: if hint lies about signB then signB*b is negative and because sign*rem is proofen to be positive the rangecheck fails
        # Note: will fail if b = 0 which is the expected behaviour
        assert [range_check_ptr+2] = signB*b - sign*rem -1	# abs(rem) < abs(b)
        let range_check_ptr = range_check_ptr + 3

        # inlined int_check(res) -- implied as res must be between a and 0 to pass the checks & a is valid (input)
        # assert [range_check_ptr+3] = res + \{show halfShift}		#Ensure: -HALF_SHIFT <= res
        # assert [range_check_ptr+4] = \{show (halfShift - 1)} - res	#Ensure: res < HALF_SHIFT

        return (res, range_check_ptr)
    end

"""
  where shift : Integer
        shift = pow2 nBits
        halfShift : Integer
        halfShift = pow2 (nBits `minus` 1)

div_intX : Nat -> (Name, Lazy CairoDef)
div_intX nBits = (name, def)
  where name : Name
        name = div_intX_name nBits
        def : CairoDef
        def = ForeignDef (MkForeignInfo True Nothing [range_check_builtin] empty (div_intX_code name nBits)) 2 1

export
[div_int] PrimFn where
    eval [_, ConstValue (I 0)] = Just $ Failure "Division by zero is not defined"
    eval [_, ConstValue (I 1)] = Just $ ArgValue 0
    eval [ConstValue (I a), ConstValue (I b)] = Just $ NewValue $ ConstValue $ I (a `div` b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [div_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (div_intX_name 64) args]

export
[div_int8] PrimFn where
    eval [_, ConstValue (I8 0)] = Just $ Failure "Division by zero is not defined"
    eval [_, ConstValue (I8 1)] = Just $ ArgValue 0
    eval [ConstValue (I8 a), ConstValue (I8 b)] = Just $ NewValue $ ConstValue $ I8 (a `div` b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [div_intX 8]
    generateCode res args implicits = Instructions [CALL [res] implicits (div_intX_name 8) args]

export
[div_int16] PrimFn where
    eval [_, ConstValue (I16 0)] = Just $ Failure "Division by zero is not defined"
    eval [_, ConstValue (I16 1)] = Just $ ArgValue 0
    eval [ConstValue (I16 a), ConstValue (I16 b)] = Just $ NewValue $ ConstValue $ I16 (a `div` b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [div_intX 16]
    generateCode res args implicits = Instructions [CALL [res] implicits (div_intX_name 16) args]

export
[div_int32] PrimFn where
    eval [_, ConstValue (I32 0)] = Just $ Failure "Division by zero is not defined"
    eval [_, ConstValue (I32 1)] = Just $ ArgValue 0
    eval [ConstValue (I32 a), ConstValue (I32 b)] = Just $ NewValue $ ConstValue $ I32 (a `div` b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [div_intX 32]
    generateCode res args implicits = Instructions [CALL [res] implicits (div_intX_name 32) args]

export
[div_int64] PrimFn where
    eval [_, ConstValue (I64 0)] = Just $ Failure "Division by zero is not defined"
    eval [_, ConstValue (I64 1)] = Just $ ArgValue 0
    eval [ConstValue (I64 a), ConstValue (I64 b)] = Just $ NewValue $ ConstValue $ I64 (a `div` b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [div_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (div_intX_name 64) args]

---------------------------------------------------------------------------------------------------
-- UINT MOD
---------------------------------------------------------------------------------------------------
mod_intX_name : Nat -> Name
mod_intX_name nBits = makeBuiltinName "mod_int\{show nBits}"

mod_intX_code : Name -> Nat -> String
mod_intX_code name nBits = """
    # Division between int.
    # Returns the remainder.
    # Conforms to Idris specifications: division by 0 yields error.
    func \{cairoName name}(range_check_ptr, a : felt, b : felt) -> (remainder : felt, range_check_ptr):
        alloc_locals
        local res
        local rem
        local sign
        local signB
        %{
            from starkware.cairo.common.math_utils import as_int
            signB = 1 if as_int(ids.b, PRIME) >= 0 else -1
            sign = signB if as_int(ids.a, PRIME) >= 0 else -signB
            ids.signB = signB % PRIME
            ids.sign = sign % PRIME
            #Note fails if b = 0, which is ok as this is expected (Hint can not find a solution)
            # Question: is it ok to fail in hint or is a check with assert 0 =  1 needed
            (res, rem) = divmod(abs(as_int(ids.a,PRIME)), abs(as_int(ids.b,PRIME)))
            ids.res = res*sign % PRIME
            ids.rem = rem*sign % PRIME
        %}

        #check that the reverse is correct (this proofs correctnes if rem is valid)
        assert res * b + signB*rem = a

        #check the sign helpers are -1 or 1
        assert sign*sign = 1
        assert signB*signB = 1

        # check that rem & res have same sign & that the sign is correct
        assert [range_check_ptr] = sign*rem 				#rem is zero or has same sign as sign
        assert [range_check_ptr+1] = sign*res 				#res is zero or has same sign as sign

        #check that rem is between b and 0 (exclusive b, inclusive 0)
        # Note: if hint lies about signB then signB*b is negative and because sign*rem is proofen to be positive the rangecheck fails
        # Note: will fail if b = 0 which is the expected behaviour
        assert [range_check_ptr+2] = signB*b - sign*rem -1	# abs(rem) < abs(b)
        let range_check_ptr = range_check_ptr + 3

        # inlined int_check(res) -- implied as res must be between a and 0 to pass the checks & a is valid (input)
        # assert [range_check_ptr+3] = res + \{show halfShift}		#Ensure: -HALF_SHIFT <= res
        # assert [range_check_ptr+4] = \{show (halfShift - 1)} - res	#Ensure: res < HALF_SHIFT

        return (rem, range_check_ptr)
    end

"""
  where shift : Integer
        shift = pow2 nBits
        halfShift : Integer
        halfShift = pow2 (nBits `minus` 1)

mod_intX : Nat -> (Name, Lazy CairoDef)
mod_intX nBits = (name, def)
  where name : Name
        name = mod_intX_name nBits
        def : CairoDef
        def = ForeignDef (MkForeignInfo True Nothing [range_check_builtin] empty (mod_intX_code name nBits)) 2 1

export
[mod_int] PrimFn where
    eval [_, ConstValue (I 1)] = Just $ NewValue $ ConstValue $ I 0
    eval [_, ConstValue (I 0)] = Just $ Failure "Division by zero is not defined"
    eval [ConstValue (I a), ConstValue (I b)] = Just $ NewValue $ ConstValue $ I (a `mod` b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [mod_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (mod_intX_name 64) args]

export
[mod_int8] PrimFn where
    eval [_, ConstValue (I8 1)] = Just $ NewValue $ ConstValue $ I8 0
    eval [_, ConstValue (I8 0)] = Just $ Failure "Division by zero is not defined"
    eval [ConstValue (I8 a), ConstValue (I8 b)] = Just $ NewValue $ ConstValue $ I8 (a `mod` b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [mod_intX 8]
    generateCode res args implicits = Instructions [CALL [res] implicits (mod_intX_name 8) args]

export
[mod_int16] PrimFn where
    eval [_, ConstValue (I16 1)] = Just $ NewValue $ ConstValue $ I16 0
    eval [_, ConstValue (I16 0)] = Just $ Failure "Division by zero is not defined"
    eval [ConstValue (I16 a), ConstValue (I16 b)] = Just $ NewValue $ ConstValue $ I16 (a `mod` b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [mod_intX 16]
    generateCode res args implicits = Instructions [CALL [res] implicits (mod_intX_name 16) args]

export
[mod_int32] PrimFn where
    eval [_, ConstValue (I32 1)] = Just $ NewValue $ ConstValue $ I32 0
    eval [_, ConstValue (I32 0)] = Just $ Failure "Division by zero is not defined"
    eval [ConstValue (I32 a), ConstValue (I32 b)] = Just $ NewValue $ ConstValue $ I32 (a `mod` b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [mod_intX 32]
    generateCode res args implicits = Instructions [CALL [res] implicits (mod_intX_name 32) args]

export
[mod_int64] PrimFn where
    eval [_, ConstValue (I64 1)] = Just $ NewValue $ ConstValue $ I64 0
    eval [_, ConstValue (I64 0)] = Just $ Failure "Division by zero is not defined"
    eval [ConstValue (I64 a), ConstValue (I64 b)] = Just $ NewValue $ ConstValue $ I64 (a `mod` b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [mod_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (mod_intX_name 64) args]

---------------------------------------------------------------------------------------------------
-- INT NEG:
---------------------------------------------------------------------------------------------------
neg_intX_name : Nat -> Name
neg_intX_name nBits = makeBuiltinName "neg_int\{show nBits}"

neg_intX_code : Name -> Nat -> String
neg_intX_code name nBits = """
    # Note that int_neg(-HALF_SHIFT) = -HALF_SHIFT
    func \{cairoName name}(a : felt) -> (res):
        if a == -\{show halfShift}:
            return (a)
        else:
            return (-a)
        end
    end
"""
  where halfShift : Integer
        halfShift = pow2 (nBits `minus` 1)

neg_intX : Nat -> (Name, Lazy CairoDef)
neg_intX nBits = (name, def)
  where name : Name
        name = neg_intX_name nBits
        def : CairoDef
        def = ForeignDef (MkForeignInfo True Nothing [] empty (neg_intX_code name nBits)) 1 1

export
[neg_int] PrimFn where
    eval [ConstValue (I a)] = Just $ NewValue $ ConstValue $ I (-a)
    eval _ = Nothing

    dependencies = fromList [neg_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (neg_intX_name 64) args]

export
[neg_int8] PrimFn where
    eval [ConstValue (I8 a)] = Just $ NewValue $ ConstValue $ I8 (-a)
    eval _ = Nothing

    dependencies = fromList [neg_intX 8]
    generateCode res args implicits = Instructions [CALL [res] implicits (neg_intX_name 8) args]

export
[neg_int16] PrimFn where
    eval [ConstValue (I16 a)] = Just $ NewValue $ ConstValue $ I16 (-a)
    eval _ = Nothing

    dependencies = fromList [neg_intX 16]
    generateCode res args implicits = Instructions [CALL [res] implicits (neg_intX_name 16) args]

export
[neg_int32] PrimFn where
    eval [ConstValue (I32 a)] = Just $ NewValue $ ConstValue $ I32 (-a)
    eval _ = Nothing

    dependencies = fromList [neg_intX 16]
    generateCode res args implicits = Instructions [CALL [res] implicits (neg_intX_name 16) args]

export
[neg_int64] PrimFn where
    eval [ConstValue (I64 a)] = Just $ NewValue $ ConstValue $ I64 (-a)
    eval _ = Nothing

    dependencies = fromList [neg_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (neg_intX_name 64) args]

---------------------------------------------------------------------------------------------------
-- INT SHL:
---------------------------------------------------------------------------------------------------
shl_intX_name : Nat -> Name
shl_intX_name nBits = makeBuiltinName "shl_int\{show nBits}"

shl_intX_code : Name -> Nat -> String
shl_intX_code name nBits = """
    ## Computes the logical left shift of an int.
    func \{cairoName name}(range_check_ptr, a, b) -> (res, range_check_ptr):
        # If b >= BIT_LENGTH - 1 , the result will be zero modulo 2**(BIT_LENGTH - 1).
        tempvar res
        %{ids.res = 1 if ids.b < \{show (nBits `minus` 1)} else 0 %}
        # b < 2**128 & a < 2**128 is given by semantics
        # b >= BIT_LENGTH - 1
        if res == 0:
            assert [range_check_ptr] = b - \{show (nBits `minus` 1)}
            let range_check_ptr = range_check_ptr + 1
            return (0, range_check_ptr)
        # b < BIT_LENGTH - 1
        else:
            assert [range_check_ptr] = \{show ((nBits `minus` 1) `minus` 1)} - b
            let range_check_ptr = range_check_ptr + 1
            let (c) = pow{range_check_ptr=range_check_ptr}(2, b)
            \{cairoName (mul_intX_name nBits)}(range_check_ptr, a, c)
            ret
        end
    end
"""

shl_intX : Nat -> (Name, Lazy CairoDef)
shl_intX nBits = (name, def)
  where name : Name
        name = shl_intX_name nBits
        def : CairoDef
        def = ForeignDef (MkForeignInfo True Nothing [range_check_builtin] empty (shl_intX_code name nBits)) 2 1

shl_intX_imports : SortedSet Import
shl_intX_imports = singleton (MkImport "starkware.cairo.common.pow" "pow")

export
[shl_int] PrimFn where
    apStable = False

    eval [_, ConstValue(I 0)] = Just $ ArgValue 0
    eval [ConstValue (I a), ConstValue (I b)] = Just $ NewValue $ ConstValue $ I (prim__shl_Int a b)
    eval _ = Nothing

    imports = shl_intX_imports
    implicits = singleton range_check_builtin
    dependencies = fromList [shl_intX 64, mul_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (shl_intX_name 64) args]

export
[shl_int8] PrimFn where
    apStable = False

    eval [_, ConstValue(I8 0)] = Just $ ArgValue 0
    eval [ConstValue (I8 a), ConstValue (I8 b)] = Just $ NewValue $ ConstValue $ I8 (prim__shl_Int8 (cast a) (cast b))
    eval _ = Nothing

    imports = shl_intX_imports
    implicits = singleton range_check_builtin
    dependencies = fromList [shl_intX 8, mul_intX 8]
    generateCode res args implicits = Instructions [CALL [res] implicits (shl_intX_name 8) args]

export
[shl_int16] PrimFn where
    apStable = False

    eval [_, ConstValue(I16 0)] = Just $ ArgValue 0
    eval [ConstValue (I16 a), ConstValue (I16 b)] = Just $ NewValue $ ConstValue $ I16 (prim__shl_Int16 (cast a) (cast b))
    eval _ = Nothing

    generateCode = generateMissingCodeError "shl_int16"

export
[shl_int32] PrimFn where
    apStable = False

    eval [_, ConstValue(I32 0)] = Just $ ArgValue 0
    eval [ConstValue (I32 a), ConstValue (I32 b)] = Just $ NewValue $ ConstValue $ I32 (prim__shl_Int32 (cast a) (cast b))
    eval _ = Nothing

    generateCode = generateMissingCodeError "shl_int32"

export
[shl_int64] PrimFn where
    apStable = False

    eval [_, ConstValue(I64 0)] = Just $ ArgValue 0
    eval [ConstValue (I64 a), ConstValue (I64 b)] = Just $ NewValue $ ConstValue $ I64 (prim__shl_Int64 (cast a) (cast b))
    eval _ = Nothing

    generateCode = generateMissingCodeError "shl_int64"

---------------------------------------------------------------------------------------------------
-- INT SHR:
---------------------------------------------------------------------------------------------------
shr_intX_name : Nat -> Name
shr_intX_name nBits = makeBuiltinName "shr_int\{show nBits}"

shr_intX_code : Name -> Nat -> String
shr_intX_code name nBits = """
    # Computes the logical right shift of an int.
    func \{cairoName name}(range_check_ptr, a : felt,  b : felt) -> (res : felt, range_check_ptr):
        #inlined int_pow2
        # If a >= BIT_LENGTH - 1, the result will be zero modulo 2**(BIT_LENGTH-1).
        tempvar res
        %{ids.res = 1 if ids.b < \{show (nBits `minus` 1)} else 0 %}
        # b < 2**128 & a < 2**128 is given by semantics
        # BIT_LENGTH <= b
        if res == 0:
            assert [range_check_ptr] = b - \{show (nBits `minus` 1)}
            let range_check_ptr = range_check_ptr + 1
            return (0, range_check_ptr)
        #b < BIT_LENGTH
        else:
            assert [range_check_ptr] = \{show ((nBits `minus` 1) `minus` 1)} - b
            let range_check_ptr = range_check_ptr + 1
            let (c) = pow{range_check_ptr=range_check_ptr}(2, b)
            \{cairoName (div_intX_name nBits)}(range_check_ptr, a, c)
            ret
        end
    end
"""

shr_intX : Nat -> (Name, Lazy CairoDef)
shr_intX nBits = (name, def)
  where name : Name
        name = shr_intX_name nBits
        def : CairoDef
        def = ForeignDef (MkForeignInfo True Nothing [range_check_builtin] empty (shr_intX_code name nBits)) 2 1

shr_intX_imports : SortedSet Import
shr_intX_imports = singleton (MkImport "starkware.cairo.common.pow" "pow")

export
[shr_int] PrimFn where
    apStable = False

    eval [_, ConstValue(I 0)] = Just $ ArgValue 0
    eval [ConstValue (I a), ConstValue (I b)] = Just $ NewValue $ ConstValue $ I (prim__shr_Int a b)
    eval _ = Nothing

    imports = shr_intX_imports
    implicits = singleton range_check_builtin
    dependencies = fromList [shr_intX 64, div_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (shr_intX_name 64) args]

export
[shr_int8] PrimFn where
    apStable = False

    eval [_, ConstValue(I8 0)] = Just $ ArgValue 0
    eval [ConstValue (I8 a), ConstValue (I8 b)] = Just $ NewValue $ ConstValue $ I8 (prim__shr_Int8 (cast a) (cast b))
    eval _ = Nothing

    imports = shr_intX_imports
    implicits = singleton range_check_builtin
    dependencies = fromList [shr_intX 8, div_intX 8]
    generateCode res args implicits = Instructions [CALL [res] implicits (shr_intX_name 8) args]

export
[shr_int16] PrimFn where
    apStable = False

    eval [_, ConstValue(I16 0)] = Just $ ArgValue 0
    eval [ConstValue (I16 a), ConstValue (I16 b)] = Just $ NewValue $ ConstValue $ I16 (prim__shr_Int16 (cast a) (cast b))
    eval _ = Nothing

    imports = shr_intX_imports
    implicits = singleton range_check_builtin
    dependencies = fromList [shr_intX 16, div_intX 16]
    generateCode res args implicits = Instructions [CALL [res] implicits (shr_intX_name 16) args]

export
[shr_int32] PrimFn where
    apStable = False

    eval [_, ConstValue(I32 0)] = Just $ ArgValue 0
    eval [ConstValue (I32 a), ConstValue (I32 b)] = Just $ NewValue $ ConstValue $ I32 (prim__shr_Int32 (cast a) (cast b))
    eval _ = Nothing

    imports = shr_intX_imports
    implicits = singleton range_check_builtin
    dependencies = fromList [shr_intX 32, div_intX 32]
    generateCode res args implicits = Instructions [CALL [res] implicits (shr_intX_name 32) args]

export
[shr_int64] PrimFn where
    apStable = False

    eval [_, ConstValue(I64 0)] = Just $ ArgValue 0
    eval [ConstValue (I64 a), ConstValue (I64 b)] = Just $ NewValue $ ConstValue $ I64 (prim__shr_Int64 (cast a) (cast b))
    eval _ = Nothing

    imports = shr_intX_imports
    implicits = singleton range_check_builtin
    dependencies = fromList [shr_intX 64, div_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (shr_intX_name 64) args]

---------------------------------------------------------------------------------------------------
-- INT AND:
---------------------------------------------------------------------------------------------------
and_intX_name : Nat -> Name
and_intX_name nBits = makeBuiltinName "and_int\{show nBits}"

and_intX_code : Name -> Nat -> String
and_intX_code name nBits = """
    func \{cairoName name}(range_check_ptr, bitwise_ptr, a : felt, b : felt) -> (res : felt, range_check_ptr, bitwise_ptr):
        alloc_locals
        local signA
        local signB
        %{
            from starkware.cairo.common.math_utils import as_int
            ids.signA = 1 if as_int(ids.a, PRIME) >= 0 else (-1 % PRIME)
            ids.signB = 1 if as_int(ids.b, PRIME) >= 0 else (-1 % PRIME)
        %}

        assert signA*signA = 1
        assert signB*signB = 1

        assert [range_check_ptr] = signA*a
        assert [range_check_ptr+1] = signB*b

        let a_offset = a - (signA - 1)*\{show halfShift} 	#from: ((signA - 1)/2)*SHIFT | ((signA - 1)/2) = -(signA == -1)
        let b_offset = b - (signB - 1)*\{show halfShift} 	#from: ((signB - 1)/2)*SHIFT | ((signB - 1)/2) = -(signB == -1)

        # inlined bitwise_and
        assert [bitwise_ptr] = a_offset
        assert [bitwise_ptr+1] = b_offset
        local res_value = [bitwise_ptr+2]
        let bitwise_ptr = bitwise_ptr + BitwiseBuiltin.SIZE

        # If the result is HALF_SHIFT or greater, this indicates a twos complement negative value
        # See int_le for explanation: note res_value < 2**127 (because max 64bits)
        local must_shift
        %{ids.must_shift = 1 if \{show halfShift} <= ids.res_value else 0 %}
        if must_shift == 0:
            assert [range_check_ptr+2] = \{show (halfShift - 1)} - res_value
            let range_check_ptr = range_check_ptr + 3
            return (res_value, range_check_ptr, bitwise_ptr )
        else:
            assert [range_check_ptr+2] = res_value - \{show halfShift}
            let range_check_ptr = range_check_ptr + 3
            return (res_value - \{show shift}, range_check_ptr, bitwise_ptr)
        end
    end

"""
  where shift : Integer
        shift = pow2 nBits
        halfShift : Integer
        halfShift = pow2 (nBits `minus` 1)

and_intX : Nat -> (Name, Lazy CairoDef)
and_intX nBits = (name, def)
  where name : Name
        name = and_intX_name nBits
        def : CairoDef
        def = ForeignDef (MkForeignInfo True Nothing [range_check_builtin, bitwise_builtin] empty (and_intX_code name nBits)) 2 1

export
[and_int] PrimFn where
    eval [_, ConstValue (I 0)] = Just $ NewValue $ ConstValue $ I 0
    eval [ConstValue (I 0), _] = Just $ NewValue $ ConstValue $ I 0
    eval [_, ConstValue (I (-1))] = Just $ ArgValue 0
    eval [ConstValue (I (-1)), _] = Just $ ArgValue 1
    eval [ConstValue (I a), ConstValue (I b)] = Just $ NewValue $ ConstValue $ I (prim__and_Int a b)
    eval _ = Nothing

    imports = singleton $ bitwiseBuiltinImport
    implicits = fromList [range_check_builtin, bitwise_builtin]
    dependencies = fromList [and_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (and_intX_name 64) args]

export
[and_int8] PrimFn where
    eval [_, ConstValue (I8 0)] = Just $ NewValue $ ConstValue $ I8 0
    eval [ConstValue (I8 0), _] = Just $ NewValue $ ConstValue $ I8 0
    eval [_, ConstValue (I8 (-1))] = Just $ ArgValue 0
    eval [ConstValue (I8 (-1)), _] = Just $ ArgValue 1
    eval [ConstValue (I8 a), ConstValue (I8 b)] = Just $ NewValue $ ConstValue $ I8 (prim__and_Int8 (cast a) (cast b))
    eval _ = Nothing

    imports = singleton $ bitwiseBuiltinImport
    implicits = fromList [range_check_builtin, bitwise_builtin]
    dependencies = fromList [and_intX 8]
    generateCode res args implicits = Instructions [CALL [res] implicits (and_intX_name 8) args]

export
[and_int16] PrimFn where
    eval [_, ConstValue (I16 0)] = Just $ NewValue $ ConstValue $ I16 0
    eval [ConstValue (I16 0), _] = Just $ NewValue $ ConstValue $ I16 0
    eval [_, ConstValue (I16 (-1))] = Just $ ArgValue 0
    eval [ConstValue (I16 (-1)), _] = Just $ ArgValue 1
    eval [ConstValue (I16 a), ConstValue (I16 b)] = Just $ NewValue $ ConstValue $ I16 (prim__and_Int16 (cast a) (cast b))
    eval _ = Nothing

    imports = singleton $ bitwiseBuiltinImport
    implicits = fromList [range_check_builtin, bitwise_builtin]
    dependencies = fromList [and_intX 16]
    generateCode res args implicits = Instructions [CALL [res] implicits (and_intX_name 16) args]

export
[and_int32] PrimFn where
    eval [_, ConstValue (I32 0)] = Just $ NewValue $ ConstValue $ I32 0
    eval [ConstValue (I32 0), _] = Just $ NewValue $ ConstValue $ I32 0
    eval [_, ConstValue (I32 (-1))] = Just $ ArgValue 0
    eval [ConstValue (I32 (-1)), _] = Just $ ArgValue 1
    eval [ConstValue (I32 a), ConstValue (I32 b)] = Just $ NewValue $ ConstValue $ I32 (prim__and_Int32 (cast a) (cast b))
    eval _ = Nothing

    imports = singleton $ bitwiseBuiltinImport
    implicits = fromList [range_check_builtin, bitwise_builtin]
    dependencies = fromList [and_intX 32]
    generateCode res args implicits = Instructions [CALL [res] implicits (and_intX_name 32) args]

export
[and_int64] PrimFn where
    eval [_, ConstValue (I64 0)] = Just $ NewValue $ ConstValue $ I64 0
    eval [ConstValue (I64 0), _] = Just $ NewValue $ ConstValue $ I64 0
    eval [_, ConstValue (I64 (-1))] = Just $ ArgValue 0
    eval [ConstValue (I64 (-1)), _] = Just $ ArgValue 1
    eval [ConstValue (I64 a), ConstValue (I64 b)] = Just $ NewValue $ ConstValue $ I64 (prim__and_Int64 (cast a) (cast b))
    eval _ = Nothing

    imports = singleton $ bitwiseBuiltinImport
    implicits = fromList [range_check_builtin, bitwise_builtin]
    dependencies = fromList [and_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (and_intX_name 64) args]

---------------------------------------------------------------------------------------------------
-- INT OR:
---------------------------------------------------------------------------------------------------
or_intX_name : Nat -> Name
or_intX_name nBits = makeBuiltinName "or_int\{show nBits}"

or_intX_code : Name -> Nat -> String
or_intX_code name nBits = """
    func \{cairoName name}(range_check_ptr, bitwise_ptr, a : felt, b : felt) -> (res : felt, range_check_ptr, bitwise_ptr):
        alloc_locals
        local signA
        local signB
        %{
            from starkware.cairo.common.math_utils import as_int
            ids.signA = 1 if as_int(ids.a, PRIME) >= 0 else (-1 % PRIME)
            ids.signB = 1 if as_int(ids.b, PRIME) >= 0 else (-1 % PRIME)
        %}

        assert signA*signA = 1
        assert signB*signB = 1

        assert [range_check_ptr] = signA*a
        assert [range_check_ptr+1] = signB*b

        let a_offset = a - (signA - 1)*\{show halfShift} 	#from: ((signA - 1)/2)*SHIFT | ((signA - 1)/2) = -(signA == -1)
        let b_offset = b - (signB - 1)*\{show halfShift} 	#from: ((signB - 1)/2)*SHIFT | ((signB - 1)/2) = -(signB == -1)

        # inlined bitwise_or
        assert [bitwise_ptr] = a_offset
        assert [bitwise_ptr+1] = b_offset
        local res_value = [bitwise_ptr+4]
        let bitwise_ptr = bitwise_ptr + BitwiseBuiltin.SIZE

        # If the result is HALF_SHIFT or greater, this indicates a twos complement negative value
        # See int_le for explanation: note res_value < 2**127 (because max 64bits)
        local must_shift
        %{ids.must_shift = 1 if \{show halfShift} <= ids.res_value else 0 %}
        if must_shift == 0:
            assert [range_check_ptr+2] = \{show (halfShift - 1)} - res_value # DK: check - is left associative
            let range_check_ptr = range_check_ptr + 3
            return (res_value, range_check_ptr, bitwise_ptr )
        else:
            assert [range_check_ptr+2] = res_value - \{show halfShift}		#-0 is for ap stability improves worse case in code gen
            let range_check_ptr = range_check_ptr + 3
            return (res_value - \{show shift}, range_check_ptr, bitwise_ptr)
        end
    end

"""
  where shift : Integer
        shift = pow2 nBits
        halfShift : Integer
        halfShift = pow2 (nBits `minus` 1)

or_intX : Nat -> (Name, Lazy CairoDef)
or_intX nBits = (name, def)
  where name : Name
        name = or_intX_name nBits
        def : CairoDef
        def = ForeignDef (MkForeignInfo True Nothing [range_check_builtin, bitwise_builtin] empty (or_intX_code name nBits)) 2 1

export
[or_int] PrimFn where
    eval [_, ConstValue (I 0)] = Just $ ArgValue 0
    eval [ConstValue (I 0), _] = Just $ ArgValue 1
    eval [_, ConstValue (I  (-1))] = Just $ NewValue $ ConstValue $ I (-1)
    eval [ConstValue (I  (-1)), _] = Just $ NewValue $ ConstValue $ I (-1)
    eval [ConstValue (I a), ConstValue (I b)] = Just $ NewValue $ ConstValue $ I (prim__or_Int a b)
    eval _ = Nothing

    imports = singleton $ bitwiseBuiltinImport
    implicits = fromList [range_check_builtin, bitwise_builtin]
    dependencies = fromList [or_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (or_intX_name 64) args]

export
[or_int8] PrimFn where
    eval [_, ConstValue (I8 0)] = Just $ ArgValue 0
    eval [ConstValue (I8 0), _] = Just $ ArgValue 1
    eval [_, ConstValue (I8 (-1))] = Just $ NewValue $ ConstValue $ I8 (-1)
    eval [ConstValue (I8 (-1)), _] = Just $ NewValue $ ConstValue $ I8 (-1)
    eval [ConstValue (I8 a), ConstValue (I8 b)] = Just $ NewValue $ ConstValue $ I8 (prim__or_Int8 (cast a) (cast b))
    eval _ = Nothing

    imports = singleton $ bitwiseBuiltinImport
    implicits = fromList [range_check_builtin, bitwise_builtin]
    dependencies = fromList [or_intX 8]
    generateCode res args implicits = Instructions [CALL [res] implicits (or_intX_name 8) args]

export
[or_int16] PrimFn where
    eval [_, ConstValue (I16 0)] = Just $ ArgValue 0
    eval [ConstValue (I16 0), _] = Just $ ArgValue 1
    eval [_, ConstValue (I16 (-1))] = Just $ NewValue $ ConstValue $ I16 (-1)
    eval [ConstValue (I16 (-1)), _] = Just $ NewValue $ ConstValue $ I16 (-1)
    eval [ConstValue (I16 a), ConstValue (I16 b)] = Just $ NewValue $ ConstValue $ I16 (prim__or_Int16 (cast a) (cast b))
    eval _ = Nothing

    imports = singleton $ bitwiseBuiltinImport
    implicits = fromList [range_check_builtin, bitwise_builtin]
    dependencies = fromList [or_intX 16]
    generateCode res args implicits = Instructions [CALL [res] implicits (or_intX_name 16) args]

export
[or_int32] PrimFn where
    eval [_, ConstValue (I32 0)] = Just $ ArgValue 0
    eval [ConstValue (I32 0), _] = Just $ ArgValue 1
    eval [_, ConstValue (I32 (-1))] = Just $ NewValue $ ConstValue $ I32 (-1)
    eval [ConstValue (I32 (-1)), _] = Just $ NewValue $ ConstValue $ I32 (-1)
    eval [ConstValue (I32 a), ConstValue (I32 b)] = Just $ NewValue $ ConstValue $ I32 (prim__or_Int32 (cast a) (cast b))
    eval _ = Nothing

    imports = singleton $ bitwiseBuiltinImport
    implicits = fromList [range_check_builtin, bitwise_builtin]
    dependencies = fromList [or_intX 32]
    generateCode res args implicits = Instructions [CALL [res] implicits (or_intX_name 32) args]

export
[or_int64] PrimFn where
    eval [_, ConstValue (I64 0)] = Just $ ArgValue 0
    eval [ConstValue (I64 0), _] = Just $ ArgValue 1
    eval [_, ConstValue (I64 (-1))] = Just $ NewValue $ ConstValue $ I64 (-1)
    eval [ConstValue (I64 (-1)), _] = Just $ NewValue $ ConstValue $ I64 (-1)
    eval [ConstValue (I64 a), ConstValue (I64 b)] = Just $ NewValue $ ConstValue $ I64 (prim__or_Int64 (cast a) (cast b))
    eval _ = Nothing

    imports = singleton $ bitwiseBuiltinImport
    implicits = fromList [range_check_builtin, bitwise_builtin]
    dependencies = fromList [or_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (or_intX_name 64) args]

---------------------------------------------------------------------------------------------------
-- INT XOR:
---------------------------------------------------------------------------------------------------
xor_intX_name : Nat -> Name
xor_intX_name nBits = makeBuiltinName "xor_int\{show nBits}"

xor_intX_code : Name -> Nat -> String
xor_intX_code name nBits = """
    func \{cairoName name}(range_check_ptr, bitwise_ptr, a : felt, b : felt) -> (res : felt, range_check_ptr, bitwise_ptr):
        alloc_locals
        let a_offset = a + \{show halfShift}  # guarantee positive number, since minimum input value is DEFAULT_PRIME - SHIFT.  e.g. 0 maps to 128, 127 maps to 255, and -1 maps to 127.  This gets applied to _both_ inputs, so intuitively, XOR doesn't notice or care.
        let b_offset = b + \{show halfShift}

        # inlined bitwise_xor
        assert [bitwise_ptr] = a_offset
        assert [bitwise_ptr+1] = b_offset
        local res_value = [bitwise_ptr+3]
        let bitwise_ptr = bitwise_ptr + BitwiseBuiltin.SIZE

        # If the result is HALF_SHIFT of greater, this indicates a twos complement negative value
        # See int_le for explanation: note res_value < 2**127 (because max 64bits)
        local must_shift
        %{ids.must_shift = 1 if \{show halfShift} <= ids.res_value else 0 %}
        if must_shift == 0:
            assert [range_check_ptr] = \{show (halfShift-1)} - res_value
            let range_check_ptr = range_check_ptr + 1
            return (res_value, range_check_ptr, bitwise_ptr)
        else:
            assert [range_check_ptr] = res_value - \{show halfShift}
            let range_check_ptr = range_check_ptr + 1
            return (res_value - \{show shift}, range_check_ptr, bitwise_ptr)
        end
    end

"""
  where shift : Integer
        shift = pow2 nBits
        halfShift : Integer
        halfShift = pow2 (nBits `minus` 1)

xor_intX : Nat -> (Name, Lazy CairoDef)
xor_intX nBits = (name, def)
  where name : Name
        name = xor_intX_name nBits
        def : CairoDef
        def = ForeignDef (MkForeignInfo True Nothing [range_check_builtin, bitwise_builtin] empty (xor_intX_code name nBits)) 2 1

export
[xor_int] PrimFn where
    eval [_, ConstValue (I 0)] = Just $ ArgValue 0
    eval [ConstValue (I 0), _] = Just $ ArgValue 1
    eval [ConstValue (I a), ConstValue (I b)] = Just $ NewValue $ ConstValue $ I (prim__xor_Int a b)
    eval _ = Nothing

    imports = singleton $ bitwiseBuiltinImport
    implicits = fromList [range_check_builtin, bitwise_builtin]
    dependencies = fromList [xor_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (xor_intX_name 64) args]

export
[xor_int8] PrimFn where
    eval [_, ConstValue (I8 0)] = Just $ ArgValue 0
    eval [ConstValue (I8 0), _] = Just $ ArgValue 1
    eval [ConstValue (I8 a), ConstValue (I8 b)] = Just $ NewValue $ ConstValue $ I8 (prim__xor_Int8 (cast a) (cast b))
    eval _ = Nothing

    imports = singleton $ bitwiseBuiltinImport
    implicits = fromList [range_check_builtin, bitwise_builtin]
    dependencies = fromList [xor_intX 8]
    generateCode res args implicits = Instructions [CALL [res] implicits (xor_intX_name 8) args]

export
[xor_int16] PrimFn where
    eval [_, ConstValue (I16 0)] = Just $ ArgValue 0
    eval [ConstValue (I16 0), _] = Just $ ArgValue 1
    eval [ConstValue (I16 a), ConstValue (I16 b)] = Just $ NewValue $ ConstValue $ I16 (prim__xor_Int16 (cast a) (cast b))
    eval _ = Nothing

    imports = singleton $ bitwiseBuiltinImport
    implicits = fromList [range_check_builtin, bitwise_builtin]
    dependencies = fromList [xor_intX 16]
    generateCode res args implicits = Instructions [CALL [res] implicits (xor_intX_name 16) args]

export
[xor_int32] PrimFn where
    eval [_, ConstValue (I32 0)] = Just $ ArgValue 0
    eval [ConstValue (I32 0), _] = Just $ ArgValue 1
    eval [ConstValue (I32 a), ConstValue (I32 b)] = Just $ NewValue $ ConstValue $ I32 (prim__xor_Int32 (cast a) (cast b))
    eval _ = Nothing

    imports = singleton $ bitwiseBuiltinImport
    implicits = fromList [range_check_builtin, bitwise_builtin]
    dependencies = fromList [xor_intX 32]
    generateCode res args implicits = Instructions [CALL [res] implicits (xor_intX_name 32) args]

export
[xor_int64] PrimFn where
    eval [_, ConstValue (I64 0)] = Just $ ArgValue 0
    eval [ConstValue (I64 0), _] = Just $ ArgValue 1
    eval [ConstValue (I64 a), ConstValue (I64 b)] = Just $ NewValue $ ConstValue $ I64 (prim__xor_Int64 (cast a) (cast b))
    eval _ = Nothing

    imports = singleton $ bitwiseBuiltinImport
    implicits = fromList [range_check_builtin, bitwise_builtin]
    dependencies = fromList [xor_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (xor_intX_name 64) args]

---------------------------------------------------------------------------------------------------
-- INT LT
---------------------------------------------------------------------------------------------------
lt_intX_name : Name
lt_intX_name = makeBuiltinName "lt_int"

lt_intX_code : Name -> String
lt_intX_code name = """
    # Returns 1 if the first integer is less than the second, otherwise returns 0.
    func \{cairoName name}(range_check_ptr, a : felt, b : felt) -> (res, range_check_ptr):
        tempvar res
        %{
            from starkware.cairo.common.math_utils import as_int
            ids.res = 1 if as_int(ids.a,PRIME) < as_int(ids.b,PRIME) else 0
        %}
        # -2**127 <= b < 2**127 & -2**127 <= a < 2**127 is given by semantics (actuall even tighter as max 64bits)
        # implies: a - b < 2**128 & b - a < 2**128 (meaning in range check upper bound)
        # b <= a
        if res == 0:
            assert [range_check_ptr] = a - b - 0 #+0 is for ap stability improves worse case in code gen
            let range_check_ptr = range_check_ptr + 1
            return (0, range_check_ptr)
        # a < b
        else:
            assert [range_check_ptr] = b - a - 1
            let range_check_ptr = range_check_ptr + 1
            return (1, range_check_ptr)
        end
    end
"""

lt_intX : (Name, Lazy CairoDef)
lt_intX = (lt_intX_name, def)
  where def : CairoDef
        def = ForeignDef (MkForeignInfo True Nothing [range_check_builtin] empty (lt_intX_code lt_intX_name)) 2 1

generateLtCallCode : CairoReg -> List CairoReg -> LinearImplicitArgs -> PrimFnCode
generateLtCallCode res args implicits = Instructions [CALL [res] implicits lt_intX_name args]

export
[lt_int] PrimFn where
    eval [ConstValue (I a), ConstValue (I b)] = Just $ NewValue $ ConstValue $ toInt (a<b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lt_intX]
    generateCode = generateLtCallCode

export
[lt_int8] PrimFn where
    eval [ConstValue (I8 a), ConstValue (I8 b)] = Just $ NewValue $ ConstValue $ toInt (a<b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lt_intX]
    generateCode = generateLtCallCode

export
[lt_int16] PrimFn where
    eval [ConstValue (I16 a), ConstValue (I16 b)] = Just $ NewValue $ ConstValue $ toInt (a<b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lt_intX]
    generateCode = generateLtCallCode

export
[lt_int32] PrimFn where
    eval [ConstValue (I32 a), ConstValue (I32 b)] = Just $ NewValue $ ConstValue $ toInt (a<b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lt_intX]
    generateCode = generateLtCallCode

export
[lt_int64] PrimFn where
    eval [ConstValue (I64 a), ConstValue (I64 b)] = Just $ NewValue $ ConstValue $ toInt (a<b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lt_intX]
    generateCode = generateLtCallCode

---------------------------------------------------------------------------------------------------
-- INT LTE
---------------------------------------------------------------------------------------------------
lte_intX_name : Name
lte_intX_name = makeBuiltinName "lte_int"

lte_intX_code : Name -> String
lte_intX_code name = """
    # Returns 1 if the first integer is less than the second, otherwise returns 0.
    func \{cairoName name}(range_check_ptr, a : felt, b : felt) -> (res, range_check_ptr):
        tempvar res
        %{
            from starkware.cairo.common.math_utils import as_int
            ids.res = 1 if as_int(ids.a,PRIME) < as_int(ids.b,PRIME) else 0
        %}
        # -2**127 <= b < 2**127 & -2**127 <= a < 2**127 is given by semantics (actuall even tighter as max 64bits)
        # implies: a - b < 2**128 & b - a < 2**128 (meaning in range check upper bound)
        # b <= a
        if res == 0:
            assert [range_check_ptr] = a - b - 0 #+0 is for ap stability improves worse case in code gen
            let range_check_ptr = range_check_ptr + 1
            return (0, range_check_ptr)
        # a < b
        else:
            assert [range_check_ptr] = b - a - 1
            let range_check_ptr = range_check_ptr + 1
            return (1, range_check_ptr)
        end
    end
"""

lte_intX : (Name, Lazy CairoDef)
lte_intX = (lte_intX_name, def)
  where def : CairoDef
        def = ForeignDef (MkForeignInfo True Nothing [range_check_builtin] empty (lte_intX_code lte_intX_name)) 2 1

generateLteCallCode : CairoReg -> List CairoReg -> LinearImplicitArgs -> PrimFnCode
generateLteCallCode res args implicits = Instructions [CALL [res] implicits lte_intX_name args]

export
[lte_int] PrimFn where
    eval [ConstValue (I a), ConstValue (I b)] = Just $ NewValue $ ConstValue $ toInt (a<=b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lte_intX]
    generateCode = generateLtCallCode

export
[lte_int8] PrimFn where
    eval [ConstValue (I8 a), ConstValue (I8 b)] = Just $ NewValue $ ConstValue $ toInt (a<=b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lte_intX]
    generateCode = generateLtCallCode

export
[lte_int16] PrimFn where
    eval [ConstValue (I16 a), ConstValue (I16 b)] = Just $ NewValue $ ConstValue $ toInt (a<=b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lte_intX]
    generateCode = generateLtCallCode

export
[lte_int32] PrimFn where
    eval [ConstValue (I32 a), ConstValue (I32 b)] = Just $ NewValue $ ConstValue $ toInt (a<=b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lte_intX]
    generateCode = generateLtCallCode

export
[lte_int64] PrimFn where
    eval [ConstValue (I64 a), ConstValue (I64 b)] = Just $ NewValue $ ConstValue $ toInt (a<=b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lte_intX]
    generateCode = generateLtCallCode

---------------------------------------------------------------------------------------------------
-- INT EQ
---------------------------------------------------------------------------------------------------
export
[eq_int] PrimFn where
    eval [ConstValue (I a), ConstValue (I b)] = Just $ NewValue $ ConstValue $ toInt (a==b)
    eval _ = Nothing

    generateCode r a i = Raw $ compileEqOp "eq_int" r a i

export
[eq_int8] PrimFn where
    eval [ConstValue (I8 a), ConstValue (I8 b)] = Just $ NewValue $ ConstValue $ toInt (a==b)
    eval _ = Nothing

    generateCode r a i = Raw $ compileEqOp "eq_int8" r a i

export
[eq_int16] PrimFn where
    eval [ConstValue (I16 a), ConstValue (I16 b)] = Just $ NewValue $ ConstValue $ toInt (a==b)
    eval _ = Nothing

    generateCode r a i = Raw $ compileEqOp "eq_int16" r a i

export
[eq_int32] PrimFn where
    eval [ConstValue (I32 a), ConstValue (I32 b)] = Just $ NewValue $ ConstValue $ toInt (a==b)
    eval _ = Nothing

    generateCode r a i = Raw $ compileEqOp "eq_int32" r a i

export
[eq_int64] PrimFn where
    eval [ConstValue (I64 a), ConstValue (I64 b)] = Just $ NewValue $ ConstValue $ toInt (a==b)
    eval _ = Nothing

    generateCode r a i = Raw $ compileEqOp "eq_int64" r a i

---------------------------------------------------------------------------------------------------
-- INT GTE
-- Implementation uses flipped arguments to lte
---------------------------------------------------------------------------------------------------

generateGteCallCode : CairoReg -> List CairoReg -> LinearImplicitArgs -> PrimFnCode
generateGteCallCode res [a,b] implicits = Instructions [CALL [res] implicits lte_intX_name [b,a]] -- Flipped arguments
generateGteCallCode _ _ _ = assert_total $ idris_crash "Bad arguments to generateCode gte_intX"

export
[gte_int] PrimFn where
    eval [ConstValue (I a), ConstValue (I b)] = Just $ NewValue $ ConstValue $ toInt (a>=b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lte_intX]
    generateCode = generateGteCallCode

export
[gte_int8] PrimFn where
    eval [ConstValue (I8 a), ConstValue (I8 b)] = Just $ NewValue $ ConstValue $ toInt (a>=b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lte_intX]
    generateCode = generateGteCallCode

export
[gte_int16] PrimFn where
    eval [ConstValue (I16 a), ConstValue (I16 b)] = Just $ NewValue $ ConstValue $ toInt (a>=b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lte_intX]
    generateCode = generateGteCallCode

export
[gte_int32] PrimFn where
    eval [ConstValue (I32 a), ConstValue (I32 b)] = Just $ NewValue $ ConstValue $ toInt (a>=b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lte_intX]
    generateCode = generateGteCallCode

export
[gte_int64] PrimFn where
    eval [ConstValue (I64 a), ConstValue (I64 b)] = Just $ NewValue $ ConstValue $ toInt (a>=b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lte_intX]
    generateCode = generateGteCallCode

---------------------------------------------------------------------------------------------------
-- INT GT
-- Implementation uses flipped arguments to lte
---------------------------------------------------------------------------------------------------

generateGtCallCode : CairoReg -> List CairoReg -> LinearImplicitArgs -> PrimFnCode
generateGtCallCode res [a,b] implicits = Instructions [CALL [res] implicits lt_intX_name [b,a]] -- Flipped arguments
generateGtCallCode _ _ _ = assert_total $ idris_crash "Bad arguments to generateCode gt_intX"

export
[gt_int] PrimFn where
    eval [ConstValue (I a), ConstValue (I b)] = Just $ NewValue $ ConstValue $ toInt (a>b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lt_intX]
    generateCode = generateGtCallCode

export
[gt_int8] PrimFn where
    eval [ConstValue (I8 a), ConstValue (I8 b)] = Just $ NewValue $ ConstValue $ toInt (a>b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lt_intX]
    generateCode = generateGtCallCode

export
[gt_int16] PrimFn where
    eval [ConstValue (I16 a), ConstValue (I16 b)] = Just $ NewValue $ ConstValue $ toInt (a>b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lt_intX]
    generateCode = generateGtCallCode

export
[gt_int32] PrimFn where
    eval [ConstValue (I32 a), ConstValue (I32 b)] = Just $ NewValue $ ConstValue $ toInt (a>b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lt_intX]
    generateCode = generateGtCallCode

export
[gt_int64] PrimFn where
    eval [ConstValue (I64 a), ConstValue (I64 b)] = Just $ NewValue $ ConstValue $ toInt (a>b)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [lt_intX]
    generateCode = generateGtCallCode

---------------------------------------------------------------------------------------------------
-- INT CASTS
---------------------------------------------------------------------------------------------------
cast_intX_name : Nat -> Name
cast_intX_name nBits = makeBuiltinName "cast_int\{show nBits}"

cast_intX_code : Name -> Nat -> String
cast_intX_code name nBits = """
    # Casts to int\{show nBits}
    # Returns the result as int\{show nBits}.
    func \{cairoName name}(range_check_ptr, a : felt) -> (res : felt, range_check_ptr):
        alloc_locals
        local res
        local offset
        %{
            from starkware.cairo.common.math_utils import as_int
            rawRes = (as_int(ids.a, PRIME) + \{show halfShift})
            (rawOffset, resRaw) = divmod(rawRes, \{show shift})
            ids.offset = rawOffset % PRIME
            ids.res = (rawRes - \{show halfShift}) % PRIME
        %}

        # Check offset can not trigger overflow / underflow
        #  Note: Max valid offset absolute = HALF_SHIFT*HALF_SHIFT/SHIFT = (SHIFT*SHIFT)/4*SHIFT = SHIFT/4
        #  For simplicity we weaken to -HALF_SHIFT <= offset <= HALF_SHIFT then no overflow (for 64bit): -2**128 <= +/-(2**63 * 2**63) +/- (2**63 * 2**64) <= +2**128 & overflow is at ~+/-PRIME/2 (with PRIME ~= 2**192)
        assert [range_check_ptr] = offset + \{show halfShift}		#Ensure: -HALF_SHIFT <= offset
        assert [range_check_ptr+1] = \{show halfShift} - offset	    #Ensure: offset <= HALF_SHIFT

        assert res = a - offset * \{show shift}

        # inlined int_check(res)
        assert [range_check_ptr+2] = res + \{show halfShift}		#Ensure: -HALF_SHIFT <= res
        assert [range_check_ptr+3] = \{show (halfShift - 1)} - res	#Ensure: res < HALF_SHIFT
        let range_check_ptr = range_check_ptr + 4

        return (res, range_check_ptr)
    end

"""
  where shift : Integer
        shift = pow2 nBits
        halfShift : Integer
        halfShift = pow2 (nBits `minus` 1)

cast_intX : Nat -> (Name, Lazy CairoDef)
cast_intX nBits = (name, def)
  where name : Name
        name = cast_intX_name nBits
        def : CairoDef
        def = ForeignDef (MkForeignInfo True Nothing [range_check_builtin] empty (cast_intX_code name nBits)) 1 1

export
[cast_to_int] PrimFn where
    eval [ConstValue c] = map (\i => NewValue $ ConstValue $ I $ cast i) (constToInteger c)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [cast_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (cast_intX_name 64) args]

export
[cast_to_int8] PrimFn where
    eval [ConstValue c] = map (\i => NewValue $ ConstValue $ I8 $ cast i) (constToInteger c)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [cast_intX 8]
    generateCode res args implicits = Instructions [CALL [res] implicits (cast_intX_name 8) args]

export
[cast_to_int16] PrimFn where
    eval [ConstValue c] = map (\i => NewValue $ ConstValue $ I16 $ cast i) (constToInteger c)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [cast_intX 16]
    generateCode res args implicits = Instructions [CALL [res] implicits (cast_intX_name 16) args]

export
[cast_to_int32] PrimFn where
    eval [ConstValue c] = map (\i => NewValue $ ConstValue $ I32 $ cast i) (constToInteger c)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [cast_intX 32]
    generateCode res args implicits = Instructions [CALL [res] implicits (cast_intX_name 32) args]

export
[cast_to_int64] PrimFn where
    eval [ConstValue c] = map (\i => NewValue $ ConstValue $ I64 $ cast i) (constToInteger c)
    eval _ = Nothing

    implicits = singleton range_check_builtin
    dependencies = fromList [cast_intX 64]
    generateCode res args implicits = Instructions [CALL [res] implicits (cast_intX_name 64) args]


foo : Either (Either () ()) () -> String
foo (Left (Left ())) = "Left Left"
foo (Left (Right ())) = "Left Right"
foo (Right ()) = "Right"

bah : List Nat -> Nat
bah [] = 0
bah (x :: xs) = x + bah xs

baz : Int -> Int
baz 10 = -1
baz _ = -2

main : IO ()
main = do
    -- primitives
    printLn $ the Int $ 10 + 11 * 92
    printLn $ the Bits32 $ 8 * 9 * 10
    printLn $ the Integer $ 19 * 19

    -- pattern matching
    putStrLn $ foo $ Left (Left ())
    putStrLn $ foo $ Right ()
    putStrLn $ foo $ Left (Right ())

    -- recursion
    printLn $ bah [0 .. 10]
    printLn $ bah [11, 2]

    -- primitive pattern matching
    printLn $ baz 10
    printLn $ baz (-100)


module PLFI.Part1.Isomorphism

import Data.Nat
import Data.Vect
import Syntax.PreorderReasoning
import Syntax.PreorderReasoning.Generic
import Control.Order
import Control.Relation

0
handwaving : {t : Type} -> (a : t) -> (b : t) -> (a = b)
handwaving a b = believe_me (the (a=a) Refl)

namespace Postulate

  export 0
  extensionality : {a,b : Type} -> {f,g : a -> b} -> ((x : a) -> f x = g x) -> f = g
  extensionality _ = handwaving f g

addX : (m,n : Nat) -> Nat
addX m Z = m
addX m (S n) = S (addX m n)

sameAdd : (m, n : Nat) -> addX m n = m + n
sameAdd m n = rewrite plusCommutative m n in helper m n
  where
    helper : (m, n : Nat) -> addX m n = n + m
    helper m 0     = Refl
    helper m (S n) = cong S (helper m n)

add : Nat -> Nat -> Nat
add = (+)

0
same : Isomorphism.addX = Isomorphism.add
same = extensionality (\p => extensionality (\q => sameAdd p q))

namespace Things

  interesting1 : (x : Nat) -> (y : Nat) -> Nat
  interesting1 x y = x + y

  interesting2 : Nat
  interesting2 = interesting1 {y=2} {x=1}

-- infix 0 _≃_
-- record _≃_ (A B : Set) : Set where
--   field
--     to   : A → B
--     from : B → A
--     from∘to : ∀ (x : A) → from (to x) ≡ x
--     to∘from : ∀ (y : B) → to (from y) ≡ y
-- open _≃_

public export
record Iso (a, b : Type) where
  constructor MkIso
  to     : a -> b
  from   : b -> a
  fromTo : (x : a) -> from (to x) === x
  toFrom : (y : b) -> to (from y) === y

-- interface IsoI a b where
--   toI   : a -> b
--   fromI : b -> a
--   fromToI : (x : a) -> fromI (toI x) = x
--   toFromI : (y : b) -> toI (fromI y) = y

-- ≃-refl : ∀ {A : Set}
--     -----
--   → A ≃ A
-- ≃-refl =
--   record
--     { to      = λ{x → x}
--     ; from    = λ{y → y}
--     ; from∘to = λ{x → refl}
--     ; to∘from = λ{y → refl}
--     }

public export
isoRefl : {a : Type} -> Iso a a
isoRefl = MkIso
  { to     = id
  , from   = id
  , fromTo = \x => Refl
  , toFrom = \x => Refl
  }

isoTo : Iso a b -> (a -> b)
isoTo i = i.to

-- ≃-sym : ∀ {A B : Set}
--   → A ≃ B
--     -----
--   → B ≃ A
-- ≃-sym A≃B =
--   record
--     { to      = from A≃B
--     ; from    = to   A≃B
--     ; from∘to = to∘from A≃B
--     ; to∘from = from∘to A≃B
--     }

public export
isoSym : Iso a b -> Iso b a
isoSym (MkIso to from fromTo toFrom) = MkIso from to toFrom fromTo

-- ≃-trans : ∀ {A B C : Set}
--   → A ≃ B
--   → B ≃ C
--     -----
--   → A ≃ C
-- ≃-trans A≃B B≃C =
--   record
--     { to       = to   B≃C ∘ to   A≃B
--     ; from     = from A≃B ∘ from B≃C
--     ; from∘to  = λ{x →
--         begin
--           (from A≃B ∘ from B≃C) ((to B≃C ∘ to A≃B) x)
--         ≡⟨⟩
--           from A≃B (from B≃C (to B≃C (to A≃B x)))
--         ≡⟨ cong (from A≃B) (from∘to B≃C (to A≃B x)) ⟩
--           from A≃B (to A≃B x)
--         ≡⟨ from∘to A≃B x ⟩
--           x
--         ∎}
--     ; to∘from = λ{y →
--         begin
--           (to B≃C ∘ to A≃B) ((from A≃B ∘ from B≃C) y)
--         ≡⟨⟩
--           to B≃C (to A≃B (from A≃B (from B≃C y)))
--         ≡⟨ cong (to B≃C) (to∘from A≃B (from B≃C y)) ⟩
--           to B≃C (from B≃C y)
--         ≡⟨ to∘from B≃C y ⟩
--           y
--         ∎}
--      }

public export
isoTrans : Iso a b -> Iso b c -> Iso a c
isoTrans ab bc = MkIso
  { to     = to bc . to ab
  , from   = from ab . from bc
  , fromTo = \x => Calc $
      |~ (from ab . from bc) ((to bc . to ab) x)
      -- ~~ (from ab . from bc) (to bc (to ab x))  ... (Refl)
      -- ~~ (from ab (from bc (to bc (to ab x))))  ... (Refl)
      ~~ (from ab (to ab x))                    ... (cong ab.from (bc.fromTo (ab.to x)))
      ~~ x                                      ... (ab.fromTo x)
  , toFrom = \y => rewrite ab.toFrom (bc.from y) in bc.toFrom y
--  , toFrom = \y => (?x (bc.toFrom y))
                    -- rewriteP (ab.toFrom (bc.from y)) (bc.toFrom y)

                    -- rewriteP : a = b -> P a -> P b
                    -- rewriteP Refl pa = pa
  }

-- module ≃-Reasoning where

--   infix  1 ≃-begin_
--   infixr 2 _≃⟨_⟩_
--   infix  3 _≃-∎

--   ≃-begin_ : ∀ {A B : Set}
--     → A ≃ B
--       -----
--     → A ≃ B
--   ≃-begin A≃B = A≃B

--   _≃⟨_⟩_ : ∀ (A : Set) {B C : Set}
--     → A ≃ B
--     → B ≃ C
--       -----
--     → A ≃ C
--   A ≃⟨ A≃B ⟩ B≃C = ≃-trans A≃B B≃C

--   _≃-∎ : ∀ (A : Set)
--       -----
--     → A ≃ A
--   A ≃-∎ = ≃-refl

-- open ≃-Reasoning

public export Reflexive  Type Iso where reflexive = isoRefl
public export Transitive Type Iso where transitive = isoTrans
public export Preorder   Type Iso where

public export
(~~) : (a, b : Type) -> Type
(~~) a b = Iso a b

-- 0
test : {a,b,c : Type} -> Iso a b -> Iso b c -> Iso a c
test ab bc = CalcWith {leq = Iso} $
  |~ a
  <~ b ... (ab)
  <~ c ... (bc)

-- [additive] Semigroup Nat where
--   a <+> b = ?h

-- [multiplicative] Semigroup Nat where
--   a <+> b = ?h2

-- infix 0 _≲_
-- record _≲_ (A B : Set) : Set where
--   field
--     to      : A → B
--     from    : B → A
--     from∘to : ∀ (x : A) → from (to x) ≡ x
-- open _≲_

-- infixr 100 <=

-- %hint -- Would be nice to have
public export
record (<=) (a, b : Type) where
  constructor MkEmb
  to     : a -> b
  from   : b -> a
  -- %hint : Would be nice to have
  fromTo : (x : a) -> from (to x) === x

%hint
fromToX : (ab : (a <= b)) -> (x : a) -> ab.from (ab.to x) === x
fromToX ab x = ab.fromTo x

-- %hint (.to)

-- (<=) : (a,b : Type) -> Type
-- (<=) a b = Emb a b

-- -- A deep question from the audience. :)
-- embIso : {a,b : Type} -> a <= b -> b <= a -> (Iso a b)
-- embIso = ?xembIso

embRefl : {a : Type} -> a <= a
embRefl = MkEmb Prelude.id Prelude.id (\x => Refl)

-- ≲-trans : ∀ {A B C : Set} → A ≲ B → B ≲ C → A ≲ C
-- ≲-trans A≲B B≲C =
--   record
--     { to      = λ{x → to   B≲C (to   A≲B x)}
--     ; from    = λ{y → from A≲B (from B≲C y)}
--     ; from∘to = λ{x →
--         begin
--           from A≲B (from B≲C (to B≲C (to A≲B x)))
--         ≡⟨ cong (from A≲B) (from∘to B≲C (to A≲B x)) ⟩
--           from A≲B (to A≲B x)
--         ≡⟨ from∘to A≲B x ⟩
--           x
--         ∎}
--      }

embTrans : {a,b,c : Type} -> a <= b -> b <= c -> a <= c
embTrans ab bc = MkEmb
  { to   = (bc.to . ab.to)
  , from = (ab.from . bc.from)
  , fromTo = \x => Calc $
      |~ ab.from (bc.from (bc.to (ab.to x)))
      ~~ ab.from (ab.to x)
          ... (cong ab.from (bc.fromTo (ab.to x)))
      ~~ x
          ... (ab.fromTo x)
--          ... (fromToX ab x)
  }

-- ≲-antisym : ∀ {A B : Set}
--   → (A≲B : A ≲ B)
--   → (B≲A : B ≲ A)
--   → (to A≲B ≡ from B≲A)
--   → (from A≲B ≡ to B≲A)
--     -------------------
--   → A ≃ B
-- ≲-antisym A≲B B≲A to≡from from≡to =
--   record
--     { to      = to A≲B
--     ; from    = from A≲B
--     ; from∘to = from∘to A≲B
--     ; to∘from = λ{y →
--         begin
--           to A≲B (from A≲B y)
--         ≡⟨ cong (to A≲B) (cong-app from≡to y) ⟩
--           to A≲B (to B≲A y)
--         ≡⟨ cong-app to≡from (to B≲A y) ⟩
--           from B≲A (to B≲A y)
--         ≡⟨ from∘to B≲A y ⟩
--           y
--         ∎}
--     }

embAntisym
  :  {a,b : Type}
  -> (ab : (a <= b)) -> (ba : (b <= a))
  -> (to ab = from ba) -> (from ab = to ba)
  -> Iso a b
embAntisym (MkEmb t f ft) (MkEmb _ _ ft1) Refl Refl = MkIso t f ft ft1

Reflexive  Type (<=) where reflexive = embRefl
Transitive Type (<=) where transitive = embTrans
Preorder   Type (<=) where

embFromIso : {a,b : Type} -> Iso a b -> (a <= b)
-- Issue in internel meta variable
-- embFromIso (MkIso to from fromTo toFrom) = MkEmb (\{arg:2022} => to arg) from fromTo
embFromIso (MkIso to from fromTo toFrom) = MkEmb to from fromTo


import Stuff

data IsSuc : Nat -> Type where
     Yes : IsSuc (S x)

pred : (x : Nat) -> IsSuc x -> Nat
pred (S x) Yes = x

test1 : Nat -> Nat
test1 x = let y = S (S x) in
              case x of
                   res => pred y ?foo

-- Propagating let binding's computational behaviour through case is
-- not supported, sorry!
{-
data Vect : Nat -> Type -> Type where
     Nil : Vect Z a
     (::) : a -> Vect k a -> Vect (S k) a

test2 : Nat -> Nat -> Nat
test2 x@(S Z) y
    = case y of
           res =>
               let fn : Vect (S Z) Nat -> Nat
                   fn ys = ?bar
               in pred x Yes
-}


module Data.Crypto.Encryption.Block.EncryptionMode

import Data.Bits

import Data.Crypto.Encryption
import Data.Crypto.Encryption.Block
import Data.Crypto.Encryption.Stream
import Data.Crypto.Util

%default total
%access public export

data ElectronicCookbook : Nat -> Type where
  ECB : ElectronicCookbook n

-- This is ECB (Electronic Cookbook) - no initialization vector
-- ECB should be considered insecure regardless of the cipher used
implementation EncryptionMode ElectronicCookbook where
  encryptBlocks key _ blocks = map (encryptBlock key) blocks
  decryptBlocks key _ blocks = map (decryptBlock key) blocks

data CipherBlockChainingMode : Nat -> Type where
  CBC : Bits n -> CipherBlockChainingMode n

implementation EncryptionMode CipherBlockChainingMode where
  encryptBlocks _ _ [] = []
  encryptBlocks key (CBC iv) (plain::rest) =
    let ciph = encryptBlock key (plain `xor` iv)
    in ciph :: encryptBlocks key (CBC ciph) rest
  decryptBlocks _ _ [] = []
  decryptBlocks key (CBC iv) (ciph::rest) =
    (decryptBlock key ciph `xor` iv) :: decryptBlocks key (CBC ciph) rest

data PropagatingCipherBlockChainingMode : Nat -> Type where
  PCBC : Bits n -> PropagatingCipherBlockChainingMode n

implementation EncryptionMode PropagatingCipherBlockChainingMode where
  encryptBlocks _ _ [] = []
  encryptBlocks key (PCBC iv) (plain::rest) =
    let ciph = encryptBlock key (plain `xor` iv)
    in ciph :: encryptBlocks key (PCBC (plain `xor` ciph)) rest
  decryptBlocks _ _ [] = []
  decryptBlocks key (PCBC iv) (ciph::rest) =
    let plain = decryptBlock key ciph `xor` iv
    in plain :: decryptBlocks key (PCBC (plain `xor` ciph)) rest

data CipherFeedbackMode : Nat -> Type where
  CFB : Bits n -> CipherFeedbackMode n

implementation EncryptionMode CipherFeedbackMode where
  encryptBlocks _ _ [] = []
  encryptBlocks key (CFB iv) (plain::rest) =
    let ciph = encryptBlock key iv `xor` plain
    in ciph :: encryptBlocks key (PCBC ciph) rest
  decryptBlocks _ _ [] = []
  decryptBlocks key (CFB iv) (ciph::rest) =
    (decryptBlock key iv `xor` ciph) :: decryptBlocks key (PCBC ciph) rest

data OutputFeedbackMode : Nat -> Type where
  OFB : Bits n -> OutputFeedbackMode n

implementation EncryptionMode OutputFeedbackMode where
  encryptBlocks _ _ [] = []
  encryptBlocks key (OFB iv) (plain::rest) =
    let newIV = encryptBlock key iv
    in (plain `xor` newIV) :: encryptBlocks key (OFB newIV) rest
  decryptBlocks _ _ [] = []
  decryptBlocks key (OFB iv) (ciph::rest) =
    let newIV = decryptBlock key iv
    in (ciph `xor` newIV) :: decryptBlocks key (OFB newIV) rest

||| `OutputFeedbackMode` allows any `BlockCipher` to be treated as a
||| `StreamCipher`.
implementation BlockCipher b bitsPerBlock _ => StreamCipher (b, OutputFeedbackMode bitsPerBlock) bitsPerBlock where
  generateKeystream (b, OFB iv) =
    let newIV = encryptBlock b iv
    in newIV :: (generateKeystream (b, OFB newIV))

||| Counter mode takes a nonce, an initial “counter” value, and a function to
||| get the next counter value.
data CounterMode : Nat -> Type where
  CTR : Bits m -> Bits n -> (Bits n -> Bits n) -> CounterMode (m + n)

||| The most common counter mode starts at 0 and moves sequentially through the
||| natural numbers.
incrementalCTR : Bits m -> CounterMode (m + n)
incrementalCTR m = CTR m (intToBits 0) (plus (intToBits 1))

implementation EncryptionMode CounterMode where
  encryptBlocks _ _ [] = []
  encryptBlocks key (CTR nonce counter f) (plain::rest) =
    (plain `xor` encryptBlock key (append nonce counter))
      :: encryptBlocks key (CTR nonce (f counter) f) rest
  decryptBlocks _ _ [] = []
  decryptBlocks key (CTR nonce counter f) (ciph::rest) =
    (ciph `xor` decryptBlock key (append nonce counter))
      :: decryptBlocks key (CTR nonce (f counter) f) rest

||| `CounterMode` allows any `BlockCipher` to be treated as a `StreamCipher`.
implementation BlockCipher b bitsPerBlock _ => StreamCipher (b, CounterMode bitsPerBlock) bitsPerBlock where
  generateKeystream (b, CTR nonce counter f) =
    encryptBlock b (append nonce counter)
      :: generateKeystream (b, CTR nonce (f counter) f)


import RealizedFunction

%default total

-- Realized incrementer
incrementer : Int -> Int
incrementer = (+1)
incrementer_attrs : RealizedAttributes
incrementer_attrs = RealizedAttributes (MkCosts 100 10 1) 1
rlz_incrementer : SimpleRealizedFunction (Int -> Int) incrementer_attrs
rlz_incrementer = MkSimpleRealizedFunction incrementer incrementer_attrs

-- Realized twicer
twicer : Int -> Int
twicer = (*2)
twicer_attrs = MkNonDTRealizedAttributes (MkCosts 200 20 2) 0.9
rlz_twicer : SimpleRealizedFunction (Int -> Int) twicer_attrs
rlz_twicer = MkSimpleRealizedFunction twicer twicer_attrs

-- Realized (twicer . incrementer).
rlz_composition_attrs = MkRealizedAttributes (MkCosts 300 30 3) 0.9
rlz_composition : RealizedFunction (Int -> Int) rlz_composition_attrs
rlz_composition = compose rlz_twicer rlz_incrementer

-- Simple test, result should be (3+1)*2 = 8
rslt : Int
rslt = apply rlz_composition 3

rsltTest : (rslt 3) = 8
rsltTest = Refl


||| http://okmij.org/ftp/Computation/free-monad.html
module Freer

%default total
%access public export

data FFree : (Type -> Type) -> Nat -> Type -> Type where
  FPure   : a -> FFree g n a
  FImpure : g x -> (x -> FFree g n a) -> FFree g (S n) a

Functor (FFree g n) where
  map f (FPure x)     = FPure (f x)
  map f (FImpure u q) = FImpure u (map f . q)

pure : a -> FFree g n a
pure = FPure

promote : {m : Nat} -> FFree g n a -> FFree g (n + m) a
promote (FPure x)     = FPure x
promote (FImpure u q) = FImpure u (promote . q)

assocS : (m : Nat) -> (n : Nat) -> m + S n = S (m + n)
assocS Z n = Refl
assocS (S m) n = cong $ assocS m n

apply : FFree g n (a -> b) -> FFree g m a -> FFree g (m + n) b
apply {n} {m} (FPure f) x =
  promote $ map f x
apply {n=S n} {m} (FImpure u q) x =
  let y = flip apply x . q
  in rewrite assocS m n in FImpure u y

bind : (a -> FFree g n b) -> FFree g m a -> FFree g (n + m) b
bind {n} {m} f (FPure x) = promote $ f x
bind {n} {m=S m} f (FImpure u f') =
  let y = bind f . f'
  in rewrite assocS n m in FImpure u y

join : FFree g m (FFree g n a) -> FFree g (n + m) a
join = bind id

--Applicative (FFree g) where
--  pure = FPure
--
--infixr 5 >=>
--
--(>=>) : Monad m => (a -> m b) -> (b -> m c) -> (a -> m c)
--f >=> g = (>>= g) . f
--
--Monad (FFree g) where
--  (FPure x) >>= f = f x
--  (FImpure u f') >>= f = FImpure u (f' >=> f)


import Data.Vect as V
import Prelude.List as L
import Language.Reflection.Utils

%language ElabReflection

--data B = T | F
%runElab
  (do declareDatatype $
        Declare `{{B}} [] `(Type)
      defineDatatype $
        DefineDatatype `{{B}}
          [ Constructor `{{T}} [] (Var `{{B}})
            , Constructor `{{F}} [] (Var `{{B}})
            ])

imp : TTName -> Raw -> FunArg
imp n t = MkFunArg n t Implicit Erased
exp : TTName -> Raw -> FunArg
exp n t = MkFunArg n t Explicit NotErased

-- equivalent to: append : Vect n a -> Vect m a -> Vect (n + m) a
-- does not define append
appendDecl : TyDecl
appendDecl =
  Declare
    `{{append}}
    [ imp `{{a}} `(Type)
      , imp `{{n}} `(Nat)
      , imp `{{m}} `(Nat)
      , exp `{{xs}}
        `(Vect ~(Var `{{n}}) ~(Var `{{a}}))
      , exp `{{ys}}
        `(Vect ~(Var `{{m}}) ~(Var `{{a}}))
    ]
    `(Vect (plus ~(Var `{{n}})
                 ~(Var `{{m}}))
                 ~(Var `{{a}}))

-- %runElab
--   (do
--     (declareType appendDecl)
-- --    the (Elab ()) debug
--     )

-- xxx : Nat
-- xxx = let vd = Declare `{{val}} [] `(Nat)
--           x = S Z
--         in %runElab (do declareType (Declare `{{val}} [] `(Nat))
--                         fill `(Z)
--                         solve)

-- xs : List Nat
-- xs = let x = S Z
--   --valDecl = Declare `{{val}} [] `(Nat)
--      in [%runElab (do declareType (Declare `{{val}} [] `(Nat))
--                       --focus `{{val}}
--                       fill `(Z)
--                       solve),
--          x + val]

-- %runElab defineFunction $ DefineFun `{{val}} [(MkFunClause (RConstant VoidType) (RConstant (BI 42)))]

-- e0 : Elab () -> Elab ()
-- e0 x = x

-- envToErrorReport : (List (TTName, Binder TT)) -> List ErrorReportPart
-- envToErrorReport xs = map ?envToErrorReport_rhs xs

-- %runElab do env <- getEnv
--             e0 $ debugMessage [(TextPart "foo")]

quicksort : List Nat -> List Nat
quicksort [] = []
quicksort (x::xs) =
  quicksort (filter (< x) xs) ++
  x :: quicksort (filter (>= x) xs)

-- processClauses : Nat -> (List (FunClause TT)) -> Elab (List (FunClause Raw, ConstructorDefn))

-- bc : (inFun, outFun, accPred : TTName) -> Elab ()
-- bc inFun outFun accPred = do
--   (inFun', Ref, inTy) <- lookupTyExact inFun
--      | _ => fail [TextPart "Not a function"]
--   (_, args, res) <- lookupArgsExact inFun'
--   declareDatatype $ Declare accPred args `(Type)
--   let accArg = MkFunArg `{{acc}}
--                 (mkApp (Var accPred)
--                    (map (Var . name) args))
--                 Explicit
--                 NotErased
--   declareType $
--     Declare outFun (args ++ [accArg]) res
--   DefineFun _ oldClauses <- lookupFunDefnExact inFun'
--   (newClauses, ctors) <-
--     Prelude.List.unzip <$> processClauses Z oldClauses
--   defineDatatype $ DefineDatatype accPred ctors
--   defineFunction $ DefineFun outFun newClauses


-- BearSSL's "adding holes" algorithm for XOR multiplication
-- Based on BearSSL ghash_ctmul64.c

module Crypto.Hash.GHash

import Utils.Bytes
import Data.Bits
import Data.List
import Data.Vect
import Data.DPair
import Data.Fin
import Data.Fin.Extra
import Data.Nat
import Data.Nat.Order.Properties
import Utils.Misc
import Crypto.Hash

HValues : Type
HValues = (Bits64, Bits64, Bits64, Bits64, Bits64, Bits64)

-- Carryless multiplication with "holes"
bmul : Bits64 -> Bits64 -> Bits64
bmul x y =
  let x0 = x .&. 0x1111111111111111
      x1 = x .&. 0x2222222222222222
      x2 = x .&. 0x4444444444444444
      x3 = x .&. 0x8888888888888888
      y0 = y .&. 0x1111111111111111
      y1 = y .&. 0x2222222222222222
      y2 = y .&. 0x4444444444444444
      y3 = y .&. 0x8888888888888888
      z0 = (x0 * y0) `xor` (x1 * y3) `xor` (x2 * y2) `xor` (x3 * y1)
      z1 = (x0 * y1) `xor` (x1 * y0) `xor` (x2 * y3) `xor` (x3 * y2)
      z2 = (x0 * y2) `xor` (x1 * y1) `xor` (x2 * y0) `xor` (x3 * y3)
      z3 = (x0 * y3) `xor` (x1 * y2) `xor` (x2 * y1) `xor` (x3 * y0)
  in (z0 .&. 0x1111111111111111) .|.
     (z1 .&. 0x2222222222222222) .|.
     (z2 .&. 0x4444444444444444) .|.
     (z3 .&. 0x8888888888888888)

rms : Bits64 -> Index {a=Bits64} -> Bits64 -> Bits64
rms m s x = (shiftL (x .&. m) s) .|. ((shiftR x s) .&. m)

rev64 : Bits64 -> Bits64
rev64 x =
  let x' =
      (rms 0x0000FFFF0000FFFF 16) $
      (rms 0x00FF00FF00FF00FF 8 ) $
      (rms 0x0F0F0F0F0F0F0F0F 4 ) $
      (rms 0x3333333333333333 2 ) $
      (rms 0x5555555555555555 1 ) x
  in (shiftL x' 32) .|. (shiftR x' 32)

gmult_core : Bits64 -> Bits64 -> HValues -> (Bits64, Bits64)
gmult_core y1 y0 (h0, h0r, h1, h1r, h2, h2r) =
  let y1r = rev64 y1
      y0r = rev64 y0
      y2 = xor y0 y1
      y2r = xor y0r y1r
      -- Karatsuba decomposition
      -- Here we decompose the 128 bit multiplication into
      -- 3 64-bits multiplication
      -- The h-suffixed variables are just multiplication for
      -- reversed bits, which is necessary because we want the
      -- high bits
      z0 = bmul y0 h0
      z1 = bmul y1 h1
      z2 = bmul y2 h2
      z0h = bmul y0r h0r
      z1h = bmul y1r h1r
      z2h = bmul y2r h2r
      z2  = xor (xor z0 z1) z2
      z2h = xor (xor z0h z1h) z2h
      z0h = shiftR (rev64 z0h) 1
      z1h = shiftR (rev64 z1h) 1
      z2h = shiftR (rev64 z2h) 1
      -- Since the operation is done in big-endian, but GHASH spec
      -- needs small endian, we flip the bits
      v0 = z0
      v1 = xor z0h z2
      v2 = xor z1 z2h
      v3 = z1h
      v3 = (shiftL v3 1) .|. (shiftR v2 63)
      v2 = (shiftL v2 1) .|. (shiftR v1 63)
      v1 = (shiftL v1 1) .|. (shiftR v0 63)
      v0 = (shiftL v0 1)
      -- Modular reduction to GF[128]
      v2 = v2 `xor` v0              `xor` (shiftR v0 1)  `xor` (shiftR v0 2) `xor` (shiftR v0 7)
      v1 = v1 `xor` (shiftL v0 63) `xor` (shiftL v0 62) `xor` (shiftL v0 57)
      v3 = v3 `xor` v1             `xor` (shiftR v1 1)  `xor` (shiftR v1 2) `xor` (shiftR v1 7)
      v2 = v2 `xor` (shiftL v1 63) `xor` (shiftL v1 62) `xor` (shiftL v1 57)
  in (v3, v2)

gcm_mult : HValues -> Vect 16 Bits8 -> Vect 16 Bits8
gcm_mult h y =
  let y1 = from_be $ take 8 y
      y0 = from_be $ drop 8 y
      (y1', y0') = gmult_core y1 y0 h
  in to_be {n=8} y1' ++ to_be y0'

mk_h_values : Vect 16 Bits8 -> HValues
mk_h_values h =
  let h1  = from_be $ take 8 h
      h0  = from_be $ drop 8 h
      h0r = rev64 h0
      h1r = rev64 h1
      h2  = xor h0 h1
      h2r = xor h0r h1r
  in (h0, h0r, h1, h1r, h2, h2r)

export
data GHash : Type where
  MkGHash : List Bits8 -> HValues -> Vect 16 Bits8 -> GHash

hash_until_done : GHash -> GHash
hash_until_done ghash@(MkGHash buffer hval state) =
  case splitAtExact 16 buffer of
    Just (a, b) => hash_until_done $ MkGHash b hval (gcm_mult hval (zipWith xor a state))
    Nothing => ghash

export
Digest GHash where
  digest_nbyte = 16
  update message (MkGHash buffer hval state) =
    hash_until_done (MkGHash (buffer <+> message) hval state)
  finalize (MkGHash buffer hval state) =
    let (MkGHash _ _ state) = hash_until_done (MkGHash (pad_zero 16 buffer) hval state)
    in state

export
MAC (Vect 16 Bits8) GHash where
  initialize_mac key = MkGHash [] (mk_h_values key) (replicate _ 0)


twice : (a -> a) -> a -> a
twice f x = f (f x)

Shape : Type
rotate : Shape -> Shape

turn_around : Shape -> Shape
turn_around = twice rotate


{-

:doc Eq

interface Eq ty where
  (==) : ty -> ty -> Bool
  (/=) : ty -> ty -> Bool
-}

occurrences : Eq ty => (item : ty) -> (values : List ty) -> Nat
occurrences item [] = 0
occurrences item (value :: values) = case value == item of
                                          False => occurrences item values
                                          True => 1 + occurrences item values

data Matter = Solid | Liquid | Gas

Eq Matter where
    (==) Solid Solid = True
    (==) Liquid Liquid = True
    (==) Gas Gas = True
    (==) _ _ = False

-- occurrences Liquid [Solid, Liquid, Liquid, Gas]

[myEq] Eq Matter where
    (==) _ _ = True

-- occurrences @{myEq} Liquid [Solid, Liquid, Liquid, Gas]


module Backprop.Node

import Backprop.CanBack
import Backprop.MathUtils
import Control.Optics
import Data.Floating
import public Data.List.Quantifiers

%hide Data.Morphisms.Op

public export
Op : (args : List Type) -> (o : Type) -> Type
Op args o = HList args -> Pair o (o -> HList args)

export
data Node : (i : Type) -> (o : Type) -> Type where
  ||| the input itself
  Input : Node a a
  ||| an operation
  Oper : {args : _} -> All CanBack args => Op args o -> All (Node i) args -> Node i o

public export
interface Functor f => NodeDistributive f where
  collect : {o : _} -> CanBack o => f (Node i o) -> Node i (f o)

public export
interface Functor f => NodeTraversable f where
  sequence : {o : _} -> CanBack o => Node i (f o) -> f (Node i o)

public export
NodeFunctor : (Type -> Type) -> Type
NodeFunctor f = (NodeTraversable f, NodeDistributive f)

export
map : {a, b : _} -> (CanBack a, CanBack b, NodeFunctor f) => (Node i a -> Node i b) -> Node i (f a) -> Node i (f b)
map f = collect . map f . sequence

export
input : Node a a
input = Input

export
op1 : {a1 : _} -> CanBack a1 => Op [a1] o -> Node i a1 -> Node i o
op1 op x = Oper op [x]

export
op2 : {a1, a2 : _} -> CanBack a1 => CanBack a2 => Op [a1, a2] o -> Node i a1 -> Node i a2 -> Node i o
op2 op x y = Oper op [x, y]

||| lift a constant into a `Node`
export
const : a -> Node i a
const x = Oper (\[] => (x, const [])) []

export
FromDouble a => FromDouble (Node i a) where
  fromDouble = const . fromDouble

export
{a : _} -> (CanBack a, Num a) => Num (Node i a) where
  (+) = op2 $ \[x, y] => (x + y, \d => [d, d])
  (*) = op2 $ \[x, y] => (x * y, \d => [d * y, x * d])
  fromInteger = const . fromInteger

export
{a : _} -> (CanBack a, Neg a) => Neg (Node i a) where
  (-) = op2 $ \[x, y] => (x - y, \d => [d, -d])
  negate = op1 $ \[x] => (-x, \d => [-d])

export
{a : _} -> (CanBack a, Ord a, Neg a, Abs a) => Abs (Node i a) where
  abs = op1 $ \[x] => (abs x, \d => [d * signum_zero x])

export
{a : _} -> (CanBack a, Neg a, Fractional a) => Fractional (Node i a) where
  (/) = op2 $ \[x, y] => (x / y, \d => [d / y, -d * x / (y * y)])
  recip = op1 $ \[x] => (recip x, \d => [-d / (x * x)])

export
{a : _} -> (CanBack a, Neg a, Floating a, Fractional (Node i a)) => Floating (Node i a) where
  pi = const pi
  sqrt = op1 $ \[x] => (sqrt x, \d => [d / (2 * sqrt x)])
  exp = op1 $ \[x] => (exp x, \d => [d * exp x])
  ln = op1 $ \[x] => (ln x, \d => [d / x])
  log base x = ln x / ln base
  pow = op2 $ \[x, y] =>
    ( pow x y
    , \d => let k = d * pow x (y - 1) in [k * y, k * x * ln x]
    )
  sin = op1 $ \[x] => (sin x, \d => [d * cos x])
  cos = op1 $ \[x] => (cos x, \d => [d * -sin x])
  tan = op1 $ \[x] => (tan x, \d => let cos' = cos x in [d / (cos' * cos')])
  asin = op1 $ \[x] => (asin x, \d => [d / sqrt (1 - x * x)])
  acos = op1 $ \[x] => (acos x, \d => [-d / sqrt (1 - x * x)])
  atan = op1 $ \[x] => (atan x, \d => [d / (1 - x * x)])
  sinh = op1 $ \[x] => (sinh x, \d => [d * cosh x])
  cosh = op1 $ \[x] => (cosh x, \d => [d * sinh x])
  tanh = op1 $ \[x] => (tanh x, \d => let cosh' = cosh x in [d / (cosh' * cosh')])

mutual
  run_all : {as : _} -> {auto prf : All CanBack as} -> CanBack i => All (Node i) as -> i -> (HList as, All (\b => b -> i -> i) as)
  run_all [] input = ([], [])
  run_all {prf = prf :: prfs} (arg :: args) input =
    let (x, y) = run arg input in bimap (x ::) (y ::) (run_all {prf = prfs} args input)

  combine_all : (backs : All (\b => b -> i -> i) as) -> HList as -> i -> i
  combine_all [] [] = id
  combine_all (back :: backs) (x :: xs) = combine_all backs xs . back x

  export
  run : {o : _} -> {auto prf : CanBack i} -> CanBack o => Node i o -> i -> (o, o -> i -> i)
  run Input input = (input, \d => add @{prf} d)
  run (Oper op arg_nodes) input =
    let
      (args, backs) = run_all arg_nodes input
      (output, back) = op args
    in
      (output, combine_all backs . back)

||| output: (the value, the gradient of inputs)
export
backprop : {o : _} -> (CanBack i, CanBack o) => Node i o -> i -> (o, i)
backprop node input = let (y, back) = run node input in (y, back one zero)

||| `backprop` but discards output value
export
grad : {o : _} -> (CanBack i, CanBack o) => Node i o -> i -> i
grad node input = snd $ backprop node input

mutual
  eval_all : All (Node i) as -> i -> HList as
  eval_all [] input = []
  eval_all (x :: xs) input = eval x input :: eval_all xs input

  ||| evaluates a `Node`, more efficient than `run` as gradient calculations are skipped
  |||
  ||| input: inputs, output node
  ||| output: output value
  export
  eval : Node i o -> i -> o
  eval Input input = input
  eval (Oper op args) input = fst $ op (eval_all args input)


module Instances

import BoundAst

%default total
%access public export

MapType : (Monotype, Monotype)
MapType =
  ( FuncType (Term "a") (Term "b")
  , FuncType
    (TypeFullyApplied TypeclassCtr (Term "a"))
    (TypeFullyApplied TypeclassCtr (Term "b"))
  )

mapImpl : String -> Method MapType
mapImpl = Method' _

FunctorTC : Typeclass Z [MapType]
FunctorTC = TyCl "Functor"

functorInstances : List (Instance FunctorTC)
functorInstances = map Inst
  [ (Appl1 (Regular 1 "Either"), [mapImpl "Either.rmap"])
  , (Appl1 (Regular 1 "Result"), [mapImpl "Result.mapError"])
  , (Regular Z "list", [mapImpl "List.map"])
  , (Regular Z "option", [mapImpl "Option.map"])
  , (Array, [mapImpl "Array.map"])
  , (Appl1 (Tuple Z), [mapImpl "Pair.map"])
  ]

BindType : (Monotype, Monotype)
BindType =
  ( FuncType (Term "a") (TypeFullyApplied TypeclassCtr (Term "b"))
  , FuncType
    (TypeFullyApplied TypeclassCtr (Term "a"))
    (TypeFullyApplied TypeclassCtr (Term "b"))
  )

ApplyType : (Monotype, Monotype)
ApplyType =
  ( TypeFullyApplied TypeclassCtr (FuncType (Term "a") (Term "b"))
  , FuncType
    (TypeFullyApplied TypeclassCtr (Term "a"))
    (TypeFullyApplied TypeclassCtr (Term "b"))
  )

--TraverseType : (Monotype, Monotype)
--TraverseType =
--  ( FuncType (Term "a") (TypeFullyApplied (TypeCtr "m") (Term "b"))
--  , FuncType
--    (TypeFullyApplied TypeclassCtr (Term "a"))
--    (TypeFullyApplied (TypeCtr "m") (NullaryTypeAppl TypeclassCtr (Term "b")))
--  )

BimapType : (Monotype, Monotype)
BimapType =
  ( FuncType (Term "a") (Term "b")
  , FuncType
    (FuncType (Term "c") (Term "d"))
    (FuncType
      (TypeFullyApplied (NAryTypeAppl TypeclassCtr (Term "a")) (Term "c"))
      (TypeFullyApplied (NAryTypeAppl TypeclassCtr (Term "b")) (Term "d"))
    )
  )

TrimapType : (Monotype, Monotype)
TrimapType =
  ( FuncType (Term "a") (Term "b")
  , FuncType
    (FuncType (Term "c") (Term "d"))
    (FuncType
        (FuncType (Term "e") (Term "f"))
        (FuncType
            (TypeFullyApplied (NAryTypeAppl (NAryTypeAppl TypeclassCtr (Term "a")) (Term "c")) (Term "e"))
            (TypeFullyApplied (NAryTypeAppl (NAryTypeAppl TypeclassCtr (Term "b")) (Term "d")) (Term "f"))
        )
    )
  )


module Main

main : IO ()
main = do
    printLn $ divNat 1809022195644390369852458 91238741987
    printLn $ div (-432642342742368327462378462387) 36473264372

    -- values at the border between C ints and bignums
    printLn $ div 1073741822 73892
    printLn $ div 1073741823 73892
    printLn $ div 1073741824 73892
    printLn $ div (-1073741823) 73892
    printLn $ div (-1073741824) 73892
    printLn $ div (-1073741825) 73892