%default covering

data Three = Vrai | Faux | Indef

getType : Three -> Type
getType Vrai = Unit
getType Faux = Unit
getType Indef = Void

swap : (t : Three) -> getType t -> Three
swap Indef _ impossible
swap Vrai () = Faux


module Sat

import Data.List
import Data.Maybe
import Data.Strings

import Data.Linear.Array

import Control.Linear
import Control.Linear.LState

%default total

public export
data Var = MkVar Int
export
Show Var where
  show (MkVar n) = show n

export
Eq Var where
  (MkVar a) == (MkVar b) = a == b
export
Ord Var where
  compare (MkVar a) (MkVar b) = compare a b

public export
data Lit = MkPos Var | MkNeg Var
export
Show Lit where
  show (MkPos v) = " " ++ show v
  show (MkNeg v) = "~" ++ show v

isPos : Lit -> Bool
isPos (MkPos _) = True
isPos (MkNeg _) = False

getVar : Lit -> Var
getVar (MkPos v) = v
getVar (MkNeg v) = v

export
data Clause : Type where
  MkClause : (l : List Lit) -> {auto ok : NonEmpty l} -> Clause

export
toClause : (l : List Lit) -> {auto ok : NonEmpty l} -> Clause
toClause l {ok} = MkClause l {ok}

export
Show Clause where
  show (MkClause l) = show l

Watchlist : Type
Watchlist = LinArray (List Clause)

litIndex : Lit -> Int
litIndex (MkPos (MkVar n)) = 2 * n
litIndex (MkNeg (MkVar n)) = 2 * n + 1

