module Compiler.LLVM.Rapid.Integer

import Data.Bits
import Data.Vect

import Compiler.LLVM.IR
import Compiler.LLVM.Instruction
import Compiler.LLVM.Rapid.Object
import Control.Codegen

GMP_LIMB_SIZE : Integer
GMP_LIMB_SIZE = 8

GMP_LIMB_BITS : Integer
GMP_LIMB_BITS = 8 * GMP_LIMB_SIZE

GMP_LIMB_BOUND : Integer
GMP_LIMB_BOUND = (1 `prim__shl_Integer` (GMP_LIMB_BITS))

-- To estimate required count of limbs (upper bound):
-- x = the number (result)
--   length of string == number of digits <= log10(x) + 1
--   required limbs: log10(x) / log10(2^LIMB_BITS)
--     LIMB_BITS=32 -> required limbs <= number of digits / 9
--     LIMB_BITS=64 -> required limbs <= number of digits / 19
GMP_ESTIMATE_DIGITS_PER_LIMB : Integer
GMP_ESTIMATE_DIGITS_PER_LIMB = 19

TARGET_SIZE_T : IRType
TARGET_SIZE_T = I64

MP_LIMB_T : IRType
MP_LIMB_T = I64

twosComplement : Num a => Bits a => a -> a
twosComplement x = 1 + (complement x)

export
cgMkConstInteger : Integer -> Codegen (IRValue IRObjPtr)
cgMkConstInteger val =
    do
      let absVal = abs val
      let (len ** limbs) = getLimbs absVal
      let len32 = cast {to=Bits32} $ cast {to=Int} len
      let lenForHeader = if (val >= 0) then len32 else (twosComplement len32)
      let newHeader = constHeader OBJECT_TYPE_ID_BIGINT lenForHeader
      let typeSignature = "{i64, [" ++ show len ++ " x %LimbT]}"
      cName <- addConstant $ "private unnamed_addr addrspace(1) constant " ++ typeSignature ++ " {" ++ toIR newHeader ++ ", [" ++ show len ++ " x %LimbT] [" ++ (getLimbsIR limbs) ++ "]}, align 8"
      pure $ SSA IRObjPtr $ "bitcast (" ++ typeSignature ++ " addrspace(1)* " ++ cName ++ " to %ObjPtr)"
  where
      getLimbs : Integer -> (n:Nat ** Vect n Integer)
      getLimbs 0 = (0 ** [])
      getLimbs x = let (n ** v) = (getLimbs (x `div` GMP_LIMB_BOUND))
                       limb = (x `mod` GMP_LIMB_BOUND) in
                       ((S n) ** (limb::v))
      getLimbsIR : Vect n Integer -> String
      getLimbsIR [] = ""
      getLimbsIR (limb::[]) = "%LimbT " ++ show limb
      getLimbsIR (limb::rest) = "%LimbT " ++ show limb ++ ", " ++ (getLimbsIR rest)

export
integer0 : Codegen (IRValue IRObjPtr)
integer0 = do
  newObj <- dynamicAllocate (Const I64 0)
  putObjectHeader newObj (constHeader OBJECT_TYPE_ID_BIGINT 0)
  pure newObj

export
cgMkIntegerSigned : IRValue I64 -> Codegen (IRValue IRObjPtr)
cgMkIntegerSigned val = do
  isNegative <- icmp "slt" val (Const I64 0)
  isZero <- icmp "eq" val (Const I64 0)
  newSize1 <- mkSelect isNegative (Const I32 (-1)) (Const I32 1)
  newSize <- mkSelect isZero (Const I32 0) newSize1
  newHeader <- mkHeader OBJECT_TYPE_ID_BIGINT newSize
  newSizeAbs <- mkAbs newSize
  allocSize <- mkMul !(mkZext newSizeAbs) (Const I64 GMP_LIMB_SIZE)
  newObj <- dynamicAllocate allocSize
  ignore $ mkIf (pure isZero) (pure (Const I1 0)) (do
       absVal <- mkAbs64 val
       putObjectSlot newObj (Const I64 0) absVal
       pure $ Const I1 0)
  putObjectHeader newObj newHeader
  pure newObj

export
cgMkIntegerUnsigned : IRValue I64 -> Codegen (IRValue IRObjPtr)
cgMkIntegerUnsigned ival = do
  isZero <- icmp "eq" (Const I64 0) ival
  newObj <- mkIf (pure isZero) integer0 (do
      newInteger <- dynamicAllocate (Const I64 GMP_LIMB_SIZE)
      putObjectHeader newInteger !(mkHeader OBJECT_TYPE_ID_BIGINT (Const I32 1))
      putObjectSlot newInteger (Const I64 0) ival
      pure newInteger
    )
  pure newObj

export
unboxIntegerUnsigned : IRValue IRObjPtr -> Codegen (IRValue I64)
unboxIntegerUnsigned integerObj = do
  isZero <- icmp "eq" (Const I32 0) !(getObjectSize integerObj)
  -- get first limb (LSB)
  mkIf (pure isZero) (pure $ Const I64 0) (getObjectSlot {t=I64} integerObj 0)

export
unboxIntegerSigned : IRValue IRObjPtr -> Codegen (IRValue I64)
unboxIntegerSigned integerObj = do
  size <- getObjectSize integerObj
  isZero <- icmp "eq" (Const I32 0) size
  let isNegative = icmp "sgt" (Const I32 0) size
  -- get first limb (LSB)
  firstLimb <- getObjectSlot {t=I64} integerObj 0
  -- TODO: this is probably wrong for 64bit
  mkIf (pure isZero) (pure $ Const I64 0) (mkIf isNegative (mkSub (Const I64 0) firstLimb) (pure firstLimb))

||| compare two BigInts `a` and `b`, return -1 if a<b, +1 if a>b, 0 otherwise
export
compareInteger : IRValue IRObjPtr -> IRValue IRObjPtr -> Codegen (IRValue I64)
compareInteger obj1 obj2 = do
  size1 <- getObjectSize obj1
  size2 <- getObjectSize obj2
  cmpResult <- mkIf (icmp "slt" size1 size2) (pure (Const I64 (-1))) (
    mkIf (icmp "sgt" size1 size2) (pure (Const I64 1)) (do
         limbs1 <- getObjectPayloadAddr {t=I64} obj1
         limbs2 <- getObjectPayloadAddr {t=I64} obj2
         absSize <- mkZext {to=I64} !(mkAbs size1)
         mpnResult <- call {t=I32} "ccc" "@__gmpn_cmp" [toIR limbs1, toIR limbs2, toIR absSize]
         sizeIsNegative <- icmp "slt" size1 (Const I32 0)
         mkSext !(mkSelect sizeIsNegative !(mkSub (Const I32 0) mpnResult) mpnResult)
         )
         )

  pure cmpResult

normaliseIntegerSize : IRValue IRObjPtr -> IRValue I32 -> IRValue I1 -> Codegen ()
normaliseIntegerSize integerObj maxSizeSigned invert = do
  maxSizeAbs <- mkAbs maxSizeSigned
  absRealNewSize <- mkTrunc {to=I32} !(call {t=I64} "ccc" "@rapid_bigint_real_size" [
    toIR !(getObjectPayloadAddr {t=I64} integerObj),
    toIR !(mkZext {to=I64} maxSizeAbs)
    ])
  isNegative <- icmp "slt" maxSizeSigned (Const I32 0)
  invertResult <- mkXOr isNegative invert
  signedNewSize <- mkSelect invertResult !(mkSub (Const I32 0) absRealNewSize) absRealNewSize
  newHeader <- mkHeader OBJECT_TYPE_ID_BIGINT signedNewSize
  putObjectHeader integerObj newHeader

export
addInteger : IRValue IRObjPtr -> IRValue IRObjPtr -> Codegen (IRValue IRObjPtr)
addInteger i1 i2 = do
      s1 <- getObjectSize i1
      s2 <- getObjectSize i2
      i1Negative <- icmp "slt" s1 (Const I32 0)
      s1a <- mkAbs s1
      s2a <- mkAbs s2
      i1longer <- icmp "ugt" s1a s2a
      -- "big" and "small" refer just to the respective limb counts
      -- it doesn't matter which number is actually bigger
      big <- mkSelect i1longer i1 i2
      small <- mkSelect i1longer i2 i1
      size1 <- mkZext {to=I64} !(mkSelect {t=I32} i1longer s1a s2a)
      size2 <- mkZext {to=I64} !(mkSelect {t=I32} i1longer s2a s1a)
      newLength <- mkAdd size1 (Const I64 1)
      newSize <- mkMul (Const I64 GMP_LIMB_SIZE) newLength
      newObj <- dynamicAllocate newSize
      carry <- call {t=I64} "ccc" "@__gmpn_add" [
        toIR !(getObjectPayloadAddr {t=I64} newObj),
        toIR !(getObjectPayloadAddr {t=I64} big),
        toIR size1,
        toIR !(getObjectPayloadAddr {t=I64} small),
        toIR size2
        ]
      putObjectSlot newObj size1 carry
      absRealNewSize <- mkAdd size1 carry
      signedNewSize <- mkSelect i1Negative !(mkSub (Const I64 0) absRealNewSize) absRealNewSize
      signedNewSize32 <- mkTrunc {to=I32} signedNewSize
      newHeader <- mkHeader OBJECT_TYPE_ID_BIGINT signedNewSize32
      putObjectHeader newObj newHeader
      pure newObj

export
subInteger : IRValue IRObjPtr -> IRValue IRObjPtr -> Codegen (IRValue IRObjPtr)
subInteger i1 i2 = do
      -- Subtract the smaller (by abs. value) from the larger (by abs. value)
      -- and use the sign of the larger (by abs. value) number as sign for the
      -- returned result.
      s1 <- getObjectSize i1
      s2 <- getObjectSize i2
      i1Negative <- icmp "slt" s1 (Const I32 0)
      s1a <- mkAbs s1
      s2a <- mkAbs s2
      i1longer <- icmp "ugt" s1a s2a
      i2longer <- icmp "ugt" s2a s1a
      i1bigger <- mkIf (pure i1longer) (pure $ Const I1 1) (mkIf (pure i2longer) (pure $ Const I1 0) (icmp "sgt" !(call "ccc" "@__gmpn_cmp" [
                                toIR !(getObjectPayloadAddr {t=I64} i1),
                                toIR !(getObjectPayloadAddr {t=I64} i2),
                                toIR !(mkZext {to=I64} s1a)
                                ]) (Const I32 0))
        )
      big <- mkSelect i1bigger i1 i2
      small <- mkSelect i1bigger i2 i1
      swapped <- mkSelect i1bigger (Const I1 0) (Const I1 1)
      bigSize <- getObjectSize big
      bigSizeAbs <- mkAbs bigSize
      smallSizeAbs <- mkAbs !(getObjectSize small)
      newSize <- mkMul (Const I64 GMP_LIMB_SIZE) !(mkZext {to=I64} bigSizeAbs)
      newObj <- dynamicAllocate newSize
      absDiff <- call {t=I64} "ccc" "@__gmpn_sub" [
        toIR !(getObjectPayloadAddr {t=I64} newObj),
        toIR !(getObjectPayloadAddr {t=I64} big),
        toIR !(mkZext {to=I64} bigSizeAbs),
        toIR !(getObjectPayloadAddr {t=I64} small),
        toIR !(mkZext {to=I64} smallSizeAbs)
        ]
      absRealNewSize <- call {t=I64} "ccc" "@rapid_bigint_real_size" [
        toIR !(getObjectPayloadAddr {t=I64} newObj),
        toIR !(mkZext {to=I64} bigSizeAbs)
        ]
      resultIsNegative <- mkXOr swapped i1Negative
      signedNewSize <- mkSelect resultIsNegative !(mkSub (Const I64 0) absRealNewSize) absRealNewSize
      signedNewSize32 <- mkTrunc {to=I32} signedNewSize
      newHeader <- mkHeader OBJECT_TYPE_ID_BIGINT signedNewSize32
      putObjectHeader newObj newHeader
      pure newObj

export
andInteger : IRValue IRObjPtr -> IRValue IRObjPtr -> Codegen (IRValue IRObjPtr)
andInteger i1 i2 = do
  -- TODO: what to do with negative numbers?
  s1 <- getObjectSize i1
  s2 <- getObjectSize i2
  zero1 <- icmp "eq" s1 (Const I32 0)
  zero2 <- icmp "eq" s2 (Const I32 0)
  resultIsZero <- mkOr zero1 zero2

  mkIf (pure resultIsZero) (mkSelect zero1 i1 i2) (do
    s1a <- mkAbs s1
    s2a <- mkAbs s2
    i1longer <- icmp "ugt" s1a s2a
    -- "long" and "short" refer just to the respective limb counts
    -- it doesn't matter which number is actually bigger
    long <- mkSelect i1longer i1 i2
    short <- mkSelect i1longer i2 i1
    size1 <- mkZext {to=I64} !(mkSelect {t=I32} i1longer s1a s2a)
    size2 <- mkZext {to=I64} !(mkSelect {t=I32} i1longer s2a s1a)
    -- result can not be longer than shortest number
    let newLength = size2
    newSize <- mkMul (Const I64 GMP_LIMB_SIZE) newLength
    newObj <- dynamicAllocate newSize
    putObjectHeader newObj !(mkHeader OBJECT_TYPE_ID_BIGINT !(mkTrunc newLength))

    newLimbs <- getObjectPayloadAddr {t=I64} newObj
    shortLimbs <- getObjectPayloadAddr {t=I64} short
    longLimbs <- getObjectPayloadAddr {t=I64} long
    voidCall "ccc" "@__gmpn_and_n" [toIR newLimbs, toIR shortLimbs, toIR longLimbs, toIR newLength]

    normaliseIntegerSize newObj !(mkTrunc newLength) (Const I1 0)

    pure newObj
    )

export
orInteger : IRValue IRObjPtr -> IRValue IRObjPtr -> Codegen (IRValue IRObjPtr)
orInteger i1 i2 = do
  -- TODO: what to do with negative numbers?
  s1 <- getObjectSize i1
  s2 <- getObjectSize i2
  zero1 <- icmp "eq" s1 (Const I32 0)
  zero2 <- icmp "eq" s2 (Const I32 0)
  resultIsZero <- mkAnd zero1 zero2

  mkIf (pure resultIsZero) (pure i1) (do
    s1a <- mkAbs s1
    s2a <- mkAbs s2
    i1longer <- icmp "ugt" s1a s2a
    -- "big" and "small" refer just to the respective limb counts
    -- it doesn't matter which number is actually bigger
    big <- mkSelect i1longer i1 i2
    small <- mkSelect i1longer i2 i1
    size1 <- mkZext {to=I64} !(mkSelect {t=I32} i1longer s1a s2a)
    size2 <- mkZext {to=I64} !(mkSelect {t=I32} i1longer s2a s1a)
    mkIf (icmp "eq" (Const I64 0) size2) (pure big) (do
      let newLength = size1
      newSize <- mkMul (Const I64 GMP_LIMB_SIZE) newLength
      newObj <- dynamicAllocate newSize
      putObjectHeader newObj !(mkHeader OBJECT_TYPE_ID_BIGINT !(mkTrunc newLength))

      newPayload <- getObjectPayloadAddr {t=I8} newObj
      bigPayload <- getObjectPayloadAddr {t=I8} big
      appendCode $ "  call void @llvm.memcpy.p1i8.p1i8.i64(" ++ toIR newPayload ++ ", " ++ toIR bigPayload ++ ", " ++ toIR newSize++ ", i1 false)"

      newLimbs <- getObjectPayloadAddr {t=I64} newObj
      smallLimbs <- getObjectPayloadAddr {t=I64} small
      voidCall "ccc" "@__gmpn_ior_n" [toIR newLimbs, toIR newLimbs, toIR smallLimbs, toIR size2]
      pure newObj
      )
    )

export
xorInteger : IRValue IRObjPtr -> IRValue IRObjPtr -> Codegen (IRValue IRObjPtr)
xorInteger i1 i2 = do
  -- TODO: what to do with negative numbers?
  s1 <- getObjectSize i1
  s2 <- getObjectSize i2
  zero1 <- icmp "eq" s1 (Const I32 0)
  zero2 <- icmp "eq" s2 (Const I32 0)
  resultIsUnchanged <- mkOr zero1 zero2

  mkIf (pure resultIsUnchanged) (mkSelect zero1 i2 i1) (do
    s1a <- mkAbs s1
    s2a <- mkAbs s2
    i1longer <- icmp "ugt" s1a s2a
    -- "long" and "short" refer just to the respective limb counts
    -- it doesn't matter which number is actually bigger
    long <- mkSelect i1longer i1 i2
    short <- mkSelect i1longer i2 i1
    size1 <- mkZext {to=I64} !(mkSelect {t=I32} i1longer s1a s2a)
    size2 <- mkZext {to=I64} !(mkSelect {t=I32} i1longer s2a s1a)
    -- result can not be longer than longest number
    let newLength = size1
    newSize <- mkMul (Const I64 GMP_LIMB_SIZE) newLength
    newObj <- dynamicAllocate newSize
    putObjectHeader newObj !(mkHeader OBJECT_TYPE_ID_BIGINT !(mkTrunc newLength))

    newPayload <- getObjectPayloadAddr {t=I8} newObj
    longPayload <- getObjectPayloadAddr {t=I8} long
    appendCode $ "  call void @llvm.memcpy.p1i8.p1i8.i64(" ++ toIR newPayload ++ ", " ++ toIR longPayload ++ ", " ++ toIR newSize ++ ", i1 false)"

    newLimbs <- getObjectPayloadAddr {t=I64} newObj
    shortLimbs <- getObjectPayloadAddr {t=I64} short
    voidCall "ccc" "@__gmpn_xor_n" [toIR newLimbs, toIR newLimbs, toIR shortLimbs, toIR size2]

    normaliseIntegerSize newObj !(mkTrunc newLength) (Const I1 0)

    pure newObj
    )

export
mulInteger : IRValue IRObjPtr -> IRValue IRObjPtr -> Codegen (IRValue IRObjPtr)
mulInteger i1 i2 = do
      s1 <- getObjectSize i1
      s2 <- getObjectSize i2
      zero1 <- icmp "eq" s1 (Const I32 0)
      zero2 <- icmp "eq" s2 (Const I32 0)
      resultIsZero <- mkOr zero1 zero2
      mkIf (pure resultIsZero) {- then -} integer0 {- else -} (do
        sx <- mkXOr s1 s2
        signsMatch <- icmp "sge" sx (Const I32 0)
        s1a <- mkAbs s1
        s2a <- mkAbs s2
        i1longer <- icmp "ugt" s1a s2a
        -- "big" and "small" refer just to the respective limb counts
        -- it doesn't matter which number is actually bigger
        big <- mkSelect i1longer i1 i2
        small <- mkSelect i1longer i2 i1
        size1 <- mkZext {to=I64} !(mkSelect {t=I32} i1longer s1a s2a)
        size2 <- mkZext {to=I64} !(mkSelect {t=I32} i1longer s2a s1a)
        newLength <- mkAdd size1 size2
        newSize <- mkMul (Const I64 GMP_LIMB_SIZE) newLength
        newObj <- dynamicAllocate newSize
        ignore $ call {t=I64} "ccc" "@__gmpn_mul" [
          toIR !(getObjectPayloadAddr {t=I64} newObj),
          toIR !(getObjectPayloadAddr {t=I64} big),
          toIR size1,
          toIR !(getObjectPayloadAddr {t=I64} small),
          toIR size2
          ]
        absRealNewSize <- call {t=I64} "ccc" "@rapid_bigint_real_size" [
          toIR !(getObjectPayloadAddr {t=I64} newObj),
          toIR newLength
          ]
        signedNewSize <- mkSelect signsMatch absRealNewSize !(mkSub (Const I64 0) absRealNewSize)
        signedNewSize32 <- mkTrunc {to=I32} signedNewSize
        newHeader <- mkHeader OBJECT_TYPE_ID_BIGINT signedNewSize32
        putObjectHeader newObj newHeader
        pure newObj)

||| divide i1 by i2, return (quotient, remainder)
divInteger : IRValue IRObjPtr -> IRValue IRObjPtr -> Codegen (IRValue IRObjPtr, IRValue IRObjPtr)
divInteger i1 i2 = do
  s1 <- getObjectSize i1
  s2 <- getObjectSize i2
  s1a <- mkZext !(mkAbs s1)
  s2a <- mkZext !(mkAbs s2)
  zero1 <- icmp "eq" s1 (Const I32 0)
  zero2 <- icmp "eq" s2 (Const I32 0)
  ignore $ mkIf (pure zero2) (do
                mkRuntimeCrash "division by 0"
                pure (Const I1 0)
                ) (pure (Const I1 0))

  retZeroLbl <- genLabel "ret0"
  checkDividendLbl <- genLabel "div_chk"
  dividendLargerLbl <- genLabel "div_lg"
  divLbl <- genLabel "div"
  endLbl <- genLabel "div_end"

  branch zero1 retZeroLbl checkDividendLbl

  beginLabel retZeroLbl
  zeroInteger <- integer0
  jump endLbl

  beginLabel checkDividendLbl
  dividendLarger <- icmp "ugt" s2a s1a
  branch dividendLarger dividendLargerLbl divLbl

  beginLabel dividendLargerLbl
  zeroQuotient <- integer0
  jump endLbl

  beginLabel divLbl
  -- i1, i2 /= 0
  sx <- mkXOr s1 s2
  signsMatch <- icmp "sge" sx (Const I32 0)

  -- remainder can not be bigger than divisor
  let maxLimbsRemainder = s2a
  remainder <- dynamicAllocate !(mkMul (Const I64 GMP_LIMB_SIZE) maxLimbsRemainder)
  -- object must have a valid header, because the next allocation might trigger a GC
  tempHeader <- mkHeader OBJECT_TYPE_ID_BIGINT !(mkTrunc maxLimbsRemainder)
  putObjectHeader remainder tempHeader

  maxLimbsQuotient <- mkMax (Const I64 1) !(mkAdd (Const I64 1) !(mkSub s1a s2a))
  quotient <- dynamicAllocate !(mkMul (Const I64 GMP_LIMB_SIZE) maxLimbsQuotient)

  voidCall "ccc" "@__gmpn_tdiv_qr" [
    toIR !(getObjectPayloadAddr {t=I64} quotient),
    toIR !(getObjectPayloadAddr {t=I64} remainder),
    toIR (Const I64 0),
    toIR !(getObjectPayloadAddr {t=I64} i1),
    toIR s1a,
    toIR !(getObjectPayloadAddr {t=I64} i2),
    toIR s2a
    ]
  qRealNewSize <- call {t=I64} "ccc" "@rapid_bigint_real_size" [
    toIR !(getObjectPayloadAddr {t=I64} quotient),
    toIR maxLimbsQuotient
    ]
  signedNewSize <- mkSelect signsMatch qRealNewSize !(mkSub (Const I64 0) qRealNewSize)
  signedNewSize32 <- mkTrunc {to=I32} signedNewSize
  newHeader <- mkHeader OBJECT_TYPE_ID_BIGINT signedNewSize32
  putObjectHeader quotient newHeader

  i1negative <- icmp "slt" s1 (Const I32 0)
  rRealNewSize <- call {t=I64} "ccc" "@rapid_bigint_real_size" [
    toIR !(getObjectPayloadAddr {t=I64} remainder),
    toIR maxLimbsRemainder
    ]
  signedNewSize <- mkSelect i1negative !(mkSub (Const I64 0) rRealNewSize) rRealNewSize
  signedNewSize32 <- mkTrunc {to=I32} signedNewSize
  newHeader <- mkHeader OBJECT_TYPE_ID_BIGINT signedNewSize32
  putObjectHeader remainder newHeader

  jump endLbl

  beginLabel endLbl
  quotient  <- phi [(zeroInteger, retZeroLbl), (zeroQuotient, dividendLargerLbl), (quotient,  divLbl)]
  remainder <- phi [(zeroInteger, retZeroLbl), (i1,           dividendLargerLbl), (remainder, divLbl)]
  pure (quotient, remainder)

export
divIntegerQuotient : IRValue IRObjPtr -> IRValue IRObjPtr -> Codegen (IRValue IRObjPtr)
divIntegerQuotient a b = fst <$> divInteger a b

export
divIntegerRemainder : IRValue IRObjPtr -> IRValue IRObjPtr -> Codegen (IRValue IRObjPtr)
divIntegerRemainder a b = snd <$> divInteger a b

export
shiftLeftInteger : IRValue IRObjPtr -> IRValue IRObjPtr -> Codegen (IRValue IRObjPtr)
shiftLeftInteger integerObj bitCountObj = do
  bitCount <- mkTrunc {to=I32} !(unboxIntegerUnsigned bitCountObj)

  size <- getObjectSize integerObj
  unchanged <- mkOr !(icmp "eq" size (Const I32 0)) !(icmp "eq" bitCount (Const I32 0))
  mkIf (pure unchanged) (do
    pure integerObj
    ) (do
    sizeAbs <- mkAbs size
    fullLimbs <- mkUDiv bitCount (Const I32 GMP_LIMB_BITS)
    maxLimbsCount <- mkAdd !(mkAdd fullLimbs sizeAbs) (Const I32 1)

    newObj <- dynamicAllocate !(mkZext !(mkMul maxLimbsCount (Const I32 GMP_LIMB_SIZE)))
    lowerLimbsAddr <- getObjectPayloadAddr {t=I8} newObj
    appendCode $ "  call void @llvm.memset.p1i8.i64(" ++ toIR lowerLimbsAddr ++ ", i8 0, " ++ toIR !(mkMul !(mkZext fullLimbs) (Const I64 8)) ++ ", i1 false)"

    restBits <- mkURem bitCount (Const I32 GMP_LIMB_BITS)
    mkIf_ (icmp "ne" (Const I32 0) restBits) (do
      srcLimbs <- getObjectPayloadAddr {t=I64} integerObj
      higherLimbsAddr <- getObjectSlotAddrVar {t=I64} newObj !(mkZext {to=I64} fullLimbs)
      msbLimb <- call {t=I64} "ccc" "@__gmpn_lshift" [
        toIR higherLimbsAddr,
        toIR srcLimbs,
        toIR !(mkZext {to=I64} sizeAbs),
        toIR restBits
        ]
      msbLimbAddr <- getObjectSlotAddrVar {t=I64} newObj !(mkZext !(mkSub maxLimbsCount (Const I32 1)))
      store msbLimb msbLimbAddr
      ) (do
      srcLimbs <- getObjectPayloadAddr {t=I8} integerObj
      higherLimbsAddr <- getObjectSlotAddrVar {t=I8} newObj !(mkZext {to=I64} fullLimbs)
      voidCall "ccc" "@llvm.memcpy.p1i8.p1i8.i64" [
        toIR higherLimbsAddr,
        toIR srcLimbs,
        toIR !(mkMul !(mkZext size) (Const I64 GMP_LIMB_SIZE)),
        "i1 false"
        ]
      )

    isNegative <- icmp "slt" size (Const I32 0)
    normaliseIntegerSize newObj maxLimbsCount isNegative
    pure newObj
    )

export
shiftRightInteger : IRValue IRObjPtr -> IRValue IRObjPtr -> Codegen (IRValue IRObjPtr)
shiftRightInteger integerObj bitCountObj = do
  bitCount <- mkTrunc {to=I32} !(unboxIntegerUnsigned bitCountObj)

  size <- getObjectSize integerObj
  unchanged <- mkOr !(icmp "eq" size (Const I32 0)) !(icmp "eq" bitCount (Const I32 0))
  mkIf (pure unchanged) (pure integerObj) (do
    sizeAbs <- mkAbs size
    fullLimbs <- mkUDiv bitCount (Const I32 GMP_LIMB_BITS)
    maxLimbsCount <- mkSub sizeAbs fullLimbs

    mkIf (icmp "sle" maxLimbsCount (Const I32 0)) integer0 (do
      newObj <- dynamicAllocate !(mkZext !(mkMul maxLimbsCount (Const I32 GMP_LIMB_SIZE)))

      restBits <- mkURem bitCount (Const I32 GMP_LIMB_BITS)
      mkIf_ (icmp "ne" (Const I32 0) restBits) (do
        srcHigherLimbs <- getObjectSlotAddrVar {t=I64} integerObj !(mkZext {to=I64} fullLimbs)
        dstLimbsAddr <- getObjectPayloadAddr {t=I64} newObj
        ignore $ call {t=I64} "ccc" "@__gmpn_rshift" [
          toIR dstLimbsAddr,
          toIR srcHigherLimbs,
          toIR !(mkZext {to=I64} maxLimbsCount),
          toIR restBits
          ]
        ) (do
        srcHigherLimbs <- getObjectSlotAddrVar {t=I8} integerObj !(mkZext {to=I64} fullLimbs)
        dstLimbsAddr <- getObjectPayloadAddr {t=I8} newObj
        voidCall "ccc" "@llvm.memcpy.p1i8.p1i8.i64" [
          toIR dstLimbsAddr,
          toIR srcHigherLimbs,
          toIR !(mkMul !(mkZext maxLimbsCount) (Const I64 GMP_LIMB_SIZE)),
          "i1 false"
          ]
        )

      isNegative <- icmp "slt" size (Const I32 0)
      normaliseIntegerSize newObj maxLimbsCount isNegative
      pure newObj
      )

    )