initWatchlist : Int -> List Clause -> Watchlist
initWatchlist maxLit clauses = newArray maxLit (initWatchlist' clauses)
  where
    initWatchlist' : List Clause -> (1 a : Watchlist) -> Watchlist
    initWatchlist' [] a = a
    initWatchlist' (c@(MkClause ls {ok})::cs) a =
      let firstLit = head ls
          idx = litIndex firstLit
          (r # arr') = mread a idx in
          initWatchlist' cs $ write arr' idx (c::(fromMaybe [] r))

getWatched : Lit -> LState Watchlist (List Clause)
getWatched lit = do
  v <- run (\wl => mread wl (litIndex lit))
  pure $ fromMaybe [] v

removeWatches : Lit -> LState Watchlist ()
removeWatches lit = do
  change $ \wl => write wl (litIndex lit) []

addWatch : Lit -> Clause -> LState Watchlist ()
addWatch l c = do
  clauses <- getWatched l
  change $ \wl => write wl (litIndex l) (c::clauses)

addWatches : List (Lit, Clause) -> LState Watchlist ()
addWatches ls = traverse_ (uncurry addWatch) ls

Assignments : Type
Assignments = LinArray (Maybe Bool)

varIndex : Var -> Int
varIndex (MkVar n) = n

value : Var -> LState Assignments (Maybe Bool)
value var = do
  val <- run $ (\as => mread as (varIndex var))
  pure $ fromMaybe Nothing val

assign : Var -> Bool -> LState Assignments ()
assign var v = change (\as => write as (varIndex var) (Just v))

unassign : Var -> LState Assignments ()
unassign var = change (\as => write as (varIndex var) Nothing)

SolverState : Type -> Type
SolverState = LState (Watchlist, Assignments)

dumpAs' : Array arr => Int -> List (Var, Maybe Bool) -> arr (Maybe Bool) -> List (Var, Maybe Bool)
dumpAs' i s as =
  let sz = size as in
      if i >= sz then reverse s
                 else case read as i of
                           Nothing => dumpAs' (assert_smaller i (i+1)) ((MkVar i, Nothing)::s) as
                           Just e => dumpAs' (assert_smaller i (i+1)) ((MkVar i, e)::s) as

dumpArr : (1 _ : Assignments) -> List (Var, Maybe Bool)
dumpArr as = toIArray as (dumpAs' 0 [])

initAssignments : Int -> Assignments
initAssignments n = newArray n (initAssignments' (integerToNat $ cast n)) where
  initAssignments' : Nat -> (1 _ : Assignments) -> Assignments
  initAssignments' Z as = as
  initAssignments' (S n) as = write (initAssignments' n as) (cast n) Nothing

updateWatchlist : Lit -> SolverState Bool
updateWatchlist assignFalse = do
  needUpdate <- liftSt $ getWatched assignFalse
  case needUpdate of
       [] => pure True
       needUpdate => do
         (Just alts) <- liftSt $ lfoldlM findAlt (Just []) needUpdate
         | Nothing => pure False
         liftSt $ removeWatches assignFalse
         liftSt $ addWatches alts
         pure True

    where

  altOk : List Lit -> LState Assignments (Maybe Lit)
  altOk [] = pure Nothing
  altOk (l::ls) = case !(value (getVar l)) of
                         Just x => if x == isPos l then pure $ Just l
                                                   else altOk ls
                         Nothing => pure $ Just l

  findAlt : (Maybe (List (Lit, Clause))) -> (_ : Clause) -> LState Assignments (Maybe (List (Lit, Clause)))
  findAlt Nothing _ = pure Nothing
  findAlt (Just acc) (MkClause ls) = do
    (Just foundAlt) <- altOk ls
    | Nothing => pure Nothing
    pure $ Just ((foundAlt, MkClause ls)::acc)

mutual
  solveTry : Bool -> Var -> SolverState Bool
  solveTry a var = do
    liftSt $ assign var a
    ok <- updateWatchlist (if a then MkNeg var else MkPos var)
    if ok then pure True
          else do liftSt $ unassign var
                  pure False

  solve : List Var -> SolverState Bool
  solve [] = pure True
  solve (var::ds) = do
    tryValues [False, True]
    where
      tryValues : List Bool -> SolverState Bool
      tryValues [] = pure False
      tryValues (val::rest) = do
        couldBe <- solveTry val var
        if couldBe
           then do
             whenIsVal <- solve ds
             if whenIsVal then pure True
                          else do
                            liftSt $ unassign var
                            if val then pure False
                                   else tryValues rest
           else if val then pure False
           else tryValues rest

prepareAndSolve : Int -> List Clause -> SolverState Bool
prepareAndSolve numVars clauses = let vars = map MkVar [0 .. (numVars - 1)] in do
  let varsFreq = map countVarNew vars
  let mostFrequentVars = reverse $ sort varsFreq
  solve $ map snd mostFrequentVars
    where
      countInClause : Int -> Var -> List Lit -> Int
      countInClause acc v [] = acc
      countInClause acc v (MkPos cv::rest) = if cv == v then countInClause (acc+1) v rest
                                                        else countInClause acc v rest
      countInClause acc v (MkNeg cv::rest) = if cv == v then countInClause (acc+1) v rest
                                                        else countInClause acc v rest

      countVarNew : Var -> (Int, Var)
      countVarNew v = (sum $ map (\(MkClause ls) => countInClause 0 v ls) clauses, v)

      countVar : Var -> SolverState (Int, Var)
      countVar v = do watchingFalse <- liftSt $ getWatched (MkPos v)
                      watchingTrue <- liftSt $ getWatched (MkNeg v)
                      let count = the Int ((cast $ length watchingFalse) + (cast $ length watchingTrue))
                      pure (count, v)

export
sat : Int -> List Clause -> Maybe (List (Var, Bool))
sat numVars clauses =
  let wl = initWatchlist (2 * numVars) clauses
      as = initAssignments numVars
      (r # (wl', as')) = runLState (wl, as) (prepareAndSolve numVars clauses) in
      if r then traverse result (dumpArr as')
           else Nothing
  where
    result : (Var, Maybe Bool) -> Maybe (Var, Bool)
    result (v, Just a) = Just (v, a)
    result (_, Nothing) = Nothing


module ImplicitCatenableDeque

import CatenableDeque

%default total
%access private

public export
interface Sized (d : Type -> Type) where
  size : d a -> Int

mutual
  export
  data ImplicitCatDeque : (Type -> Type) -> Type -> Type where
    Shallow : d a -> ImplicitCatDeque d a
    Deep : d a -> Lazy (ImplicitCatDeque d (CmpdElem d a)) -> d a
        -> Lazy (ImplicitCatDeque d (CmpdElem d a)) -> d a
        -> ImplicitCatDeque d a

  data CmpdElem : (Type -> Type) -> Type -> Type where
    Simple : d a -> CmpdElem d a
    Cmpd : d a -> Lazy (ImplicitCatDeque d (CmpdElem d a)) -> d a
        -> CmpdElem d a

dappendL : Deque d => d a {- size < 4 -} -> d a -> d a
dappendL d1 d2 =
  if isEmpty d1 then d2 else dappendL (assert_smaller d1 $ init d1) (cons (last d1) d2)

dappendR : Deque d => d a -> d a {- size < 4 -} -> d a
dappendR d1 d2 =
  if isEmpty d2 then d1 else dappendR (snoc d1 (head d2)) (assert_smaller d2 $ tail d2)

replaceHead : Deque d => a -> ImplicitCatDeque d a -> ImplicitCatDeque d a
replaceHead x (Shallow d) = Shallow (cons x (tail d))
replaceHead x (Deep f a m b r) = Deep (cons x (tail f)) a m b r

replaceLast : Deque d => ImplicitCatDeque d a -> a -> ImplicitCatDeque d a
replaceLast (Shallow d) x = Shallow (snoc (init d) x)
replaceLast (Deep f a m b r) x = Deep f a m b (snoc (init r) x)

share : Deque d => d a -> d a -> (d a, d a, d a)
share f r = let m = cons (last f) (cons (head r) empty) in (init f, m, tail r)

mutual
  export
  (Deque d, Sized d) => Deque (ImplicitCatDeque d) where
    empty = Shallow empty
    isEmpty (Shallow d) = isEmpty d
    isEmpty _ = False

    cons x (Shallow d) = Shallow (cons x d)
    cons x (Deep f a m b r) = Deep (cons x f) a m b r

    head (Shallow d) = head d
    head (Deep f a m b r) = head f

    tail (Shallow d) = Shallow (tail d)
    tail (Deep f a m b r) =
      if size f > 3 then Deep (tail f) a m b r
      else if not (isEmpty a) then
        case head a of
          (Simple d) => let f' = dappendL (tail f) d
                        in Deep f' (tail a) m b r
          (Cmpd f' c' r') => let f'' = dappendL (tail f) f'
                                 a'' = c' ++ replaceHead (Simple r') a
                             in Deep f'' a'' m b r
      else if not (isEmpty b) then
        case head b of
          (Simple d) => let f' = dappendL (tail f) m
                        in Deep f' empty d (tail b) r
          (Cmpd f' c' r') => let f'' = dappendL (tail f) m
                                 a'' = cons (Simple f') c'
                             in Deep f'' a'' r' (tail b) r
      else Shallow (dappendL (tail f) m) ++ Shallow r

    snoc (Shallow d) x = Shallow (snoc d x)
    snoc (Deep f a m b r) x = Deep f a m b (snoc r x)

    last (Shallow d) = last d
    last (Deep f a m b r) = last r

    init (Shallow d) = Shallow (init d)
    init (Deep f a m b r) =
      if size r > 3 then Deep f a m b (init r)
      else if not (isEmpty b) then
        case last b of
          (Simple d) => let r' = dappendR d (init r)
                        in Deep f a m (init b) r'
          (Cmpd f' c' r') => let r'' = dappendR r' (init r)
                                 b'' = replaceLast b (Simple f') ++ c'
                             in Deep f a m b'' r''
      else if not (isEmpty a) then
        case last a of
          (Simple d) => let r' = dappendR m (init r)
                        in Deep f (init a) d empty r'
          (Cmpd f' c' r') => let r'' = dappendR m (init r)
                                 b'' = snoc c' (Simple r')
                             in Deep f (init a) f' b'' r''
      else Shallow f ++ Shallow (dappendR m (init r))

  export
  (Deque d, Sized d) => CatenableDeque (ImplicitCatDeque d) where
    (Shallow d1) ++ (Shallow d2) =
      if size d1 < 4 then Shallow (dappendL d1 d2)
      else if size d2 < 4 then Shallow (dappendR d1 d2)
      else let (f, m, r) = share d1 d2 in Deep f empty m empty r
    (Shallow d) ++ (Deep f a m b r) =
      if size d < 4 then Deep (dappendL d f) a m b r
      else Deep d (cons (Simple f) a) m b r
    (Deep f a m b r) ++ (Shallow d) =
      if size d < 4 then Deep f a m b (dappendR r d)
      else Deep f a m (snoc b (Simple r)) d
    (Deep f1 a1 m1 b1 r1) ++ (Deep f2 a2 m2 b2 r2) =
      let (r1', m, f2') = share r1 f2
          a1' = snoc a1 (Cmpd m1 b1 r1')
          b2' = cons (Cmpd f2' a2 m2) b2
      in Deep f1 a1' m b2' r2


module TypedContainers.RBTree.Balance

import Data.DPair
import Data.List
import TypedContainers.In
import TypedContainers.LawfulOrd
import TypedContainers.RBTree.Base
import TypedContainers.Util.Filter

%default total

0 lemma0 : LawfulOrd kt => {keys : List kt} -> (xk : kt) -> (yk : kt) -> In (yk ==) (filter (xk <) keys) -> (compare xk yk = LT)
lemma0 xk yk p0 =
  let MkDPair v (p3, p4) = extractIn p0 in
  let p5 = equality2 _ _ xk (convEQ _ _ p3) in
  let p6 = convLT xk v p4 in
  trans p5 p6

0 lemma1 : LawfulOrd kt => {keys : List kt} -> (xk : kt) -> (yk : kt) -> In (yk ==) (filter (xk >) keys) -> (compare xk yk = GT)
lemma1 xk yk p0 =
  let MkDPair v (p3, p4) = extractIn p0 in
  let p5 = equality2 _ _ xk (convEQ _ _ p3) in
  let p6 = convGT xk v p4 in
  trans p5 p6

public export
balanceLeft :
  (kord : LawfulOrd kt) =>
  (zk : kt) ->
  {0 zkp : In (zk ==) keys} ->
  vt zk ->
  BadTree {height, kt, kord, keys = filter (zk >) keys, vt} ->
  GoodTree {height, color = colorRight, kt, kord, keys = filter (zk <) keys, vt} ->
  Exists \color : Color => GoodTree {height = S height, color, kt, kord, keys = keys, vt}
balanceLeft zk zv (BadRedNode {kp = xkp} xk xv a (RedNode yk yv b c {kp = ykp})) d =
  let
    0 p0 : (arg : kt) -> (xk == arg = True) -> (yk > arg = True)
    p0 arg prr =
      let pp0 = reversion1 xk yk $ lemma0 xk yk ykp in
      let pp1 = equality2 xk arg yk (convEQ xk arg prr) in
      let pp2 = trans (sym pp1) pp0 in
      cong (\case { GT => True; _ => False }) pp2

    0 p0' : (arg : kt) -> (zk == arg = True) -> (yk < arg = True)
    p0' arg prr =
      let pp0 = lemma1 zk yk (outFilter ykp) in
      let pp1 = equality1 zk arg yk (convEQ zk arg prr) in
      let pp2 = trans (sym pp1) pp0 in
      cong (\case { LT => True; _ => False }) $ reversion2 arg yk pp2

    0 p1 : (arg : kt) -> (xk > arg = True) -> (yk > arg = True)
    p1 arg pp0 =
      let pp1 : (compare arg xk = LT) = reversion2 xk arg $ convGT _ _ pp0 in
      let pp2 : (compare xk yk = LT) = lemma0 xk yk ykp in
      let pp3 : (compare arg yk = LT) = transitivity arg xk yk pp1 pp2 in
      cong (\case { GT => True; _ => False }) $ reversion1 arg yk pp3

    0 p1' : (arg : kt) -> (yk < arg = True) -> (xk < arg = True)
    p1' arg pp0 =
      let pp1 : (compare yk arg = LT) = convLT _ _ pp0 in
      let pp2 : (compare xk arg = LT) = transitivity xk yk arg (lemma0 xk yk ykp) pp1 in
      cong (\case { LT => True; _ => False }) pp2

    0 p2 : (arg : kt) -> (yk > arg = True) -> (zk > arg = True)
    p2 arg pp0 =
      let pp1 : (compare arg yk = LT) = reversion2 yk arg $ convGT _ _ pp0 in
      let pp2 : (compare yk zk = LT) = reversion2 zk yk $ lemma1 zk yk (outFilter ykp) in
      let pp3 : (compare arg zk = LT) = transitivity arg yk zk pp1 pp2 in
      cong (\case { GT => True; _ => False }) $ reversion1 arg zk pp3

    0 p2' : (arg : kt) -> (zk < arg = True) -> (yk < arg = True)
    p2' arg pp0 =
      let pp1 : (compare zk arg = LT) = convLT _ _ pp0 in
      let pp2 : (compare yk arg = LT) = transitivity yk zk arg (reversion2 zk yk $ lemma1 zk yk (outFilter ykp)) pp1 in
      cong (\case { LT => True; _ => False }) pp2

    0 p3 : ((filter (xk >) $ filter (zk >) keys) === (filter (xk >) $ filter (yk >) keys))
    p3 =
      let pp0 : ((filter (yk >) $ filter (xk >) $ filter (zk >) keys) === (filter (xk >) $ filter (zk >) keys))
          = filterImplication {g = (xk >), f = (yk >)} p1
      in
      let pp1 : ((filter (xk >) $ filter (yk >) $ filter (zk >) keys) === (filter (yk >) $ filter (xk >) $ filter (zk >) keys))
          = filterCommutative
      in
      let pp2 : ((filter (xk >) $ filter (zk >) $ filter (yk >) keys) === (filter (xk >) $ filter (yk >) $ filter (zk >) keys))
          = cong (filter (xk >)) filterCommutative
      in
      let pp3 : ((filter (xk >) $ filter (yk >) keys) === (filter (xk >) $ filter (zk >) $ filter (yk >) keys))
          = cong (filter (xk >)) $ sym $ filterImplication p2
      in
      sym $ trans pp3 $ trans pp2 $ trans pp1 pp0

    0 p4 : ((filter (yk >) $ filter (xk <) $ filter (zk >) keys) === (filter (xk <) $ filter (yk >) keys))
    p4 =
      let pp0 : ((filter (xk <) $ filter (yk >) $ filter (zk >) keys) === (filter (yk >) $ filter (xk <) $ filter (zk >) keys))
          = filterCommutative
      in
      let pp1 : ((filter (xk <) $ filter (zk >) $ filter (yk >) keys) === (filter (xk <) $ filter (yk >) $ filter (zk >) keys))
          = cong (filter (xk <)) filterCommutative
      in
      let pp2 : ((filter (xk <) $ filter (yk >) keys) === (filter (xk <) $ filter (zk >) $ filter (yk >) keys))
          = cong (filter (xk <)) $ sym $ filterImplication p2
      in
      sym $ trans pp2 $ trans pp1 pp0

    0 p5 : ((filter (yk <) $ filter (xk <) $ filter (zk >) keys) === (filter (zk >) $ filter (yk <) keys))
    p5 =
      let pp0 : ((filter (xk <) $ filter (yk <) $ filter (zk >) keys) === (filter (yk <) $ filter (xk <) $ filter (zk >) keys))
          = filterCommutative
      in
      let pp1 : ((filter (yk <) $ filter (zk >) keys) === (filter (xk <) $ filter (yk <) $ filter (zk >) keys))
          = sym $ filterImplication p1'
      in
      let pp2 : ((filter (zk >) $ filter (yk <) keys) === (filter (yk <) $ filter (zk >) keys))
          = filterCommutative
      in
      sym $ trans pp2 $ trans pp1 pp0

    0 p6 : (filter (zk <) keys) === (filter (zk <) $ filter (yk <) keys)
    p6 =
      let pp0 : ((filter (yk <) $ filter (zk <) keys) === (filter (zk <) keys))
          = filterImplication p2'
      in
      let pp1 : ((filter (zk <) $ filter (yk <) keys) === (filter (yk <) $ filter (zk <) keys))
          = filterCommutative
      in
      sym $ trans pp1 pp0
  in
  let a' : GoodTree {keys = filter (xk >) $ filter (yk >) keys} = rewrite sym p3 in a in
  let b' : GoodTree {keys = filter (xk <) $ filter (yk >) keys} = rewrite sym p4 in b in
  let c' : GoodTree {keys = filter (zk >) $ filter (yk <) keys} = rewrite sym p5 in c in
  let d' : GoodTree {keys = filter (zk <) $ filter (yk <) keys} = rewrite sym p6 in d in
  let l' = BlackNode xk xv a' b' {kp = inFilter (yk >) (xk ==) p0 _ (outFilter xkp)} in
  let r' = BlackNode zk zv c' d' {kp = inFilter (yk <) (zk ==) p0' _ zkp} in
  Evidence Red $ RedNode yk yv l' r' {kp = outFilter $ outFilter ykp}
balanceLeft zk zv (BadRedNode {kp = ykp} yk yv (RedNode xk xv a b {kp = xkp}) c) d =
  let
    0 p0 : (arg : kt) -> (zk == arg = True) -> (yk < arg = True)
    p0 arg prr =
      let pp0 = lemma1 zk yk ykp in
      let pp1 = equality1 zk arg yk (convEQ zk arg prr) in
      let pp2 = trans (sym pp1) pp0 in
      cong (\case { LT => True; _ => False }) $ reversion2 arg yk pp2

    0 p1 : (In (xk ==) (filter (yk >) keys))
    p1 = outFilter $ replace {p = In (xk ==)} filterCommutative xkp

    0 p2 : (arg : kt) -> (yk > arg = True) -> (zk > arg = True)
    p2 arg pp0 =
      let pp1 : (compare arg yk = LT) = reversion2 yk arg $ convGT _ _ pp0 in
      let pp2 : (compare yk zk = LT) = reversion2 zk yk $ lemma1 zk yk ykp in
      let pp3 : (compare arg zk = LT) = transitivity arg yk zk pp1 pp2 in
      cong (\case { GT => True; _ => False }) $ reversion1 arg zk pp3

    0 p2' : (arg : kt) -> (zk < arg = True) -> (yk < arg = True)
    p2' arg pp0 =
      let pp1 : (compare zk arg = LT) = convLT _ _ pp0 in
      let pp2 : (compare yk arg = LT) = transitivity yk zk arg (reversion2 zk yk $ lemma1 zk yk ykp) pp1 in
      cong (\case { LT => True; _ => False }) pp2

    0 p3 : ((filter (xk >) $ filter (yk >) $ filter (zk >) keys) === (filter (xk >) $ filter (yk >) keys))
    p3 =
      let pp2 : ((filter (xk >) $ filter (zk >) $ filter (yk >) keys) === (filter (xk >) $ filter (yk >) $ filter (zk >) keys))
          = cong (filter (xk >)) filterCommutative
      in
      let pp3 : ((filter (xk >) $ filter (yk >) keys) === (filter (xk >) $ filter (zk >) $ filter (yk >) keys))
          = cong (filter (xk >)) $ sym $ filterImplication p2
      in
      sym $ trans pp3 pp2

    0 p4 : ((filter (xk <) $ filter (yk >) $ filter (zk >) keys) === (filter (xk <) $ filter (yk >) keys))
    p4 =
      let pp1 : ((filter (xk <) $ filter (zk >) $ filter (yk >) keys) === (filter (xk <) $ filter (yk >) $ filter (zk >) keys))
          = cong (filter (xk <)) filterCommutative
      in
      let pp2 : ((filter (xk <) $ filter (yk >) keys) === (filter (xk <) $ filter (zk >) $ filter (yk >) keys))
          = cong (filter (xk <)) $ sym $ filterImplication p2
      in
      sym $ trans pp2 pp1

    0 p5 : ((filter (yk <) $ filter (zk >) keys) === (filter (zk >) $ filter (yk <) keys))
    p5 = filterCommutative

    0 p6 : (filter (zk <) keys) === (filter (zk <) $ filter (yk <) keys)
    p6 =
      let pp0 : ((filter (yk <) $ filter (zk <) keys) === (filter (zk <) keys))
          = filterImplication p2'
      in
      let pp1 : ((filter (zk <) $ filter (yk <) keys) === (filter (yk <) $ filter (zk <) keys))
          = filterCommutative
      in
      sym $ trans pp1 pp0
  in
  let a' : GoodTree {keys = filter (xk >) $ filter (yk >) keys} = rewrite sym p3 in a in
  let b' : GoodTree {keys = filter (xk <) $ filter (yk >) keys} = rewrite sym p4 in b in
  let c' : GoodTree {keys = filter (zk >) $ filter (yk <) keys} = rewrite sym p5 in c in
  let d' : GoodTree {keys = filter (zk <) $ filter (yk <) keys} = rewrite sym p6 in d in
  let l' = BlackNode xk xv a' b' {kp = p1} in
  let r' = BlackNode zk zv c' d' {kp = inFilter (yk <) (zk ==) p0 _ zkp} in
  Evidence Red $ RedNode yk yv l' r' {kp = outFilter ykp}
balanceLeft zk zv (BadRedNode {kp = ykp} yk yv a@(BlackNode _ _ _ _) b@(BlackNode _ _ _ _)) c =
  Evidence Black $ BlackNode {kp = zkp} zk zv (RedNode {kp = ykp} yk yv a b) c
balanceLeft zk zv (BadRedNode {kp = ykp} yk yv a@(BlackNode _ _ _ _) b@(Empty _)) c impossible
balanceLeft zk zv (BadRedNode {kp = ykp} yk yv a@(Empty _) b@(BlackNode _ _ _ _)) c impossible
balanceLeft zk zv (BadRedNode {kp = ykp} yk yv a@(Empty _) b@(Empty _)) c =
  Evidence Black $ BlackNode {kp = zkp} zk zv (RedNode {kp = ykp} yk yv a b) c

public export
balanceRight :
  (kord : LawfulOrd kt) =>
  (xk : kt) ->
  {0 xkp : In (xk ==) keys} ->
  vt xk ->
  GoodTree {height, color = colorLeft, kt, kord, keys = filter (xk >) keys, vt} ->
  BadTree {height, kt, kord, keys = filter (xk <) keys, vt} ->
  Exists \color : Color => GoodTree {height = S height, color, kt, kord, keys, vt}
balanceRight xk xv a (BadRedNode {kp = zkp} zk zv (RedNode yk yv b c {kp = ykp}) d) =
  let
    0 p0 : (arg : kt) -> (yk < arg = True) -> (xk < arg = True)
    p0 arg pp0 =
      let pp1 : (compare yk arg = LT) = convLT _ _ pp0 in
      let pp2 : (compare xk arg = LT) = transitivity xk yk arg (lemma0 xk yk (outFilter ykp)) pp1 in
      cong (\case { LT => True; _ => False }) pp2

    0 p0' : (arg : kt) -> (zk < arg = True) -> (yk < arg = True)
    p0' arg pp0 =
      let pp1 : (compare zk arg = LT) = convLT _ _ pp0 in
      let pp2 : (compare yk arg = LT) = transitivity yk zk arg (reversion2 zk yk $ lemma1 zk yk ykp) pp1 in
      cong (\case { LT => True; _ => False }) pp2

    0 p1 : (arg : kt) -> (yk > arg = True) -> (zk > arg = True)
    p1 arg pp0 =
      let pp1 : (compare arg yk = LT) = reversion2 yk arg $ convGT _ _ pp0 in
      let pp2 : (compare yk zk = LT) = reversion2 zk yk $ lemma1 zk yk ykp in
      let pp3 : (compare arg zk = LT) = transitivity arg yk zk pp1 pp2 in
      cong (\case { GT => True; _ => False }) $ reversion1 arg zk pp3

    0 p2 : (arg : kt) -> (xk > arg = True) -> (yk > arg = True)
    p2 arg pp0 =
      let pp1 : (compare arg xk = LT) = reversion2 xk arg $ convGT _ _ pp0 in
      let pp2 : (compare xk yk = LT) = lemma0 xk yk (outFilter ykp) in
      let pp3 : (compare arg yk = LT) = transitivity arg xk yk pp1 pp2 in
      cong (\case { GT => True; _ => False }) $ reversion1 arg yk pp3

    0 p3 : (filter (xk >) $ filter (yk >) keys) === (filter (xk >) keys)
    p3 = trans filterCommutative $ filterImplication p2

    0 p4 : (filter (xk <) $ filter (yk >) keys) === (filter (yk >) $ filter (zk >) $ filter (xk <) keys)
    p4 = trans filterCommutative $ trans (sym $ filterImplication p1) filterCommutative

    0 p5 : (filter (zk >) $ filter (yk <) keys) === (filter (yk <) $ filter (zk >) $ filter (xk <) keys)
    p5 = trans (cong (filter (zk >))
      $ sym $ filterImplication p0)
      $ trans (cong (filter (zk >)) filterCommutative) filterCommutative

    0 p6 : (filter (zk <) $ filter (yk <) keys) === (filter (zk <) $ filter (xk <) keys)
    p6 =
      let
        pp0 : (filter (yk <) $ filter (zk <) $ filter (xk <) keys) === (filter (zk <) $ filter (xk <) keys)
        pp0 = filterImplication p0'
        pp1 : (filter (yk <) $ filter (xk <) $ filter (zk <) keys) === (filter (yk <) $ filter (zk <) $ filter (xk <) keys)
        pp1 = cong (filter (yk <)) filterCommutative
        pp2 : (filter (xk <) $ filter (yk <) $ filter (zk <) keys) === (filter (yk <) $ filter (xk <) $ filter (zk <) keys)
        pp2 = filterCommutative
        pp3 : (filter (yk <) $ filter (zk <) keys) === (filter (xk <) $ filter (yk <) $ filter (zk <) keys)
        pp3 = sym $ filterImplication p0
        pp4 : (filter (zk <) $ filter (yk <) keys) === (filter (yk <) $ filter (zk <) keys)
        pp4 = filterCommutative
      in
      trans pp4 $ trans pp3 $ trans pp2 $ trans pp1 pp0

    0 p7 : (arg : kt) -> (xk == arg = True) -> (yk > arg = True)
    p7 arg pp0 =
      let pp1 = lemma0 xk yk (outFilter ykp) in
      let pp2 = trans (sym pp1) $ equality1 xk arg yk (convEQ xk arg pp0) in
      cong (\case { GT => True; _ => False }) $ reversion1 arg yk $ sym pp2

    0 p8 : (arg : kt) -> (zk == arg = True) -> (yk < arg = True)
    p8 arg pp0 =
      let pp1 = lemma1 zk yk ykp in
      let pp2 = trans (sym pp1) $ equality1 zk arg yk (convEQ zk arg pp0) in
      cong (\case { LT => True; _ => False }) $ reversion2 arg yk $ sym pp2
  in
  let a' : GoodTree {keys = filter (xk >) $ filter (yk >) keys} = rewrite p3 in a in
  let b' : GoodTree {keys = filter (xk <) $ filter (yk >) keys} = rewrite p4 in b in
  let c' : GoodTree {keys = filter (zk >) $ filter (yk <) keys} = rewrite p5 in c in
  let d' : GoodTree {keys = filter (zk <) $ filter (yk <) keys} = rewrite p6 in d in
  let l' = BlackNode xk xv a' b' {kp = inFilter (yk >) (xk ==) p7 _ xkp} in
  let r' = BlackNode zk zv c' d' {kp = inFilter (yk <) (zk ==) p8 _ (outFilter zkp)} in
  Evidence Red $ RedNode yk yv l' r' {kp = outFilter $ outFilter ykp}
balanceRight xk xv a (BadRedNode {kp = ykp} yk yv b (RedNode zk zv c d {kp = zkp})) =
  let
    0 p1 : (arg : kt) -> (yk < arg = True) -> (xk < arg = True)
    p1 arg pp0 =
      let pp1 : (compare yk arg = LT) = convLT _ _ pp0 in
      let pp2 : (compare xk yk = LT) = lemma0 xk yk ykp in
      let pp3 : (compare xk arg = LT) = transitivity xk yk arg pp2 pp1 in
      cong (\case { LT => True; _ => False }) pp3

    0 p2 : (arg : kt) -> (xk > arg = True) -> (yk > arg = True)
    p2 arg pp0 =
      let pp1 : (compare arg xk = LT) = reversion2 xk arg $ convGT _ _ pp0 in
      let pp2 : (compare xk yk = LT) = lemma0 xk yk ykp in
      let pp3 : (compare arg yk = LT) = transitivity arg xk yk pp1 pp2 in
      cong (\case { GT => True; _ => False }) $ reversion1 arg yk pp3

    0 p3 : (filter (xk >) $ filter (yk >) keys) === (filter (xk >) keys)
    p3 = trans filterCommutative $ filterImplication p2

    0 p4 : (filter (xk <) $ filter (yk >) keys) === (filter (yk >) $ filter (xk <) keys)
    p4 = filterCommutative

    0 p5 : (filter (zk >) $ filter (yk <) keys) === (filter (zk >) $ filter (yk <) $ filter (xk <) keys)
    p5 =
      let
        pp0 : (filter (zk >) $ filter (xk <) $ filter (yk <) keys) === (filter (zk >) $ filter (yk <) $ filter (xk <) keys)
        pp0 = cong (filter (zk >)) filterCommutative
        pp1 : (filter (zk >) $ filter (yk <) keys) === (filter (zk >) $ filter (xk <) $ filter (yk <) keys)
        pp1 = cong (filter (zk >)) $ sym $ filterImplication p1
      in
      trans pp1 pp0

    0 p6 : (filter (zk <) $ filter (yk <) keys) === (filter (zk <) $ filter (yk <) $ filter (xk <) keys)
    p6 =
      let
        pp0 : (filter (zk <) $ filter (xk <) $ filter (yk <) keys) === (filter (zk <) $ filter (yk <) $ filter (xk <) keys)
        pp0 = cong (filter (zk <)) filterCommutative
        pp1 : (filter (zk <) $ filter (yk <) keys) === (filter (zk <) $ filter (xk <) $ filter (yk <) keys)
        pp1 = cong (filter (zk <)) $ sym $ filterImplication p1
      in
      trans pp1 pp0

    0 p7 : (arg : kt) -> (xk == arg = True) -> (yk > arg = True)
    p7 arg pp0 =
      let pp1 = lemma0 xk yk ykp in
      let pp2 = trans (sym pp1) $ equality1 xk arg yk (convEQ xk arg pp0) in
      cong (\case { GT => True; _ => False }) $ reversion1 arg yk $ sym pp2
  in
  let a' : GoodTree {keys = filter (xk >) $ filter (yk >) keys} = rewrite p3 in a in
  let b' : GoodTree {keys = filter (xk <) $ filter (yk >) keys} = rewrite p4 in b in
  let c' : GoodTree {keys = filter (zk >) $ filter (yk <) keys} = rewrite p5 in c in
  let d' : GoodTree {keys = filter (zk <) $ filter (yk <) keys} = rewrite p6 in d in
  let l' = BlackNode xk xv a' b' {kp = inFilter (yk >) (xk ==) p7 _ xkp} in
  let r' = BlackNode zk zv c' d' {kp = replace {p = In (zk ==)} (trans filterCommutative $ filterImplication p1) zkp} in
  Evidence Red $ RedNode yk yv l' r' {kp = outFilter ykp}
balanceRight _ _ _ (BadRedNode _ _ (Empty _) (BlackNode _ _ _ _)) impossible
balanceRight _ _ _ (BadRedNode _ _ (BlackNode _ _ _ _) (Empty _)) impossible
balanceRight xk xv a (BadRedNode {kp = ykp} yk yv b@(BlackNode _ _ _ _) c@(BlackNode _ _ _ _)) =
  Evidence Black $ BlackNode {kp = xkp} xk xv a (RedNode {kp = ykp} yk yv b c)
balanceRight xk xv a (BadRedNode {kp = ykp} yk yv b@(Empty _) c@(Empty _)) =
  Evidence Black $ BlackNode {kp = xkp} xk xv a (RedNode {kp = ykp} yk yv b c)


module Builtin

%default total

-- The most primitive data types; things which are used by desugaring

-- Totality assertions

||| Assert to the totality checker that the given expression will always
||| terminate.
|||
||| The multiplicity of its argument is 1, so `assert_total` won't affect how
||| many times variables are used. If you're not writing a linear function,
||| this doesn't make a difference.
|||
||| Note: assert_total can reduce at compile time, if required for unification,
||| which might mean that it's no longer guarded a subexpression. Therefore,
||| it is best to use it around the smallest possible subexpression.
%inline
public export
assert_total : (1 _ : a) -> a
assert_total x = x

||| Assert to the totality checker that y is always structurally smaller than x
||| (which is typically a pattern argument, and *must* be in normal form for
||| this to work).
|||
||| The multiplicity of x is 0, so in a linear function, you can pass values to
||| x even if they have already been used.
||| The multiplicity of y is 1, so `assert_smaller` won't affect how many times
||| its y argument is used.
||| If you're not writing a linear function, the multiplicities don't make a
||| difference.
|||
||| @ x the larger value (typically a pattern argument)
||| @ y the smaller value (typically an argument to a recursive call)
%inline
public export
assert_smaller : (0 x : a) -> (1 y : b) -> b
assert_smaller x y = y

-- Unit type and pairs

||| The canonical single-element type, also known as the trivially true
||| proposition.
public export
data Unit =
  ||| The trivial constructor for `()`.
  MkUnit

||| The non-dependent pair type, also known as conjunction.
public export
data Pair : Type -> Type -> Type where
  ||| A pair of elements.
  ||| @ a the left element of the pair
  ||| @ b the right element of the pair
  MkPair : {0 a, b : Type} -> (x : a) -> (y : b) -> Pair a b

||| Return the first element of a pair.
public export
fst : {0 a, b : Type} -> (a, b) -> a
fst (x, y) = x

||| Return the second element of a pair.
public export
snd : {0 a, b : Type} -> (a, b) -> b
snd (x, y) = y

-- This directive tells auto implicit search what to use to look inside pairs
%pair Pair fst snd

infix 5 #

||| A pair type where each component is linear
public export
data LPair : Type -> Type -> Type where
  ||| A linear pair of elements.
  ||| If you take one copy of the linear pair apart
  ||| then you only get one copy of its left and right elements.
  ||| @ a the left element of the pair
  ||| @ b the right element of the pair
  (#) : (1 _ : a) -> (1 _ : b) -> LPair a b

namespace DPair
  ||| Dependent pairs aid in the construction of dependent types by providing
  ||| evidence that some value resides in the type.
  |||
  ||| Formally, speaking, dependent pairs represent existential quantification -
  ||| they consist of a witness for the existential claim and a proof that the
  ||| property holds for it.
  |||
  ||| @ a the value to place in the type.
  ||| @ p the dependent type that requires the value.
  public export
  record DPair a (p : a -> Type) where
    constructor MkDPair
    fst : a
    snd : p fst

  ||| A dependent variant of LPair, pairing a result value with a resource
  ||| that depends on the result value
  public export
  data Res : (a : Type) -> (a -> Type) -> Type where
       (#) : (val : a) -> (1 r : t val) -> Res a t

-- The empty type

||| The empty type, also known as the trivially false proposition.
|||
||| Use `void` or `absurd` to prove anything if you have a variable of type
||| `Void` in scope.
public export
data Void : Type where

-- Equality

public export
data Equal : forall a, b . a -> b -> Type where
     [search a b]
     Refl : {0 x : a} -> Equal x x

%name Equal prf

infix 6 ===, ~=~

-- An equality type for when you want to assert that each side of the
-- equality has the same type, but there's not other evidence available
-- to help with unification
public export
(===) : (x : a) -> (y : a) -> Type
(===) = Equal

||| Explicit heterogeneous ("John Major") equality.  Use this when Idris
||| incorrectly chooses homogeneous equality for `(=)`.
||| @ a the type of the left side
||| @ b the type of the right side
||| @ x the left side
||| @ y the right side
public export
(~=~) : (x : a) -> (y : b) -> Type
(~=~) = Equal

||| Perform substitution in a term according to some equality.
|||
||| Like `replace`, but with an explicit predicate, and applying the rewrite in
||| the other direction, which puts it in a form usable by the `rewrite` tactic
||| and term.
%inline
public export
rewrite__impl : {0 x, y : a} -> (0 p : _) ->
                (0 rule : x = y) -> (val : p y) -> p x
rewrite__impl p Refl prf = prf

%rewrite Equal rewrite__impl

||| Perform substitution in a term according to some equality.
%inline
public export
replace : forall x, y, p . (0 rule : x = y) -> p x -> p y
replace Refl prf = prf

||| Symmetry of propositional equality.
%inline
public export
sym : (0 rule : x = y) -> y = x
sym Refl = Refl

||| Transitivity of propositional equality.
%inline
public export
trans : forall a, b, c . (0 l : a = b) -> (0 r : b = c) -> a = c
trans Refl Refl = Refl

||| Injectivity of MkDPair (first components)
export
mkDPairInjectiveFst : MkDPair a pa === MkDPair b qb -> a === b
mkDPairInjectiveFst Refl = Refl

||| Injectivity of MkDPair (snd components)
export
mkDPairInjectiveSnd : MkDPair a pa === MkDPair a qa -> pa === qa
mkDPairInjectiveSnd Refl = Refl

||| Subvert the type checker.  This function is abstract, so it will not reduce
||| in the type checker.  Use it with care - it can result in segfaults or
||| worse!
public export
believe_me : a -> b
believe_me = prim__believe_me _ _

||| Assert to the usage checker that the given function uses its argument linearly.
public export
assert_linear : (1 f : a -> b) -> (1 val : a) -> b
assert_linear = believe_me id
  where
    id : (1 f : a -> b) -> a -> b
    id f = f

export partial
idris_crash : String -> a
idris_crash = prim__crash _

public export %inline
delay : a -> Lazy a
delay x = Delay x

public export %inline
force : Lazy a -> a
force x = Force x

%stringLit fromString

||| Interface for types that can be constructed from string literals.
public export
interface FromString ty where
  constructor MkFromString
  ||| Conversion from String.
  fromString : String -> ty

%allow_overloads fromString

%inline
public export
FromString String where
  fromString s = s

%defaulthint
%inline
public export
defaultString : FromString String
defaultString = %search

%charLit fromChar

||| Interface for types that can be constructed from char literals.
public export
interface FromChar ty where
  constructor MkFromChar
  ||| Conversion from Char.
  fromChar : Char -> ty

%allow_overloads fromChar

%inline
public export
FromChar Char where
  fromChar s = s

%defaulthint
%inline
public export
defaultChar : FromChar Char
defaultChar = %search

%doubleLit fromDouble

||| Interface for types that can be constructed from double literals.
public export
interface FromDouble ty where
  constructor MkFromDouble
  ||| Conversion from Double.
  fromDouble : Double -> ty

%allow_overloads fromDouble

%inline
public export
FromDouble Double where
  fromDouble s = s

%defaulthint
%inline
public export
defaultDouble : FromDouble Double
defaultDouble = %search


module Network.TLS.AEAD

import Data.Stream
import Data.Bits
import Data.Vect
import Data.List
import Utils.Misc
import Utils.Bytes
import Crypto.AES.Big
import Crypto.AES.Common
import Crypto.Hash
import Crypto.Hash.GHash
import Crypto.Hash.Poly1305
import Crypto.ChaCha

public export
interface AEAD (0 a : Type) where
  ||| IV generated during key exchange
  fixed_iv_length : Nat
  enc_key_length : Nat
  mac_length : Nat
  mac_key_length : Nat
  ||| Part of IV that is sent along with the ciphertext, should always be 0 in TLS 1.3
  record_iv_length : Nat

  encrypt : Vect enc_key_length Bits8 -> Vect fixed_iv_length Bits8 -> Vect mac_key_length Bits8 -> Nat ->
            (plaintext : List Bits8) -> (aad : List Bits8) -> (Vect record_iv_length Bits8, List Bits8, Vect mac_length Bits8)
  decrypt : Vect enc_key_length Bits8 -> Vect fixed_iv_length Bits8 -> Vect record_iv_length Bits8 -> Vect mac_key_length Bits8 -> Nat ->
            (ciphertext : List Bits8) -> (plaintext_to_aad : List Bits8 -> List Bits8) -> (mac_tag : List Bits8) -> (List Bits8, Bool)

aes_pad_iv_block : {iv : Nat} -> Vect iv Bits8 -> Stream (Vect (iv+4) Bits8)
aes_pad_iv_block iv = map ((iv ++) . to_be . (cast {to=Bits32})) $ drop 2 nats

aes_keystream : (mode : Mode) -> Vect ((get_n_k mode) * 4) Bits8 -> Vect 12 Bits8 -> Stream Bits8
aes_keystream mode key iv =
  stream_concat $ map (toList . encrypt_block mode key) (aes_pad_iv_block iv)

aes_gcm_create_aad : (mode : Mode) -> Vect ((get_n_k mode) * 4) Bits8 -> Vect 12 Bits8 -> List Bits8 -> List Bits8 -> Vect 16 Bits8
aes_gcm_create_aad mode key iv aad ciphertext =
  let a = toList $ to_be {n=8} $ cast {to=Bits64} $ 8 * (length aad)
      c = toList $ to_be {n=8} $ cast {to=Bits64} $ 8 * (length ciphertext)
      input = pad_zero 16 aad <+> pad_zero 16 ciphertext <+> a <+> c
      h = encrypt_block mode key (replicate _ 0)
      output = mac GHash h input
      j0 = encrypt_block mode key (iv ++ (to_be $ the Bits32 1))
  in zipWith xor j0 output

public export
data TLS13_AES_128_GCM : Type where

public export
AEAD TLS13_AES_128_GCM where
  fixed_iv_length = 12
  enc_key_length = 16
  mac_length = 16
  mac_key_length = 0
  record_iv_length = 0

  encrypt key iv mac_key seq_no plaintext aad =
    let iv' = zipWith xor iv $ integer_to_be _ $ natToInteger seq_no
        ciphertext = zipWith xor plaintext (toList $ Stream.take (length plaintext) $ aes_keystream AES128 key iv')
        mac_tag = aes_gcm_create_aad AES128 key iv' aad ciphertext
    in ([], ciphertext, mac_tag)

  decrypt key iv [] mac_key seq_no ciphertext aadf mac_tag' =
    let iv' = zipWith xor iv $ integer_to_be _ $ natToInteger seq_no
        plaintext = zipWith xor ciphertext (toList $ Stream.take (length ciphertext) $ aes_keystream AES128 key iv')
        mac_tag = aes_gcm_create_aad AES128 key iv' (aadf plaintext) ciphertext
    in (plaintext, s_eq' (toList mac_tag) mac_tag')

public export
data TLS12_AES_128_GCM : Type where

public export
AEAD TLS12_AES_128_GCM where
  fixed_iv_length = 4
  enc_key_length = 16
  mac_length = 16
  mac_key_length = 0
  record_iv_length = 8

  encrypt key iv mac_key seq_no plaintext aad =
    let explicit_iv = to_be {n=8} $ cast {to=Bits64} seq_no
        iv' = iv ++ explicit_iv
        ciphertext = zipWith xor plaintext (toList $ Stream.take (length plaintext) $ aes_keystream AES128 key iv')
        mac_tag = aes_gcm_create_aad AES128 key iv' aad ciphertext
    in (explicit_iv, ciphertext, mac_tag)

  decrypt key iv explicit_iv mac_key seq_no ciphertext aadf mac_tag' =
    let iv' = iv ++ explicit_iv
        plaintext = zipWith xor ciphertext (toList $ Stream.take (length ciphertext) $ aes_keystream AES128 key iv')
        mac_tag = aes_gcm_create_aad AES128 key iv' (aadf plaintext) ciphertext
    in (plaintext, s_eq' (toList mac_tag) mac_tag')

public export
data TLS13_AES_256_GCM : Type where

public export
AEAD TLS13_AES_256_GCM where
  fixed_iv_length = 12
  enc_key_length = 32
  mac_length = 16
  mac_key_length = 0
  record_iv_length = 0

  encrypt key iv mac_key seq_no plaintext aad =
    let iv' = zipWith xor iv $ integer_to_be _ $ natToInteger seq_no
        ciphertext = zipWith xor plaintext (toList $ Stream.take (length plaintext) $ aes_keystream AES256 key iv')
        mac_tag = aes_gcm_create_aad AES256 key iv' aad ciphertext
    in ([], ciphertext, mac_tag)

  decrypt key iv [] mac_key seq_no ciphertext aadf mac_tag' =
    let iv' = zipWith xor iv $ integer_to_be _ $ natToInteger seq_no
        plaintext = zipWith xor ciphertext (toList $ Stream.take (length ciphertext) $ aes_keystream AES256 key iv')
        mac_tag = aes_gcm_create_aad AES256 key iv' (aadf plaintext) ciphertext
    in (plaintext, s_eq' (toList mac_tag) mac_tag')

public export
data TLS12_AES_256_GCM : Type where

public export
AEAD TLS12_AES_256_GCM where
  fixed_iv_length = 4
  enc_key_length = 32
  mac_length = 16
  mac_key_length = 0
  record_iv_length = 8

  encrypt key iv mac_key seq_no plaintext aad =
    let explicit_iv = to_be {n=8} $ cast {to=Bits64} seq_no
        iv' = iv ++ explicit_iv
        ciphertext = zipWith xor plaintext (toList $ Stream.take (length plaintext) $ aes_keystream AES256 key iv')
        mac_tag = aes_gcm_create_aad AES256 key iv' aad ciphertext
    in (explicit_iv, ciphertext, mac_tag)

  decrypt key iv explicit_iv mac_key seq_no ciphertext aadf mac_tag' =
    let iv' = iv ++ explicit_iv
        plaintext = zipWith xor ciphertext (toList $ Stream.take (length ciphertext) $ aes_keystream AES256 key iv')
        mac_tag = aes_gcm_create_aad AES256 key iv' (aadf plaintext) ciphertext
    in (plaintext, s_eq' (toList mac_tag) mac_tag')

chacha_create_aad : Vect 64 Bits8 -> List Bits8 -> List Bits8 -> Vect 16 Bits8
chacha_create_aad polykey aad ciphertext =
  let key = take 32 polykey
      length_aad = toList $ to_le {n=8} $ cast {to=Bits64} $ length aad
      length_ciphertext = toList $ to_le {n=8} $ cast {to=Bits64} $ length ciphertext
      input = pad_zero 16 aad ++ pad_zero 16 ciphertext ++ length_aad ++ length_ciphertext
  in mac Poly1305 key input

public export
data TLS1213_ChaCha20_Poly1305 : Type where

public export
AEAD TLS1213_ChaCha20_Poly1305 where
  fixed_iv_length = 12
  enc_key_length = 32
  mac_length = 16
  mac_key_length = 0
  record_iv_length = 0

  encrypt key iv [] seq_no plaintext aad =
    let k' = from_le {n=4} <$> group 8 4 key
        iv' = zipWith xor iv $ integer_to_be _ $ natToInteger seq_no
        i' = from_le {n=4} <$> group 3 4 iv'
        (polykey :: keystream) = map (\c => chacha_rfc8439_block 10 (cast c) k' i') nats
        ciphertext = zipWith xor plaintext (toList $ Stream.take (length plaintext) $ stream_concat keystream)
        auth_tag = chacha_create_aad polykey aad ciphertext
    in ([], ciphertext, auth_tag)

  decrypt key iv [] [] seq_no ciphertext aadf mac_tag' =
    let k' = from_le {n=4} <$> group 8 4 key
        iv' = zipWith xor iv $ integer_to_be _ $ natToInteger seq_no
        i' = from_le {n=4} <$> group 3 4 iv'
        (polykey :: keystream) = map (\c => chacha_rfc8439_block 10 (cast c) k' i') nats
        plaintext = zipWith xor ciphertext (toList $ Stream.take (length ciphertext) $ stream_concat keystream)
        auth_tag = chacha_create_aad polykey (aadf plaintext) ciphertext
    in (plaintext, toList auth_tag `s_eq'` mac_tag')


module Compiler.Opts.ConstantFold

import Compiler.CompileExpr
import Core.Context
import Core.Context.Log
import Core.Primitives
import Core.Value
import Data.List
import Data.Vect

findConstAlt : Constant -> List (CConstAlt vars) ->
               Maybe (CExp vars) -> Maybe (CExp vars)
findConstAlt c [] def = def
findConstAlt c (MkConstAlt c' exp :: alts) def = if c == c'
    then Just exp
    else findConstAlt c alts def

foldableOp : PrimFn ar -> Bool
foldableOp BelieveMe = False
foldableOp (Cast from to)
  = fromMaybe False [| intKind from `smaller` intKind to |]
  where
    smaller : IntKind -> IntKind -> Bool
    smaller (Signed x) (Signed y) = x <= y
    smaller (Unsigned x) (Unsigned y) = x <= y
    smaller (Signed x) (Unsigned y) = x < P y
    smaller (Unsigned x) (Signed y) = P x < y
foldableOp _ = True

data Subst : List Name -> List Name -> Type where
  Nil  : Subst [] vars
  (::) : CExp vars -> Subst ds vars -> Subst (d :: ds) vars
  Wk   : SizeOf ws -> Subst ds vars -> Subst (ws ++ ds) (ws ++ vars)

initSubst : (vars : List Name) -> Subst vars vars
initSubst vars
  = rewrite sym $ appendNilRightNeutral vars in
    Wk (mkSizeOf vars) []

wk : SizeOf out -> Subst ds vars -> Subst (out ++ ds) (out ++ vars)
wk sout (Wk {ws, ds, vars} sws rho)
  = rewrite appendAssociative out ws ds in
    rewrite appendAssociative out ws vars in
    Wk (sout + sws) rho
wk ws rho = Wk ws rho

record WkCExp (vars : List Name) where
  constructor MkWkCExp
  {0 outer, supp : List Name}
  size : SizeOf outer
  0 prf : vars === outer ++ supp
  expr : CExp supp

Weaken WkCExp where
  weaken (MkWkCExp s Refl e) = MkWkCExp (suc s) Refl e

lookup : FC -> Var ds -> Subst ds vars -> CExp vars
lookup fc (MkVar p) rho = case go p rho of
    Left (MkVar p') => CLocal fc p'
    Right (MkWkCExp s Refl e) => weakenNs s e

  where

  go : {i : Nat} -> {0 ds, vars : _} -> (0 _ : IsVar n i ds) ->
       Subst ds vars -> Either (Var vars) (WkCExp vars)
  go First     (val :: rho) = Right (MkWkCExp zero Refl val)
  go (Later p) (val :: rho) = go p rho
  go p         (Wk ws  rho) = case sizedView ws of
    Z => go p rho
    S ws' => case i of
      Z => Left (MkVar First)
      S i' => bimap later weaken (go (dropLater p) (Wk ws' rho))

-- constant folding of primitive functions
-- if a primitive function is applied to only constant
-- then replace with the result
-- if there's only 1 constant argument to a commutative function
-- move the constant to the right. This simplifies Compiler.Opts.Identity
constFold : {vars' : _} ->
            Subst vars vars' ->
            CExp vars -> CExp vars'
constFold rho (CLocal fc p) = lookup fc (MkVar p) rho
constFold rho e@(CRef fc x) = CRef fc x
constFold rho (CLam fc x y)
  = CLam fc x $ constFold (wk (mkSizeOf [x]) rho) y
constFold rho (CLet fc x inlineOK y z) =
    let val = constFold rho y
     in case val of
        CPrimVal _ _ => if inlineOK
            then constFold (val :: rho) z
            else CLet fc x inlineOK val (constFold (wk (mkSizeOf [x]) rho) z)
        _ => CLet fc x inlineOK val (constFold (wk (mkSizeOf [x]) rho) z)
constFold rho (CApp fc x xs) = CApp fc (constFold rho x) (constFold rho <$> xs)
constFold rho (CCon fc x y tag xs) = CCon fc x y tag $ constFold rho <$> xs
constFold rho (COp {arity} fc fn xs) =
    let xs' = map (constFold rho) xs
        e = constRight fc fn xs'
     in fromMaybe e $ do
          guard (foldableOp fn)
          nfs <- traverse toNF xs'
          nf <- getOp fn nfs
          fromNF nf
  where
    toNF : CExp vars' -> Maybe (NF vars')
    -- Don't fold `Int` and `Double` because they have varying widths
    toNF (CPrimVal fc (I _)) = Nothing
    toNF (CPrimVal fc (Db _)) = Nothing
    -- Fold the rest
    toNF (CPrimVal fc c) = Just $ NPrimVal fc c
    toNF _ = Nothing

    fromNF : NF vars' -> Maybe (CExp vars')
    fromNF (NPrimVal fc c) = Just $ CPrimVal fc c
    fromNF _ = Nothing

    commutative : Constant -> Bool
    commutative DoubleType = False
    commutative _ = True

    constRight : {ar : _} -> FC -> PrimFn ar ->
                 Vect ar (CExp vars') -> CExp vars'
    constRight fc (Add ty) [x@(CPrimVal _ _), y] =
        if commutative ty
            then COp fc (Add ty) [y, x]
            else COp fc (Add ty) [x, y]
    constRight fc (Mul ty) [x@(CPrimVal _ _), y] =
        if commutative ty
            then COp fc (Mul ty) [y, x]
            else COp fc (Mul ty) [x, y]
    constRight fc fn args = COp fc fn args

constFold rho (CExtPrim fc p xs) = CExtPrim fc p $ constFold rho <$> xs
constFold rho (CForce fc x y) = CForce fc x $ constFold rho y
constFold rho (CDelay fc x y) = CDelay fc x $ constFold rho y
constFold rho (CConCase fc sc xs x)
  = CConCase fc (constFold rho sc) (foldAlt <$> xs) (constFold rho <$> x)
  where
    foldAlt : CConAlt vars -> CConAlt vars'
    foldAlt (MkConAlt n ci t xs e)
      = MkConAlt n ci t xs $ constFold (wk (mkSizeOf xs) rho) e

constFold rho (CConstCase fc sc xs x) =
    let sc' = constFold rho sc
     in case sc' of
        CPrimVal _ val => case findConstAlt val xs x of
            Just exp => constFold rho exp
            Nothing => CConstCase fc (constFold rho sc) (foldAlt <$> xs) (constFold rho <$> x)
        _ => CConstCase fc (constFold rho sc) (foldAlt <$> xs) (constFold rho <$> x)
  where
    foldAlt : CConstAlt vars -> CConstAlt vars'
    foldAlt (MkConstAlt c e) = MkConstAlt c $ constFold rho e
constFold rho (CPrimVal fc v) = CPrimVal fc v
constFold rho (CErased fc) = CErased fc
constFold rho (CCrash fc err) = CCrash fc err

constFoldCDef : CDef -> Maybe CDef
constFoldCDef (MkFun args exp)
  = Just $ MkFun args $ constFold (initSubst args) exp
constFoldCDef _ = Nothing

export
constantFold : Ref Ctxt Defs => Name -> Core ()
constantFold fn = do
    defs <- get Ctxt
    Just (fnIdx, gdef) <- lookupCtxtExactI fn defs.gamma
        | Nothing => pure ()
    let Just cdef = gdef.compexpr
        | Nothing => pure ()
    let Just cdef' = constFoldCDef cdef
        | Nothing => pure ()
    log "compiler.const-fold" 50 $ "constant folding " ++ show !(getFullName fn)
                                 ++ "\n\told def: " ++ show cdef
                                 ++ "\n\tnew def: " ++ show cdef'
    setCompiled (Resolved fnIdx) cdef'


module Lens.Uninhabited

import Lens
import Map

%default total

public export
Uninhabited (a = b) => Uninhabited (Boolean a = Boolean b) where
    uninhabited Refl @{ab} = uninhabited @{ab} Refl

public export
Uninhabited (Boolean a = Number b) where
    uninhabited Refl impossible

public export
Uninhabited (Boolean a = Text b) where
    uninhabited Refl impossible

public export
Uninhabited (a = b) => Uninhabited (Number a = Number b) where
    uninhabited Refl @{ab} = uninhabited @{ab} Refl

public export
Uninhabited (Number a = Boolean b) where
    uninhabited Refl impossible

public export
Uninhabited (Number a = Text b) where
    uninhabited Refl impossible

public export
Uninhabited (a = b) => Uninhabited (Text a = Text b) where
    uninhabited Refl @{ab} = uninhabited @{ab} Refl

public export
Uninhabited (Text a = Boolean b) where
    uninhabited Refl impossible

public export
Uninhabited (Text a = Number b) where
    uninhabited Refl impossible

public export
Uninhabited (KBoolean = KNumber) where
    uninhabited Refl impossible

public export
Uninhabited (KBoolean = KText) where
    uninhabited Refl impossible

public export
Uninhabited (KNumber = KBoolean) where
    uninhabited Refl impossible

public export
Uninhabited (KNumber = KText) where
    uninhabited Refl impossible

public export
Uninhabited (KText = KBoolean) where
    uninhabited Refl impossible

public export
Uninhabited (KText = KNumber) where
    uninhabited Refl impossible

public export
Uninhabited (Null = Primitive _) where
    uninhabited Refl impossible

public export
Uninhabited (Null = Array _) where
    uninhabited Refl impossible

public export
Uninhabited (Null = Object _) where
    uninhabited Refl impossible

public export
Uninhabited (Primitive _ = Null) where
    uninhabited Refl impossible

public export
Uninhabited (a = b) => Uninhabited (Primitive a = Primitive b) where
    uninhabited Refl @{ab} = uninhabited @{ab} Refl

public export
Uninhabited (Primitive _ = Array _) where
    uninhabited Refl impossible

public export
Uninhabited (Primitive _ = Object _) where
    uninhabited Refl impossible

public export
Uninhabited (Array _ = Null) where
    uninhabited Refl impossible

public export
Uninhabited (Array _ = Primitive _) where
    uninhabited Refl impossible

public export
Uninhabited (a = b) => Uninhabited (Array a = Array b) where
    uninhabited Refl @{ab} = uninhabited @{ab} Refl

public export
Uninhabited (Array _ = Object _) where
    uninhabited Refl impossible

public export
Uninhabited (Object _ = Null) where
    uninhabited Refl impossible

public export
Uninhabited (Object _ = Primitive _) where
    uninhabited Refl impossible

public export
Uninhabited (Object _ = Array _) where
    uninhabited Refl impossible

public export
Uninhabited (a = b) => Uninhabited (Object a = Object b) where
    uninhabited Refl @{ab} = uninhabited @{ab} Refl

public export
Uninhabited (SNull = SBoolean) where
    uninhabited Refl impossible

public export
Uninhabited (SNull = SNumber) where
    uninhabited Refl impossible

public export
Uninhabited (SNull = SText) where
    uninhabited Refl impossible

public export
Uninhabited (SNull = SArray _ _) where
    uninhabited Refl impossible

public export
Uninhabited (SNull = SObject _) where
    uninhabited Refl impossible

public export
Uninhabited (SBoolean = SNull) where
    uninhabited Refl impossible

public export
Uninhabited (SBoolean = SNumber) where
    uninhabited Refl impossible

public export
Uninhabited (SBoolean = SText) where
    uninhabited Refl impossible

public export
Uninhabited (SBoolean = SArray _ _) where
    uninhabited Refl impossible

public export
Uninhabited (SBoolean = SObject _) where
    uninhabited Refl impossible

public export
Uninhabited (SNumber = SNull) where
    uninhabited Refl impossible

public export
Uninhabited (SNumber = SBoolean) where
    uninhabited Refl impossible

public export
Uninhabited (SNumber = SText) where
    uninhabited Refl impossible

public export
Uninhabited (SNumber = SArray _ _) where
    uninhabited Refl impossible

public export
Uninhabited (SNumber = SObject _) where
    uninhabited Refl impossible

public export
Uninhabited (SText = SNull) where
    uninhabited Refl impossible

public export
Uninhabited (SText = SBoolean) where
    uninhabited Refl impossible

public export
Uninhabited (SText = SNumber) where
    uninhabited Refl impossible

public export
Uninhabited (SText = SArray _ _) where
    uninhabited Refl impossible

public export
Uninhabited (SText = SObject _) where
    uninhabited Refl impossible

public export
Uninhabited (SArray _ _ = SNull) where
    uninhabited Refl impossible

public export
Uninhabited (SArray _ _ = SBoolean) where
    uninhabited Refl impossible

public export
Uninhabited (SArray _ _ = SNumber) where
    uninhabited Refl impossible

public export
Uninhabited (SArray _ _ = SText) where
    uninhabited Refl impossible

public export
Uninhabited (a = b) => Uninhabited (SArray a _ = SArray b _) where
    uninhabited Refl @{ab} = uninhabited @{ab} Refl

public export
Uninhabited (a = b) => Uninhabited (SArray _ a = SArray _ b) where
    uninhabited Refl @{ab} = uninhabited @{ab} Refl

public export
Uninhabited (SArray _ _ = SObject _) where
    uninhabited Refl impossible

public export
Uninhabited (SObject _ = SNull) where
    uninhabited Refl impossible

public export
Uninhabited (SObject _ = SBoolean) where
    uninhabited Refl impossible

public export
Uninhabited (SObject _ = SNumber) where
    uninhabited Refl impossible

public export
Uninhabited (SObject _ = SText) where
    uninhabited Refl impossible

public export
Uninhabited (SObject _ = SArray _ _) where
    uninhabited Refl impossible

public export
Uninhabited (a = b) => Uninhabited (SObject a = SObject b) where
    uninhabited Refl @{ab} = uninhabited @{ab} Refl


module Sec15b_ProcessList

import Part3.Sec15a_ProcessLib

data ListAction : Type where
     Length : List elem -> ListAction
     Append : List elem -> List elem -> ListAction

||| Note: this is not directly representable in Haskell
||| ListType can be represented as TypeFamily but TypeFamily is not
||| first class (cannot be used with partially applied type variables in type signatures).
ListType : ListAction -> Type
ListType (Length xs) = Nat
ListType (Append {elem} xs ys) = List elem

-- Definitions:
-- Service iface a = Process iface a Ready Looping
-- Client a = Process NoRecv a Ready Ready

total
procList : Service ListType ()
procList = do Respond (\msg => case msg of
                                    Length xs => Pure (length xs)
                                    Append xs ys => Pure (xs ++ ys))
              Loop procList

procMain : Client ()
procMain = do Just list <- Spawn procList
                      | Nothing => Action (putStrLn "Spawn failed")
              len <- Request list (Length [1,2,3])
              Action (printLn len)

              app <- Request list (Append [1,2,3] [4,5,6])
              Action (printLn app)


-- ------------------------------------------------------------ [ DeBruijn.idr ]
-- Module    : DeBruijn.idr
-- Copyright : (c) 2015,2016 See CONTRIBUTORS.md
-- License   : see LICENSE
-- --------------------------------------------------------------------- [ EOH ]
||| Data structure to compute de Bruijn indices.
|||
||| Thanks to christiansen's Galois tutorials for the accessor and
||| mutator functions for environments/object store.
module Data.DeBruijn

import Data.DList

||| A de Bruijn Index.
|||
||| @T    The type of type's collected.
||| @ctxt The collection of types.
||| @t    The element collected at the current position.
public export
data Index : (T : Type)
          -> (ctxt : List T)
          -> (t : T)
          -> Type where
  ||| The first element in the index.
  First : Index ty (t::ts) t

  ||| The next element in the index.
  Next  : Index ty ts t -> Index ty (t'::ts) t

||| Sometimes it is bettern to think that we have this thing called an
||| environment and not a `DList`.
|||
||| @t    The Type for Types in our environment.
||| @obj  How we interpret the types in our DSL. Either this is a
|||       dependent type or a function that computes a type.
||| @ctxt The typing context.
public export
Env : (t : Type) -> (obj : t -> Type) -> (ctxt : List t) -> Type
Env ty obj ctxt = DList ty (\x => obj x) ctxt

||| Read an object from our typing environment.
|||
||| @idx The typing context.
||| @env The typing environment.
export
read : (idx : Index ty ctxt t)
    -> (env : Env ty e ctxt)
    -> e t
read First    (obj::store) = obj
read (Next x) (obj::store) = read x store

||| Add an object to our typnig environment.
|||
||| @idx The typing context.
||| @obj The object to add.
||| @env The environment to which the object is added.
export
write : (idx : Index ty ctxt t )
     -> (obj : e t)
     -> (env : Env ty e ctxt)
     -> Env ty e ctxt
write First    obj (_    :: store) = obj  :: store
write (Next x) obj (obj' :: store) = obj' :: write x obj store

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


module LightClick.Types.Compatibility

import public Toolkit.Decidable.Equality.Indexed

import public Toolkit.Decidable.Informative

import Toolkit.Data.Rig
import Toolkit.Data.Vect.Extra

import LightClick.Types.Meta
import LightClick.Types.Direction
import LightClick.Types.Sensitivity
import LightClick.Types.WireType

import LightClick.Types
import LightClick.Types.Equality

%default total

namespace Data

  mutual
    namespace Fields
      public export
      data Compatible : (kxs : Vect n (Pair String (Ty DATA)))
                     -> (kys : Vect m (Pair String (Ty DATA)))
                                 -> Type
        where
          Empty : Compatible Nil Nil
          Cons : (compatLabels :                     i        =  j)
              -> (compatTypes  : Data.Compatible       a           b)
              -> (rest         : Fields.Compatible         xs          ys)
                              -> Fields.Compatible ((i,a)::xs) ((j,b)::ys)

    public export
    data Compatible : (x,y : Ty DATA) -> Type where
      CompatLogic : Data.Compatible TyLogic TyLogic
      CompatArray : (sameLength : i = j)
                 -> (types      : Data.Compatible a b)
                               -> Data.Compatible (TyArray a i)
                                                  (TyArray b j)
      CompatStruct : (compatFields : Fields.Compatible xs ys)
                                  -> Data.Compatible (TyStruct xs) (TyStruct ys)
      CompatUnion : (compatFields : Fields.Compatible xs ys)
                                 -> Data.Compatible (TyUnion xs) (TyUnion ys)

    export
    posFlip : Data.Compatible l r -> Data.Compatible r l
    posFlip (CompatArray Refl prf) = CompatArray Refl (posFlip prf)
    posFlip (CompatStruct prfs) = CompatStruct (posFlips prfs)
    posFlip (CompatUnion  prfs) = CompatUnion (posFlips prfs)
    posFlip CompatLogic = CompatLogic

    export
    negFlip : (Data.Compatible l r -> Void) -> Data.Compatible r l -> Void
    negFlip f (CompatArray Refl prf) = f (CompatArray Refl (posFlip prf))
    negFlip f (CompatStruct prfs) = f (CompatStruct (posFlips prfs))
    negFlip f (CompatUnion  prfs) = f (CompatUnion (posFlips prfs))
    negFlip f CompatLogic = f CompatLogic

    export
    posFlips : (Fields.Compatible ls rs) -> Fields.Compatible rs ls
    posFlips Empty = Empty
    posFlips (Cons Refl prf rest) = Cons Refl (posFlip prf) (posFlips rest)

namespace Logic
  export
  logicNotCompatWithArray : Data.Compatible TyLogic (TyArray l t) -> Void
  logicNotCompatWithArray (CompatArray _ _) impossible
  logicNotCompatWithArray (CompatStruct _) impossible
  logicNotCompatWithArray (CompatUnion _) impossible
  logicNotCompatWithArray CompatLogic impossible

  export
  logicNotCompatWithStruct : Data.Compatible TyLogic (TyStruct xs) -> Void
  logicNotCompatWithStruct (CompatArray _ _) impossible
  logicNotCompatWithStruct (CompatStruct _) impossible
  logicNotCompatWithStruct (CompatUnion _) impossible
  logicNotCompatWithStruct CompatLogic impossible

  export
  logicNotCompatWithUnion : Data.Compatible TyLogic (TyUnion xs) -> Void
  logicNotCompatWithUnion (CompatArray _ _) impossible
  logicNotCompatWithUnion (CompatStruct _) impossible
  logicNotCompatWithUnion (CompatUnion _) impossible
  logicNotCompatWithUnion CompatLogic impossible

namespace Array

  export
  arraysNotSameLength : (contra : (k = j) -> Void)
                     -> Data.Compatible (TyArray x k) (TyArray y j)
                     -> Void
  arraysNotSameLength contra (CompatArray prf x) = contra prf

  export
  arraysNotSameType : (f : Data.Compatible x y -> Void)
                   -> Data.Compatible (TyArray x j) (TyArray y j)
                   -> Void
  arraysNotSameType f (CompatArray Refl x) = f x

  export
  arrayNotCompatWithStruct : Data.Compatible (TyArray x k) (TyStruct xs) -> Void
  arrayNotCompatWithStruct (CompatArray _ _) impossible
  arrayNotCompatWithStruct (CompatStruct _) impossible
  arrayNotCompatWithStruct (CompatUnion _) impossible
  arrayNotCompatWithStruct CompatLogic impossible

  export
  arrayNotCompatWithUnion : Data.Compatible (TyArray x k) (TyUnion xs) -> Void
  arrayNotCompatWithUnion (CompatArray _ _) impossible
  arrayNotCompatWithUnion (CompatStruct _) impossible
  arrayNotCompatWithUnion (CompatUnion _ ) impossible
  arrayNotCompatWithUnion CompatLogic impossible

namespace Struct

  export
  structNotCompatWithUnion : Data.Compatible (TyStruct xs) (TyUnion ys) -> Void
  structNotCompatWithUnion (CompatArray _ _) impossible
  structNotCompatWithUnion (CompatStruct _) impossible
  structNotCompatWithUnion (CompatUnion _ ) impossible
  structNotCompatWithUnion CompatLogic impossible

  export
  structsAreNotCompat : (Compatible xs ys -> Void) -> Compatible (TyStruct xs) (TyStruct ys) -> Void
  structsAreNotCompat contra (CompatStruct prf) = contra prf

namespace Union

  export
  unionsAreNotCompat : (Compatible xs ys -> Void) -> Compatible (TyUnion xs) (TyUnion ys) -> Void
  unionsAreNotCompat contra (CompatUnion prf) = contra prf

namespace Data
  public export
  data Error : Type where
    Mismatch : Data.Error
    MismatchArrayLength : Data.Error
    MismatchArrayType   : (error : Data.Error) -> Data.Error
    MismatchStructureFieldType  : (position : Nat) -> (error : Data.Error) -> Data.Error
    MismatchStructureFieldLabel  : (position : Nat) -> (x,y : String) -> Data.Error
    MismatchStructureLength : Data.Error

-- [ Function Declaration ]

export
compatible : (x : Ty DATA)
          -> (y : Ty DATA)
                 -> DecInfo Data.Error
                            (Data.Compatible x y)

namespace Fields
  fieldMismatchLabel : (xl = yl -> Void)
                    -> Fields.Compatible ((xl, xt) :: xs) ((yl, yt) :: ys)
                    -> Void
  fieldMismatchLabel contra (Cons Refl prf rest) = contra Refl

  fieldMismatchType : (Compatible xt yt -> Void)
                   -> Compatible ((xl, xt) :: xs) ((xl, yt) :: ys)
                   -> Void
  fieldMismatchType contra (Cons Refl prf rest) = contra prf

  restNotCompatible : (Compatible xs ys -> Void)
                   -> Compatible ((xl, xt) :: xs) ((xl, yt) :: ys)
                   -> Void
  restNotCompatible contra (Cons Refl prf rest) = contra rest

  fieldsLeftIsLarger : Compatible (x :: xs) [] -> Void
  fieldsLeftIsLarger _ impossible

  fieldsRightIsLarger : Compatible [] (y :: ys) -> Void
  fieldsRightIsLarger _ impossible

  export
  compatible : (kxs : Vect n (Pair String (Ty DATA)))
            -> (kys : Vect m (Pair String (Ty DATA)))
                        -> DecInfo Compatibility.Data.Error
                                   (Fields.Compatible kxs kys)

  compatible {n} xs {m} ys with (shape xs ys)
    compatible {n = 0} [] {m = 0} [] | Empty = Yes Empty

    compatible {n = (S len)} (x :: xs) {m = 0} [] | LH = No MismatchStructureLength fieldsLeftIsLarger
    compatible {n = 0} [] {m = (S len)} (y :: ys) | RH = No MismatchStructureLength fieldsRightIsLarger

    compatible {n = (S len)} ((xl, xt) :: xs) {m = (S len)} ((yl, yt) :: ys) | Both with (decEq xl yl)
      compatible {n = (S len)} ((xl, xt) :: xs) {m = (S len)} ((xl, yt) :: ys) | Both | (Yes Refl) with (compatible xt yt)
        compatible {n = (S len)} ((xl, xt) :: xs) {m = (S len)} ((xl, yt) :: ys) | Both | (Yes Refl) | (Yes prf) with (compatible xs ys)
          compatible {n = (S len)} ((xl, xt) :: xs) {m = (S len)} ((xl, yt) :: ys) | Both | (Yes Refl) | (Yes prf) | (Yes rest)
            = Yes (Cons Refl prf rest)
            -- End of happy path

          -- Rest not happy somewhere, and need to update position when it is recorded.
          compatible {n = (S len)} ((xl, xt) :: xs) {m = (S len)} ((xl, yt) :: ys) | Both | (Yes Refl) | (Yes prf) | (No (MismatchStructureFieldType position error) contra)
            = No (MismatchStructureFieldType (S position) error) (restNotCompatible contra)
          compatible {n = (S len)} ((xl, xt) :: xs) {m = (S len)} ((xl, yt) :: ys) | Both | (Yes Refl) | (Yes prf) | (No (MismatchStructureFieldLabel position x y) contra)
            = No (MismatchStructureFieldLabel (S position) x y) (restNotCompatible contra)
          compatible {n = (S len)} ((xl, xt) :: xs) {m = (S len)} ((xl, yt) :: ys) | Both | (Yes Refl) | (Yes prf) | (No msg contra)
            = No msg (restNotCompatible contra)

        -- Type mismatch
        compatible {n = (S len)} ((xl, xt) :: xs) {m = (S len)} ((xl, yt) :: ys) | Both | (Yes Refl) | (No msg contra)
          = No (MismatchStructureFieldType Z msg) (fieldMismatchType contra)
      -- Label mismatch
      compatible {n = (S len)} ((xl, xt) :: xs) {m = (S len)} ((yl, yt) :: ys) | Both | (No contra)
        = No (MismatchStructureFieldLabel Z xl yl) (fieldMismatchLabel contra)

-- [ Definition Starts Here]

-- [ Logic vs X ]

compatible TyLogic TyLogic = Yes CompatLogic
compatible TyLogic (TyArray l t) = No Mismatch logicNotCompatWithArray
compatible TyLogic (TyStruct kvs) = No Mismatch logicNotCompatWithStruct
compatible TyLogic (TyUnion kvs) = No Mismatch logicNotCompatWithUnion

-- [ Array vs X ]

compatible (TyArray type length) TyLogic = No Mismatch (negFlip logicNotCompatWithArray)

compatible (TyArray typeA lengthA) (TyArray typeB lengthB) with (decEq lengthA lengthB)
  compatible (TyArray typeA lengthA) (TyArray typeB lengthA) | Yes Refl with (compatible typeA typeB)
    compatible (TyArray typeA lengthA) (TyArray typeB lengthA) | Yes Refl | Yes prf
      = Yes (CompatArray Refl prf)

    compatible (TyArray typeA lengthA) (TyArray typeB lengthA) | Yes Refl | No msg contra
      = No (MismatchArrayType msg) (arraysNotSameType contra)
  compatible (TyArray typeA lengthA) (TyArray typeB lengthB) | No contra
    = No MismatchArrayLength (arraysNotSameLength contra)

compatible (TyArray type length) (TyStruct kvs) = No Mismatch arrayNotCompatWithStruct
compatible (TyArray type length) (TyUnion kvs) = No Mismatch arrayNotCompatWithUnion

-- [ Struct vs X ]

compatible (TyStruct kvs) TyLogic
  = No Mismatch (negFlip logicNotCompatWithStruct)

compatible (TyStruct kvs) (TyArray type length)
  = No Mismatch (negFlip arrayNotCompatWithStruct)

compatible (TyStruct xs) (TyStruct ys) with (Fields.compatible xs ys)
  compatible (TyStruct xs) (TyStruct ys) | (Yes prf) = Yes (CompatStruct prf)
  compatible (TyStruct xs) (TyStruct ys) | (No msg contra) = No msg (structsAreNotCompat contra)

compatible (TyStruct kvs) (TyUnion xs) = No Mismatch structNotCompatWithUnion

-- [ Union vs X ]

compatible (TyUnion kvs) TyLogic = No Mismatch (negFlip logicNotCompatWithUnion)

compatible (TyUnion kvs) (TyArray type length) = No Mismatch (negFlip arrayNotCompatWithUnion)

compatible (TyUnion xs) (TyStruct ys) = No Mismatch (negFlip structNotCompatWithUnion)

compatible (TyUnion xs) (TyUnion ys) with (Fields.compatible xs ys)
  compatible (TyUnion xs) (TyUnion ys) | (Yes prf) = Yes (CompatUnion prf)
  compatible (TyUnion xs) (TyUnion ys) | (No msg contra) = No msg (unionsAreNotCompat contra)

-- [ EOF ]


module SqlDsl

import Data.Vect

--%access public export

data SqlType : Type where
  SqlInt : SqlType
  SqlBool : SqlType
  SqlString : SqlType

data SqlExpr : SqlType -> Type where
  IntC : Int -> SqlExpr SqlInt
  BoolC : Bool -> SqlExpr SqlBool
  StringC : String -> SqlExpr SqlString
  PlusC : SqlExpr SqlInt -> SqlExpr SqlInt -> SqlExpr SqlInt
  ConcatC : SqlExpr SqlString -> SqlExpr SqlString -> SqlExpr SqlString
  EqualC : SqlExpr a -> SqlExpr a -> SqlExpr SqlBool
  NotC : SqlExpr SqlBool -> SqlExpr SqlBool
  Field : {t : SqlType} -> (alias : String) -> (name : String) -> SqlExpr t

int : Int -> SqlExpr SqlInt
int x = IntC x

bool : Bool -> SqlExpr SqlBool
bool x = BoolC x

str : String -> SqlExpr SqlString
str x = StringC x

(==) : SqlExpr a -> SqlExpr a -> SqlExpr SqlBool
(==) x y = EqualC x y

data Tuple : {a : Type} -> (f : a -> Type) -> Vect n a -> Type where
  Nil : Tuple f []
  (::) : f t -> Tuple f ts -> Tuple f (t :: ts)

Table : {n : Nat} -> Type
Table {n} = (String, Vect n (String, SqlType))

index : (i : Fin k) -> Tuple f ts -> f (index i ts)
index FZ (x::xs) = x
index (FS j) (x::xs) = index j xs

get : (t : (String, SqlType)) -> {ts : Vect n (String, SqlType)} -> Tuple (\t => SqlExpr (snd t)) ts -> {auto p : Elem t ts} -> SqlExpr (snd t)
get t (x :: xs) {p = Here} = x
get t (x :: xs) {p = There p'} = get t {p = p'} xs

data SqlQuery : Vect n (String, SqlType) -> Type where
  From : (t : Table) -> SqlQuery (snd t)
  Product : SqlQuery ts -> SqlQuery ts' -> (Tuple (\t => SqlExpr (snd t)) ts -> Tuple (\t => SqlExpr (snd t)) ts' -> Tuple (\t => SqlExpr (snd t)) ts'') -> SqlQuery ts''
  Where : SqlQuery ts -> (Tuple (\t => SqlExpr (snd t)) ts -> SqlExpr SqlBool) -> SqlQuery ts
  Select : SqlQuery ts -> (Tuple (\t => SqlExpr (snd t)) ts -> Tuple (\t => SqlExpr (snd t)) ts') -> SqlQuery ts'

from : (t : Table) -> SqlQuery (snd t)
from t = From t

product : (Tuple (\t => SqlExpr (snd t)) ts -> Tuple (\t => SqlExpr (snd t)) ts' -> Tuple (\t => SqlExpr (snd t)) ts'') -> SqlQuery ts -> SqlQuery ts' -> SqlQuery ts''
product proj q1 q2 = Product q1 q2 proj

where' : (Tuple (\t => SqlExpr (snd t)) ts -> SqlExpr SqlBool) -> SqlQuery ts -> SqlQuery ts
where' pred query = Where query pred

select : (Tuple (\t => SqlExpr (snd t)) ts -> Tuple (\t => SqlExpr (snd t)) ts') -> SqlQuery ts -> SqlQuery ts'
select f query = Select query f

compileExpr : SqlExpr t -> String
compileExpr (IntC x) = show x
compileExpr (BoolC x) = show x
compileExpr (StringC x) = "\"" ++ x ++ "\""
compileExpr (PlusC x y) = "(" ++ compileExpr x ++ " + " ++ compileExpr y ++ ")"
compileExpr (ConcatC x y) = "(" ++ compileExpr x ++ " ++ " ++ compileExpr y ++ ")"
compileExpr (EqualC x y) = "(" ++ compileExpr x ++ " = " ++ compileExpr y ++ ")"
compileExpr (NotC x) = "NOT" ++ compileExpr x
compileExpr (Field alias name) = alias ++ "." ++ name

mapToTuple : String -> (ts : Vect n (String, SqlType)) -> Tuple (\t => SqlExpr (snd t)) ts
mapToTuple alias [] = []
mapToTuple alias ((name, t) :: ts) = Field {t = t} alias name :: mapToTuple alias ts

tupleToString : {ts : Vect n (String, SqlType)} -> Tuple (\t => SqlExpr (snd t)) ts -> String
tupleToString {ts = []} [] = ""
tupleToString {ts = [t]} [e] = compileExpr e ++ " AS " ++ fst t
tupleToString {ts = t :: ts} (e :: es) = compileExpr e ++ " AS " ++ fst t ++ ", " ++ tupleToString es

compile' : {ts : Vect n (String, SqlType)} -> SqlQuery ts -> Int -> (Int -> String -> String) -> String
compile' (Select {ts} (From (tableName, _)) f) i k =
  let alias = "c" ++ show i in
  let tuple' = f $ mapToTuple alias ts in
  let sql = "SELECT " ++ tupleToString tuple' ++ " FROM " ++ tableName ++ " AS " ++ alias in
  k (i + 1) sql
compile' (Select {ts} (Where (From (tableName, _)) pred) f) i k =
  let alias = "c" ++ show i in
  let tuple' = f $ mapToTuple alias ts in
  let sql = "SELECT " ++ tupleToString tuple' ++ " FROM " ++ tableName ++ " AS " ++ alias ++ " WHERE " ++ (compileExpr $ pred $ mapToTuple alias ts) in
  k (i + 1) sql
compile' (From (tableName, ts)) i k =
  let alias = "c" ++ show i in
  let sql = "SELECT * FROM " ++ tableName ++ " AS " ++ alias in
  k (i + 1) sql
compile' {ts} (Where query pred) i k =
  compile' query i (\i, sql =>
    let alias = "c" ++ show i in
    let sql' = "SELECT * FROM (" ++ sql ++ ") AS " ++ alias ++ " WHERE " ++ (compileExpr $ pred $ mapToTuple alias ts) in
    k (i + 1) sql')
compile' (Select {ts} query f) i k =
  compile' query i (\i, sql =>
    let alias = "c" ++ show i in
    let tuple' = f $ mapToTuple alias ts in
    let sql' = "SELECT " ++ tupleToString tuple' ++ " FROM (" ++ sql ++ ") AS " ++ alias in
    k (i + 1) sql')
compile' (Product {ts} {ts'} q1 q2 f) i k =
  compile' q1 i (\i, sql =>
    compile' q2 i (\i, sql' =>
      let alias = "c" ++ show i in
      let i' = i + 1 in
      let alias' = "c" ++ show i' in
      let tuple' = f (mapToTuple alias ts) (mapToTuple alias' ts') in
      let sql'' = "SELECT " ++ tupleToString tuple' ++ " FROM (" ++ sql ++ ") AS " ++ alias ++ ", (" ++ sql' ++ ") AS " ++ alias' in
      k (i' + 1) sql''))

compile : SqlQuery ts -> String
compile query = compile' query 0 (\_, sql => sql)

customer : Table {n = 3}
customer = ("Customer", [("Id", SqlInt),
                         ("Name", SqlString),
                         ("Age", SqlInt)])

--
example0 : SqlQuery [("Id", SqlInt), ("Name", SqlString), ("Age", SqlInt)]
example0 = from customer

example1 : SqlQuery [("Name", SqlString)]
example1 = select (\ta  => [get ("Name", SqlString) {ts = snd customer} ta]) $ from customer

example2 : SqlQuery [("Id", SqlInt)]
example2 = select (\ta => [index 0 ta]) $ from customer

example3 : SqlQuery [("Id", SqlInt), ("Age", SqlInt)]
example3 = select (\ta => [index 0 ta, index 2 ta]) $ where' (\ta => index 1 ta == str "Nick") $ from customer


module Main

-- base
import Data.List
import Data.Strings
import System.File

-- contrib
import Data.SortedMap
import Text.Lexer
import Text.Lexer.Core
import Text.Parser
import Text.Parser.Core
import Text.Token

----------------------------------------------------------------------------------------------------

namespace Rules

  data Terminal = A
                | B

  Eq Terminal where
    A == A = True
    B == B = True
    _ == _ = False

  Show Terminal where
    show A = "A"
    show B = "B"

  data Ruleset : Type where
    Match : Terminal -> Ruleset
    References : (references : List Nat ** NonEmpty references) -> Ruleset
    Option : Ruleset -> Ruleset -> Ruleset

  Show Ruleset where
    show (Match terminal) = "Match " ++ show terminal
    show (References (ns**_)) = "References " ++ show ns
    show (Option x y) = "Option " ++ show x ++ " " ++ show y

  record Rule where
    constructor MkRule
    number : Nat
    ruleset : Ruleset

  [EqNumber] Eq Rule where
    (MkRule number _) == (MkRule number' _) = number == number'

  [OrdNumber] Ord Rule using EqNumber where
    compare (MkRule number _) (MkRule number' _) = compare number number'

  Show Rule where
    show (MkRule number ruleset) = "Rule " ++ show number ++ " " ++ show ruleset

----------------------------------------------------------------------------------------------------

namespace RulesParsing

  data RulesKind = RKNumber
                 | RKCharA
                 | RKCharB
                 | RKColon
                 | RKSpace
                 | RKPipe
                 | RKQuotationMarks

  Eq RulesKind where
    RKNumber == RKNumber = True
    RKCharA == RKCharA = True
    RKCharB == RKCharB = True
    RKColon == RKColon = True
    RKSpace == RKSpace = True
    RKPipe == RKPipe = True
    RKQuotationMarks == RKQuotationMarks = True
    _ == _ = False

  TokenKind RulesKind where
    TokType RKNumber = Nat
    TokType _ = ()
    tokValue RKNumber x = integerToNat $ cast x
    tokValue RKCharA _ = ()
    tokValue RKCharB _ = ()
    tokValue RKColon _ = ()
    tokValue RKSpace _ = ()
    tokValue RKPipe _ = ()
    tokValue RKQuotationMarks _ = ()

  tokenMap : TokenMap $ Token RulesKind
  tokenMap = toTokenMap
    [ (digits, RKNumber)
    , (is 'a', RKCharA)
    , (is 'b', RKCharB)
    , (is ':', RKColon)
    , (spaces, RKSpace)
    , (is '|', RKPipe)
    , (is '"', RKQuotationMarks)
    ]

  grammarRuleset : Grammar (Token RulesKind) True Ruleset
  grammarRuleset = grammarOption <|> grammarReferences <|> grammarMatch where
    grammarMatch : Grammar (Token RulesKind) True Ruleset
    grammarMatch = do
      match RKQuotationMarks
      terminal <- (match RKCharA *> pure A) <|> (match RKCharB *> pure B)
      match RKQuotationMarks
      pure $ Match terminal
    grammarReferences : Grammar (Token RulesKind) True Ruleset
    grammarReferences = References <$> sepBy1' (match RKSpace) (match RKNumber)
    grammarOption : Grammar (Token RulesKind) True Ruleset
    grammarOption = do
      left <- (grammarMatch <|> grammarReferences)
      match RKSpace
      match RKPipe
      match RKSpace
      right <- (grammarMatch <|> grammarReferences)
      pure $ Option left right

  grammarRule : Grammar (Token RulesKind) True Rule
  grammarRule = do
    number <- match RKNumber
    match RKColon
    match RKSpace
    ruleset <- grammarRuleset
    eof
    pure $ MkRule number ruleset

  export
  readRule : String -> Maybe Rule
  readRule str = do
    let (tokens, debugLexed) = lex tokenMap str
        parsed = parse grammarRule $ tok <$> tokens
    case parsed of
         Left _ => Nothing
         Right (rule, _) => Just rule

----------------------------------------------------------------------------------------------------

namespace ParserCombination

  ruleMap : List Rule -> SortedMap Nat Ruleset
  ruleMap rules = fromList $ extract <$> rules where
    extract : Rule -> (Nat, Ruleset)
    extract (MkRule number ruleset) = (number, ruleset)

  public export
  TestGrammar : Type
  TestGrammar = Grammar Terminal True ()

  TestGrammarMap : Type
  TestGrammarMap = SortedMap Nat TestGrammar

  lookupTestGrammar : TestGrammarMap -> Nat -> Maybe TestGrammar
  lookupTestGrammar = flip lookup

  mapOnDPair : (f : a -> b) -> (xs : List a ** NonEmpty xs) -> (ys : List b ** NonEmpty ys)
  mapOnDPair f ([x] ** IsNonEmpty) = ([f x] ** IsNonEmpty)
  mapOnDPair f ((x :: x' :: xs) ** IsNonEmpty) with (mapOnDPair f ((x' :: xs) ** IsNonEmpty))
    mapOnDPair f ((x :: x' :: xs) ** IsNonEmpty) | ((y :: ys) ** IsNonEmpty) = ((f x :: y :: ys) ** IsNonEmpty)

  lookupTestGrammars' : TestGrammarMap -> (references : List Nat ** NonEmpty references) ->
                                          (grammars : List (Maybe TestGrammar) ** NonEmpty grammars)
  lookupTestGrammars' map refs = mapOnDPair (lookupTestGrammar map) refs

  takeMaybeOutProof : (xs : List (Maybe a) ** NonEmpty xs) -> Maybe (ys : List a ** NonEmpty ys)
  takeMaybeOutProof ([ Nothing ] ** IsNonEmpty) = Nothing
  takeMaybeOutProof ([ Just x ] ** IsNonEmpty) = Just ([ x ] ** IsNonEmpty)
  takeMaybeOutProof ((Nothing :: x' :: xs) ** IsNonEmpty) = Nothing
  takeMaybeOutProof ((Just x :: x' :: xs) ** IsNonEmpty) with (takeMaybeOutProof ((x' :: xs) ** IsNonEmpty))
    takeMaybeOutProof ((Just x :: x' :: xs) ** IsNonEmpty) | Nothing = Nothing
    takeMaybeOutProof ((Just x :: x' :: xs) ** IsNonEmpty) | (Just dpair) = Just $ f x (g dpair) where
      f : a -> List a -> (bs : List a ** NonEmpty bs)
      f q qs = ((q :: qs) ** IsNonEmpty)
      g : (x : List a ** NonEmpty x) -> List a
      g (x ** y) = x

  lookupTestGrammars : TestGrammarMap -> (references : List Nat ** NonEmpty references) ->
                                         Maybe (grammars : List TestGrammar ** NonEmpty grammars)
  lookupTestGrammars map refs = takeMaybeOutProof $ lookupTestGrammars' map refs

  combineTestGrammars : TestGrammar -> TestGrammar -> TestGrammar
  combineTestGrammars x y = x *> y

  grammar : (ruleset : Ruleset) -> (grammars : TestGrammarMap) -> Maybe TestGrammar
  grammar (Match char) grammars =
    Just $ terminal "Failed to parse Terminal" test where
      test : Terminal -> Maybe ()
      test x = if x == char
                  then Just ()
                  else Nothing
  grammar (References references) grammars = do
    let testGrammarsStuff = lookupTestGrammars grammars references
    case testGrammarsStuff of
         Nothing => Nothing
         Just (testGrammars ** nonEmptyTestGrammarProof) =>
           Just $ foldl1 combineTestGrammars testGrammars
  grammar (Option leftRuleset rightRuleset) grammars =
    [| grammar leftRuleset grammars <|> grammar rightRuleset grammars |]

  grammarMap' : (rules : List Rule) -> TestGrammarMap -> TestGrammarMap
  grammarMap' [] acc = acc
  grammarMap' (rule@(MkRule number ruleset) :: rules) acc with (grammar ruleset acc)
    grammarMap' (rule@(MkRule number ruleset) :: rules) acc | Nothing = grammarMap' (rules ++ [rule]) acc
    grammarMap' (rule@(MkRule number ruleset) :: rules) acc | Just g = grammarMap' rules $ insert number g acc

  grammarMap : (rules : List Rule) -> TestGrammarMap
  grammarMap rules = grammarMap' rules empty

  export
  grammar0 : (rules : List Rule) -> Maybe TestGrammar
  grammar0 rules = lookup 0 $ grammarMap rules

----------------------------------------------------------------------------------------------------

  newGrammar8 : TestGrammarMap -> Maybe TestGrammar
  newGrammar8 gmap = do
    grammar42 <- lookupTestGrammar gmap 42
    let g : TestGrammar
        g = do
          grammar42
          optional g
          pure ()
    Just $ g

  newGrammar11 : TestGrammarMap -> Maybe TestGrammar
  newGrammar11 gmap = do
    grammar42 <- lookupTestGrammar gmap 42
    grammar31 <- lookupTestGrammar gmap 31
    let g : TestGrammar
        g = do
          grammar42
          optional g
          grammar31
          pure ()
    Just g

  contains8Or11 : List Nat -> Bool
  contains8Or11 [] = False
  contains8Or11 (8 :: _) = True
  contains8Or11 (11 :: _) = True
  contains8Or11 (_ :: ns) = contains8Or11 ns

  fromDepPair : (xs : List a ** NonEmpty xs) -> List a
  fromDepPair (xs ** _) = xs

  contains8Or11' : Ruleset -> Bool
  contains8Or11' (References refs) = contains8Or11 $ fromDepPair refs
  contains8Or11' _ = False

  grammar' : (ruleset : Rule) -> (grammars : TestGrammarMap) -> Maybe TestGrammar
  grammar' (MkRule 8 _) grammars = newGrammar8 grammars
  grammar' (MkRule 11 _) grammars = newGrammar11 grammars
  grammar' (MkRule _ (Match char)) grammars =
    Just $ terminal "Failed to parse Terminal" test where
      test : Terminal -> Maybe ()
      test x = if x == char
                  then Just ()
                  else Nothing
  grammar' (MkRule _ (References references)) grammars = do
    let testGrammarsStuff = lookupTestGrammars grammars references
    case testGrammarsStuff of
         Nothing => Nothing
         Just (testGrammars ** nonEmptyTestGrammarProof) =>
           Just $ foldl1 combineTestGrammars testGrammars
  grammar' (MkRule _ (Option leftRuleset rightRuleset)) grammars =
    if contains8Or11' leftRuleset
       then [| grammar' (MkRule 0 rightRuleset) grammars <|> grammar' (MkRule 0 leftRuleset) grammars |]
       else [| grammar' (MkRule 0 leftRuleset) grammars <|> grammar' (MkRule 0 rightRuleset) grammars |]

  grammar'Map' : (rules : List Rule) -> TestGrammarMap -> TestGrammarMap
  grammar'Map' [] acc = acc
  grammar'Map' (rule@(MkRule number ruleset) :: rules) acc with (grammar' rule acc)
    grammar'Map' (rule@(MkRule number ruleset) :: rules) acc | Nothing = grammar'Map' (rules ++ [rule]) acc
    grammar'Map' (rule@(MkRule number ruleset) :: rules) acc | Just g = grammar'Map' rules $ insert number g acc

  grammar'Map : (rules : List Rule) -> TestGrammarMap
  grammar'Map rules = grammar'Map' rules empty

  export
  grammar0' : (rules : List Rule) -> Maybe TestGrammar
  grammar0' rules = lookup 0 $ grammar'Map rules

----------------------------------------------------------------------------------------------------

namespace MessagesParsing

  toTerminal : Char -> Maybe Terminal
  toTerminal 'a' = Just A
  toTerminal 'b' = Just B
  toTerminal _ = Nothing

  export
  toTerminals : String -> Maybe $ List Terminal
  toTerminals message = traverse toTerminal $ unpack message

  export
  testMessage : TestGrammar -> List Terminal -> Bool
  testMessage grammar word =
    case parse (grammar <* eof) word of
         Left _ => False
         Right _ => True

----------------------------------------------------------------------------------------------------

main : IO ()
main = do
  rulesData <- readFile "./rules"
  messagesData <- readFile "./messages"
  case (lines <$> rulesData, lines <$> messagesData) of
       (Right rulesData, Right messagesData) => do
         let mbRules = sort @{OrdNumber} <$> sequence (readRule <$> rulesData)
             mbMessages = traverse toTerminals messagesData
         case (mbRules, mbMessages) of
              (Just rules, Just messages) => do
                --case grammar0 rules of
                --     Nothing => putStrLn "Can't construct grammar0"
                --     Just grammar => do
                --       let tests = testMessage grammar <$> messages
                --           countParsed = length $ filter id tests
                --       print countParsed
                case grammar0' rules of
                     Nothing => putStrLn "Can't construct grammar0'"
                     Just grammar => do
                       let tests = testMessage grammar <$> messages
                           countParsed = length $ filter id tests
                       print countParsed
              _ => pure ()
         pure ()
       _ => pure ()
  pure ()


module NoCo

nats : Stream Nat
nats = Z :: map S nats

record Foo : Type where
  MkFoo : (bar : Bool) -> (baz : if bar then Nat else Bool) -> Foo


module Crypto.AES.Big

import Utils.ConstantTable
import Crypto.AES.Common
import Data.Bits
import Utils.Bytes
import Data.List
import Data.Fin
import Data.Fin.Extra
import Data.Vect
import Utils.Misc

ssm0 : ConstantTable 256 Bits32
ssm0 = from_vect
  [ 0xC66363A5, 0xF87C7C84, 0xEE777799, 0xF67B7B8D, 0xFFF2F20D, 0xD66B6BBD
  , 0xDE6F6FB1, 0x91C5C554, 0x60303050, 0x02010103, 0xCE6767A9, 0x562B2B7D
  , 0xE7FEFE19, 0xB5D7D762, 0x4DABABE6, 0xEC76769A, 0x8FCACA45, 0x1F82829D
  , 0x89C9C940, 0xFA7D7D87, 0xEFFAFA15, 0xB25959EB, 0x8E4747C9, 0xFBF0F00B
  , 0x41ADADEC, 0xB3D4D467, 0x5FA2A2FD, 0x45AFAFEA, 0x239C9CBF, 0x53A4A4F7
  , 0xE4727296, 0x9BC0C05B, 0x75B7B7C2, 0xE1FDFD1C, 0x3D9393AE, 0x4C26266A
  , 0x6C36365A, 0x7E3F3F41, 0xF5F7F702, 0x83CCCC4F, 0x6834345C, 0x51A5A5F4
  , 0xD1E5E534, 0xF9F1F108, 0xE2717193, 0xABD8D873, 0x62313153, 0x2A15153F
  , 0x0804040C, 0x95C7C752, 0x46232365, 0x9DC3C35E, 0x30181828, 0x379696A1
  , 0x0A05050F, 0x2F9A9AB5, 0x0E070709, 0x24121236, 0x1B80809B, 0xDFE2E23D
  , 0xCDEBEB26, 0x4E272769, 0x7FB2B2CD, 0xEA75759F, 0x1209091B, 0x1D83839E
  , 0x582C2C74, 0x341A1A2E, 0x361B1B2D, 0xDC6E6EB2, 0xB45A5AEE, 0x5BA0A0FB
  , 0xA45252F6, 0x763B3B4D, 0xB7D6D661, 0x7DB3B3CE, 0x5229297B, 0xDDE3E33E
  , 0x5E2F2F71, 0x13848497, 0xA65353F5, 0xB9D1D168, 0x00000000, 0xC1EDED2C
  , 0x40202060, 0xE3FCFC1F, 0x79B1B1C8, 0xB65B5BED, 0xD46A6ABE, 0x8DCBCB46
  , 0x67BEBED9, 0x7239394B, 0x944A4ADE, 0x984C4CD4, 0xB05858E8, 0x85CFCF4A
  , 0xBBD0D06B, 0xC5EFEF2A, 0x4FAAAAE5, 0xEDFBFB16, 0x864343C5, 0x9A4D4DD7
  , 0x66333355, 0x11858594, 0x8A4545CF, 0xE9F9F910, 0x04020206, 0xFE7F7F81
  , 0xA05050F0, 0x783C3C44, 0x259F9FBA, 0x4BA8A8E3, 0xA25151F3, 0x5DA3A3FE
  , 0x804040C0, 0x058F8F8A, 0x3F9292AD, 0x219D9DBC, 0x70383848, 0xF1F5F504
  , 0x63BCBCDF, 0x77B6B6C1, 0xAFDADA75, 0x42212163, 0x20101030, 0xE5FFFF1A
  , 0xFDF3F30E, 0xBFD2D26D, 0x81CDCD4C, 0x180C0C14, 0x26131335, 0xC3ECEC2F
  , 0xBE5F5FE1, 0x359797A2, 0x884444CC, 0x2E171739, 0x93C4C457, 0x55A7A7F2
  , 0xFC7E7E82, 0x7A3D3D47, 0xC86464AC, 0xBA5D5DE7, 0x3219192B, 0xE6737395
  , 0xC06060A0, 0x19818198, 0x9E4F4FD1, 0xA3DCDC7F, 0x44222266, 0x542A2A7E
  , 0x3B9090AB, 0x0B888883, 0x8C4646CA, 0xC7EEEE29, 0x6BB8B8D3, 0x2814143C
  , 0xA7DEDE79, 0xBC5E5EE2, 0x160B0B1D, 0xADDBDB76, 0xDBE0E03B, 0x64323256
  , 0x743A3A4E, 0x140A0A1E, 0x924949DB, 0x0C06060A, 0x4824246C, 0xB85C5CE4
  , 0x9FC2C25D, 0xBDD3D36E, 0x43ACACEF, 0xC46262A6, 0x399191A8, 0x319595A4
  , 0xD3E4E437, 0xF279798B, 0xD5E7E732, 0x8BC8C843, 0x6E373759, 0xDA6D6DB7
  , 0x018D8D8C, 0xB1D5D564, 0x9C4E4ED2, 0x49A9A9E0, 0xD86C6CB4, 0xAC5656FA
  , 0xF3F4F407, 0xCFEAEA25, 0xCA6565AF, 0xF47A7A8E, 0x47AEAEE9, 0x10080818
  , 0x6FBABAD5, 0xF0787888, 0x4A25256F, 0x5C2E2E72, 0x381C1C24, 0x57A6A6F1
  , 0x73B4B4C7, 0x97C6C651, 0xCBE8E823, 0xA1DDDD7C, 0xE874749C, 0x3E1F1F21
  , 0x964B4BDD, 0x61BDBDDC, 0x0D8B8B86, 0x0F8A8A85, 0xE0707090, 0x7C3E3E42
  , 0x71B5B5C4, 0xCC6666AA, 0x904848D8, 0x06030305, 0xF7F6F601, 0x1C0E0E12
  , 0xC26161A3, 0x6A35355F, 0xAE5757F9, 0x69B9B9D0, 0x17868691, 0x99C1C158
  , 0x3A1D1D27, 0x279E9EB9, 0xD9E1E138, 0xEBF8F813, 0x2B9898B3, 0x22111133
  , 0xD26969BB, 0xA9D9D970, 0x078E8E89, 0x339494A7, 0x2D9B9BB6, 0x3C1E1E22
  , 0x15878792, 0xC9E9E920, 0x87CECE49, 0xAA5555FF, 0x50282878, 0xA5DFDF7A
  , 0x038C8C8F, 0x59A1A1F8, 0x09898980, 0x1A0D0D17, 0x65BFBFDA, 0xD7E6E631
  , 0x844242C6, 0xD06868B8, 0x824141C3, 0x299999B0, 0x5A2D2D77, 0x1E0F0F11
  , 0x7BB0B0CB, 0xA85454FC, 0x6DBBBBD6, 0x2C16163A ]

lookup_ssm0 : Bits32 -> Bits32
lookup_ssm0 x = index_bits8 (cast x) ssm0

lookup_sbox : Bits32 -> Bits32 -> Bits32 -> Bits32 -> Bits32
lookup_sbox s0 s1 s2 s3 =
  let a = cast $ index_bits8 (cast $ shiftR s0 24) sbox
      b = cast $ index_bits8 (cast $ shiftR s1 16) sbox
      c = cast $ index_bits8 (cast $ shiftR s2 8) sbox
      d = cast $ index_bits8 (cast s3) sbox
  in (shiftL a 24) .|. (shiftL b 16) .|. (shiftL c 8) .|. d

sbox_ext : Fin 32 -> Bits32 -> Bits32
sbox_ext i = rotr i . lookup_ssm0

expand_key_stream : {n_k : _} -> {auto 0 ok : NonZero n_k} -> Vect n_k (Vect 4 Bits8) -> Stream Bits32
expand_key_stream key = map from_be $ prepend (toList key) $ expand_key key

encrypt_block' : Nat -> Stream Bits32 -> Vect 4 Bits32 -> Vect 16 Bits8

encrypt_block' Z skey [s0, s1, s2, s3] =
  concat $ map (to_be {n=4}) $ vecxor (Stream.take 4 skey)
  [ lookup_sbox s0 s1 s2 s3
  , lookup_sbox s1 s2 s3 s0
  , lookup_sbox s2 s3 s0 s1
  , lookup_sbox s3 s0 s1 s2 ]

encrypt_block' (S n) skey [s0, s1, s2, s3] =
  let (s, skey) = splitAt 4 skey
  in encrypt_block' n skey $ vecxor s
     [ (sbox_ext 0 $ shiftR s0 24) `xor` (sbox_ext 8 $ shiftR s1 16) `xor` (sbox_ext 16 $ shiftR s2 8) `xor` (sbox_ext 24 s3)
     , (sbox_ext 0 $ shiftR s1 24) `xor` (sbox_ext 8 $ shiftR s2 16) `xor` (sbox_ext 16 $ shiftR s3 8) `xor` (sbox_ext 24 s0)
     , (sbox_ext 0 $ shiftR s2 24) `xor` (sbox_ext 8 $ shiftR s3 16) `xor` (sbox_ext 16 $ shiftR s0 8) `xor` (sbox_ext 24 s1)
     , (sbox_ext 0 $ shiftR s3 24) `xor` (sbox_ext 8 $ shiftR s0 16) `xor` (sbox_ext 16 $ shiftR s1 8) `xor` (sbox_ext 24 s2)
     ]

export
encrypt_block : (mode : Mode) -> (key : Vect ((get_n_k mode) * 4) Bits8) -> Vect 16 Bits8 -> Vect 16 Bits8
encrypt_block mode key plaintext =
  let buf = map (from_be {n=4,a=Bits32}) $ group 4 4 plaintext
      (s, skey) = splitAt 4 (expand_key_stream {ok = n_k_never_Z mode} $ group _ _ key)
  in encrypt_block' (get_main_rounds mode) skey $ vecxor buf s


-- -------------------------------------------------------------- [ Lens.idr ]
-- Description : Idris port of Control.Lens
-- Copyright   : (c) Huw Campbell
-- --------------------------------------------------------------------- [ EOH ]
module Control.Lens.Types
import Control.Lens.Const
import Data.Contravariant
import Control.Lens.First
import Control.Monad.Identity
import Data.Bifunctor
import Data.Profunctor
import Data.Tagged

%default total

public export
Optic : (Type -> Type -> Type) -> (Type -> Type) -> Type -> Type -> Type -> Type -> Type
Optic p f s t a b = p a (f b) -> p s (f t)

public export
Simple : (Type -> Type -> Type -> Type -> Type) -> Type -> Type -> Type
Simple p s a = p s s a a

public export
Optic' : (Type -> Type -> Type) -> (Type -> Type) -> Type -> Type -> Type
Optic' p f = Simple (Optic p f)

public export
LensLike : (Type -> Type) -> Type -> Type -> Type -> Type -> Type
LensLike = Optic Morphism

public export
LensLike' : (Type -> Type) -> Type -> Type -> Type
LensLike' f = Simple (LensLike f)

-- type Lens s t a b = forall f. Functor f => (a -> f b) -> s -> f t
public export
Lens : Type -> Type -> Type -> Type -> Type
Lens s t a b  = { f : Type -> Type } -> Functor f => LensLike f s t a b

-- type Lens' s a = Lens s s a a
public export
Lens' : Type -> Type -> Type
Lens' = Simple Lens

-- type Getting r s a = (a -> Const r a) -> s -> Const r s
public export
Getting : Type -> Type -> Type -> Type
Getting r = LensLike' (Const r)

-- type Getter s a = forall f. (Contravariant f, Functor f) => (a -> f a) -> s -> f s
public export
Getter : Type -> Type -> Type
Getter s a = { f : Type -> Type } -> (Contravariant f, Functor f) => LensLike' f s a

-- type ASetter s t a b = (a -> Identity b) -> s -> Identity t
public export
Setter : Type -> Type -> Type -> Type -> Type
Setter = LensLike Identity

public export
Setter' : Type -> Type -> Type
Setter' = Simple Setter

-- type Traversal s t a b = forall f. Applicative f => (a -> f b) -> s -> f t
public export
Traversal : Type -> Type -> Type -> Type -> Type
Traversal s t a b  = { f : Type -> Type } -> Applicative f => LensLike f s t a b

-- type Traversal' s a = Traversal s s a a
public export
Traversal' : Type -> Type -> Type
Traversal' = Simple Traversal

-- type Prism s t a b = forall p f. (Choice p, Applicative f) => p a (f b) -> p s (f t)
-- The choice of `Choice p` for a valid `Traversal` is the `(->)` implementation. Which is used
-- when using a prism as a getter or setter.

public export
Prism : Type -> Type -> Type -> Type -> Type
Prism s t a b  = {p : Type -> Type -> Type} -> { f : Type -> Type } ->
                 (Choice p,Applicative f) => Optic p f s t a b

public export
Prism' : Type -> Type -> Type
Prism' = Simple Prism

public export
Iso : Type -> Type -> Type -> Type -> Type
Iso s t a b = {p : Type -> Type -> Type} -> {f : Type -> Type} ->
              (Profunctor p,Functor f) => Optic p f s t a b

public export
Iso' : Type -> Type -> Type
Iso' = Simple Iso

-- type Review t b = forall p f. (Choice p, Bifunctor p, Settable f) => Optic' p f t b
public export
Review : Type -> Type -> Type -> Type -> Type
Review s t a b = { p : Type -> Type -> Type } -> { f : Type -> Type } ->
                 (Choice p, Bifunctor p) => Optic p Identity s t a b

-- type AReview t b = Optic' Tagged Identity t b
public export
AReview : Type -> Type -> Type
AReview t b = Optic' Tagged Identity t b

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