IEEE_DOUBLE_MASK_EXP  : Bits64
IEEE_DOUBLE_MASK_EXP  = 0x7ff0000000000000
IEEE_DOUBLE_MASK_FRAC : Bits64
IEEE_DOUBLE_MASK_FRAC = 0x000fffffffffffff
IEEE_DOUBLE_MASK_SIGN : Bits64
IEEE_DOUBLE_MASK_SIGN = 0x8000000000000000
IEEE_DOUBLE_INF_POS   : Bits64
IEEE_DOUBLE_INF_POS   = 0x7ff0000000000000
IEEE_DOUBLE_INF_NEG   : Bits64
IEEE_DOUBLE_INF_NEG   = 0xfff0000000000000

export
castDoubleToInteger : IRValue IRObjPtr -> Codegen (IRValue IRObjPtr)
castDoubleToInteger floatObj = do
  floatBitsAsI64 <- getObjectSlot {t=I64} floatObj 0
  exponent <- mkShiftR !(mkAnd (Const I64 $ cast IEEE_DOUBLE_MASK_EXP) floatBitsAsI64) (Const I64 52)
  -- NaN and infinity will be returned as "0"
  isInfOrNaN <- icmp "eq" exponent (Const I64 0x7ff)
  -- absolute values < 1.0 will be returned as "0"
  isSmallerThanOne <- icmp "ult" exponent (Const I64 1023)
  returnZero <- mkOr isInfOrNaN isSmallerThanOne

  mkIf (pure returnZero) {- then -} integer0 {- else -} (do
    fraction <- mkAnd (Const I64 $ cast IEEE_DOUBLE_MASK_FRAC) floatBitsAsI64
    -- highest bit is sign bit in both, Double and I64, so we can just use that one
    isNegative <- icmp "slt" floatBitsAsI64 (Const I64 0)
    initial <- mkOr fraction (Const I64 0x10000000000000)
    toShift <- mkSub exponent (Const I64 1075)
    shiftLeft <- icmp "sgt" toShift (Const I64 0)
    mkIf (pure shiftLeft) (do
      let maxLimbCount = Const I64 17
      payloadSize <- mkMul maxLimbCount (Const I64 GMP_LIMB_SIZE)
      -- requiredBits <- (exponent - 1022)
      -- requiredLimbs <- (requiredBits+63) / 64
      -- newObj <- allocObject (requiredLimbs * 8)
      newObj <- dynamicAllocate payloadSize
      payloadAddr <- getObjectPayloadAddr {t=I8} newObj
      appendCode $ "  call void @llvm.memset.p1i8.i64(" ++ toIR payloadAddr ++ ", i8 0, " ++ toIR payloadSize ++ ", i1 false)"
      putObjectSlot newObj (Const I64 0) initial
      absRealNewSize <- call {t=I64} "ccc" "@rapid_bigint_lshift_inplace" [
        toIR !(getObjectPayloadAddr {t=I64} newObj),
        toIR maxLimbCount,
        toIR !(mkTrunc {to=I32} toShift)
      ]
      signedNewSize <- mkSelect isNegative !(mkSub (Const I64 0) absRealNewSize) absRealNewSize
      size32 <- mkTrunc {to=I32} signedNewSize
      newHeader <- mkHeader OBJECT_TYPE_ID_BIGINT size32
      putObjectHeader newObj newHeader
      pure newObj
      ) (do
        newObj <- dynamicAllocate (Const I64 8)
        toShiftRight <- mkSub (Const I64 0) toShift
        shifted <- mkShiftR initial toShiftRight
        putObjectSlot newObj (Const I64 0) shifted
        signedNewSize32 <- mkSelect isNegative (Const I32 (-1)) (Const I32 1)
        newHeader <- mkHeader OBJECT_TYPE_ID_BIGINT signedNewSize32
        putObjectHeader newObj newHeader
        pure newObj
      )
    )

export
castIntegerToDouble : IRValue IRObjPtr -> Codegen (IRValue IRObjPtr)
castIntegerToDouble intObj = do
  size <- getObjectSize intObj
  isZero <- icmp "eq" (Const I32 0) size

  mkIf (pure isZero) (cgMkConstDouble 0.0) (do
    sizeAbs <- mkAbs size
    highestLimbIndex <- mkZext {to=I64} !(mkSub sizeAbs (Const I32 1))
    msbLimbAddr <- getObjectSlotAddrVar {t=I64} intObj highestLimbIndex
    msbLimb <- load msbLimbAddr
    countLeadingZeros <- SSA I64 <$> assignSSA ("  call ccc i64 @llvm.ctlz.i64(" ++ toIR msbLimb ++ ", i1 1)")
    exponentOffset <- mkAdd (Const I64 (1023 + 63)) !(mkMul highestLimbIndex (Const I64 64))
    exponent <- mkSub exponentOffset countLeadingZeros

    isInfinity <- icmp "ugt" exponent (Const I64 2046)
    doubleAsBits <- mkIf (pure isInfinity) (do
        pure (Const I64 $ cast IEEE_DOUBLE_INF_POS)
      ) (do
        fracShiftLeft <- icmp "uge" countLeadingZeros (Const I64 12)
        shiftedFraction <- mkIf (pure fracShiftLeft) (do
            mkShiftL msbLimb !(mkSub countLeadingZeros (Const I64 11))
          ) (do
            mkShiftR msbLimb !(mkSub (Const I64 11) countLeadingZeros)
          )
        fraction <- mkIf (icmp "eq" sizeAbs (Const I32 1)) (do
            pure shiftedFraction
          ) (do
            mkIf (pure fracShiftLeft) (do
                secondMsbLimbAddr <- getObjectSlotAddrVar {t=I64} intObj !(mkSub highestLimbIndex (Const I64 1))
                secondMsbLimb <- load secondMsbLimbAddr
                fractionLowerPart <- mkShiftR secondMsbLimb !(mkSub (Const I64 (64 + 11)) countLeadingZeros)
                mkOr shiftedFraction fractionLowerPart
              ) (do
                pure shiftedFraction
              )
          )
        shiftedExponent <- mkShiftL exponent (Const I64 52)
        maskedFraction <- mkAnd (Const I64 $ cast IEEE_DOUBLE_MASK_FRAC) fraction
        mkOr shiftedExponent maskedFraction
      )
    isNegative <- icmp "slt" size (Const I32 0)
    sign <- mkSelect isNegative (Const I64 $ cast IEEE_DOUBLE_MASK_SIGN) (Const I64 0)
    signedDoubleAsBits <- mkOr sign doubleAsBits
    cgMkDoubleFromBits signedDoubleAsBits
    )

export
castStringToInteger : IRValue IRObjPtr -> Codegen (IRValue IRObjPtr)
castStringToInteger strObj = do
  strLength <- getObjectSize strObj
  maxLimbsCount <- mkUDiv strLength (Const I32 GMP_ESTIMATE_DIGITS_PER_LIMB)
  -- GMP requires 1 limb scratch space
  maxLimbsCountPlus1 <- mkAdd maxLimbsCount (Const I32 1)

  newObj <- dynamicAllocate !(mkZext !(mkMul maxLimbsCountPlus1 (Const I32 GMP_LIMB_SIZE)))
  putObjectHeader newObj !(mkHeader OBJECT_TYPE_ID_BIGINT maxLimbsCountPlus1)
  ignore $ call {t=I32} "ccc" "@rapid_bigint_set_str" [
    toIR newObj,
    toIR strObj
    ]
  pure newObj

export
castIntegerToString : IRValue IRObjPtr -> Codegen (IRValue IRObjPtr)
castIntegerToString i1 = do
  s1 <- getObjectSize i1
  u1 <- mkZext {to=I64} !(mkAbs s1)

  isZero <- icmp "eq" s1 (Const I32 0)

  mkIf (pure isZero) (mkStr "0") (do
      maxDigits <- call {t=TARGET_SIZE_T} "ccc" "@__gmpn_sizeinbase" [toIR !(getObjectPayloadAddr {t=MP_LIMB_T} i1), toIR u1, "i32 10"]
      isNegative <- icmp "slt" s1 (Const I32 0)

      -- we need to add one extra byte of "scratch space" for mpn_get_str
      -- if the number is negative we need one character more for the leading minus
      needsSign <- mkSelect isNegative (Const I64 2) (Const I64 1)
      maxDigitsWithSign <- mkAdd maxDigits needsSign

      newStr <- dynamicAllocate maxDigitsWithSign
      newHeader <- mkHeader OBJECT_TYPE_ID_STR !(mkTrunc maxDigitsWithSign)
      putObjectHeader newStr newHeader

      actualDigits <- call {t=I64} "ccc" "@rapid_bigint_get_str" [toIR newStr, toIR i1, "i32 10"]
      actualLengthHeader <- mkHeader OBJECT_TYPE_ID_STR !(mkTrunc actualDigits)
      putObjectHeader newStr actualLengthHeader

      pure newStr
    )


module Examples.Declassification.DatingSystem

import Pruviloj.Disjoint

import DepSec.Labeled
import DepSec.DIO
import DepSec.Declassification

import Examples.AuthorityLattice

%language ElabReflection

%default total

data Gender : Type where
 Unknown : Gender
 M : Gender
 F : Gender

implementation DecEq Gender where
  decEq Unknown Unknown = Yes Refl
  decEq Unknown M = No (%runElab disjoint)
  decEq Unknown F = No (%runElab disjoint)
  decEq M Unknown = No (%runElab disjoint)
  decEq M M = Yes Refl
  decEq M F = No (%runElab disjoint)
  decEq F Unknown = No (%runElab disjoint)
  decEq F M = No (%runElab disjoint)
  decEq F F = Yes Refl

record ProfileInfo (Owner : Principal) where
  constructor MkProfileInfo
  name      : Labeled Owner String
  gender    : Labeled Owner Gender
  birthdate : Labeled Owner String
  interests : Labeled Owner String
  city      : Labeled Owner String

DiscoveryResult : Principal -> Type
DiscoveryResult = ProfileInfo

Discoverer : Type
Discoverer = {bOwner : Principal}
          -> (b : ProfileInfo bOwner)
          -> DIO bOwner (Maybe (DiscoveryResult bOwner))

record Profile (Owner : Principal) where
  constructor MkProfile
  info       : ProfileInfo Owner
  discoverer : Discoverer

data GenProfile : (Principal : Type) -> Type where
  MkGenProfile : {t : Principal} -> Profile t -> GenProfile Principal

data Authority : a -> Type where
  ||| TCB
  MkAuthority : (t : a) -> Authority t

authorityHatch : { l, l' : Principal }
              -> (s ** (l = s, Authority s)) -> Labeled l a -> Labeled l' a
authorityHatch {l} = tokenHatch (\s => (l=s, Authority s))

--------------------------------------------------------------------------------
-- Convenience
--------------------------------------------------------------------------------
%hint
hint1 : {a : Principal} -> leq a a
hint1 {a} = reflexive a

implicit
genderToString : Gender -> String
genderToString M = "Male"
genderToString F = "Female"
genderToString Unknown = "Unknown"

implicit
profileToString : (p : DiscoveryResult viewer) -> DIO viewer String
profileToString (MkProfileInfo name gender birthdate likes location) =
   pure ("NAME: "      ++ !(unlabel name) ++ "\n" ++
         "BIRTHDATE: " ++ !(unlabel birthdate) ++ "\n" ++
         "LIKES: "     ++ !(unlabel likes) ++ "\n" ++
         "LOCATION: "  ++ !(unlabel location) ++ "\n")

implicit
toLabeled : Poset labelType => {l : labelType} -> String -> Labeled l String
toLabeled x = label x

implicit
toLabeledG : Poset labelType => {l : labelType} -> Gender -> Labeled l Gender
toLabeledG x = label x

--------------------------------------------------------------------------------
-- UTCB
--------------------------------------------------------------------------------
sampleDiscoverer : {A, B : Principal}
                -> Authority A
                -> (a : ProfileInfo A)
                -> (b : ProfileInfo B)
                -> DIO B (Maybe (DiscoveryResult B))
sampleDiscoverer {A} {B} auth a b =
  do bGender <- unlabel $ gender b
     aGender <- unlabel $ authorityHatch (A ** (Refl, auth)) (gender a)
     aName <- unlabel $ authorityHatch (A ** (Refl, auth)) (name a)
     case decEq bGender aGender of
       Yes prf => pure Nothing
       No contra =>
         let result : DiscoveryResult B = MkProfileInfo (label aName) (label Unknown) "" "" ""
         in pure $ Just result

--------------------------------------------------------------------------------
-- TCB
--------------------------------------------------------------------------------
profileBob : Profile Bob
profileBob =
  let profileInfo = MkProfileInfo "Bob" M "1994" "Nothing" "Aarhus"
  in MkProfile profileInfo (sampleDiscoverer (MkAuthority Bob) profileInfo)

profileAlice : Profile Alice
profileAlice =
  let profileInfo = MkProfileInfo "Alice" F "1993" "Hunting" "Aarhus"
  in MkProfile profileInfo (sampleDiscoverer (MkAuthority Alice) profileInfo)

profileChuck : Profile Chuck
profileChuck =
  let profileInfo = MkProfileInfo "Chuck" M "1992" "Swimming" "Ribe"
  in  MkProfile profileInfo (sampleDiscoverer (MkAuthority Chuck) profileInfo)

profileQueueAlice : List (GenProfile Principal)
profileQueueAlice = [MkGenProfile profileBob, MkGenProfile profileChuck]

printAlice : String -> DIO Alice ()
printAlice x = lift $ putStrLn x

runThroughQueueAlice : (acc : DIO Alice ()) -> List (GenProfile Principal) -> DIO Alice ()
runThroughQueueAlice acc [] = acc
runThroughQueueAlice acc ((MkGenProfile p) :: xs) =
  let thisProfilePrint : DIO Alice ()
    = do acc
         result <- (discoverer p) (info profileAlice)
         case result of
           Nothing => printAlice "Did not wanna be seen"
           Just result' => printAlice !result'
  in runThroughQueueAlice thisProfilePrint xs

main : IO ()
main =
  do putStrLn "########## Alice's queue is checked ##########\n"
     run $ runThroughQueueAlice (pure ()) profileQueueAlice
     putStrLn "########## Done checking ##########"
     pure ()


module ArgsWrongOrder

import IdrisPrelude

%hide Prelude.Nat.Nat

myFun : (a : Type) -> a -> Nat -> Nat
myFun _ x y = y

myApp : Nat
myApp = myFun _ Zero (IdrisPrelude.Nil Nat)


module Tools.AstF

public export
interface AstF (idx : Type) (f : (idx -> Type) -> (idx -> Type)) where
  localMap : (f1 : idx -> Type) -> (f2 : idx -> Type) -> ((i : idx) -> f1 i -> f2 i)
          -> (i : idx) -> f f1 i -> f f2 i
  -- TODO find a food localFold
  -- localFold : Monoid m => (i : idx) -> f (\_ => m) i -> m


%foreign "scheme:my-mul"
myMul : Int -> Int -> Int

main : IO ()
main = do
  putStrLn $ "6 * 7 = " ++ show (myMul 6 7)


module Circuits.NetList.Linear.Usage.TermType.Use.Channel

import Decidable.Equality

import Data.Nat
import Data.List.Elem
import Data.List.Quantifiers

import Data.Vect
import Data.Vect.AtIndex
import Data.Vect.Quantifiers

import Data.String

import Toolkit.Data.Whole
import Toolkit.Data.DList
import Toolkit.Data.DList.Elem

import Toolkit.Data.Vect.Extra
import Toolkit.Data.DVect

import Circuits.Common

import Circuits.NetList.Types

import Circuits.NetList.Linear.Usage.DataType
import Circuits.NetList.Linear.Usage.DataType.Use.Full
import Circuits.NetList.Linear.Usage.DataType.Use.Partial

import Circuits.NetList.Linear.Usage.TermType

%default total

public export
data CanProject : (how  : Project dir)
               -> (type : Ty)
               -> (u    : Usage type)
                       -> Type
  where
    ProjectWrite : (prf : IsFree                   type          ein)
                       -> CanProject WRITE (TyChan type) (TyChan ein eout)

    ProjectRead : (prf : IsFree                  type              eout)
                      -> CanProject READ (TyChan type) (TyChan ein eout)

Uninhabited (CanProject how TyGate u) where
  uninhabited (ProjectWrite prf) impossible
  uninhabited (ProjectRead prf) impossible

Uninhabited (CanProject how TyUnit u) where
  uninhabited (ProjectWrite prf) impossible
  uninhabited (ProjectRead prf) impossible

Uninhabited (CanProject how (TyPort (f,dtype)) u) where
  uninhabited (ProjectWrite prf) impossible
  uninhabited (ProjectRead prf) impossible

cannotWriteTo : (IsFree type ein -> Void) -> CanProject WRITE (TyChan type) (TyChan ein eout) -> Void
cannotWriteTo f (ProjectWrite prf) = f prf

cannotReadFrom : (IsFree type eout -> Void) -> CanProject READ (TyChan type) (TyChan ein eout) -> Void
cannotReadFrom f (ProjectRead prf) = f prf

export
canProject : {type : Ty}
          -> (how : Project dir)
          -> (u   : Usage type)
                 -> Dec (CanProject how type u)
canProject how (TyPort u) = No absurd

canProject WRITE (TyChan ein eout) with (isFree ein)
  canProject WRITE (TyChan ein eout) | (Yes prf)
    = Yes (ProjectWrite prf)

  canProject WRITE (TyChan ein eout) | (No contra)
    = No (cannotWriteTo contra)

canProject READ (TyChan ein eout) with (isFree eout)
  canProject READ (TyChan ein eout) | (Yes prf)
    = Yes (ProjectRead prf)

  canProject READ (TyChan ein eout) | (No contra)
    = No (cannotReadFrom contra)

canProject how TyGate = No absurd
canProject how TyUnit = No absurd

public export
data Use : (how  : Project dir)
        -> (type : Ty)
        -> (old  : Usage type)
        -> (prf  : CanProject how type old)
        -> (new  : Usage type)
                -> Type

  where
    IsWrite : Use               type          ein                     free          newU       prf
           -> Use WRITE (TyChan type) (TyChan ein eout) (ProjectWrite free) (TyChan newU eout)

    IsRead : Use              type             eout               free              newU  prf
          -> Use READ (TyChan type)(TyChan ein eout) (ProjectRead free) (TyChan ein newU)

public export
data Result : (how  : Project dir)
           -> (type : Ty)
           -> (old  : Usage type)
           -> (prf  : CanProject how type old)
                   -> Type
  where
    R : (new : Usage  (TyChan type))
     -> (prf : Use    how (TyChan type) old proj new)
            -> Result how (TyChan type) old proj

export
use : {type : Ty}
   -> {how  : Project dir}
   -> {old  : Usage type}
   -> (prf  : CanProject how type old)
           -> Result     how type old prf
use (ProjectWrite prf) with (use prf)
  use (ProjectWrite prf) | (R u used result)
    = R (TyChan u _) (IsWrite result)
use (ProjectRead prf) with (use prf)
  use (ProjectRead prf) | (R u used result)
    = R (TyChan _ u) (IsRead result)

-- [ EOF ]


module Core.SchemeEval.Builtins

import Core.Context
import Core.SchemeEval.ToScheme
import Core.TT

import Data.Vect
import Libraries.Utils.Scheme

-- Integers are wrapped, so unwrap then wrap again
add : Maybe IntKind -> SchemeObj Write -> SchemeObj Write -> SchemeObj Write
add (Just (Signed (P n))) x y = Apply (Var "ct-s+") [x, y, toScheme (n-1)]
add (Just (Unsigned n)) x y = Apply (Var "ct-u+") [x, y, toScheme n]
add _ x y = Apply (Var "ct+") [x, y]

sub : Maybe IntKind -> SchemeObj Write -> SchemeObj Write -> SchemeObj Write
sub (Just (Signed (P n))) x y = Apply (Var "ct-s-") [x, y, toScheme (n-1)]
sub (Just (Unsigned n)) x y = Apply (Var "ct-u-") [x, y, toScheme n]
sub _ x y = Apply (Var "ct-") [x, y]

mul : Maybe IntKind -> SchemeObj Write -> SchemeObj Write -> SchemeObj Write
mul (Just (Signed (P n))) x y = Apply (Var "ct-s*") [x, y, toScheme (n-1)]
mul (Just (Unsigned n)) x y = Apply (Var "ct-u*") [x, y, toScheme n]
mul _ x y = Apply (Var "ct*") [x, y]

div : Maybe IntKind -> SchemeObj Write -> SchemeObj Write -> SchemeObj Write
div (Just (Signed Unlimited)) x y = Apply (Var "ct/") [x, y]
div (Just (Signed (P n))) x y = Apply (Var "ct-s/") [x, y, toScheme (n-1)]
div (Just (Unsigned n)) x y = Apply (Var "ct-u/") [x, y, toScheme n]
div _ x y = Apply (Var "ct/") [x, y]

mod : SchemeObj Write -> SchemeObj Write -> SchemeObj Write
mod x y = Apply (Var "ct-mod") [x, y]

shl : Maybe IntKind -> SchemeObj Write -> SchemeObj Write -> SchemeObj Write
shl (Just (Signed (P n))) x y = Apply (Var "ct-bits-shl-signed") [x, y, toScheme (n-1)]
shl (Just (Unsigned n)) x y = Apply (Var "ct-bits-shl") [x, y, toScheme n]
shl _ x y = Apply (Var "ct-shl") [x, y]

shr : Maybe IntKind -> SchemeObj Write -> SchemeObj Write -> SchemeObj Write
shr _ x y = Apply (Var "ct-shr") [x, y]

-- Doubles don't need wrapping, since there's only one double type
addDbl : SchemeObj Write -> SchemeObj Write -> SchemeObj Write
addDbl x y = Apply (Var "+") [x, y]

subDbl : SchemeObj Write -> SchemeObj Write -> SchemeObj Write
subDbl x y = Apply (Var "-") [x, y]

mulDbl : SchemeObj Write -> SchemeObj Write -> SchemeObj Write
mulDbl x y = Apply (Var "*") [x, y]

divDbl : SchemeObj Write -> SchemeObj Write -> SchemeObj Write
divDbl x y = Apply (Var "/") [x, y]

-- Check necessary arguments are in canonical form before applying the
-- operator, otherwise return the blocked form
-- Current assumption is that all primitives that we can evaluate at
-- compile time work on constants, if they do anything in Scheme at all.
canonical : SchemeObj Write ->
            Vect n (SchemeObj Write) -> SchemeObj Write -> SchemeObj Write
canonical blk [] body = body
canonical blk (n :: ns) body
   = If (Apply (Var "ct-isConstant") [n]) (canonical blk ns body) blk

-- Return blocked application if a partial operator is given an input
-- on which it's undefined
testPartial : SchemeObj Write -> SchemeObj Write -> SchemeObj Write
testPartial blk res
    = Let "p-0" res $
          (If (Apply (Var "ct-isConstant") [Var "p-0"])
              (Var "p-0")
              blk)

unaryOp : SchemeObj Write -> String ->
          SchemeObj Write -> SchemeObj Write
unaryOp blk op x = canonical blk [x] $ Apply (Var op) [x]

binOp : SchemeObj Write -> String ->
        SchemeObj Write -> SchemeObj Write -> SchemeObj Write
binOp blk op x y = canonical blk [x, y] $ Apply (Var op) [x, y]

ternaryOp : SchemeObj Write -> String ->
            SchemeObj Write -> SchemeObj Write -> SchemeObj Write ->
            SchemeObj Write
ternaryOp blk op x y z = canonical blk [x, y, z] $ Apply (Var op) [x, y, z]

int : SchemeObj Write -> SchemeObj Write
int obj = Vector (-100) [obj]

int8 : SchemeObj Write -> SchemeObj Write
int8 obj = Vector (-101) [obj]

int16 : SchemeObj Write -> SchemeObj Write
int16 obj = Vector (-102) [obj]

int32 : SchemeObj Write -> SchemeObj Write
int32 obj = Vector (-103) [obj]

int64 : SchemeObj Write -> SchemeObj Write
int64 obj = Vector (-104) [obj]

integer : SchemeObj Write -> SchemeObj Write
integer obj = Vector (-105) [obj]

bits8 : SchemeObj Write -> SchemeObj Write
bits8 obj = Vector (-106) [obj]

bits16 : SchemeObj Write -> SchemeObj Write
bits16 obj = Vector (-107) [obj]

bits32 : SchemeObj Write -> SchemeObj Write
bits32 obj = Vector (-108) [obj]

bits64 : SchemeObj Write -> SchemeObj Write
bits64 obj = Vector (-109) [obj]

wrap : IntKind -> SchemeObj Write -> SchemeObj Write
wrap (Signed Unlimited) = integer
wrap (Signed (P 8)) = int8
wrap (Signed (P 16)) = int16
wrap (Signed (P 32)) = int32
wrap (Signed (P 64)) = int64
wrap (Unsigned 8) = bits8
wrap (Unsigned 16) = bits16
wrap (Unsigned 32) = bits32
wrap (Unsigned 64) = bits64
wrap _ = integer

-- Result has to be wrapped in Int, which is Vector (-100)
boolOp : SchemeObj Write -> String ->
         SchemeObj Write -> SchemeObj Write -> SchemeObj Write
boolOp blk op x y
    = canonical blk [x, y] $
         int $
             Apply (Var "or")
                [Apply (Var "and") [Apply (Var op) [x, y],
                                    IntegerVal 1],
                 IntegerVal 0]

applyIntCast : IntKind -> IntKind -> SchemeObj Write -> SchemeObj Write
applyIntCast _ (Signed Unlimited) x = x
applyIntCast (Signed m) k@(Signed (P n)) x
    = if P n >= m
         then x
         else wrap k $ Apply (Var "ct-cast-signed") [x, toScheme (n - 1)]
applyIntCast (Unsigned m) k@(Signed (P n)) x
    = if n > m
         then x
         else wrap k $ Apply (Var "ct-cast-signed") [x, toScheme (n - 1)]
applyIntCast (Signed _) k@(Unsigned n) x
    = wrap k $ Apply (Var "ct-cast-unsigned") [x, toScheme n]
applyIntCast (Unsigned m) (Unsigned n) x
    = if n >= m
         then x
         else Apply (Var "ct-cast-unsigned") [x, toScheme n]

applyCast : SchemeObj Write ->
            Constant -> Constant ->
            SchemeObj Write -> SchemeObj Write
applyCast blk CharType to x
    = canonical blk [x] $
        case intKind to of
           Nothing =>
              case to of
                   StringType => Apply (Var "string") [x]
                   _ => blk
           Just (Signed Unlimited) => integer $ Apply (Var "char->integer") [x]
           Just k@(Signed (P n)) => wrap k $ Apply (Var "ct-cast-char-boundedInt") [x, toScheme (n - 1)]
           Just k@(Unsigned n) => wrap k $ Apply (Var "ct-cast-char-boundedUInt") [x, toScheme n]
applyCast blk from CharType x
    = canonical blk [x] $
        case intKind from of
           Nothing => blk
           Just k => Apply (Var "ct-cast-int-char") [x]
applyCast blk StringType to x
    = canonical blk [x] $
        case intKind to of
           Nothing => case to of
                           DoubleType => Apply (Var "ct-cast-string-double") [x]
                           _ => blk
           Just (Signed Unlimited) => integer $ Apply (Var "ct-cast-string-int") [x]
           Just k@(Signed (P n)) => wrap k $ Apply (Var "ct-cast-string-boundedInt") [x, toScheme (n - 1)]
           Just k@(Unsigned n) => wrap k $ Apply (Var "ct-cast-string-boundedUInt") [x, toScheme n]
applyCast blk from StringType x
    = canonical blk [x] $
        case intKind from of
           Nothing => case from of
                           DoubleType => Apply (Var "number->string") [x]
                           _ => blk
           Just k => Apply (Var "ct-cast-number-string") [x]
applyCast blk DoubleType to x
    = canonical blk [x] $
        case intKind to of
           Nothing => case to of
                           StringType => Apply (Var "number->string") [x]
                           _ => blk
           Just (Signed Unlimited) => integer $ Apply (Var "ct-exact-truncate") [x]
           Just k@(Signed (P n)) => wrap k $ Apply (Var "ct-exact-truncate-boundedInt") [x, toScheme (n - 1)]
           Just k@(Unsigned n) => wrap k $ Apply (Var "ct-exact-truncate-boundedUInt") [x, toScheme n]
applyCast blk from DoubleType x
    = canonical blk [x] $
        case intKind from of
           Nothing => case from of
                           StringType => Apply (Var "ct-cast-string-double") [x]
                           _ => blk
           Just k => Apply (Var "ct-int-double") [x]
applyCast blk from to x
    = canonical blk [x] $
        case (intKind from, intKind to) of
           (Just f, Just t) => applyIntCast f t x
           _ => blk

applyOp : SchemeObj Write -> -- if we don't have arguments in canonical form
          PrimFn n -> Vect n (SchemeObj Write) ->
          SchemeObj Write
applyOp blk (Add DoubleType) [x, y] = binOp blk "+" x y
applyOp blk (Sub DoubleType) [x, y] = binOp blk "-" x y
applyOp blk (Mul DoubleType) [x, y] = binOp blk "*" x y
applyOp blk (Div DoubleType) [x, y] = binOp blk "/" x y
applyOp blk (Neg DoubleType) [x] = unaryOp blk "-" x
applyOp blk (Add ty) [x, y] = canonical blk [x, y] $ add (intKind ty) x y
applyOp blk (Sub ty) [x, y] = canonical blk [x, y] $ sub (intKind ty) x y
applyOp blk (Mul ty) [x, y] = canonical blk [x, y] $ mul (intKind ty) x y
applyOp blk (Div ty) [x, y] = canonical blk [x, y] $ div (intKind ty) x y
applyOp blk (Mod ty) [x, y] = canonical blk [x, y] $ mod x y
applyOp blk (Neg ty) [x] = canonical blk [x] $ Apply (Var "ct-neg") [x]
applyOp blk (ShiftL ty) [x, y] = canonical blk [x, y] $ shl (intKind ty) x y
applyOp blk (ShiftR ty) [x, y] = canonical blk [x, y] $ shr (intKind ty) x y
applyOp blk (BAnd ty) [x, y] = binOp blk "ct-and" x y
applyOp blk (BOr ty) [x, y] = binOp blk "ct-or" x y
applyOp blk (BXOr ty) [x, y] = binOp blk "ct-xor" x y
applyOp blk (LT CharType) [x, y] = boolOp blk "char<?" x y
applyOp blk (LTE CharType) [x, y] = boolOp blk "char<=?" x y
applyOp blk (EQ CharType) [x, y] = boolOp blk "char=?" x y
applyOp blk (GTE CharType) [x, y] = boolOp blk "char>=?" x y
applyOp blk (GT CharType) [x, y] = boolOp blk "char>?" x y
applyOp blk (LT StringType) [x, y] = boolOp blk "string<?" x y
applyOp blk (LTE StringType) [x, y] = boolOp blk "string<=?" x y
applyOp blk (EQ StringType) [x, y] = boolOp blk "string=?" x y
applyOp blk (GTE StringType) [x, y] = boolOp blk "string>=?" x y
applyOp blk (GT StringType) [x, y] = boolOp blk "string>?" x y
applyOp blk (LT DoubleType) [x, y] = boolOp blk "<" x y
applyOp blk (LTE DoubleType) [x, y] = boolOp blk "<=" x y
applyOp blk (EQ DoubleType) [x, y] = boolOp blk "=" x y
applyOp blk (GTE DoubleType) [x, y] = boolOp blk ">=" x y
applyOp blk (GT DoubleType) [x, y] = boolOp blk ">" x y
applyOp blk (LT ty) [x, y] = boolOp blk "ct<" x y
applyOp blk (LTE ty) [x, y] = boolOp blk "ct<=" x y
applyOp blk (EQ ty) [x, y] = boolOp blk "ct=" x y
applyOp blk (GTE ty) [x, y] = boolOp blk "ct>=" x y
applyOp blk (GT ty) [x, y] = boolOp blk "ct>" x y
applyOp blk StrLength [x]
    = canonical blk [x] $ Vector (-100) [Apply (Var "string-length") [x]]
applyOp blk StrHead [x]
    = canonical blk [x] $ Apply (Var "string-ref")
                                [x, IntegerVal 0]
applyOp blk StrTail [x]
    = canonical blk [x] $ Apply (Var "substring")
                                [x, IntegerVal 1,
                                 Apply (Var "string-length") [x]]
applyOp blk StrIndex [x, y]
    = canonical blk [x, y] $ testPartial blk $
         Apply (Var "ct-string-ref") [x, y]
applyOp blk StrCons [x, y]
    = canonical blk [x, y] $ Apply (Var "ct-string-cons") [x, y]
applyOp blk StrAppend [x, y]
    = canonical blk [x, y] $ Apply (Var "string-append") [x, y]
applyOp blk StrReverse [x]
    = canonical blk [x] $ Apply (Var "ct-string-reverse") [x]
applyOp blk StrSubstr [x, y, z]
    = canonical blk [x, y, z] $ Apply (Var "ct-string-substr") [x]

applyOp blk DoubleExp [x] = unaryOp blk "flexp" x
applyOp blk DoubleLog [x] = unaryOp blk "fllog" x
applyOp blk DoublePow [x, y] = binOp blk "expt" x y
applyOp blk DoubleSin [x] = unaryOp blk "flsin" x
applyOp blk DoubleCos [x] = unaryOp blk "flcos" x
applyOp blk DoubleTan [x] = unaryOp blk "fltan" x
applyOp blk DoubleASin [x] = unaryOp blk "flasin" x
applyOp blk DoubleACos [x] = unaryOp blk "flacos" x
applyOp blk DoubleATan [x] = unaryOp blk "flatan" x
applyOp blk DoubleSqrt [x] = unaryOp blk "flsqrt" x
applyOp blk DoubleFloor [x] = unaryOp blk "flfloor" x
applyOp blk DoubleCeiling [x] = unaryOp blk "flceiling" x

applyOp blk (Cast from to) [x] = applyCast blk from to x
applyOp blk BelieveMe [_, _, x] = x
applyOp blk Crash [_, msg] = blk

mkArgList : Int -> (n : Nat) -> Vect n String
mkArgList i Z = []
mkArgList i (S k) = ("x-" ++ show i) :: mkArgList (i + 1) k

export
compileBuiltin : {farity : Nat} ->
                Name -> PrimFn farity -> SchemeObj Write
compileBuiltin nm fn
    = let args = mkArgList 0 farity in
          bindArgs args [] args
  where
    makeBlockedApp : Vect n String -> SchemeObj Write
    makeBlockedApp args = Vector (-2) [toScheme nm, vars args]
      where
        vars : forall n . Vect n String -> SchemeObj Write
        vars [] = Null
        vars (x :: xs) = Cons (Var x) (vars xs)

    bindArgs : Vect n String -> Vect n' String ->
               Vect farity String -> SchemeObj Write
    bindArgs [] done args = applyOp (makeBlockedApp args) fn (map Var args)
    bindArgs (x :: xs) done args
        = Vector (-9) [makeBlockedApp (reverse done),
                       Lambda [x] (bindArgs xs (x :: done) args)]


module QStateT

import Control.Monad.State
import LinearTypes
import Control.Linear.LIO

%default total

public export
R : Type -> Type
R = L IO {use = Linear}

--Inspired by the indexed state monad in Haskell, and for linear types
public export
data QStateT : (initialType : Type) -> (finalType : Type) -> (returnType : Type) -> Type where
  MkQST : (1 _ : (1 _ : initialType) -> R (LPair finalType returnType)) -> QStateT initialType finalType returnType

||| Unwrap and apply a QStateT monad computation.
public export
runQStateT : (1 _ : initialType) -> (1 _ : QStateT initialType finalType returnType) -> R (LPair finalType returnType)
runQStateT i (MkQST f) = f i

public export
pure : (1 _ : a) -> QStateT t t a
pure x = MkQST (\st => pure1 (st # x))

public export
(>>=) : (1 _ : QStateT i m a) -> (1 _ : ((1 _ : a) -> QStateT m o b)) -> QStateT i o b
v >>= f = MkQST $ \i => do
                         (a # m) <- runQStateT i v
                         runQStateT a (f m)

public export
modify : ((1 _ : i) -> o) -> QStateT i o ()
modify f = MkQST $ \i => pure1 (f i # ())


module Node.Free.FileSystem

import Control.Monad.Free

%access public export

FilePath : Type
FilePath = String

Contents : Type
Contents = String

public export
data FileSystemF next
  = WriteFile FilePath Contents next
  | ReadFile FilePath (Contents -> next)

FileSystem : Type -> Type
FileSystem = Free FileSystemF

Functor FileSystemF where
  map f (WriteFile filePath contents next) = WriteFile filePath contents (f next)
  map f (ReadFile filePath nextFn) = ReadFile filePath (f . nextFn)

writeFile : FilePath -> Contents -> FileSystem Unit
writeFile filePath contents = liftFree $ WriteFile filePath contents ()

readFile : FilePath -> FileSystem String
readFile filePath = liftFree $ ReadFile filePath id


module Toolkit.Data.Location

import Toolkit.Data.Nat

%default total

public export
record Location where
  constructor MkLoc
  source : Maybe String
  line   : Nat
  col    : Nat

public export
record FileContext where
  constructor MkFC
  source : Maybe String
  start  : Location
  end    : Location

public export
record Ref where
  constructor MkRef
  span : FileContext
  get  : String

public export
FC : Type
FC = FileContext

export
newFC : Maybe String -> Location -> Location -> FileContext
newFC n s e = MkFC n (record {source = n} s) (record {source = n} e)

namespace FromCoords
  export
  newLoc : Maybe String -> (Nat, Nat) -> Location
  newLoc n (l,c) = MkLoc n l c

  export
  newFC : Maybe String -> (Nat, Nat) -> (Nat, Nat) -> FileContext
  newFC n s e = newFC n (newLoc n s) (newLoc n e)

  namespace Int
    export
    newLoc : Maybe String -> (Int, Int) -> Location
    newLoc n (l,c) = newLoc n (toNat l, toNat c)

    export
    newFC : Maybe String -> (Int , Int) -> (Int, Int) -> FileContext
    newFC n s e = newFC n (newLoc n s) (newLoc n e)

namespace Anon

  export
  newLoc : (Nat, Nat) -> Location
  newLoc (l,c) = MkLoc Nothing l c

  export
  newFC : Location -> Location -> FileContext
  newFC s e = newFC Nothing s e

export
emptyFC : FileContext
emptyFC = newFC Nothing (Z,Z) (Z,Z)

export
setSource : String -> FileContext -> FileContext
setSource str fc
  = record { source       = Just str
           , start.source = Just str
           , end.source   = Just str
           } fc

namespace Ref

  export
  setSource : String -> Ref -> Ref
  setSource new ref = record {span $= setSource new} ref

export
Show Location where
  show (MkLoc Nothing  l c) = with List concat [show l, ":", show c]
  show (MkLoc (Just n) l c) = with List concat [n, ":", show l, ":", show c]

export
Show FileContext where
  show (MkFC Nothing (MkLoc _ l scol) (MkLoc _ _ ecol)) = with List concat ["global:", show l, ":", show scol, "-", show ecol, ":"]
  show (MkFC (Just x) (MkLoc _ l scol) (MkLoc _ _ ecol)) = with List concat [x, ":", show l, ":", show scol, "-", show ecol, ":"]

-- [ EOF ]


-- SPDX-FileCopyrightText: 2022 The toml-idr developers
--
-- SPDX-License-Identifier: MPL-2.0

module Language.TOML.Processing

import Data.List.Elem
import Decidable.Equality
import Language.TOML.Value

infixr 2 `And`

mutual
    public export
    data ValueTy
        = TString
        | TInteger
        | TFloat
        | TBoolean
        | TEnum (List String)
        | TArray ValueTy
        | TTable TableTy

    public export
    record FieldTy where
        constructor MkFieldTy
        name : String
        optional : Bool
        ty : ValueTy

    public export
    data TableTy : Type where
        Emp : TableTy
        Ext : (ft : FieldTy) -> (FieldOf ft -> TableTy) -> TableTy

    public export
    And : FieldTy -> TableTy -> TableTy
    And x y = Ext x (const y)

    public export
    ValueOf : ValueTy -> Type
    ValueOf TString    = String
    ValueOf TInteger   = Integer
    ValueOf TFloat     = Double
    ValueOf TBoolean   = Bool
    ValueOf (TEnum names) = (name ** Elem name names)
    ValueOf (TArray t) = List (ValueOf t)
    ValueOf (TTable x) = TableOf x

    public export
    FieldOf : FieldTy -> Type
    FieldOf x = (if x.optional then Maybe else id) (ValueOf x.ty)

    public export
    data TableOf : TableTy -> Type where
        Nil : TableOf Emp
        (::) : (x : FieldOf ft) -> TableOf (f x) -> TableOf (Ext ft f)

mutual
    public export
    data ValueError
        = ExpectedType ValueTy
        | InsideArray ValueError
        | InsideTable TableError
        | BadEnumVal String

    public export
    data TableError
        = FieldError String ValueError
        | ExpectedField String
        | UnexpectedFields (List String)

mutual
    public export
    processTable : (ty : TableTy) -> Table -> Either TableError (TableOf ty)
    processTable Emp x with (keys x)
      _ | [] = Right []
      _ | ks = Left $ UnexpectedFields ks
    processTable (Ext (MkFieldTy name optional ty) fields) x with (lookup name x)
      _ | Nothing = case optional of
                         True => Right $ Nothing :: !(processTable (fields Nothing) x)
                         False => Left $ ExpectedField name
      _ | Just val = do
              val' <- bimap (FieldError name) id $ processValue ty val
              let val'' = case optional of
                               True => Just val'
                               False => val'
              fields' <- processTable (fields val'') (delete name x)
              Right (val'' :: fields')

    public export
    processValue : (ty : ValueTy) -> Value -> Either ValueError (ValueOf ty)
    processValue (TEnum xs) (VString x) with (isElem x xs)
      _ | Yes el = Right (x ** el)
      _ | No nel = Left $ BadEnumVal x
    processValue TString    (VString x)  = Right x
    processValue TInteger   (VInteger x) = Right x
    processValue TFloat     (VFloat x)   = Right x
    processValue TBoolean   (VBoolean x) = Right x
    processValue (TArray t) (VArray xs)  = traverse (bimap InsideArray id . processValue t) xs
    processValue (TTable t) (VTable x)   = bimap InsideTable id $ processTable t x
    processValue t          v            = Left $ ExpectedType t


module CompoInt2Float

import Int2Str

-- This function is used to fill the hole that the registry is
-- supposed fill by returning the service and procedure implementing a
-- given type signature.
export
registry_hole_filler : a -> b

-- int2float is defined using the registry hole filler.  Basically
-- assuming (Registry.retrieve (String -> Double)) would return a
-- function with such type signature, then int2float can be
-- implemented as
--
-- (Registry.retrieve (String -> Double)) . int2str
export
int2float : Int -> Double
int2float = registry_hole_filler . Int2Str.int2str


module Data.HashMap

import Data.Array16
import public Data.Hashable
import Data.HashMap.Internal as Internal
import Data.List
import Data.String

||| Hash map implemented with Hash-array mapped tries
export
data HashMap : Type -> Type -> Type where
    MkHM : (eq : k -> k -> Bool) -> (hws : Bits64 -> k -> Bits64) -> HashArrayMapTrie k v -> HashMap k v

||| Unsafely get the internals of a HashMap.
export
unsafeFoldHashMap : ((eq : k -> k -> Bool) -> (hws : Bits64 -> k -> Bits64) -> HashArrayMapTrie k v -> r) -> HashMap k v -> r
unsafeFoldHashMap f (MkHM eq hws m) = f eq hws m

||| Lookup the value at `k` if it exists.
export
lookup : k -> HashMap k v -> Maybe v
lookup k (MkHM eq hws m) = Internal.lookup eq hws bitMask0 (hws defaultSalt k) k m

||| Insert a value at a key, replacing it if it already exists.
export
insert : k -> v -> HashMap k v -> HashMap k v
insert k v (MkHM eq hws m) = MkHM eq hws $ Internal.insert eq hws defaultSalt bitMask0 (hws defaultSalt k) k v m

||| Delete the value at a key.
export
delete : k -> HashMap k v -> HashMap k v
delete k (MkHM eq hws m) = MkHM eq hws $ Internal.delete eq hws bitMask0 (hws defaultSalt k) k m

||| Alter the value at a key.
-- TODO: Make faster.
export
alter : k -> (Maybe v -> Maybe v) -> HashMap k v -> HashMap k v
alter k f m = case f (lookup k m) of
    Nothing => delete k m
    Just v => insert k v m

||| Empty HashMap.
export
empty : Eq k => Hashable k => HashMap k v
empty = MkHM (==) hashWithSalt Internal.empty

||| Construct a HashMap from a single key and value.
export
singleton : Eq k => Hashable k => k -> v -> HashMap k v
singleton k v = MkHM (==) hashWithSalt $ Internal.singleton (hashWithSalt defaultSalt k) k v

||| Construct a HashMap from a `Foldable` structure of key-value pairs.
export
fromFoldable : Foldable f => Eq k => Hashable k => f (k, v) -> HashMap k v
fromFoldable kvs =
    MkHM (==) hashWithSalt
        $ foldr (\(k, v) =>
            Internal.insert (==) hashWithSalt defaultSalt bitMask0 (hashWithSalt defaultSalt k) k v)
            Internal.empty kvs

||| Construct a HashMap from a List of key-value pairs.
||| Specialised and inlined for speed.
export
fromList : Eq k => Hashable k => List (k, v) -> HashMap k v
fromList kvs = MkHM (==) hashWithSalt (go kvs Internal.empty)
  where
    go : List (k, v) -> HashArrayMapTrie k v -> HashArrayMapTrie k v
    go [] m = m
    go ((k, v) :: kvs) m =
        go kvs
            $ Internal.insert (==) hashWithSalt defaultSalt bitMask0
                (hashWithSalt defaultSalt k) k v m

||| Fold a hashmap with the key.
||| Note: Due to hashing keys and values are stored in an arbitrary order.
export
foldWithKey : (k -> v -> acc -> acc) -> acc -> HashMap k v -> acc
foldWithKey f z (MkHM _ _ m) = Internal.foldWithKey f z m

-- Should this exist given the lack of order?
export
Foldable (HashMap k) where
    foldl f = foldWithKey (const $ flip f)
    foldr f = foldWithKey (const f)

||| Convert a HashMap to a List of key-value pairs.
export
toList : HashMap k v -> List (k, v)
toList = foldWithKey (\k, v, acc => (k, v) :: acc) []

export
Show k => Show v => Show (HashMap k v) where
    showPrec d hm = showCon d "fromList" $ showArg $ Data.HashMap.toList hm

||| Merge 2 hashmaps with a preference for entries in the left map.
export
mergeL : HashMap k v -> HashMap k v -> HashMap k v
mergeL m0 m1 = foldWithKey insert m1 m0

||| Merge 2 hashmaps with a preference for entries in the right map.
export
mergeR : HashMap k v -> HashMap k v -> HashMap k v
mergeR m0 m1 = foldWithKey insert m0 m1

||| Merges 2 hashmaps with a preference for entries in the left map.
Semigroup (HashMap k v) where
    (<+>) = mergeL

||| Merges 2 hashmaps with a preference for entries in the right map.
[HMRightSemigroup] Semigroup (HashMap k v) where
    (<+>) = mergeR

Eq k => Hashable k => Monoid (HashMap k v) where
    neutral = empty

[HMRightMonoid] Eq k => Hashable k => Monoid (HashMap k v) using HMRightSemigroup where
    neutral = empty


-- exercises in "Type-Driven Development with Idris"
-- chapter 11, section 3, questions 2 and 3

-- check that all functions are total
%default total

data LazyFuel = LDry | LMore (Lazy LazyFuel)

partial
forever : LazyFuel
forever = LMore forever

data Command : Type -> Type where
  PutStr : String -> Command ()
  GetLine : Command String
  Pure : ty -> Command ty
  Bind : Command a -> (a -> Command b) -> Command b
  ReadFile : String -> Command (Either FileError String)
  WriteFile : String -> String -> Command (Either FileError ())

namespace Command
  (>>=) : Command a -> (a -> Command b) -> Command b
  (>>=) = Bind

runCommand : Command a -> IO a
runCommand (PutStr x) = putStr x
runCommand GetLine = getLine
runCommand (Pure x) = pure x
runCommand (Bind x f) = do res <- runCommand x
                           runCommand (f res)
runCommand (ReadFile f) = readFile f
runCommand (WriteFile f x) = writeFile f x

data ConsoleIO : Type -> Type where
  Quit : a -> ConsoleIO a
  Do : Command a -> (a -> Inf (ConsoleIO b)) -> ConsoleIO b

namespace Console
  (>>=) : Command a -> (a -> Inf (ConsoleIO b)) -> ConsoleIO b
  (>>=) = Do

run : LazyFuel -> ConsoleIO a -> IO (Maybe a)
run _ (Quit v) = pure $ Just v
run LDry (Do act gen) = pure Nothing
run (LMore fuel) (Do act gen) = do res <- runCommand act
                                   run fuel $ gen res

data ShellInput = Cat String
                | Copy String String
                | Exit
                | Error String

readInput : (prompt : String) -> Command ShellInput
readInput prompt
  = do PutStr prompt
       input <- GetLine
       parts <- Pure $ words input
       case parts of
         ["cat", filename] => Pure $ Cat filename
         ["copy", source, destination] => Pure $ Copy source destination
         ["exit"] => Pure $ Exit
         _ => Pure $ Error "Unknown command\n"

cat : String -> Command ()
cat filename = do res <- ReadFile filename
                  case res of
                    Right contents => PutStr contents
                    Left error => PutStr $ show error ++ "\n"

copy : String -> String -> Command ()
copy source destination =
  do res1 <- ReadFile source
     case res1 of
       Left error => PutStr $ show error ++ " (source) \n"
       Right contents => do res2 <- WriteFile destination contents
                            case res2 of
                              Left error => PutStr $ show error ++ " (destination)\n"
                              Right () => PutStr "File copied\n"

shell : ConsoleIO ()
shell = do cmd <- readInput "> "
           case cmd of
             Cat filename => do cat filename
                                shell
             Copy source destination => do copy source destination
                                           shell
             Exit => Quit ()
             Error error => do PutStr error
                               shell

-- :exec
-- > cat input.txt
-- (...) file
-- > copy input.txt output.txt
-- File copied
-- > cat output.txt
-- (...) file
partial
main : IO ()
main = do run forever shell
          pure ()


import Data.Primitives.Views
import Data.Bits
import Data.String
import System

%default total

export
data Command : Type -> Type where
     PutStr : String -> Command ()
     GetLine : Command String

     Pure : ty -> Command ty
     Bind : Command a -> (a -> Command b) -> Command b

export
data ConsoleIO : Type -> Type where
     Quit : a -> ConsoleIO a
     Do : Command a -> (a -> Inf (ConsoleIO b)) -> ConsoleIO b
     Seq : Command () -> Inf (ConsoleIO a) -> ConsoleIO a

namespace CommandDo
  export
  (>>=) : Command a -> (a -> Command b) -> Command b
  (>>=) = Bind

  export
  (>>) : Command () -> Command b -> Command b
  ma >> mb = Bind ma (\ _ => mb)

namespace ConsoleDo
  export
  (>>=) : Command a -> (a -> Inf (ConsoleIO b)) -> ConsoleIO b
  (>>=) = Do

  export
  (>>) : Command () -> Inf (ConsoleIO b) -> ConsoleIO b
  (>>) = Seq

data Fuel = Dry | More (Lazy Fuel)

runCommand : Command a -> IO a
runCommand (PutStr x) = putStr x
runCommand GetLine = getLine
runCommand (Pure val) = pure val
runCommand (Bind c f) = do res <- runCommand c
                           runCommand (f res)

run : Fuel -> ConsoleIO a -> IO (Maybe a)
run fuel (Quit val) = do pure (Just val)
run (More fuel) (Do c f) = do res <- runCommand c
                              run fuel (f res)
run (More fuel) (Seq c d) = do runCommand c; run fuel d
run Dry p = pure Nothing

randoms : Int -> Stream Int
randoms seed = let seed' = 1664525 * seed + 1013904223 in
                   (seed' `shiftR` fromNat 2) :: randoms seed'

arithInputs : Int -> Stream Int
arithInputs seed = map bound (randoms seed)
  where
    bound : Int -> Int
    bound x with (divides x 12)
      bound ((12 * div) + rem) | (DivBy div rem prf) = abs rem + 1

mutual
  correct : Stream Int -> (score : Nat) -> ConsoleIO Nat
  correct nums score
          = do PutStr "Correct!\n"
               quiz nums (score + 1)

  wrong : Stream Int -> Int -> (score : Nat) -> ConsoleIO Nat
  wrong nums ans score
        = do PutStr ("Wrong, the answer is " ++ show ans ++ "\n")
             quiz nums score

  data Input = Answer Int
             | QuitCmd

  readInput : (prompt : String) -> Command Input
  readInput prompt
     = do PutStr prompt
          answer <- GetLine
          if toLower answer == "quit"
             then Pure QuitCmd
             else Pure (Answer (cast answer))

  quiz : Stream Int -> (score : Nat) -> ConsoleIO Nat
  quiz (num1 :: num2 :: nums) score
     = do PutStr ("Score so far: " ++ show score ++ "\n")
          input <- readInput (show num1 ++ " * " ++ show num2 ++ "? ")
          case input of
               Answer answer => if answer == num1 * num2
                                   then correct nums score
                                   else wrong nums (num1 * num2) score
               QuitCmd => Quit score

partial
forever : Fuel
forever = More forever

partial
main : IO ()
main = do seed <- time
          Just score <- run forever (quiz (arithInputs (fromInteger seed)) 0)
               | Nothing => putStrLn "Ran out of fuel"
          putStrLn ("Final score: " ++ show score)


impure_op : Bits64 -> IO Bits64
impure_op foo = pure $ foo + 1

impure_int : IO Int
impure_int = pure 41

main : IO ()
main = impure_int >>= impure_op . prim__zextInt_B64 >>= printLn


module Flexidisc.Record.Update

import        Flexidisc.RecordContent
import public Flexidisc.Record.Type

%default total
%access export

||| Replace a row, can change its type
|||
||| Complexity is _O(n)_
|||
||| @ xs  the record
||| @ loc the proof that the row is in it
||| @ new the new value for the row
setByLabel : (xs : RecordM m k header) -> (loc : Label query header) ->
             (new : m ty) ->
             RecordM m k (changeType header loc ty)
setByLabel (Rec xs nub) (L loc) new =
  Rec (set xs loc new) (changeValuePreservesNub nub)

||| Update a row, the update can change the row type.
|||
||| Complexity is _O(n)_
|||
||| @ key the row name
||| @ xs    the record
||| @ loc   the proof that the row is in it
||| @ new   the new value for the row
set : (key : k) -> (new : m ty) -> (xs : RecordM m k header) ->
      {auto loc : Label key header} ->
      RecordM m k (changeType header loc ty)
set _ new xs {loc} = setByLabel xs loc new

||| Update a row, the update can change the row type.
|||
||| Complexity is _O(n)_
|||
||| @ query the row name
||| @ xs    the record
||| @ loc   the proof that the row is in it
||| @ new   the new value for the row
setA : Applicative m =>
       (query : k) -> (new : ty) -> (xs : RecordM m k header) ->
       {auto loc : Label query header} ->
       RecordM m k (changeType header loc ty)
setA loc = set loc . pure

||| Update a row at a given `Label`, can change its type.
|||
||| Complexity is _O(n)_
|||
||| @ xs     the record
||| @ loc    the proof that the row is in it
||| @ f      the update function
updateByLabel : (xs : RecordM m k header) -> (loc : Row query a header) ->
                (f : m a -> m b) ->
                RecordM m k (changeType header loc b)
updateByLabel (Rec xs nub) (R loc) f =
  Rec (update xs loc f) (changeValuePreservesNub nub)

||| Update a row at the given key, can change its type.
|||
||| Complexity is _O(n)_
|||
||| @ query  the row name
||| @ xs     the record
||| @ loc    the proof that the row is in it
||| @ f      the update function
update : (query : k) -> (f : m a -> m b) -> (xs : RecordM m k header) ->
         {auto loc : Row query a header} ->
         RecordM m k (changeType header loc b)
update _ f xs {loc} = updateByLabel xs loc f

||| Map at a given `Row`, can change its type.
|||
||| Complexity is _O(n)_
|||
||| @ query  the row name
||| @ xs     the record
||| @ loc    the proof that the row is in it
||| @ f      the update function
mapRow : Functor m =>
         (query : k) -> (f : a -> b) -> (xs : RecordM m k header) ->
         {auto loc : Row query a header} ->
         RecordM m k (changeType header loc b)
mapRow _ f xs {loc} = updateByLabel xs loc (map f)


-- ----------------------------------------------------------------- [ ABE.idr ]
-- Module      : ABE
-- Description : Types for Attribute Based Encryption
-- Copyright   : (c) Jan de Muijnck-Hughes
-- License     : see LICENSE
-- --------------------------------------------------------------------- [ EOF ]
module Crypto.Adv.ABE

import Crypto.Adv.Key

data ABEKeyType = CipherTextPolicy | KeyPolicy


module Channel

import Data.List.Views
import System.Concurrency.Process
import public Network.Socket

%access public export

data Request msg = WriteLock
                 | Read
                 | ReleaseLock
                 | Pure msg

data ResultType msg = Fatal   String
                    | Warning String
                    | Correct msg

data Channel msg = MkChan (ProcID (Request msg))

data ChannelState = Unlocked | Locked

data ChannelCmd : (resultType : Type) -> (sharedResTyp : Type) -> (stateIn : ChannelState) ->
                  (stateOut : ChannelState) -> Type where
     ReadOnlyCmd : (chan : Channel typ) -> ChannelCmd (ResultType typ) typ Unlocked Unlocked
     ReadCmd : (chan : Channel typ) -> ChannelCmd (ResultType typ) typ Unlocked Locked
     WriteCmd : (chan : Channel typ) -> (newRes : typ) -> ChannelCmd Bool typ Locked Unlocked
     ReleaseLockCmd : (chan : Channel typ) -> ChannelCmd Bool typ Locked Unlocked

     PureCmd : (res : ty) -> ChannelCmd ty typ state state
     (>>=) : ChannelCmd a typ state1 state2 -> (a -> ChannelCmd b typ state2 state3) ->
             ChannelCmd b typ state1 state3

mutual
  write_inner : (req : Request msgType) -> (state : msgType) ->
                (procId : ProcID (Request msgType)) -> Process (Request msgType) ()
  write_inner req state procId =
    do send procId (Pure state)
       Just newReq <- recvFrom procId
        | Nothing => channelLoop state
       case newReq of
            (Pure newState) => channelLoop newState
            (ReleaseLock) => channelLoop state
            _ =>  channelLoop state

  channelLoop : (state : msgType) -> Process (Request msgType) ()
  channelLoop state = do (procId, req) <- recvWithSender
                         case req of
                               WriteLock => write_inner req state procId
                               ReleaseLock => channelLoop state
                               Read => do send procId (Pure state)
                                          channelLoop state

makeChan : (init : ty) -> IO (Channel (ty))
makeChan init = do procId <- run (create (channelLoop init))
                   pure (MkChan procId)

read_inner : Channel msg -> (req : Request msg) -> IO (ResultType msg)
read_inner (MkChan procId) req =
    do ok <- run (send procId req)
       Just newReq <- run (recvFrom procId)
        | Nothing => pure (Fatal "Channel thread died")
       case newReq of
            (Pure val) => pure (Correct val)
            _ => pure (Warning "Did not recieve correct formatted state from channel thread")

read : Channel msg -> IO (ResultType msg)
read chan = read_inner chan WriteLock

readOnly : Channel msg -> IO (ResultType msg)
readOnly chan = read_inner chan Read

write : Channel msg -> msg -> IO Bool
write (MkChan procId) msg = run (send procId (Pure msg))

run : ChannelCmd res typ inState outState -> IO res
run (ReadOnlyCmd chan) = readOnly chan
run (ReadCmd chan) = read chan
run (WriteCmd chan newRes) = write chan newRes
run (ReleaseLockCmd (MkChan procId)) = run (send procId ReleaseLock)
run (PureCmd res) = pure res
run (x >>= f) = do x' <- run x
                   run (f x')


module Common.Memory
import Common

export
data CairoPtr : Type -> Type where [external]

export
data Memory : Type where [external]

%foreign
    "apStable:True"
    "untupled:(_,_)"
    "imports:starkware.cairo.common.alloc alloc"
    """
    code:
    func Common_Memory_createMemoryPrim(world) -> (world, res):
        let (mem) = alloc()
        return (world, cast(mem,felt))
    end
    """
createMemoryPrim : PrimCairo (CairoPtr Memory)

public export %inline
createMemory : Cairo (CairoPtr Memory)
createMemory = fromPrimCairo createMemoryPrim

%foreign
    "apStable:True"
    """
    code:
    func Common_Memory_writeMemoryPrim(offset, value, ptr, world) -> (world):
        assert [ptr+offset] = value
        return (world)
    end
    """
writeMemoryPrim : (offset:Felt) -> (value: Felt) -> (CairoPtr Memory) -> PrimCairoUnit

public export %inline
writeMemory : Felt -> Felt -> (CairoPtr Memory) -> Cairo ()
writeMemory offset value ptr = fromPrimCairoUnit $ writeMemoryPrim offset value ptr

%foreign
    "apStable:True"
    "untupled:(_,_)"
    """
    code:
    func Common_Memory_readMemoryPrim(offset, ptr, world) -> (world,res):
        let res = [ptr+offset]
        return (world, res)
    end
    """
readMemoryPrim : (offset:Felt) -> (CairoPtr Memory) -> PrimCairo Felt

public export  %inline
readMemory : Felt -> (CairoPtr Memory) -> Cairo Felt
readMemory offset ptr = fromPrimCairo $ readMemoryPrim offset ptr


-- ------------------------------------------------------------- [ Process.idr ]
--
-- Implementation of the echo protocol for process communication.
--
-- --------------------------------------------------------------------- [ EOH ]
module TCPHandshake

import Effects
import Effect.StdIO
import Effect.Default
import Effect.Msg

import System.Protocol

import RFC.Handshake

covering
bob : Nat -> (client : PID)
      -> Process (handshake) 'Bob ['Alice := client] [STDIO] ()
bob seqno client = do
    (msg, x) <- recvFrom 'Alice
    let x' = S x
    putStrLn $ "Bob -> Alice: SynAck, " ++ show seqno ++ ", " ++ show x'
    sendTo 'Alice (SynAck, seqno, (S x ** Refl))
    (msg1, no') <- recvFrom 'Alice
    putStrLn $ "Received from Alice: " ++ show (getWitness no')
    pure ()

covering
alice : Nat -> (server : PID)
      -> Process (handshake) 'Alice ['Bob := server] [STDIO] ()
alice seqno server = do
    putStrLn $ "Alice -> Bob: Syn, " ++ show seqno
    sendTo 'Bob (Syn, seqno)
    (msg, y, x) <- recvFrom 'Bob
    putStrLn $ "Alice -> Bob: Ack " ++ show (S y)
    sendTo 'Bob (Ack, (S y ** Refl))
    pure ()

covering
doHandshake : IO ()
doHandshake = runConc [()] doShake
  where
    doShake : Process (handshake) 'Alice [] [STDIO] ()
    doShake = do
       server <- spawn (bob 2) [()]
       setChan 'Bob server
       alice 100 server
       dropChan 'Bob
       pure ()

-- -------------------------------------------------------------------- [ Main ]
namespace Main
  main : IO ()
  main = doHandshake

-- --------------------------------------------------------------------- [ EOF ]


import Control.ST
import Control.ST.ImplicitCall
import Network.Socket

data SocketState = Ready | Bound | Listening | Open | Closed

data CloseOK : SocketState -> Type where
    CloseOpen : CloseOK Open
    CloseListening : CloseOK Listening

interface Sockets (m : Type -> Type) where
    Sock : SocketState -> Type

    socket : SocketType -> ST m (Either () Var) [addIfRight (Sock Ready)]

    bind : (sock : Var) -> (addr : Maybe SocketAddress) -> (port : Port) ->
            ST m (Either () ()) [sock ::: Sock Ready :-> (Sock Closed `or` Sock Bound)]
    listen : (sock : Var) -> ST m (Either () ())
            [sock ::: Sock Bound :-> (Sock Closed `or` Sock Listening)]
    accept : (sock : Var) -> ST m (Either () Var)
            [sock ::: Sock Listening, addIfRight $ Sock Open]

    connect : (sock : Var) -> SocketAddress -> Port ->
            ST m (Either () ()) [sock ::: Sock Ready :-> (Sock Closed `or` Sock Open)]

    send : (sock : Var) -> String ->
            ST m (Either () ()) [sock ::: Sock Open :-> (Sock Closed `or` Sock Open)]
    recv : (sock : Var) -> ST m (Either () String)
            [sock ::: Sock Open :-> (Sock Closed `or` Sock Open)]

    close : (sock : Var) -> {auto prf : CloseOK st} -> ST m () [sock ::: Sock st :-> Sock Closed]
    remove : (sock : Var) -> ST m () [Remove sock (Sock Closed)]

echoServer : (ConsoleIO m, Sockets m) => (sock : Var) ->
    ST m () [remove sock (Sock {m} Listening)]
echoServer sock =
    do
        Right new <- accept sock
            | Left err => do close sock; remove sock
        Right msg <- recv new
            | Left err => do close sock; remove sock; remove new
        Right ok <- send new ("You said " ++ msg)
            | Left err => do remove new; close sock; remove sock
        close new; remove new; echoServer sock

startServer : (ConsoleIO m, Sockets m) => ST m () []
startServer =
    do
        Right sock <- socket Stream
            | Left err => pure ()
        Right ok <- bind sock Nothing 9442
            | Left err => remove sock
        Right ok <- listen sock
            | Left err => remove sock
        echoServer sock

implementation Sockets IO where
    Sock _ = State Socket

    socket ty =
        do
            Right sock <- lift $ Socket.socket AF_INET ty 0
                | Left err => pure (Left ())
            lbl <- new sock
            pure (Right lbl)

    bind sock addr port =
        do
            ok <- lift $ bind !(read sock) addr port
            if ok /= 0
                then pure (Left ())
                else pure (Right ())

    listen sock =
        do
            ok <- lift $ listen !(read sock)
            if ok /= 0
                then pure (Left ())
                else pure (Right ())

    accept sock =
        do
            Right (conn, addr) <- lift $ accept !(read sock)
                | Left err => pure (Left ())
            lbl <- new conn
            returning (Right lbl) (toEnd lbl)

    connect sock addr port =
        do
            ok <- lift $ connect !(read sock) addr port
            if ok /= 0
                then pure (Left ())
                else pure (Right ())

    close sock =
        do
            lift $ close !(read sock)
            pure ()

    remove sock = delete sock

    send sock msg =
        do
            Right _ <- lift $ send !(read sock) msg
                | Left _ => pure (Left ())
            pure (Right ())
    recv sock =
        do
            Right (msg, len) <- lift $ recv !(read sock) 1024
                | Left _ => pure (Left ())
            pure (Right msg)


module Data.SortedMap.Dependent

import Data.DPair

import Decidable.Equality

-- TODO: write split

-------------------------------
--- Internal representation ---
-------------------------------

private
data Tree : Nat -> (k : Type) -> (v : k -> Type) -> Ord k -> Type where
  Leaf : (x : k) -> v x -> Tree Z k v o
  Branch2 : Tree n k v o -> k -> Tree n k v o -> Tree (S n) k v o
  Branch3 : Tree n k v o -> k -> Tree n k v o -> k -> Tree n k v o -> Tree (S n) k v o

branch4 :
  Tree n k v o -> k ->
  Tree n k v o -> k ->
  Tree n k v o -> k ->
  Tree n k v o ->
  Tree (S (S n)) k v o
branch4 a b c d e f g =
  Branch2 (Branch2 a b c) d (Branch2 e f g)

branch5 :
  Tree n k v o -> k ->
  Tree n k v o -> k ->
  Tree n k v o -> k ->
  Tree n k v o -> k ->
  Tree n k v o ->
  Tree (S (S n)) k v o
branch5 a b c d e f g h i =
  Branch2 (Branch2 a b c) d (Branch3 e f g h i)

branch6 :
  Tree n k v o -> k ->
  Tree n k v o -> k ->
  Tree n k v o -> k ->
  Tree n k v o -> k ->
  Tree n k v o -> k ->
  Tree n k v o ->
  Tree (S (S n)) k v o
branch6 a b c d e f g h i j k =
  Branch3 (Branch2 a b c) d (Branch2 e f g) h (Branch2 i j k)

branch7 :
  Tree n k v o -> k ->
  Tree n k v o -> k ->
  Tree n k v o -> k ->
  Tree n k v o -> k ->
  Tree n k v o -> k ->
  Tree n k v o -> k ->
  Tree n k v o ->
  Tree (S (S n)) k v o
branch7 a b c d e f g h i j k l m =
  Branch3 (Branch3 a b c d e) f (Branch2 g h i) j (Branch2 k l m)

merge1 : Tree n k v o -> k -> Tree (S n) k v o -> k -> Tree (S n) k v o -> Tree (S (S n)) k v o
merge1 a b (Branch2 c d e) f (Branch2 g h i) = branch5 a b c d e f g h i
merge1 a b (Branch2 c d e) f (Branch3 g h i j k) = branch6 a b c d e f g h i j k
merge1 a b (Branch3 c d e f g) h (Branch2 i j k) = branch6 a b c d e f g h i j k
merge1 a b (Branch3 c d e f g) h (Branch3 i j k l m) = branch7 a b c d e f g h i j k l m

merge2 : Tree (S n) k v o -> k -> Tree n k v o -> k -> Tree (S n) k v o -> Tree (S (S n)) k v o
merge2 (Branch2 a b c) d e f (Branch2 g h i) = branch5 a b c d e f g h i
merge2 (Branch2 a b c) d e f (Branch3 g h i j k) = branch6 a b c d e f g h i j k
merge2 (Branch3 a b c d e) f g h (Branch2 i j k) = branch6 a b c d e f g h i j k
merge2 (Branch3 a b c d e) f g h (Branch3 i j k l m) = branch7 a b c d e f g h i j k l m

merge3 : Tree (S n) k v o -> k -> Tree (S n) k v o -> k -> Tree n k v o -> Tree (S (S n)) k v o
merge3 (Branch2 a b c) d (Branch2 e f g) h i = branch5 a b c d e f g h i
merge3 (Branch2 a b c) d (Branch3 e f g h i) j k = branch6 a b c d e f g h i j k
merge3 (Branch3 a b c d e) f (Branch2 g h i) j k = branch6 a b c d e f g h i j k
merge3 (Branch3 a b c d e) f (Branch3 g h i j k) l m = branch7 a b c d e f g h i j k l m

treeLookup : Ord k => (x : k) -> Tree n k v o -> Maybe (y : k ** v y) -- may also return an erased `So (x == y)`
treeLookup k (Leaf k' v) =
  if k == k' then
    Just (k' ** v)
  else
    Nothing
treeLookup k (Branch2 t1 k' t2) =
  if k <= k' then
    treeLookup k t1
  else
    treeLookup k t2
treeLookup k (Branch3 t1 k1 t2 k2 t3) =
  if k <= k1 then
    treeLookup k t1
  else if k <= k2 then
    treeLookup k t2
  else
    treeLookup k t3

treeInsert' : Ord k => (x : k) -> v x -> Tree n k v o -> Either (Tree n k v o) (Tree n k v o, k, Tree n k v o)
treeInsert' k v (Leaf k' v') =
  case compare k k' of
    LT => Right (Leaf k v, k, Leaf k' v')
    EQ => Left (Leaf k v)
    GT => Right (Leaf k' v', k', Leaf k v)
treeInsert' k v (Branch2 t1 k' t2) =
  if k <= k' then
    case treeInsert' k v t1 of
      Left t1' => Left (Branch2 t1' k' t2)
      Right (a, b, c) => Left (Branch3 a b c k' t2)
  else
    case treeInsert' k v t2 of
      Left t2' => Left (Branch2 t1 k' t2')
      Right (a, b, c) => Left (Branch3 t1 k' a b c)
treeInsert' k v (Branch3 t1 k1 t2 k2 t3) =
  if k <= k1 then
    case treeInsert' k v t1 of
      Left t1' => Left (Branch3 t1' k1 t2 k2 t3)
      Right (a, b, c) => Right (Branch2 a b c, k1, Branch2 t2 k2 t3)
  else
    if k <= k2 then
      case treeInsert' k v t2 of
        Left t2' => Left (Branch3 t1 k1 t2' k2 t3)
        Right (a, b, c) => Right (Branch2 t1 k1 a, b, Branch2 c k2 t3)
    else
      case treeInsert' k v t3 of
        Left t3' => Left (Branch3 t1 k1 t2 k2 t3')
        Right (a, b, c) => Right (Branch2 t1 k1 t2, k2, Branch2 a b c)

treeInsert : Ord k => (x : k) -> v x -> Tree n k v o -> Either (Tree n k v o) (Tree (S n) k v o)
treeInsert k v t =
  case treeInsert' k v t of
    Left t' => Left t'
    Right (a, b, c) => Right (Branch2 a b c)

delType : Nat -> (k : Type) -> (v : k -> Type) -> Ord k -> Type
delType Z k v o = ()
delType (S n) k v o = Tree n k v o

treeDelete : Ord k => (n : Nat) -> k -> Tree n k v o -> Either (Tree n k v o) (delType n k v o)
treeDelete _ k (Leaf k' v) =
  if k == k' then
    Right ()
  else
    Left (Leaf k' v)
treeDelete (S Z) k (Branch2 t1 k' t2) =
  if k <= k' then
    case treeDelete Z k t1 of
      Left t1' => Left (Branch2 t1' k' t2)
      Right () => Right t2
  else
    case treeDelete Z k t2 of
      Left t2' => Left (Branch2 t1 k' t2')
      Right () => Right t1
treeDelete (S Z) k (Branch3 t1 k1 t2 k2 t3) =
  if k <= k1 then
    case treeDelete Z k t1 of
      Left t1' => Left (Branch3 t1' k1 t2 k2 t3)
      Right () => Left (Branch2 t2 k2 t3)
  else if k <= k2 then
    case treeDelete Z k t2 of
      Left t2' => Left (Branch3 t1 k1 t2' k2 t3)
      Right () => Left (Branch2 t1 k1 t3)
  else
    case treeDelete Z k t3 of
      Left t3' => Left (Branch3 t1 k1 t2 k2 t3')
      Right () => Left (Branch2 t1 k1 t2)
treeDelete (S (S n)) k (Branch2 t1 k' t2) =
  if k <= k' then
    case treeDelete (S n) k t1 of
      Left t1' => Left (Branch2 t1' k' t2)
      Right t1' =>
        case t2 of
          Branch2 a b c => Right (Branch3 t1' k' a b c)
          Branch3 a b c d e => Left (branch4 t1' k' a b c d e)
  else
    case treeDelete (S n) k t2 of
      Left t2' => Left (Branch2 t1 k' t2')
      Right t2' =>
        case t1 of
          Branch2 a b c => Right (Branch3 a b c k' t2')
          Branch3 a b c d e => Left (branch4 a b c d e k' t2')
treeDelete (S (S n)) k (Branch3 t1 k1 t2 k2 t3) =
  if k <= k1 then
    case treeDelete (S n) k t1 of
      Left t1' => Left (Branch3 t1' k1 t2 k2 t3)
      Right t1' => Left (merge1 t1' k1 t2 k2 t3)
  else if k <= k2 then
    case treeDelete (S n) k t2 of
      Left t2' => Left (Branch3 t1 k1 t2' k2 t3)
      Right t2' => Left (merge2 t1 k1 t2' k2 t3)
  else
    case treeDelete (S n) k t3 of
      Left t3' => Left (Branch3 t1 k1 t2 k2 t3')
      Right t3' => Left (merge3 t1 k1 t2 k2 t3')

treeToList : Tree n k v o -> List (x : k ** v x)
treeToList = treeToList' (:: [])
  where
    -- explicit quantification to avoid conflation with {n} from treeToList
    treeToList' : {0 n : Nat} -> ((x : k ** v x) -> List (x : k ** v x)) -> Tree n k v o -> List (x : k ** v x)
    treeToList' cont (Leaf k v) = cont (k ** v)
    treeToList' cont (Branch2 t1 _ t2) = treeToList' (:: treeToList' cont t2) t1
    treeToList' cont (Branch3 t1 _ t2 _ t3) = treeToList' (:: treeToList' (:: treeToList' cont t3) t2) t1

----------------------
--- User interface ---
----------------------

export
data SortedDMap : (k : Type) -> (v : k -> Type) -> Type where
  Empty : Ord k => SortedDMap k v
  M : (o : Ord k) => (n : Nat) -> Tree n k v o -> SortedDMap k v

export
empty : Ord k => SortedDMap k v
empty = Empty

export
lookup : (x : k) -> SortedDMap k v -> Maybe (y : k ** v y) -- could return also `So (x == y)`
lookup _ Empty = Nothing
lookup k (M _ t) = treeLookup k t

export
lookupPrecise : DecEq k => (x : k) -> SortedDMap k v -> Maybe (v x)
lookupPrecise x = lookup x >=> \(y ** v) =>
  case decEq x y of
    Yes Refl => Just v
    No _     => Nothing

export
insert : (x : k) -> v x -> SortedDMap k v -> SortedDMap k v
insert k v Empty = M Z (Leaf k v)
insert k v (M _ t) =
  case treeInsert k v t of
    Left t' => (M _ t')
    Right t' => (M _ t')

export
singleton : Ord k => (x : k) -> v x -> SortedDMap k v
singleton k v = insert k v empty

export
insertFrom : Foldable f => f (x : k ** v x) -> SortedDMap k v -> SortedDMap k v
insertFrom = flip $ foldl $ flip $ uncurry insert

export
delete : k -> SortedDMap k v -> SortedDMap k v
delete _ Empty = Empty
delete k (M Z t) =
  case treeDelete Z k t of
    Left t' => (M _ t')
    Right () => Empty
delete k (M (S n) t) =
  case treeDelete (S n) k t of
    Left t' => (M _ t')
    Right t' => (M _ t')

export
fromList : Ord k => List (x : k ** v x) -> SortedDMap k v
fromList = foldl (flip (uncurry insert)) empty

export
toList : SortedDMap k v -> List (x : k ** v x)
toList Empty = []
toList (M _ t) = treeToList t

||| Gets the keys of the map.
export
keys : SortedDMap k v -> List k
keys = map fst . toList

export
values : SortedDMap k v -> List (x : k ** v x)
values = toList

treeMap : ({x : k} -> a x -> b x) -> Tree n k a o -> Tree n k b o
treeMap f (Leaf k v) = Leaf k (f v)
treeMap f (Branch2 t1 k t2) = Branch2 (treeMap f t1) k (treeMap f t2)
treeMap f (Branch3 t1 k1 t2 k2 t3)
    = Branch3 (treeMap f t1) k1 (treeMap f t2) k2 (treeMap f t3)

treeTraverse : Applicative f => ({x : k} -> a x -> f (b x)) -> Tree n k a o -> f (Tree n k b o)
treeTraverse f (Leaf k v) = Leaf k <$> f v
treeTraverse f (Branch2 t1 k t2) =
  Branch2
    <$> treeTraverse f t1
    <*> pure k
    <*> treeTraverse f t2
treeTraverse f (Branch3 t1 k1 t2 k2 t3) =
  Branch3
    <$> treeTraverse f t1
    <*> pure k1
    <*> treeTraverse f t2
    <*> pure k2
    <*> treeTraverse f t3

export
map : ({x : k} -> v x -> w x) -> SortedDMap k v -> SortedDMap k w
map _ Empty = Empty
map f (M _ t) = M _ $ treeMap f t

export
foldl : (acc -> (x : k ** v x) -> acc) -> acc -> SortedDMap k v -> acc
foldl f i = foldl f i . values

export
foldr : ((x : k ** v x) -> acc -> acc) -> acc -> SortedDMap k v -> acc
foldr f i = foldr f i . values

export
foldlM : Monad m => (acc -> (x : k ** v x) -> m acc) -> acc -> SortedDMap k v -> m acc
foldlM f i = foldl (\ma, b => ma >>= flip f b) (pure i)

export
foldMap : Monoid m => (f : (x : k) -> v x -> m) -> SortedDMap k v -> m
foldMap f = foldr ((<+>) . uncurry f) neutral

export
null : SortedDMap k v -> Bool
null Empty   = True
null (M _ _) = False

export
traverse : Applicative f => ({x : k} -> v x -> f (w x)) -> SortedDMap k v -> f (SortedDMap k w)
traverse _ Empty = pure Empty
traverse f (M _ t) = M _ <$> treeTraverse f t

||| Merge two maps. When encountering duplicate keys, using a function to combine the values.
||| Uses the ordering of the first map given.
export
mergeWith : DecEq k => ({x : k} -> v x -> v x -> v x) -> SortedDMap k v -> SortedDMap k v -> SortedDMap k v
mergeWith f x y = insertFrom inserted x where
  inserted : List (x : k ** v x)
  inserted = do
    (k ** v) <- toList y
    let v' = (maybe id f $ lookupPrecise k x) v
    pure (k ** v')

||| Merge two maps using the Semigroup (and by extension, Monoid) operation.
||| Uses mergeWith internally, so the ordering of the left map is kept.
export
merge : DecEq k => ({x : k} -> Semigroup (v x)) => SortedDMap k v -> SortedDMap k v -> SortedDMap k v
merge = mergeWith (<+>)

||| Left-biased merge, also keeps the ordering specified  by the left map.
export
mergeLeft : DecEq k => SortedDMap k v -> SortedDMap k v -> SortedDMap k v
mergeLeft = mergeWith const

treeLeftMost : Tree n k v o -> (x : k ** v x)
treeLeftMost (Leaf x y) = (x ** y)
treeLeftMost (Branch2 x _ _) = treeLeftMost x
treeLeftMost (Branch3 x _ _ _ _) = treeLeftMost x

treeRightMost : Tree n k v o -> (x : k ** v x)
treeRightMost (Leaf x y) = (x ** y)
treeRightMost (Branch2 _ _ x) = treeRightMost x
treeRightMost (Branch3 _ _ _ _ x) = treeRightMost x

treeLookupBetween : Ord k => k -> Tree n k v o -> (Maybe (x : k ** v x), Maybe (x : k ** v x))
treeLookupBetween k (Leaf k' v) with (k < k')
  treeLookupBetween k (Leaf k' v) | True = (Nothing, Just (k' ** v))
  treeLookupBetween k (Leaf k' v) | False = (Just (k' ** v), Nothing)
treeLookupBetween k (Branch2 t1 k' t2) with (k < k')
  treeLookupBetween k (Branch2 t1 k' t2) | True = -- k < k'
    let (lower, upper) = treeLookupBetween k t1 in
    (lower, upper <|> pure (treeLeftMost t2))
  treeLookupBetween k (Branch2 t1 k' t2) | False = -- k >= k'
    let (lower, upper) = treeLookupBetween k t2 in
    (lower <|> pure (treeRightMost t1), upper)
treeLookupBetween k (Branch3 t1 k1 t2 k2 t3) with (k < k1)
  treeLookupBetween k (Branch3 t1 k1 t2 k2 t3) | True = treeLookupBetween k (Branch2 t1 k1 t2)
  treeLookupBetween k (Branch3 t1 k1 t2 k2 t3) | False with (k < k2)
    treeLookupBetween k (Branch3 t1 k1 t2 k2 t3) | False | False = treeLookupBetween k (Branch2 t2 k2 t3)
    treeLookupBetween k (Branch3 t1 k1 t2 k2 t3) | False | True = --k1 <= k < k2
      let (lower, upper) = treeLookupBetween k (Branch2 t1 k1 t2) in
      (lower, upper <|> pure (treeLeftMost t3))

||| looks up a key in map, returning the left and right closest values, so that
|||  k1 <= k < k2. If at the end of the beginning and/or end of the sorted map, returns
||| nothing appropriately
export
lookupBetween : k -> SortedDMap k v -> (Maybe (x : k ** v x), Maybe (x : k ** v x))
lookupBetween k Empty = (Nothing, Nothing)
lookupBetween k (M _ t) = treeLookupBetween k t

||| Returns the leftmost (least) key and value
export
leftMost : SortedDMap k v -> Maybe (x : k ** v x)
leftMost Empty = Nothing
leftMost (M _ t) = Just $ treeLeftMost t

||| Returns the rightmost (greatest) key and value
export
rightMost : SortedDMap k v -> Maybe (x : k ** v x)
rightMost Empty = Nothing
rightMost (M _ t) = Just $ treeRightMost t

export
(Show k, {x : k} -> Show (v x)) => Show (SortedDMap k v) where
   show m = "fromList " ++ (show $ toList m)

export
(DecEq k, {x : k} -> Eq (v x)) => Eq (SortedDMap k v) where
  (==) = (==) `on` toList

export
strictSubmap : DecEq k => ({x : k} -> Eq (v x)) => (sub : SortedDMap k v) -> (sup : SortedDMap k v) -> Bool
strictSubmap sub sup = all (\(k ** v) => Just v == lookupPrecise k sup) $ toList sub

-- TODO: is this the right variant of merge to use for this? I think it is, but
-- I could also see the advantages of using `mergeLeft`. The current approach is
-- strictly more powerful I believe, because `mergeLeft` can be emulated with
-- the `First` monoid. However, this does require more code to do the same
-- thing.
export
DecEq k => ({x : k} -> Semigroup (v x)) => Semigroup (SortedDMap k v) where
  (<+>) = merge

||| For `neutral <+> y`, y is rebuilt in `Ord k`, so this is not a "strict" Monoid.
||| However, semantically, it should be equal.
export
DecEq k => Ord k => ({x : k} -> Semigroup (v x)) => Monoid (SortedDMap k v) where
  neutral = empty