module Typing.Equate

import Typing.Common
import Typing.Util

-- 8YYYY _YYY_ 0   0  o8o  YY8YY Y8Y _YYY_ 8o  0
-- 8___  0 _ 8 0   0 8   8   0    0  0   0 8Yo 8
-- 8"""  0 "%8 0   0 8YYY8   0    0  0   0 8 Yo8
-- 8oooo "oooo "ooo" 0   0   0   o8o "ooo" 0   8

infix 4 ~~
public export
data Equation = (~~) Term Term

-- _YYY_  oYYo  oYYo 0   0 8YYYo oYYYo
-- 0   0 0   " 0   " 0   0 8___P %___
-- 0   0 0   , 0   , 0   0 8""Yo `"""p
-- "ooo"  YooY  YooY "ooo" 0   0 YoooY

export
occurs : MetaVar -> Term -> Bool
occurs x (Meta y)         = x == y
occurs _ (Var _)          = False
occurs _ (DottedVar _ _)  = False
occurs _ (AbsoluteDottedVar _ _)
                          = False
occurs x (Ann m t)        = occurs x m || occurs x t
occurs _ Star             = False
occurs x (Pi _ a sc)      = occurs x a || occurs x (descope (Var . Name) sc)
occurs x (Lam _ sc)       = occurs x (descope (Var . Name) sc)
occurs x (App _ f a)      = occurs x f || occurs x a
occurs x (Con _ as)       = any (occurs x . snd) as
occurs x (Case as mot cs) = any (occurs x) as || occursCaseMotive mot || any occursClause cs
where
	occursCaseMotive : CaseMotive -> Bool
	occursClause : Clause -> Bool
	occursPattern : Pattern -> Bool

	occursCaseMotive (CaseMotiveNil m) = occurs x m
	occursCaseMotive (CaseMotiveCons a sc)
		= occurs x a || occursCaseMotive (descope (Var . Name) sc)

	occursClause (NewClause psc sc) = any occursPattern (descope Name psc) || occurs x (descope (Var . Name) sc)

	occursPattern (VarPat _) = False
	occursPattern (ConPat _ xs) = any (occursPattern . snd) xs
	occursPattern (AssertionPat m) = occurs x m
	occursPattern MakeMeta = False
occurs x (RecordType tele)  = occursTelescope tele
where
	occursTelescope : Telescope -> Bool
	occursTelescope TelescopeNil = False
	occursTelescope (TelescopeCons t sc) = occurs x t || occursTelescope (descope (Var . Name) sc)
occurs x (RecordCon fields) = any (occurs x . snd) fields
occurs x (Project m _)    = occurs x m

export
evalWithSubs : Substitution -> Equation -> TypeChecker Equation
evalWithSubs subs (l ~~ r) = do
	l' <- evaluate (instantiateMetas subs l)
	r' <- evaluate (instantiateMetas subs r)
	pure (l' ~~ r')

--  oYYo _YYY_ 8_   _8 8YYYo  o8o  8YYYo 8YYYY  oYYo  o8o  oYYYo 8YYYY 8_   _8 _YYY_ YY8YY Y8Y 0   0 8YYYY
-- 0   " 0   0 8"o_o"8 8___P 8   8 8___P 8___  0   " 8   8 %___  8___  8"o_o"8 0   0   0    0  0   0 8___
-- 0   , 0   0 0  8  0 8"""  8YYY8 8""Yo 8"""  0   , 8YYY8 `"""p 8"""  0  8  0 0   0   0    0  "o o" 8"""
--  YooY "ooo" 0     0 0     0   0 0   0 8oooo  YooY 0   0 YoooY 8oooo 0     0 "ooo"   0   o8o  "8"  8oooo

export
compareCaseMotive : CaseMotive -> CaseMotive -> TypeChecker (List Equation)
compareCaseMotive (CaseMotiveNil a1) (CaseMotiveNil a2) = pure [(a1 ~~ a2)]
compareCaseMotive (CaseMotiveCons a1 sc1) (CaseMotiveCons a2 sc2) = do
	i <- newName {s=TCState}
	case (head' (names sc1), head' (names sc2)) of
		(Just h1, Just h2) => do
			reqs <- compareCaseMotive
				(instantiate sc1 [Var (Generated h1 i)])
				(instantiate sc2 [Var (Generated h2 i)])
			pure ((a1 ~~ a2) :: reqs)
		_ => throw "Should not happen"
compareCaseMotive mot mot' = throw $ "Motives not equal: " ++ show mot ++ " and " ++ show mot'

--  oYYo _YYY_ 8_   _8 8YYYo  o8o  8YYYo 8YYYY     oYYo 0     o8o  0   0 oYYYo 8YYYY
-- 0   " 0   0 8"o_o"8 8___P 8   8 8___P 8___     0   " 0    8   8 0   0 %___  8___
-- 0   , 0   0 0  8  0 8"""  8YYY8 8""Yo 8"""     0   , 0    8YYY8 0   0 `"""p 8"""
--  YooY "ooo" 0     0 0     0   0 0   0 8oooo     YooY 8ooo 0   0 "ooo" YoooY 8oooo

export
compareClauses : List Clause -> List Clause -> TypeChecker (List Equation)
export
comparePattern : Pattern -> Pattern -> TypeChecker (List Equation)
export
comparePatterns : List (Plict,Pattern) -> List (Plict,Pattern) -> TypeChecker (List Equation)

compareClauses [] [] = pure []
compareClauses (NewClause psc1 sc1::cs1) (NewClause psc2 sc2::cs2) = do
	unless (length (names psc1) == length (names psc2))
		$ throw "Patterns bind different numbers of arguments."
	is <- replicateM (length (names psc1)) $ newName
	let xs1 = zipWith Generated (names psc1) is
	let xs1' = map Var (removeByDummies (names psc1) xs1)
	let xs2 = zipWith Generated (names psc2) is
	let xs2' = map Var (removeByDummies (names psc2) xs2)
	reqss <- zipWithM comparePattern (instantiate psc1 xs1) (instantiate psc2 xs2)
	reqs' <- compareClauses cs1 cs2
	pure ((instantiate sc1 xs1' ~~ instantiate sc2 xs2') :: concat reqss ++ reqs')
compareClauses _ _ = throw $ "Mismatching number of clauses."

comparePattern (VarPat x) (VarPat x') = do
	unless (x == x')
		$ throw $ "Variable patters not equal: " ++ show x ++ " and " ++ show x'
	pure []
comparePattern (ConPat c ps) (ConPat c' ps') with (c == c')
	| True  = comparePatterns ps ps'
	| False = throw $ "Mismatching constructors " ++ show c ++ " and " ++ show c'
comparePattern _ _ = throw "Patterns not equal."

comparePatterns [] [] = pure []
comparePatterns ((plic1,p1)::ps1) ((plic2,p2)::ps2) = do
	unless (plic1 == plic2)
		$ throw "Mismatched plicities when trying to unify constructor sequences."
	eqs <- comparePattern p1 p2
	eqs' <- comparePatterns ps1 ps2
	pure $ eqs ++ eqs'
comparePatterns _ _ = throw "Patterns not equal."

--  oYYo _YYY_ 8_   _8 8YYYo  o8o  8YYYo 8YYYY YY8YY 8YYYY 0    8YYYY oYYYo  oYYo _YYY_ 8YYYo 8YYYY
-- 0   " 0   0 8"o_o"8 8___P 8   8 8___P 8___    0   8___  0    8___  %___  0   " 0   0 8___P 8___
-- 0   , 0   0 0  8  0 8"""  8YYY8 8""Yo 8"""    0   8"""  0    8"""  `"""p 0   , 0   0 8"""  8"""
--  YooY "ooo" 0     0 0     0   0 0   0 8oooo   0   8oooo 8ooo 8oooo YoooY  YooY "ooo" 0     8oooo

export
compareTelescope : Telescope -> Telescope -> TypeChecker (List Equation)
compareTelescope TelescopeNil TelescopeNil = pure []
compareTelescope (TelescopeCons t1 sc1) (TelescopeCons t2 sc2) = do
	unless (names sc1 == names sc2)
		$ throw $ "Mismatching record field names: " ++ unwords (names sc1) ++ " and " ++ unwords (names sc2)
	eqs' <- compareTelescope (descope (Var . Name) sc1) (descope (Var . Name) sc2)
	pure ((t1 ~~ t2) :: eqs')
compareTelescope _ _ = throw "Mismatched number of record fields."

--  oYYo _YYY_ 8_   _8 8YYYo  o8o  8YYYo 8YYYY 8YYYY Y8Y 8YYYY 0    8888_ oYYYo
-- 0   " 0   0 8"o_o"8 8___P 8   8 8___P 8___  8___   0  8___  0    0   0 %___
-- 0   , 0   0 0  8  0 8"""  8YYY8 8""Yo 8"""  8"""   0  8"""  0    0   0 `"""p
--  YooY "ooo" 0     0 0     0   0 0   0 8oooo 0     o8o 8oooo 8ooo 8oooY YoooY

export
compareFields : List (String,Term) -> List (String,Term) -> TypeChecker (List Equation)
compareFields [] [] = pure []
compareFields (_ :: _) [] = throw "Mismatching number of record fields."
compareFields [] (_ :: _) = throw "Mismatching number of record fields."
compareFields ((x,m)::xms) xms' = case extract (\(x',_) => x == x') xms' of
	Nothing      => throw "Mismatching record field names."
	Just ((_,m'),xms'') => do
		eqs' <- compareFields xms xms''
		pure ((m ~~ m') :: eqs')

-- 8YYYY _YYY_ 0   0  o8o  YY8YY 8YYYY
-- 8___  0 _ 8 0   0 8   8   0   8___
-- 8"""  0 "%8 0   0 8YYY8   0   8"""
-- 8oooo "oooo "ooo" 0   0   0   8oooo

export
equate : List Equation -> Substitution -> TypeChecker ()
equate eqs0 subs = go eqs0
where
	go : List Equation -> TypeChecker ()

	go [] = pure ()

	go ((Meta x ~~ Meta y) :: eqs) with (x == y)
		| True = go eqs
		| False = do
			eqs' <- traverse (evalWithSubs subs) eqs
			go eqs'

	go ((Meta x ~~ t2) :: eqs) = do
		unless (not (occurs x t2))
			$ throw $ "Cannot unify because " ++ show (Meta x)
				++ " occurs in " ++ show t2
		eqs' <- traverse (evalWithSubs subs) eqs
		go eqs'

	go ((t1 ~~ Meta y) :: eqs) = do
		unless (not (occurs y t1))
			$ throw $ "Cannot unify because " ++ show (Meta y)
				++ " occurs in " ++ show t1
		eqs' <- traverse (evalWithSubs subs) eqs
		go eqs'

	go ((Var x ~~ Var y) :: eqs) = do
		unless (x == y)
			$ throw $ "Cannot unify variables " ++ show (Var x)
				++ " and " ++ show (Var y)
		go eqs

	go ((DottedVar m1 x1 ~~ DottedVar m2 x2) :: eqs) = do
		unless (m1 == m2 && x1 == x2)
			$ throw $ "Cannot unify symbols " ++ show (DottedVar m1 x1)
				++ " and " ++ show (DottedVar m2 x2)
		go eqs

	go ((AbsoluteDottedVar m1 x1 ~~ AbsoluteDottedVar m2 x2) :: eqs) = do
		unless (m1 == m2 && x1 == x2)
			$ throw $ "Cannot unify symbols " ++ show (AbsoluteDottedVar m1 x1)
				++ " and " ++ show (AbsoluteDottedVar m2 x2)
		go eqs

	go ((Ann m1 t1 ~~ Ann m2 t2) :: eqs) = go ((m1 ~~ m2) :: (t1 ~~ t2) :: eqs)

	go ((Star ~~ Star) :: eqs) = go eqs

	go ((Pi plic1 a1 sc1 ~~ Pi plic2 a2 sc2) :: eqs) = do
		unless (plic1 == plic2)
			$ throw $ "Mismatched plicities when trying to unify "
						++ show (Pi plic1 a1 sc1) ++ " with "
						++ show (Pi plic2 a2 sc2)
		i <- newName
		case (head' (names sc1), head' (names sc2)) of
			(Just h1, Just h2) => do
				let b1 = instantiate sc1 [Var (Generated h1 i)]
				let b2 = instantiate sc2 [Var (Generated h2 i)]
				go ((a1 ~~ a2) :: (b1 ~~ b2) :: eqs)
			_ => throw "should not happen"

	go ((Lam plic1 sc1 ~~ Lam plic2 sc2) :: eqs) = do
		unless (plic1 == plic2)
			$ throw $ "Mismatched plicities when trying to unify "
						++ show (Lam plic1 sc1) ++ " with "
						++ show (Lam plic2 sc2)
		i <- newName
		case (head' (names sc1), head' (names sc2)) of
			(Just h1, Just h2) => do
				let b1 = instantiate sc1 [Var (Generated h1 i)]
				let b2 = instantiate sc2 [Var (Generated h2 i)]
				go ((b1 ~~ b2) :: eqs)
			_ => throw "should not happen"

	go ((App plic1 f1 a1 ~~ App plic2 f2 a2) :: eqs) = do
		unless (plic1 == plic2)
			$ throw $ "Mismatched plicities when trying to unify "
						++ show (App plic1 f1 a1) ++ " with "
						++ show (App plic2 f2 a2)
		go ((f1 ~~ f2) :: (a1 ~~ a2) :: eqs)

	go ((Con c1 as1 ~~ Con c2 as2) :: eqs) = do
		unless (c1 == c2) $ throw $ "Mismatching constructors " ++ show c1 ++ " and " ++ show c2
		unless (length as1 == length as2) $ throw $ "Mismatching number of arguments in "
			++ show (Con c1 as1) ++ " and "
			++ show (Con c2 as2)
		eqs' <- zipConArgs as1 as2
		go (eqs' ++ eqs)
	where
		zipConArgs : List (Plict, Term) -> List (Plict, Term) -> TypeChecker (List Equation)
		zipConArgs [] [] = pure []
		zipConArgs ((plic1',a1')::as1') ((plic2',a2')::as2') = if plic1' == plic2'
				then do
					eqs' <- zipConArgs as1' as2'
					pure $ (a1' ~~ a2') :: eqs'
				else
					throw "Mismatching plicity."
		zipConArgs _ _ = throw "Mismatching number of arguments."

	go ((Case as1 mot1 cs1 ~~ Case as2 mot2 cs2) :: eqs) = do
		unless(length as1 == length as2)
			$ throw $ "Mismatching number of case arguments in "
						++ show (Case as1 mot1 cs1) ++ " and "
						++ show (Case as2 mot2 cs2)
		unless (length cs1 == length cs2)
			$ throw $ "Mismatching number of clauses in "
						++ show (Case as1 mot1 cs1) ++ " and "
						++ show (Case as2 mot2 cs2)
		let argEqs = zipWith (~~) as1 as2
		motEqs <- compareCaseMotive mot1 mot2
		clauseEqs <- compareClauses cs1 cs2
		go (argEqs ++ motEqs ++ clauseEqs ++ eqs)

	go ((RecordType tele1 ~~ RecordType tele2) :: eqs) = do
		eqs' <- compareTelescope tele1 tele2
		go (eqs' ++ eqs)

	go ((RecordCon fields1 ~~ RecordCon fields2) :: eqs) = do
		eqs' <- compareFields fields1 fields2
		go (eqs' ++ eqs)

	go ((Project m1 x1 ~~ Project m2 x2) :: eqs) = do
		unless (x1 == x2) $ throw $ "Field names not equal: " ++ x1 ++ " and " ++ x2
		go ((m1 ~~ m2) :: eqs)

	go ((m ~~ m') :: _) = throw $ "Terms not equal: " ++ show m ++ " and " ++ show m'


module ProgramSynthSMT

import Data.Vect
import Data.String
import src.SMTLib
import src.Tensor

record VarPos (size : Nat) where
  constructor MkVar
  pos : Expr (BitVecT size)
  val : Expr (BitVecT size)

record Arg (size : Nat) where
  constructor MkArg
  isVar : Expr BoolT
  varPos : Expr (BitVecT size)
  instrPos : Expr (BitVecT size)
  val : Expr (BitVecT size)

record ArgName (size : Nat) where
  constructor MkArgName
  isVar : String
  varPos : String
  instrPos : String

record Instr (argsN : Nat) (size : Nat) where
  constructor MkInstr
  pos : Expr (BitVecT size)
  isConst : Expr BoolT
  constVal : Expr (BitVecT size)
  val : Expr (BitVecT size)
  op : Expr (BitVecT size)
  args : Vect argsN (Arg size)

record InstrName (argsN : Nat) (size : Nat) where
    constructor MkInstrName
    pos : Nat
    isConst : String
    constVal : String
    op : String
    args : Vect argsN (ArgName size)

record Op (argsN : Nat) (size : Nat) where
  constructor MkOp
  id : Integer
  func : Expr (BitVecT size) -> Vect argsN (Expr (BitVecT size)) -> Expr BoolT
  str : Vect argsN String -> String

createVarPos : Vect n (Expr (BitVecT size)) -> Vect n (VarPos size)
createVarPos {n} {size} vs = tabulate (\i => MkVar (bv (finToInteger i) size) (index i vs))

lookupVar : Expr (BitVecT size) -> Expr (BitVecT size) -> Vect n (VarPos size) -> Expr BoolT
lookupVar p v vs = or [ite' (p == pos vp) (v == val vp) (bool False) | vp <- vs]

createArgs : Vect n (Vect argsN (Expr BoolT)) -> Vect n (Vect argsN (Expr (BitVecT size))) -> Vect n (Vect argsN (Expr (BitVecT size))) ->  Vect n (Vect argsN (Expr (BitVecT size))) -> Vect n (Vect argsN (Arg size))
createArgs isVars vars instrs vals = tabulate (\i0 => tabulate (\i1 => MkArg (index i0 i1 isVars) (index i0 i1 vars) (index i0 i1 instrs) (index i0 i1 vals)))

createInstrs : Vect n (Expr BoolT) -> Vect n (Expr (BitVecT size)) -> Vect n (Expr (BitVecT size)) -> Vect n (Expr (BitVecT size)) -> Vect n (Vect argsN (Arg size)) -> Vect n (Instr argsN size)
createInstrs {size} isConsts constVals vals ops args = tabulate (\i => MkInstr (bv (finToInteger i) size) (index i isConsts) (index i constVals) (index i vals) (index i ops) (index i args))

lookupInstr : Expr (BitVecT size) -> Expr (BitVecT size) -> Vect n (Instr argsN size) -> Expr BoolT
lookupInstr p v instrs = or [ite' (p == pos instr) (v == val instr) (bool False) | instr <- instrs]

varNames : String -> Tensor [n] String
varNames name = toTensor $ tabulate (\i => name ++ "_" ++ show (finToNat i))

varNames' : String -> Tensor [n, m] String
varNames' name = toTensor $ tabulate (\i0 => tabulate (\i1 => name ++ "_" ++ show (finToNat i0) ++ "_" ++ show (finToNat i1)))

lookupOp : Expr (BitVecT size) -> Expr (BitVecT size) -> Vect argsN (Expr (BitVecT size)) -> Vect opsN (Op argsN size) -> Expr BoolT
lookupOp {size} opPos instrVal args ops = or [ite' (opPos == (bv (id op) size)) (func op instrVal args) (bool False) | op <- ops]

checkInstr : Vect n (Instr argsN size) -> Expr BoolT
checkInstr [] = bool True
checkInstr [instr] = and $ [constVal instr `bvule` bv 65000 size] ++ [isVar arg == bool True | arg <- args instr]
checkInstr (instr :: instrs) = and $ [constVal instr `bvule` bv 65000 size] ++ [or [instrPos arg == pos instr' | instr' <- instrs] | arg <- args instr] ++ [checkInstr instrs]

evalInstr : Vect opsN (Op argsN size) -> Vect varsN (VarPos size) -> Vect instrsN (Instr argsN size) -> Expr BoolT
evalInstr {argsN} ops vars instrs = and [ite' (isConst instr) (val instr == constVal instr) (and [assocArg vars instrs instr, lookupOp (op instr) (val instr) (map val $ args instr) ops]) | instr <- instrs]
  where
    assocArg : Vect varsN (VarPos size) -> Vect instrsN (Instr argsN size) -> Instr argsN size -> Expr BoolT
    assocArg vars instrs instr = and $ tabulate (\i => let arg = index i (args instr) in ite' (isVar arg) (lookupVar (varPos arg) (val arg) vars)
                                                                                (lookupInstr (instrPos arg) (val arg) instrs))

VarsN : Nat
VarsN = 1

ArgsN : Nat
ArgsN = 3

OpsN : Nat
OpsN = 17

BitSize : Nat
BitSize = 32

ops : Vect OpsN (Op ArgsN BitSize)
ops = let ops = the (Vect OpsN (Op ArgsN BitSize))
                [MkOp 0 (\r, arg => r == bvand (index 0 arg) (index 1 arg)) (\args => (index 0 args) ++ " & " ++ (index 1 args)),
                 MkOp 0 (\r, arg => r == bvor (index 0 arg) (index 1 arg)) (\args => (index 0 args) ++ " | " ++ (index 1 args)),
                 MkOp 0 (\r, arg => r == bvxor (index 0 arg) (index 1 arg)) (\args => (index 0 args) ++ " ^ " ++ (index 1 args)),
                 MkOp 0 (\r, arg => r == bvnot (index 0 arg)) (\args => "~" ++ (index 0 args)),
                 MkOp 0 (\r, arg => ite' ((index 0 arg) == (index 1 arg)) (r == (bv 1 BitSize)) (r == (bv 0 BitSize))) (\args => (index 0 args) ++ " == " ++ (index 1 args) ++ " ? (uint)1 : 0"),
                 MkOp 0 (\r, arg => r == bvsub (index 0 arg) (index 1 arg)) (\args => (index 0 args) ++ " - " ++ (index 1 args)),
                 MkOp 0 (\r, arg => r == bvadd (index 0 arg) (index 1 arg)) (\args => (index 0 args) ++ " + " ++ (index 1 args)),
                 MkOp 0 (\r, arg => r == bvmul (index 0 arg) (index 1 arg)) (\args => (index 0 args) ++ " * " ++ (index 1 args)),
                 MkOp 0 (\r, arg => ite' ((index 0 arg) == bv 1 BitSize) (r == (index 1 arg)) (r == (index 2 arg))) (\args => (index 0 args) ++ " == 1 ? " ++ (index 1 args) ++ " : " ++ (index 2 args)),
                 MkOp 0 (\r, arg => ite' ((index 1 arg) == bv 0 BitSize) (r == bv 1 BitSize) (r == bvudiv (index 0 arg) (index 1 arg))) (\args => (index 1 args) ++ " == 0 ? (uint)1 : " ++ (index 0 args) ++ " / " ++ (index 1 args)),
                 --MkOp 0 (\r, arg => ite' ((index 1 arg) == bv 0 BitSize) (r == index 0 arg) (r == bvurem (index 0 arg) (index 1 arg))) (\args => (index 1 args) ++ " == 0 ? (uint)1 : " ++ (index 0 args) ++ " % " ++ (index 1 args)),
                 MkOp 0 (\r, arg => r == bvshl (index 0 arg) (index 1 arg)) (\args => (index 0 args) ++ " << " ++ (index 1 args)),
                 MkOp 0 (\r, arg => r == bvlshr (index 0 arg) (index 1 arg)) (\args => (index 0 args) ++ " >> " ++ (index 1 args)),
                 MkOp 0 (\r, arg => r == bvneg (index 0 arg)) (\args => "-" ++ (index 0 args)),
                 MkOp 0 (\r, arg => ite' (bvule (index 0 arg) (index 1 arg)) (r == (bv 1 BitSize)) (r == (bv 0 BitSize)))  (\args => (index 0 args) ++ " <= " ++ (index 1 args) ++ " ? (uint)1 : 0"),
                 MkOp 0 (\r, arg => ite' (bvult (index 0 arg) (index 1 arg)) (r == (bv 1 BitSize)) (r == (bv 0 BitSize))) (\args => (index 0 args) ++ " < " ++ (index 1 args) ++ " ? (uint)1 : 0"),
                 MkOp 0 (\r, arg => ite' (bvuge (index 0 arg) (index 1 arg)) (r == (bv 1 BitSize)) (r == (bv 0 BitSize))) (\args => (index 0 args) ++ " >= " ++ (index 1 args) ++ " ? (uint)1 : 0"),
                 MkOp 0 (\r, arg => ite' (bvugt (index 0 arg) (index 1 arg)) (r == (bv 1 BitSize)) (r == (bv 0 BitSize))) (\args => (index 0 args) ++ " > " ++ (index 1 args) ++ " ? (uint)1 : 0")] in
                 --MkOp 0 (\r, arg => ite' (bvsle (index 0 arg) (index 1 arg)) (r == (bv 1 BitSize)) (r == (bv 0 BitSize))) (\args => (index 0 args) ++ " <= " ++ (index 1 args) ++ " ? 1 : 0"),
                 -- MkOp 0 (\r, arg => ite' (bvslt (index 0 arg) (index 1 arg)) (r == (bv 1 BitSize)) (r == (bv 0 BitSize))) (\args => (index 0 args) ++ " < " ++ (index 1 args) ++ " ? 1 : 0"),
                 -- MkOp 0 (\r, arg => ite' (bvsge (index 0 arg) (index 1 arg)) (r == (bv 1 BitSize)) (r == (bv 0 BitSize))) (\args => (index 0 args) ++ " >= " ++ (index 1 args) ++ " ? 1 : 0"),
                 -- MkOp 0 (\r, arg => ite' (bvsgt (index 0 arg) (index 1 arg)) (r == (bv 1 BitSize)) (r == (bv 0 BitSize))) (\args => (index 0 args) ++ " > " ++ (index 1 args) ++ " ? 1 : 0")] in
      zipWith (\op, i =>  record { id = cast i } op) ops (fromList [0..OpsN - 1])

varPosVal : Int -> Tensor [VarsN] String
varPosVal n = varNames $ "varPosVal" ++ show n

argIsVar : Tensor [ArgsN, instrsN] String
argIsVar = varNames' "argIsVar"

argVarPos : Tensor [ArgsN, instrsN] String
argVarPos = varNames' "argVarPos"

argInstrPos : Tensor [ArgsN, instrsN] String
argInstrPos = varNames' "argInstrPos"

argVal : Int -> Tensor [ArgsN, instrsN] String
argVal n = varNames' $ "argVal" ++ show n

instrVal : Int -> Tensor [instrsN] String
instrVal n = varNames $ "instrVal" ++ show n

instrOp : Tensor [instrsN] String
instrOp = varNames "instrOp"

instrIsConst : Tensor [instrsN] String
instrIsConst = varNames "isConst"

instrConstVal : Tensor [instrsN] String
instrConstVal = varNames "constVal"

instrNames : Vect instrsN (InstrName ArgsN BitSize)
instrNames = tabulate (\i => MkInstrName (finToNat i)
                                         (index i (toVect instrIsConst))
                                         (index i (toVect instrConstVal))
                                         (index i (toVect instrOp))
                                         (tabulate (\j => MkArgName (index i j (toVect argIsVar))
                                                                    (index i j (toVect argVarPos))
                                                                    (index i j (toVect argInstrPos)))))

eval : {prf : GTE instrsN 1} -> Int -> List (Expr (BitVecT BitSize), Expr (BitVecT BitSize) -> Expr BoolT) ->
       Tensor [ArgsN, instrsN] (Expr BoolT) -> Tensor [ArgsN, instrsN] (Expr (BitVecT BitSize)) ->
       Tensor [ArgsN, instrsN] (Expr (BitVecT BitSize)) -> Tensor [instrsN] (Expr (BitVecT BitSize)) ->
       Tensor [instrsN] (Expr BoolT) -> Tensor [instrsN] (Expr (BitVecT BitSize)) -> Smt ()
eval {instrsN} {prf} n [] argIsVar argVarPos argInstrPos instrOp instrIsConst instrConstVal = pure ()
eval {instrsN = S instrsN'} {prf = (LTESucc LTEZero)} n ((inp, outf) :: xs) argIsVar argVarPos argInstrPos instrOp instrIsConst instrConstVal =
    do vars <- declareVars (varPosVal n) (BitVecT BitSize)
       argVal <- declareVars (argVal {instrsN = S instrsN'} n) (BitVecT BitSize)
       instrVal <- declareVars (instrVal {instrsN = S instrsN'} n) (BitVecT BitSize)
       let vps = createVarPos (toVect vars)
       let args = createArgs (toVect argIsVar) (toVect argVarPos) (toVect argInstrPos) (toVect argVal)
       let instrs = createInstrs (toVect instrIsConst) (toVect instrConstVal) (toVect instrVal) (toVect instrOp) args
       assert $ evalInstr ops vps instrs
       assert $ val (index 0 vps) == inp
       r <- declareVar ("output" ++ show n) (BitVecT BitSize)
       assert $ outf r
       assert $ val (index 0 instrs) == r
       eval {prf = (LTESucc LTEZero)} (n + 1) xs argIsVar argVarPos argInstrPos instrOp instrIsConst instrConstVal
       pure ()

solver : {instrsN : Nat} -> {prf : GTE instrsN 1} -> List (Expr (BitVecT BitSize), Expr (BitVecT BitSize) -> Expr BoolT) -> Smt ()
solver {instrsN} {prf} xs =
         do setLogic "QF_BV"
            setOption ":timeout 600000"
            setOption ":parallel.enable true"
            setOption ":model_validate true"
            setOption ":pp.bv-literals false"
            let n = 0
            argIsVar <- declareVars (argIsVar {instrsN = instrsN}) BoolT
            argVarPos <- declareVars (argVarPos {instrsN = instrsN}) (BitVecT BitSize)
            argInstrPos <- declareVars (argInstrPos {instrsN = instrsN}) (BitVecT BitSize)
            instrOp <- declareVars (instrOp {instrsN = instrsN}) (BitVecT BitSize)
            instrIsConst <- declareVars (instrIsConst {instrsN = instrsN}) BoolT
            instrConstVal <- declareVars (instrConstVal {instrsN = instrsN}) (BitVecT BitSize)
            vars <- declareVars (varPosVal n) (BitVecT BitSize)
            argVal <- declareVars (argVal {instrsN = instrsN} n) (BitVecT BitSize)
            instrVal <- declareVars (instrVal {instrsN = instrsN} n) (BitVecT BitSize)
            let vps = createVarPos (toVect vars)
            let args = createArgs (toVect argIsVar) (toVect argVarPos) (toVect argInstrPos) (toVect argVal)
            let instrs = createInstrs (toVect instrIsConst) (toVect instrConstVal) (toVect instrVal) (toVect instrOp) args
            assert $ checkInstr instrs
            eval {prf = prf} (n + 1) xs argIsVar argVarPos argInstrPos instrOp instrIsConst instrConstVal
            checkSat
            getModel
            end

parseBool : String -> Maybe Bool
parseBool "true" = Just True
parseBool "false" = Just False
parseBool _ = Nothing

parseArgs : Model -> Vect n (ArgName size) -> Maybe (Vect n String)
parseArgs {n = Z} model args = Just []
parseArgs {n = (S n)} model ((MkArgName isVar varPos instrPos) :: args) =
  do isVarVal <- lookup isVar model
     isVar <- parseBool isVarVal
     varPosVal <- lookup varPos model
     instrPosVal <- lookup instrPos model
     args <- parseArgs model args
     pure $ [if isVar then "var" ++ varPosVal else "instr" ++ instrPosVal] ++ args

parseInstrs : Model -> Vect opsN (Op argsN size) -> Vect n (InstrName argsN size) -> Maybe (List String)
parseInstrs model ops [] = Just []
parseInstrs model ops ((MkInstrName pos isConst constVal opStr args) :: instrs) =
    do isConst <- lookup isConst model
       isConst <- parseBool isConst
       constVal <- lookup constVal model
       case isConst of
         True =>
          do args <- parseArgs model args
             instrs <- parseInstrs model ops instrs
             pure $ ["uint instr" ++ show pos ++ " = " ++ constVal ++ ";"] ++ instrs
         False =>
           do opVal <- lookup opStr model
              opId <- parseInteger {a = Integer} opVal
              op <- find (\op => id op == opId) ops
              args <- parseArgs model args
              instrs <- parseInstrs model ops instrs
              pure $ ["uint instr" ++ show pos ++ " = " ++ (str op args) ++ ";"] ++ instrs

equiv : Vect n (Expr a) -> Expr a -> Expr a -> (Expr a -> Vect n (Expr a) -> Expr BoolT) -> (Expr a -> Vect n (Expr a) -> Expr BoolT) -> Expr BoolT
equiv xs r r' f g = and [f r xs, g r' xs, not $ r == r']

testEquiv : Smt ()
testEquiv = let size = 64 in
            do setLogic "QF_BV"
               setOption ":pp.bv-literals false"
               setOption ":model_validate true"
               xs <- declareVars {f = Vect 3} ["a", "b", "c"] (BitVecT size)
               r <- declareVar "r" (BitVecT size)
               r' <- declareVar "_r" (BitVecT size)
               assert $ equiv xs r r' (\r, xs => r == ((((index 0 xs) `bvmul` (index 0 xs)) `bvmul` (index 0 xs)) `bvadd` (((index 1 xs) `bvmul` (index 1 xs)) `bvmul` (index 1 xs))))
                                      (\r, xs => r == (((index 2 xs) `bvmul` (index 2 xs)) `bvmul` (index 2 xs)))
               checkSat
               getModel
               end

runEquiv : IO ()
runEquiv = do r <- sat testEquiv
              case r of
                Nothing => do putStrLn "Error parsing result"
                Just (Sat, model) =>
                   do
                      putStrLn "sat"
                      putStrLn $ show model
                      pure ()
                Just (UnSat, _) => putStrLn "unsat"
                Just (Unknown, _) => putStrLn "unknown"

data' : List (Expr (BitVecT BitSize), Expr (BitVecT BitSize) -> Expr BoolT)
data' = let inp = example in
        map (\(a, b) => (bv a BitSize, \r => r == bv b BitSize)) inp
  where
    isPowerOfTwo : List (Integer, Integer)
    isPowerOfTwo = [(2, 1), (3, 0), (4, 1), (5, 0), (6, 0), (8, 1)]
    isEven : List (Integer, Integer)
    isEven = [(0, 1), (1, 0), (2, 1), (3, 0), (4, 1), (5, 0), (6, 1), (7, 0)]
    primes : List (Integer, Integer)
    primes = [(2, 3), (3, 5), (5, 7), (7, 11), (11, 13), (17, 19), (19, 23)]
    example : List (Integer, Integer)
    example = [(0, 2646464646)]

runSolver : (instrsN : Nat) -> {prf : GTE instrsN 1} -> IO ()
runSolver instrsN {prf} = do r <- sat $ solver {prf = prf} data'
                             --putStrLn $ "model: " ++ show r
                             case r of
                               Nothing => do putStrLn "Error parsing result"
                               Just (Sat, model) =>
                                  case parseInstrs {n = instrsN} model ops instrNames of
                                    Just instrs =>
                                      do putStrLn $ unlines $ reverse $ instrs
                                         pure ()
                                    Nothing => do putStrLn "Error parsing instrs"
                               Just (UnSat, _) => putStrLn "unsat"
                               Just (Unknown, _) => putStrLn "unknown"

run : (n : Nat) -> (instrsN : Nat) -> {auto prf : GTE instrsN 1} -> IO ()
run Z instrsN {prf = prf} = pure ()
run (S n) instrsN {prf = prf} = do putStrLn $ "n instrs: " ++ show instrsN
                                   runSolver {prf = prf} instrsN
                                   run n (S instrsN)


module Libraries.Text.Parser

import Data.Bool
import Data.List
import Data.Nat
import public Data.List1

import public Libraries.Text.Parser.Core
import public Libraries.Text.Quantity
import public Libraries.Text.Token

%default total

||| Parse a terminal based on a kind of token.
export
match : (Eq k, TokenKind k) =>
        (kind : k) ->
        Grammar state (Token k) True (TokType kind)
match k = terminal "Unrecognised input" $
    \t => if t.kind == k
             then Just $ tokValue k t.text
             else Nothing

||| Optionally parse a thing, with a default value if the grammar doesn't
||| match. May match the empty input.
export
option : {c : Bool} ->
         (def : a) -> (p : Grammar state tok c a) ->
         Grammar state tok False a
option {c = False} def p = p <|> pure def
option {c = True} def p = p <|> pure def

||| Optionally parse a thing.
||| To provide a default value, use `option` instead.
export
optional : {c : Bool} ->
           (p : Grammar state tok c a) ->
           Grammar state tok False (Maybe a)
optional p = option Nothing (map Just p)

||| Try to parse one thing or the other, producing a value that indicates
||| which option succeeded. If both would succeed, the left option
||| takes priority.
export
choose : {c1, c2 : Bool} ->
         (l : Grammar state tok c1 a) ->
         (r : Grammar state tok c2 b) ->
         Grammar state tok (c1 && c2) (Either a b)
choose l r = map Left l <|> map Right r

||| Produce a grammar by applying a function to each element of a container,
||| then try each resulting grammar until the first one succeeds. Fails if the
||| container is empty.
export
choiceMap : {c : Bool} ->
            (a -> Grammar state tok c b) ->
            Foldable t => t a ->
            Grammar state tok c b
choiceMap {c} f xs = foldr (\x, acc => rewrite sym (andSameNeutral c) in
                                               f x <|> acc)
                           (fail "No more options") xs

%hide Prelude.(>>=)

||| Try each grammar in a container until the first one succeeds.
||| Fails if the container is empty.
export
choice : Foldable t =>
         {c : Bool} ->
         t (Grammar state tok c a) ->
         Grammar state tok c a
choice = choiceMap id

mutual
  ||| Parse one or more things
  export
  some : Grammar state tok True a ->
         Grammar state tok True (List1 a)
  some p = pure (!p ::: !(many p))

  ||| Parse zero or more things (may match the empty input)
  export
  many : Grammar state tok True a ->
         Grammar state tok False (List a)
  many p = option [] (forget <$> some p)

mutual
  private
  count1 : (q : Quantity) ->
           (p : Grammar state tok True a) ->
           Grammar state tok True (List a)
  count1 q p = do x <- p
                  seq (count q p)
                      (\xs => pure (x :: xs))

  ||| Parse `p`, repeated as specified by `q`, returning the list of values.
  export
  count : (q : Quantity) ->
          (p : Grammar state tok True a) ->
          Grammar state tok (isSucc (min q)) (List a)
  count (Qty Z Nothing) p = many p
  count (Qty Z (Just Z)) _ = pure []
  count (Qty Z (Just (S max))) p = option [] $ count1 (atMost max) p
  count (Qty (S min) Nothing) p = count1 (atLeast min) p
  count (Qty (S min) (Just Z)) _ = fail "Quantity out of order"
  count (Qty (S min) (Just (S max))) p = count1 (between (S min) max) p

mutual
  ||| Parse one or more instances of `p` until `end` succeeds, returning the
  ||| list of values from `p`. Guaranteed to consume input.
  export
  someTill : {c : Bool} ->
             (end : Grammar state tok c e) ->
             (p : Grammar state tok True a) ->
             Grammar state tok True (List1 a)
  someTill {c} end p = do x <- p
                          seq (manyTill end p)
                              (\xs => pure (x ::: xs))

  ||| Parse zero or more instances of `p` until `end` succeeds, returning the
  ||| list of values from `p`. Guaranteed to consume input if `end` consumes.
  export
  manyTill : {c : Bool} ->
             (end : Grammar state tok c e) ->
             (p : Grammar state tok True a) ->
             Grammar state tok c (List a)
  manyTill {c} end p = rewrite sym (andTrueNeutral c) in
                               map (const []) end <|> (forget <$> someTill end p)

mutual
  ||| Parse one or more instance of `skip` until `p` is encountered,
  ||| returning its value.
  export
  afterSome : {c : Bool} ->
              (skip : Grammar state tok True s) ->
              (p : Grammar state tok c a) ->
              Grammar state tok True a
  afterSome skip p = do ignore $ skip
                        afterMany skip p

  ||| Parse zero or more instance of `skip` until `p` is encountered,
  ||| returning its value.
  export
  afterMany : {c : Bool} ->
              (skip : Grammar state tok True s) ->
              (p : Grammar state tok c a) ->
              Grammar state tok c a
  afterMany {c} skip p = rewrite sym (andTrueNeutral c) in
                                 p <|> afterSome skip p

||| Parse one or more things, each separated by another thing.
export
sepBy1 : {c : Bool} ->
         (sep : Grammar state tok True s) ->
         (p : Grammar state tok c a) ->
         Grammar state tok c (List1 a)
sepBy1 {c} sep p = rewrite sym (orFalseNeutral c) in
                           [| p ::: many (sep *> p) |]

||| Parse zero or more things, each separated by another thing. May
||| match the empty input.
export
sepBy : {c : Bool} ->
        (sep : Grammar state tok True s) ->
        (p : Grammar state tok c a) ->
        Grammar state tok False (List a)
sepBy sep p = option [] $ forget <$> sepBy1 sep p

||| Parse one or more instances of `p` separated by and optionally terminated by
||| `sep`.
export
sepEndBy1 : {c : Bool} ->
            (sep : Grammar state tok True s) ->
            (p : Grammar state tok c a) ->
            Grammar state tok c (List1 a)
sepEndBy1 {c} sep p = rewrite sym (orFalseNeutral c) in
                              sepBy1 sep p <* optional sep

||| Parse zero or more instances of `p`, separated by and optionally terminated
||| by `sep`. Will not match a separator by itself.
export
sepEndBy : {c : Bool} ->
           (sep : Grammar state tok True s) ->
           (p : Grammar state tok c a) ->
           Grammar state tok False (List a)
sepEndBy sep p = option [] $ forget <$> sepEndBy1 sep p

||| Parse one or more instances of `p`, separated and terminated by `sep`.
export
endBy1 : {c : Bool} ->
         (sep : Grammar state tok True s) ->
         (p : Grammar state tok c a) ->
         Grammar state tok True (List1 a)
endBy1 {c} sep p = some $ rewrite sym (orTrueTrue c) in
                                  p <* sep

export
endBy : {c : Bool} ->
        (sep : Grammar state tok True s) ->
        (p : Grammar state tok c a) ->
        Grammar state tok False (List a)
endBy sep p = option [] $ forget <$> endBy1 sep p

||| Parse an instance of `p` that is between `left` and `right`.
export
between : {c : Bool} ->
          (left : Grammar state tok True l) ->
          (right : Grammar state tok True r) ->
          (p : Grammar state tok c a) ->
          Grammar state tok True a
between left right contents = left *> contents <* right

export
location : Grammar state token False (Int, Int)
location = startBounds <$> position

export
endLocation : Grammar state token False (Int, Int)
endLocation = endBounds <$> position

export
column : Grammar state token False Int
column = snd <$> location

public export
when : Bool -> Lazy (Grammar state token False ()) -> Grammar state token False ()
when True y = y
when False y = pure ()

public export
%inline
unless : Bool -> Lazy (Grammar state token False ()) -> Grammar state token False ()
unless = when . not


module Lens

import Data.List
import Data.Maybe
import Map

%default total

data PrimitiveValue =
      Boolean Bool
    | Number Nat
    | Text (List Char)

%name PrimitiveValue val, v, v1, v2

Eq PrimitiveValue where
    Boolean b1 == Boolean b2 = b1 == b2
    Number n1 == Number n2 = n1 == n2
    Text t1 == Text t2 = t1 == t2
    _ == _ = False

boolean : (v : PrimitiveValue) -> Maybe Bool
boolean (Boolean b) = Just b
boolean _ = Nothing

number : (v : PrimitiveValue) -> Maybe Nat
number (Number n) = Just n
number _ = Nothing

text : (v : PrimitiveValue) -> Maybe (List Char)
text (Text t) = Just t
text _ = Nothing

data PrimitiveKind =
      KBoolean
    | KNumber
    | KText

%name PrimitiveKind k, k1, k2

Eq PrimitiveKind where
    KBoolean == KBoolean = True
    KNumber == KNumber = True
    KText == KText = True
    _ == _ = False

prim_kind_of : PrimitiveValue -> PrimitiveKind
prim_kind_of (Boolean x) = KBoolean
prim_kind_of (Number x) = KNumber
prim_kind_of (Text x) = KText

data Value =
      Null
    | Primitive PrimitiveValue
    | Array (List Value)
    | Object (Map Value)

%name Value val, v, v1, v2

Eq Value where
    Null == Null = True
    Primitive v1 == Primitive v2 = v1 == v2
    Array vs1 == Array vs2 = assert_total (vs1 == vs2)
    Object ps1 == Object ps2 = assert_total (ps1 == ps2)
    _ == _ = False

prim : Value -> Maybe PrimitiveValue
prim (Primitive v) = Just v
prim _ = Nothing

array : Value -> Maybe (List Value)
array (Array xs) = Just xs
array _ = Nothing

object : Value -> Maybe (Map Value)
object (Object m) = Just m
object _ = Nothing

data Kind =
      KNull
    | KPrimitive PrimitiveKind
    | KArray
    | KObject

%name Kind k, k1, k2

Eq Kind where
    KNull == KNull = True
    KPrimitive k1 == KPrimitive k2 = k1 == k2
    KArray == KArray = True
    KObject == KObject = True
    _ == _ = False

data Schema =
      SNull
    | SBoolean
    | SNumber
    | SText
    | SArray Bool Schema
    | SObject (Map Schema)

%name Schema s, s1, s2

allow_empty : Schema -> Maybe Bool
allow_empty (SArray e _) = Just e
allow_empty _ = Nothing

sarray : Schema -> Maybe Schema
sarray (SArray _ s) = Just s
sarray _ = Nothing

mutual
    validate_properties' : List Key -> Map Value -> Map Schema -> Bool
    validate_properties' [] _ _ = True
    validate_properties' (k :: ks) vm sm with (get k vm, get k sm)
        validate_properties' (k :: ks) vm sm | (Just v, Just s) =
             assert_total (validate s v) && validate_properties' ks vm sm
        validate_properties' (k :: ks) vm sm | (_, _) = False

    validate_properties : Map Value -> Map Schema -> Bool
    validate_properties vm sm = validate_properties' (keys vm) vm sm

    validate : Schema -> Value -> Bool
    validate SNull Null = True
    validate SNull _ = False
    validate SBoolean (Primitive (Boolean _)) = True
    validate SBoolean _ = False
    validate SNumber (Primitive (Number _)) = True
    validate SNumber _ = False
    validate SText (Primitive (Text _)) = True
    validate SText _ = False
    validate (SArray e _) (Array []) = e
    validate (SArray _ s) (Array (x :: xs)) =
        assert_total (validate s) x && assert_total (validate (SArray True s) (Array xs))
    validate (SArray _ _) _ = False
    validate (SObject sm) (Object vm) =
        half_keys_eq sm vm && validate_properties vm sm
    validate (SObject _) _ = False

data Lens =
      Make Kind
    | Destroy Kind
    | AddProperty Key
    | RemoveProperty Key
    | RenameProperty Key Key
    | HoistProperty Key Key
    | PlungeProperty Key Key
    | Wrap
    | Head
    | LensIn Key Lens
    | LensMap Lens
    | Convert PrimitiveKind PrimitiveKind (List (PrimitiveValue, PrimitiveValue))

%name Lens lens, l, l1, l2

Eq Lens where
    Make k1 == Make k2 = k1 == k2
    Destroy k1 == Destroy k2 = k1 == k2
    AddProperty p1 == AddProperty p2 = p1 == p2
    RemoveProperty p1 == RemoveProperty p2 = p1 == p2
    RenameProperty a1 b1 == RenameProperty a2 b2 = a1 == a2 && b1 == b2
    HoistProperty h1 p1 == HoistProperty h2 p2 = h1 == h2 && p1 == p2
    PlungeProperty h1 p1 == PlungeProperty h2 p2 = h1 == h2 && p1 == p2
    Wrap == Wrap = True
    Head == Head = True
    LensIn p1 l1 == LensIn p2 l2 = p1 == p2 && l1 == l2
    LensMap l1 == LensMap l2 = l1 == l2
    Convert k11 k21 f1 == Convert k12 k22 f2 = k11 == k12 && k22 == k22 && f1 == f2
    _ == _ = False

flip_map : List (PrimitiveValue, PrimitiveValue) -> List (PrimitiveValue, PrimitiveValue)
flip_map [] = []
flip_map ((x, y) :: xs) = (y, x) :: (flip_map xs)

validate_map : PrimitiveKind -> PrimitiveKind -> List (PrimitiveValue, PrimitiveValue) -> Bool
validate_map _ _ [] = True
validate_map kx ky ((x, y) :: map) =
    prim_kind_of x == kx && prim_kind_of y == ky && validate_map kx ky map

convert_prim : PrimitiveValue -> List (PrimitiveValue, PrimitiveValue) -> Maybe PrimitiveValue
convert_prim _ [] = Nothing
convert_prim key ((k, v) :: xs) = if key == k then Just v else convert_prim key xs

reverse_lens : Lens -> Lens
reverse_lens (Make k) = Destroy k
reverse_lens (Destroy k) = Make k
reverse_lens (AddProperty x) = RemoveProperty x
reverse_lens (RemoveProperty x) = AddProperty x
reverse_lens (RenameProperty x y) = RenameProperty y x
reverse_lens (HoistProperty x y) = PlungeProperty x y
reverse_lens (PlungeProperty x y) = HoistProperty x y
reverse_lens Wrap = Head
reverse_lens Head = Wrap
reverse_lens (LensIn x y) = LensIn x (reverse_lens y)
reverse_lens (LensMap x) = LensMap (reverse_lens x)
reverse_lens (Convert a b m) = Convert b a (flip_map m)

strip_postfix : Eq a => List a -> List a -> (List a, List a)
strip_postfix (a::as) (b::bs) =
    if a == b
    then strip_postfix as bs
    else (a::as, b::bs)
strip_postfix a b = (a, b)

transform_lenses : List Lens -> List Lens -> List Lens
transform_lenses a b =
    let
        a = reverse a
        b = reverse b
        (a, b) = strip_postfix a b
        a = reverse a
        a = map reverse_lens a
    in a ++ b

transform_schema : Lens -> Schema -> Maybe Schema
transform_schema (Make (KPrimitive KBoolean)) SNull = Just SBoolean
transform_schema (Make (KPrimitive KNumber)) SNull = Just SNumber
transform_schema (Make (KPrimitive KText)) SNull = Just SText
transform_schema (Make KArray) SNull = Just (SArray True SNull)
transform_schema (Make KObject) SNull = Just (SObject Empty)
transform_schema (Make _) _ = Nothing
transform_schema (Destroy (KPrimitive KBoolean)) SBoolean = Just SNull
transform_schema (Destroy (KPrimitive KNumber)) SNumber = Just SNull
transform_schema (Destroy (KPrimitive KText)) SText = Just SNull
transform_schema (Destroy KArray) (SArray True SNull) = Just SNull
transform_schema (Destroy KObject) (SObject Empty) = Just SNull
transform_schema (Destroy KObject) (SObject _) = Nothing
transform_schema (Destroy _) _ = Nothing
transform_schema (AddProperty key) (SObject smap) =
    case get key smap of
        Just _ => Nothing
        Nothing => Just (SObject (insert key SNull smap))
transform_schema (AddProperty _) _ = Nothing
transform_schema (RemoveProperty key) (SObject smap) =
    case get key smap of
        Just SNull => Just (SObject (remove key smap))
        _ => Nothing
transform_schema (RemoveProperty _) _ = Nothing
transform_schema (RenameProperty x y) (SObject smap) =
    case (get x smap, get y smap) of
        (Just p, Nothing) =>
            let
                smap = (remove x smap)
                smap = (insert y p smap)
            in Just (SObject smap)
        _ => Nothing
transform_schema (RenameProperty _ _) _ = Nothing
transform_schema (HoistProperty h t) (SObject m) =
    case (get h m, get t m) of
        (Just (SObject hm), Nothing) =>
            case get t hm of
                Just p =>
                    let
                        hm = remove t hm
                        m = insert t p m
                        m = insert h (SObject hm) m
                     in Just (SObject m)
                Nothing => Nothing
        _ => Nothing
transform_schema (HoistProperty _ _) _ = Nothing
transform_schema (PlungeProperty h t) (SObject m) =
    case (get t m, get h m, t == h) of
        (Just p, Just (SObject hm), False) =>
            case get t hm of
                Nothing =>
                    let
                        hm = insert t p hm
                        m = remove t m
                        m = insert h (SObject hm) m
                    in Just (SObject m)
                _ => Nothing
        _ => Nothing
transform_schema (PlungeProperty _ _) _ = Nothing
transform_schema Wrap schema = Just (SArray False schema)
transform_schema Head (SArray False schema) = Just schema
transform_schema Head _ = Nothing
transform_schema (LensIn key lens) (SObject smap) =
    case get key smap of
        Just schema =>
            case transform_schema lens schema of
                Just schema => Just (SObject (insert key schema smap))
                Nothing => Nothing
        Nothing => Nothing
transform_schema (LensIn _ _) _ = Nothing
transform_schema (LensMap lens) (SArray allowEmpty schema) =
    case transform_schema lens schema of
        Just schema' => Just (SArray allowEmpty schema')
        Nothing => Nothing
transform_schema (LensMap _) _ = Nothing
transform_schema (Convert a b map) s with (validate_map a b map)
    transform_schema (Convert KBoolean KBoolean _) SBoolean | True = Just SBoolean
    transform_schema (Convert KBoolean KNumber _) SBoolean | True = Just SNumber
    transform_schema (Convert KBoolean KText _) SBoolean | True = Just SText
    transform_schema (Convert KNumber KBoolean _) SNumber | True = Just SBoolean
    transform_schema (Convert KNumber KNumber _) SNumber | True = Just SNumber
    transform_schema (Convert KNumber KText _) SNumber | True = Just SText
    transform_schema (Convert KText KBoolean _) SText | True = Just SBoolean
    transform_schema (Convert KText KNumber _) SText | True = Just SNumber
    transform_schema (Convert KText KText _) SText | True = Just SText
    transform_schema (Convert _ _ _) _ | _ = Nothing

lenses_to_schema : List Lens -> Maybe Schema
lenses_to_schema [] = Just SNull
lenses_to_schema (l::ls) =
    case lenses_to_schema ls of
        Just s => transform_schema l s
        Nothing => Nothing

transform_value : Lens -> Value -> Value
transform_value (Make KNull) _ = Null
transform_value (Make (KPrimitive KBoolean)) _ = Primitive (Boolean False)
transform_value (Make (KPrimitive KNumber)) _ = Primitive (Number 0)
transform_value (Make (KPrimitive KText)) _ = Primitive (Text [])
transform_value (Make KArray) _ = Array []
transform_value (Make KObject) _ = Object Empty
transform_value (Destroy _) _ = Null
transform_value (AddProperty k) (Object m) = Object (insert k Null m)
transform_value (AddProperty _) v = v
transform_value (RemoveProperty k) (Object m) = Object (remove k m)
transform_value (RemoveProperty _) v = v
transform_value (RenameProperty k1 k2) (Object m) =
    case get k1 m of
        Just v => Object (insert k2 v (remove k1 m))
        Nothing => Object m
transform_value (RenameProperty _ _) v = v
transform_value (HoistProperty h t) (Object m) =
    case get h m of
        Just (Object hm) =>
            case get t hm of
                Just v => Object (insert h (Object (remove t hm)) $ insert t v m)
                Nothing => Object m
        _ => Object m
transform_value (HoistProperty _ _) v = v
transform_value (PlungeProperty h t) (Object m) =
    case (get h m, get t m) of
        (Just (Object hm), Just v) => Object (insert h (Object (insert t v hm)) $ remove t m)
        _ => Object m
transform_value (PlungeProperty _ _) v = v
transform_value Wrap v = Array [v]
transform_value Head (Array (x :: _)) = x
transform_value Head v = v
transform_value (LensIn k l) (Object vm) =
    case get k vm of
        Just v => Object (insert k (transform_value l v) vm)
        Nothing => Object vm
transform_value (LensIn _ _) v = v
transform_value (LensMap l) (Array xs) = Array (map (transform_value l) xs)
transform_value (LensMap _) v = v
transform_value (Convert _ b m) (Primitive v) =
    case convert_prim v m of
        Just v => (Primitive v)
        Nothing => case b of
            KBoolean => Primitive (Boolean False)
            KNumber => Primitive (Number 0)
            KText => Primitive (Text [])
transform_value (Convert _ _ _) v = v


module Flexidisc.TaggedValue

%default total

infixl 8 :=

public export
data TaggedValue : (key : k) -> (v : Type) -> Type where
  (:=) : (key : k) -> v -> TaggedValue key v


-- --------------------------------------------------------- [ Description.idr ]
-- Module    : Description.idr
-- Copyright : (c) Jan de Muijnck-Hughes
-- License   : see LICENSE
-- --------------------------------------------------------------------- [ EOH ]
||| Data types for modelling abstract descriptions of interfaces.
module Cordial.Model.Specification

import        Data.Vect
import        Data.DList
import public Cordial.Model.Common

%default total
%access public export

||| Individual port in an interface.
|||
record PortTy annot (labelTy : Type) where
  constructor PortDesc
  info : annot
  kind  : PortShape
  label : Name nKind labelTy
  wkind : WireTy
  flow  : Flow
  width : PortWidthDesc kind

  sens  : Sensitivity

  required : OptionalTy
  origin   : Origin

  valueOrigin  : ValueOrigin width
  doc          : DocString

data OptionalPortType : (endpoint : Endpoint) -> (type : PortTy annot labelTy) -> Type where
  IsOptionalPortType : (prf : OptionalKind endpoint optTy OPT)
                    -> OptionalPortType endpoint (PortDesc i k l ty f w s optTy o vo doc)

data SystemPortType : (type : PortTy annot labelTy) -> Type where
   IsSystemPortType : (name : String) -> SystemPortType (PortDesc i k l ty f w s optTy (SYSTEM name) vo doc)

data IPPortType : (type : PortTy annot labelTy) -> Type where
   IsIPPortType : IPPortType (PortDesc i k l ty f w s optTy IP vo doc)

portIsSystem : IPPortType (PortDesc i k l wk f w s r (SYSTEM x) v d)
            -> Void
portIsSystem IsIPPortType impossible

portIsIP : SystemPortType (PortDesc i k l wk f w s r IP v d)
        -> Void
portIsIP (IsSystemPortType _) impossible

isPortTypeSystem : (type : PortTy annot labelTy)
                -> Dec (SystemPortType type)
isPortTypeSystem (PortDesc info kind label wkind flow width sens required (SYSTEM x) valueOrigin doc) = Yes (IsSystemPortType x)
isPortTypeSystem (PortDesc info kind label wkind flow width sens required IP valueOrigin doc) = No portIsIP

isPortTypeIP : (type : PortTy annot labelTy)
           -> Dec (IPPortType type)
isPortTypeIP (PortDesc info kind label wkind flow width sens required (SYSTEM x) valueOrigin doc) = No portIsSystem
isPortTypeIP (PortDesc info kind label wkind flow width sens required IP valueOrigin doc) = Yes IsIPPortType

||| Collection of Ports
PortGroupTy : (annot : Type) -> (labelTy : Type) -> Type
PortGroupTy annot labelTy = List (PortTy annot labelTy)

record Channel (annot : Type) (labelTy : Type) where
  constructor MkChan
  name : String
  info : annot
  ports : PortGroupTy annot labelTy

flatten : List (Channel annot labelTy) -> PortGroupTy annot labelTy
flatten Nil     = Nil
flatten (c::cs) = ports c ++ flatten cs

||| An interface is a list of ports with a given number of producers
||| and consumers, and a communication style.
|||
record InterfaceTy (annot : Type) (labelTy : Type) where
   constructor MkInterfaceTy
   name         : VLNV
   info         : annot
   cstyle       : CommStyle
   maxProducers : Whole
   maxConsumers : Whole
   portGroupTy  : PortGroupTy annot labelTy
   channels     : List (Channel annot labelTy)
   doc          : DocString

fullPortGroupTy : InterfaceTy annot labelTy -> PortGroupTy annot labelTy
fullPortGroupTy iface = (portGroupTy iface) ++ (flatten $ channels iface)

||| Collection of interfaces
data InterfaceTys : (annot : Type) -> (labelTys : List Type)-> Type where
  Nil  : InterfaceTys  annot Nil
  (::) : InterfaceTy   annot labelTy
      -> InterfaceTys  annot hs
      -> InterfaceTys  annot (labelTy::hs)

data IsBroadcast : InterfaceTy annot labelTy -> Type where
  YesIsBroadcast : IsBroadcast (MkInterfaceTy n i Broadcast p c ps cs doc)
-- --------------------------------------------------------------------- [ EOF ]


module Utils.Parser

import Data.List
import Data.List.Elem
import Data.List1
import Data.Vect
import Data.Void
import Decidable.Equality
import Syntax.WithProof

public export
interface Cons a s | s where
  singleton : a -> s
  uncons : s -> Maybe (a, s)

public export
Cons Char String where
  singleton = cast
  uncons = strUncons

public export
Cons a (List a) where
  singleton = pure
  uncons [] = Nothing
  uncons (x :: xs) = Just (x, xs)

||| a simple incremental parser
public export
data Parser : (input : Type) -> (error : Type) -> (a : Type) -> Type where
  Fail : e -> Parser i e a
  Pure : (leftover : i) -> a -> Parser i e a
  More : (on_feed : (i -> Parser i e a)) -> Parser i e a
  Alt : Parser i e a -> Lazy (Parser i e a) -> Parser i e a

public export
Functor (Parser i e) where
  map f (Fail e) = Fail e
  map f (Pure leftover x) = Pure leftover (f x)
  map f (More on_feed) = More (map f . on_feed)
  map f (Alt p1 p2) = Alt (map f p1) (map f p2)

||| maps over the errors of the parser
public export
map_error : (e -> e') -> Parser i e a -> Parser i e' a
map_error f (Fail e) = Fail (f e)
map_error f (Pure leftover x) = Pure leftover x
map_error f (More on_feed) = More (map_error f . on_feed)
map_error f (Alt p1 p2) = Alt (map_error f p1) (map_error f p2)

public export
(<|>) : Semigroup e => Parser i e a -> Lazy (Parser i e a) -> Parser i e a
Fail e <|> p = map_error (e <+>) p
Pure leftover x <|> p = Pure leftover x
p <|> q = Alt p q

||| fail with an error
public export
fail : e -> Parser i e a
fail = Fail

||| feed input into the parser incrementally
public export
feed : (Semigroup e, Semigroup i) => i -> Parser i e a -> Parser i e a
feed input (Fail e) = Fail e
feed input (Pure leftover x) = Pure (leftover <+> input) x
feed input (More on_feed) = on_feed input
feed input (Alt p1 p2) = feed input p1 <|> feed input p2

apply : (Semigroup e, Semigroup i) => (Parser i e a -> Parser i e b) -> Parser i e a -> Parser i e b
apply f (Fail msg) = Fail msg
apply f (Alt p1 p2) = Alt (f p1) (f p2)
apply f (More on_feed) = More (f . on_feed)
apply f parser = More (\input => f $ feed input parser)

public export
pure : Monoid i => a -> Parser i e a
pure x = Pure neutral x

public export
(<*>) : (Semigroup e, Monoid i) => Parser i e (a -> b) -> Lazy (Parser i e a) -> Parser i e b
Pure leftover f <*> p = map f $ feed leftover p
p1 <*> p2 = apply (<*> p2) p1

public export
(<*) : (Semigroup e, Monoid i) => Parser i e a -> Parser i e b -> Parser i e a
x <* y = map const x <*> y

public export
(*>) : (Semigroup e, Monoid i) => Parser i e a -> Parser i e b -> Parser i e b
x *> y = map (const id) x <*> y

public export
(>>=) : (Semigroup e, Monoid i) => Parser i e a -> (a -> Parser i e b) -> Parser i e b
(>>=) (Pure leftover x) f = feed leftover $ f x
(>>=) p f = apply (>>= f) p

public export
more : (i -> Parser i e a) -> Parser i e a
more = More

||| peek into the next token without consuming it
public export
peek : Cons c i => Parser i e c
peek = more $ \input =>
  case uncons input of
    Just (x, _) => Pure input x
    Nothing => peek

||| reads the next token
public export
token : Cons c i => Parser i e c
token = more $ \input =>
  case uncons input of
    Just (x, xs) => Pure xs x
    Nothing => token

||| run `p` `k` times and collect the results
public export
count : (Semigroup e, Monoid i, Cons c i) => (k : Nat) -> (p : Parser i e a) -> Parser i e (Vect k a)
count Z parser = pure []
count (S k) parser = pure $ !parser :: !(count k parser)

||| return the result of `p` if it succeeds, otherwise return `x`
public export
option : (Semigroup e, Monoid i) => (x : a) -> (p : Parser i e a) -> Parser i e a
option x p = p <|> pure x

mutual
  public export
  some : (Semigroup e, Monoid i) => Parser i e a -> Parser i e (List1 a)
  some p = pure $ !p ::: !(many p)

  public export
  many : (Semigroup e, Monoid i) => Parser i e a -> Parser i e (List a)
  many p = option [] (forget <$> some p)

namespace Error
  ||| example: `SimpleError String`
  public export
  data SimpleError : Type -> Type where
    Msg : a -> SimpleError a
    Alt : SimpleError a -> SimpleError a -> SimpleError a
    Under : a -> SimpleError a -> SimpleError a

  public export
  Semigroup (SimpleError a) where
    (<+>) = Alt

  public export
  Show a => Show (SimpleError a) where
    show (Msg x) = show x
    show (Alt a b) = "(" <+> show a <+> " <|> " <+> show b <+> ")"
    show (Under x a) = "(" <+> show x <+> ": " <+> show a <+> ")"

  public export
  msg : e -> SimpleError e
  msg = Msg

  public export
  under : e -> Parser i (SimpleError e) a -> Parser i (SimpleError e) a
  under = map_error . Under


module Ziria

-- Based on https://github.com/dimitriv/ziria-sem/blob/master/Haskell/ZirStepper.hs

data ZirS : (s : Type) -> (i : Type) -> (o : Type) -> (v : Type) -> Type where
  Yield : s -> o -> ZirS s i o v
  Skip : s -> ZirS s i o v
  Done : v -> ZirS s i o v
  NeedInput : (i -> s) -> ZirS s i o v

data Zir : (i : Type) -> (o : Type) -> (v : Type) -> Type where
  ZirC : s -> (s -> ZirS s i o v) -> Zir i o v

return : v -> Zir i o v
return v = ZirC () (\() => Done v)

cast : (s -> s') -> ZirS s i o v -> ZirS s' i o v
cast f (Yield s o) = Yield (f s) o
cast f (Skip s) = Skip (f s)
cast f (Done x) = Done x
cast f (NeedInput g) = NeedInput (f . g)

(>>=) : Zir i o v -> (v -> Zir i o w) -> Zir i o w
(>>=) (ZirC x step1) f = ZirC (Left x) (step step1)
  where
    s2step : Zir i o w -> ZirS (Either s (Zir i o w)) i o w
    s2step (ZirC s2 step2) = cast (\x => Right (ZirC x step2)) (step2 s2)
    step : (s -> ZirS s i o v) -> Either s (Zir i o w) -> ZirS (Either s (Zir i o w)) i o w
    step step1 (Right zir) = s2step zir
    step step1 (Left l) with (step1 l)
      step step1 (Left l) | (Yield s1' z) = Yield (Left s1') z
      step step1 (Left l) | (Skip s1') = Skip (Left s1')
      step step1 (Left l) | (Done y) = s2step (f y)
      step step1 (Left l) | (NeedInput g) = NeedInput (\i => Left (g i))

step' : (s1 -> ZirS s1 i m v) -> (s2 -> ZirS s2 m o v) -> (s1, s2) -> ZirS (s1, s2) i o v
step' step1 step2 (s1, s2) with (step2 s2)
  step' step1 step2 (s1, s2) | (Yield s2' ot) = Yield (s1, s2') ot
  step' step1 step2 (s1, s2) | (Skip s2') = Skip (s1, s2')
  step' step1 step2 (s1, s2) | (Done x) = Done x
  step' step1 step2 (s1, s2) | (NeedInput f) with (step1 s1)
    step' step1 step2 (s1, s2) | (NeedInput f) | (Yield s1' ot) = Skip (s1', f ot)
    step' step1 step2 (s1, s2) | (NeedInput f) | (Skip s1') = Skip (s1', s2)
    step' step1 step2 (s1, s2) | (NeedInput f) | (Done x) = Done x
    step' step1 step2 (s1, s2) | (NeedInput f) | (NeedInput g) = NeedInput (\i => ((g i), s2))

infixl 9 >>>
(>>>) : Zir i m v -> Zir m o v -> Zir i o v
(>>>) (ZirC s1 step1) (ZirC s2 step2) = ZirC (s1, s2) (step' step1 step2)

emit : o -> Zir i o ()
emit x = ZirC True aux
  where
    aux : Bool -> ZirS Bool i o ()
    aux False = Done ()
    aux True = Yield False x

take : Zir i o i
take = ZirC Nothing aux
  where
    aux : Maybe i -> ZirS (Maybe i) i o i
    aux Nothing = NeedInput Just
    aux (Just x) = Done x

run : Zir i o v -> Stream i -> (v, List o)
run (ZirC st step) input = loop step [] input (step st)
  where
    loop : (s -> ZirS s i o v) -> List o -> Stream i -> ZirS s i o v -> (v, List o)
    loop step acc input (Yield st ot) = loop step (ot :: acc) input (step st)
    loop step acc input (Skip st) = loop step acc input (step st)
    loop step acc input (Done v) = (v, reverse acc)
    loop step acc (x :: xs) (NeedInput f) = loop step acc xs (step (f x))

example : Zir Int Int Int
example = do x <- take
             emit (x + 1)
             y <- take
             emit (y + 2)
             return (x + y)

test : (Int, List Int)
test = run example [1..]


||| Tail-call optimization.
|||
||| Here is a lengthy explanation how this works at the
||| moment. Assume the following call graph of functions f1,f2,f3,f4 all
||| calling each other in tail call position:
|||
||| ```
|||       ------------ f2 ---- f4 (result)
|||      /          /     \
||| f1 ---- f1     /       -- f1
|||      \        /
|||       -- f3 --
||| ```
|||
||| First, a directed graph of all toplevel function calls
||| (incoming and outgoing) in tail-call position is created:
|||
||| ```idris
||| MkCallGraph $ fromList [(f1,[f1,f2,f3]),(f2,[f1,f4]),(f3,[f2])]
|||             $ fromList [(f1,[f1,f2]),(f2,[f1,f3]),(f3,[f1]),(f4,[f2])]
||| ```
|||
||| Mutually tail-recursive functions form a strongly connected
||| component in such a call graph: There is a (directed) path from every function
||| to every other function. Tarjan's algorithm is used to identify
||| these strongly connected components and grouping them in
||| a `List` of `List1`s.
|||
||| A tail-recursive group of functions is now converted to an imperative
||| loop as follows: Let `obj={h:_, a1:_, a2:_, ...}`
||| be a Javascript object consisting
||| of a tag `h` and arguments `a1`,`a2`,... . `h` indicates, whether `obj.a1`
||| contains the result of the computation (`h = 0`) or describes
||| a continuation indicating the next function to be invoked, in which
||| case fields `a1`,`a2`,... are the function's arguments.
||| together with the necessary arguments. The group of mutually
||| recursive functions is now converted to a single switch statement
||| where each branch corresponds to one of the function.
||| Each function will be changed in such a way that instead of
||| (recursively) calling another function in its group it will return
||| a new object `{h:_, a1:_, ...}` with `h` indicating the next
||| function to call (the next branch to choose in the `switch`
||| statement and `a1`,`a2`,... being the next function's set of
||| arguments. The function and initial argument object will then
||| be passed to toplevel function `__tailRec`, which loops
||| until the object signals that we have arrived at a result.
|||
||| Here is an example of two mutually tail-recursive functions
||| together with the generated tail-call optimized code.
|||
||| Original version:
|||
||| ```javascript
||| function isEven(arg){
|||   switch (arg) {
|||     case 0  : return 1;
|||     default : return isOdd(arg - 1);
|||   }
||| }
|||
||| function isOdd(arg){
|||   switch (arg) {
|||     case 0  : return 0;
|||     default : return isEven(arg - 1);
|||   }
||| }
||| ```
|||
||| The above gets converted to code similar to
||| the following.
|||
||| ```javascript
||| function tcOpt(arg) {
|||   switch(arg.h) {
|||   // former function isEven
|||   case 1: {
|||     switch (arg.a1) {
|||       case 0  : return {h: 0, a1: 1};
|||       default : return {h: 2, a1: arg.a1 - 1};
|||     }
|||   }
|||   // former function isOdd
|||   case 2: {
|||     switch (a1) {
|||       case 0  : return {h: 0, a1: 0};
|||       default : return {h: 1, a1: arg.a1 - 1};
|||     }
|||   }
||| }
|||
||| function isEven(arg){
|||   return __tailRec(tcOpt,{h: 1, a1: arg})
||| }
|||
||| function isOdd(arg){
|||   return __tailRec(tcOpt,{h: 2, a1: arg})
||| }
||| ```
|||
||| Finally, `__tailRec` is implemented as follows:
|||
||| ```javascript
|||   function __tailRec(f,ini) {
|||     let obj = ini;
|||     while(true){
|||       switch(obj.h){
|||         case 0: return obj.a1;
|||         default: obj = f(obj);
|||       }
|||     }
|||   }
||| ```
|||
||| While the above example is in Javascript, this module operates
||| on `NamedCExp` exclusively, so it might be used with any backend
||| where the things described above can be expressed.
module Compiler.ES.TailRec

import Data.Maybe
import Data.List
import Data.List1
import Data.String
import Libraries.Data.Graph
import Libraries.Data.SortedSet
import Libraries.Data.List.Extra as L
import Libraries.Data.SortedMap as M
import Core.Name
import Core.CompileExpr
import Core.Context

--------------------------------------------------------------------------------
--          Utilities
--------------------------------------------------------------------------------

-- indices of a list starting at 1
indices : List a -> List Int
indices as = [1 .. cast (length as)]

zipWithIndices : List a -> List (Int,a)
zipWithIndices as = zip (indices as) as

--------------------------------------------------------------------------------
--          Tailcall Graph
--------------------------------------------------------------------------------

||| A (toplevel) function in a group of mutually tail recursive functions.
public export
record TcFunction where
  constructor MkTcFunction
  ||| Function's name
  name  : Name

  ||| Function's index in its tail call group
  ||| This is used to decide on which branch to choose in
  ||| the next iteration
  index : Int

  ||| Argument list
  args  : List Name

  ||| Function's definition
  exp   : NamedCExp

||| A group of mutually tail recursive toplevel functions.
public export
record TcGroup where
  constructor MkTcGroup
  ||| Index of the group. This is used to generate a uniquely
  ||| named tail call optimized toplevel function.
  index     : Int

  ||| Set of mutually recursive functions.
  functions : SortedMap Name TcFunction

-- tail calls made from an expression
tailCalls : NamedCExp -> SortedSet Name
tailCalls (NmLet _ _ _ z) = tailCalls z
tailCalls (NmApp _ (NmRef _ x) _) = singleton x
tailCalls (NmConCase fc sc xs x) =
  concatMap (\(MkNConAlt _ _ _ _ x) => tailCalls x) xs <+>
  concatMap tailCalls x
tailCalls (NmConstCase fc sc xs x) =
  concatMap (\(MkNConstAlt _ x) => tailCalls x) xs <+>
  concatMap tailCalls x
tailCalls _ = empty

-- Checks if a `List1` of functions actually has any tail recursive
-- function calls and needs to be optimized.
-- In case of a larger group (more than one element)
-- the group contains tailcalls by construction. In case
-- of a single function, we need to check that at least one
-- outgoing tailcall goes back to the function itself.
-- We use the given mapping from `Name` to set of names
-- called in tail position to verify this.
hasTailCalls : SortedMap Name (SortedSet Name) -> List1 Name -> Bool
hasTailCalls g (x ::: Nil) = maybe False (contains x) $ lookup x g
hasTailCalls _ _           = True

-- Given a strongly connected group of functions, plus
-- a unique index, convert them to a list of
-- consisting of the function names plus the `TcGroup`,
-- to which they belong.
toGroup :  SortedMap Name (Name,List Name,NamedCExp)
        -> (Int,List1 Name)
        -> List (Name,TcGroup)
toGroup funMap (groupIndex,functions) =
  let ns    = zipWithIndices $ forget functions
      group = MkTcGroup groupIndex . fromList $ mapMaybe fun ns
   in (,group) <$> forget functions
  where fun : (Int,Name) -> Maybe (Name,TcFunction)
        fun (fx, n) = do
          (_,args,exp) <- lookup n funMap
          pure (n,MkTcFunction n fx args exp)

-- Returns the connected components of the tail call graph
-- of a set of toplevel function definitions.
-- Every function name that is part of a tail-call group
-- (a set of mutually tail-recursive functions)
-- points to its corresponding group.
tailCallGroups :  List (Name,List Name,NamedCExp)
               -> SortedMap Name TcGroup
tailCallGroups funs =
  let funMap = M.fromList $ map (\t => (fst t,t)) funs
      graph  = map (\(_,_,x) => tailCalls x) funMap
      groups = filter (hasTailCalls graph) $ tarjan graph
   in fromList $ concatMap (toGroup funMap) (zipWithIndices groups)

--------------------------------------------------------------------------------
--          Converting tail call groups to expressions
--------------------------------------------------------------------------------

public export
record Function where
  constructor MkFunction
  name : Name
  args : List Name
  body : NamedCExp

tcFunction : Int -> Name
tcFunction = MN "$tcOpt"

tcArgName : Name
tcArgName = MN "$a" 0

tcContinueName : (groupIndex : Int) -> (functionIndex : Int) -> Name
tcContinueName gi fi = MN ("TcContinue" ++ show gi) fi

tcDoneName : (groupIndex : Int) -> Name
tcDoneName gi = MN "TcDone" gi

-- Converts a single function in a mutually tail-recursive
-- group of functions to a single branch in a pattern match.
-- Recursive function calls will be replaced with an
-- applied data constructor whose tag indicates the
-- branch in the pattern match to use next, and whose values
-- will be used as the arguments for the next function.
conAlt : TcGroup -> TcFunction -> NamedConAlt
conAlt (MkTcGroup tcIx funs) (MkTcFunction n ix args exp) =
  let name = tcContinueName tcIx ix
   in MkNConAlt name DATACON (Just ix) args (toTc exp)

   where mutual
     -- this is returned in case we arrived at a result
     -- (an expression not corresponding to a recursive
     -- call in tail position).
     tcDone : NamedCExp -> NamedCExp
     tcDone x = NmCon EmptyFC (tcDoneName tcIx) DATACON (Just 0) [x]

     -- this is returned in case we arrived at a resursive call
     -- in tail position. The index indicates the next "function"
     -- to call, the list of expressions are the function's
     -- arguments.
     tcContinue : (index : Int) -> List NamedCExp -> NamedCExp
     tcContinue ix = NmCon EmptyFC (tcContinueName tcIx ix) DATACON (Just ix)

     -- recursively converts an expression. Only the `NmApp` case is
     -- interesting, the rest is pretty much boiler plate.
     toTc : NamedCExp -> NamedCExp
     toTc (NmLet fc x y z) = NmLet fc x y $ toTc z
     toTc x@(NmApp _ (NmRef _ n) xs) =
       case lookup n funs of
         Just v  => tcContinue v.index xs
         Nothing => tcDone x
     toTc (NmConCase fc sc a d) = NmConCase fc sc (map con a) (map toTc d)
     toTc (NmConstCase fc sc a d) = NmConstCase fc sc (map const a) (map toTc d)
     toTc x@(NmCrash _ _) = x
     toTc x = tcDone x

     con : NamedConAlt -> NamedConAlt
     con (MkNConAlt x y tag xs z) = MkNConAlt x y tag xs (toTc z)

     const : NamedConstAlt -> NamedConstAlt
     const (MkNConstAlt x y) = MkNConstAlt x (toTc y)

-- Converts a group of mutually tail recursive functions
-- to a list of toplevel function declarations. `tailRecLoopName`
-- is the name of the toplevel function that does the
-- infinite looping (function `__tailRec` in the example at
-- the top of this module).
convertTcGroup : (tailRecLoopName : Name) -> TcGroup -> List Function
convertTcGroup loop g@(MkTcGroup gindex fs) =
  let functions = sortBy (comparing index) $ values fs
      branches  = map (conAlt g) functions
      switch    = NmConCase EmptyFC (local tcArgName) branches Nothing
   in MkFunction tcFun [tcArgName] switch :: map toFun functions

  where tcFun : Name
        tcFun = tcFunction gindex

        local : Name -> NamedCExp
        local = NmLocal EmptyFC

        toFun : TcFunction -> Function
        toFun (MkTcFunction n ix args x) =
          let exps  = map local args
              tcArg = NmCon EmptyFC (tcContinueName gindex ix)
                        DATACON (Just ix) exps
              tcFun = NmRef EmptyFC tcFun
              body  = NmApp EmptyFC (NmRef EmptyFC loop) [tcFun,tcArg]
           in MkFunction n args body

-- Tail recursion optimizations: Converts all groups of
-- mutually tail recursive functions to an imperative loop.
tailRecOptim :  SortedMap Name TcGroup
             -> (tcLoopName : Name)
             -> List (Name,List Name,NamedCExp)
             -> List Function
tailRecOptim groups loop ts =
  let regular = mapMaybe toFun ts
      tailOpt = concatMap (convertTcGroup loop) $ values groups
   in tailOpt ++ regular

  where toFun : (Name,List Name,NamedCExp) -> Maybe Function
        toFun (n,args,exp) = case lookup n groups of
          Just _  => Nothing
          Nothing => Just $ MkFunction n args exp

||| Converts a list of toplevel definitions (potentially
||| several groups of mutually tail-recursive functions)
||| to a new set of tail-call optimized function definitions.
||| `MkNmFun`s are converted. Other constructors of `NamedDef`
||| are ignored and silently dropped.
export
functions :  (tcLoopName : Name)
          -> List (Name,FC,NamedDef)
          -> List Function
functions loop dfs =
  let ts = mapMaybe def dfs
   in tailRecOptim (tailCallGroups ts) loop ts
   where def : (Name,FC,NamedDef) -> Maybe (Name,List Name,NamedCExp)
         def (n,_,MkNmFun args x) = Just (n,args,x)
         def _                    = Nothing


module LightClick.IR.ChannelCentric

import Data.List
import Data.String
import Data.Vect

import Toolkit.Data.DList
import Toolkit.Data.DVect

import Language.SystemVerilog.Gates

import LightClick.Types
import LightClick.Terms
import LightClick.Error

import LightClick.IR.ModuleCentric

%default total

public export
data ChannelIR : TyIR -> Type where
  CRef  : String -> (type : TyIR) -> ChannelIR type
  CLet  : {term : TyIR}
       -> (bind : String)
       -> (this : ChannelIR term)
       -> (inThis : ChannelIR expr)
                 -> ChannelIR expr

  CSeq  : {a,b : TyIR}
       -> ChannelIR a
       -> ChannelIR b
       -> ChannelIR b

  CEnd : ChannelIR UNIT

  CPort : String
       -> Direction
       -> ChannelIR DATA
       -> ChannelIR PORT

  CModule : {n : Nat} -> Vect (S n) (ChannelIR PORT) -> ChannelIR MODULE

  CDataLogic : ChannelIR DATA
  CDataArray : ChannelIR DATA -> Nat -> ChannelIR DATA
  CDataStruct : {n : Nat} -> Vect (S n) (Pair String (ChannelIR DATA)) -> ChannelIR DATA
  CDataUnion  : {n : Nat} -> Vect (S n) (Pair String (ChannelIR DATA)) -> ChannelIR DATA

  CIDX : (label : String) -- don't exist in the normal form, required for mapping MIDX to something for coverage...
      -> ChannelIR MODULE
      -> ChannelIR PORT
      -> ChannelIR PORT

  CChan : ChannelIR DATA -> ChannelIR CHAN

  CModuleInst : {n : Nat}
              -> (mname : ChannelIR MODULE)
              -> Vect (S n)
                      (Pair String (ChannelIR CHAN))
              -> ChannelIR CONN

  CNot : ChannelIR CHAN
      -> ChannelIR CHAN
      -> ChannelIR GATE

  CGate : {n : Nat}
       -> TyGateComb
       -> ChannelIR CHAN
       -> Vect (S (S n)) (ChannelIR CHAN)
       -> ChannelIR GATE

covering
showC : ChannelIR type -> String
showC (CRef name type) =
  "(CRef " <+> show name <+> " " <+> show type <+> ")"

showC (CLet x y z) =
   "(CLet "
     <+> show x <+> " "
     <+> showC y <+> "\n"
     <+> showC z
     <+> "\n)"

showC (CSeq x y) =
    "(CSeq "
      <+> showC x <+> " "
      <+> showC y
      <+> "\n)"
showC CEnd = "(CEnd)"

showC (CPort x y z) =
    "(CPort "
      <+> show x <+> " "
      <+> show y <+> " "
      <+> showC z <+> " "
      <+> ")"

showC (CModule x) =
    "(CModule "
      <+> show (map showC x)
      <+> ")"

showC CDataLogic = "(CTyLogic)"

showC (CDataArray x k) =
  "(CTyArray "
    <+> showC x <+> " "
    <+> show k
    <+> ")"

showC (CDataStruct {n} xs) = "(CTyStruct " <+> show ps <+> ")"
  where
    covering
    ps : Vect (S n) String
    ps = map (\(l,c) => "(" <+> show l <+> " " <+> showC c <+> ")") xs

showC (CDataUnion {n} xs) = "(TyUnion " <+> show ps <+> ")"
  where
    covering
    ps : Vect (S n) String
    ps = map (\(l,c) => "(" <+> show l <+> " " <+> showC c <+> ")") xs

showC (CChan x) = "(CChan " <+> showC x <+> ")"

showC (CIDX l m p) =
  "(CIDX "
     <+> show l  <+> " "
     <+> showC m <+> " "
     <+> ")"
showC (CModuleInst m {n} params) =
    "(CModuleInst "
      <+> showC m <+> " "
      <+> show ps
      <+> ")"
  where
    covering
    ps : Vect (S n) String
    ps = map (\(l,c) => show l <+> " " <+> showC c) params

showC (CNot o i)
  = unwords ["(CNot", showC o, showC i, ")"]

showC (CGate ty o ins)
    = unwords ["(CGate", show ty, showC o, ins', ")"]

  where
    covering
    ins' : String
    ins' = unwords $ toList $ map (\c => "(" <+> showC c <+> ")") ins

export
Show (ChannelIR type) where
  show expr = assert_total $ showC expr

public export
Convert : Type -> Type
Convert = Either LightClick.Error

covering
convert : ModuleIR type -> Convert (ChannelIR type)

covering
convPort : ModuleIR PORT -> Convert (ChannelIR MODULE, String)
convPort (MRef name PORT) = Left (NotSupposedToHappen (Just "convPort CIR Ref"))
convPort (MLet name beThis inThis) = Left (NotSupposedToHappen (Just "convPort CIR Let"))
convPort (MSeq doThis thenThis) = Left (NotSupposedToHappen (Just "convPort CIR Seq"))
convPort (MPort label dir type) = Left (NotSupposedToHappen (Just "convPort CIR Port"))
convPort (MIDX label x z)
    = do m' <- convert x
         m'' <- convModule m'
         pure (m'',label)

  where

    convModule : ChannelIR MODULE -> Convert (ChannelIR MODULE)
    convModule m@(CRef x MODULE) = Right m
    convModule (CLet x z w)
      = Left (NotSupposedToHappen (Just "convModule CIR let"))

    convModule (CSeq x z)
      = Left (NotSupposedToHappen (Just "convModule CIR CSeq"))

    convModule (CModule xs)
      = Left (NotSupposedToHappen (Just "convModule CIR CMod"))

covering
mkConn : (p : ModuleIR PORT)
      -> (c : ChannelIR CHAN)
           -> Convert (ChannelIR CONN)
mkConn p c =
  do (m,l) <- convPort p
     pure (CModuleInst m [(l,c)])

covering
convPorts : (c    : ChannelIR CHAN)
         -> (outs : Vect (S n) (ModuleIR PORT))
                 -> Convert (ChannelIR CONN)
convPorts c (x :: []) = mkConn x c
convPorts c (x :: (z :: xs)) =
    do con <- mkConn x c
       rest <- convPorts c (z::xs)
       pure (CSeq con rest)

-- [ Definition ]
convert (MRef x t) = Right $ CRef x t
convert (MLet x y z) =
  do y' <- convert y
     z' <- convert z
     pure $ CLet x y' z'
convert (MSeq x y) =
  do x' <- convert x
     y' <- convert y
     pure $ CSeq x' y'
convert MEnd = pure CEnd

convert (MPort x y z) =
  do z' <- convert z
     pure $ CPort x y z'

convert (MModule {n} x) =
    do xs <- traverse convert x
       pure $ CModule xs

convert MDataLogic = Right $ CDataLogic

convert (MDataArray x k) =
  do x' <- convert x
     pure $ CDataArray x' k
convert (MDataStruct xs) =
    do xs' <- traverse (\(l,x) => do {x' <- convert x; pure (l, x')}) xs
       pure $ CDataStruct xs'

convert (MDataUnion xs)  =
    do xs' <- traverse (\(l,x) => do {x' <- convert x; pure (l,x')}) xs
       pure $ CDataUnion xs'

convert (MChan x) =
  do x' <- convert x
     pure (CChan x')

convert (MIDX x s y) =
  do s' <- convert s
     y' <- convert y
     pure (CIDX x s' y')

convert (MConn cname y ps) =
    do c <- convert cname
       i <- mkConn y c
       rest <- convPorts c ps
       pure (CSeq i rest)

convert (MNot o i)
  = do o' <- convert o
       i' <- convert i
       pure (CNot o' i')

convert (MGate ty o ins)
  = do o' <- convert o
       rest <- traverse convert ins
       pure (CGate ty o' rest)

convert (MConnG c idx)
  = do c' <- convert c
       p  <- mkConn idx c'
       pure p

export
covering
runConvert : ModuleIR type -> Convert (ChannelIR type)
runConvert = convert

-- [ EOF ]


import Data.Maybe
import Data.Nat

namespace SuperDog
  public export
  record SuperDog where
     constructor MkDog
     supername : String
     age : Int
     weight : Int

namespace OrdinaryDog
  public export
  record OrdinaryDog where
    constructor MkDog
    name : String
    age : Int
    weight : Int

record Other a where
  constructor MkOther
  {imp : String}
  fieldA : a
  fieldB : b

------ Using new application syntax as sugar for data instantiation (be it data/record/interface) -------

myDog : OrdinaryDog
myDog = MkDog { age = 4
              , weight = 12
              , name = "Sam" }

mySuperDog : SuperDog
mySuperDog = MkDog { age = 3
                   , weight = 10
                   , supername = "Super-Sam" }

other : ?
other = MkOther {fieldB = the Int 1, fieldA = "hi", imp = "Secret string"}

otherOk1 : Main.other.fieldB = (the Int 1)
otherOk1 = Refl

otherOk2 : Main.other.fieldA = "hi"
otherOk2 = Refl

otherOk3 : Main.other.imp = "Secret string"
otherOk3 = Refl

same : MkDog {age = 2, name = "Rex", weight = 10} = (the OrdinaryDog $ MkDog "Rex" 2 10)
same = Refl

namespace R1

  public export
  record R1 where
    constructor MkR
    field : a

namespace R2

  public export
  record R2 where
    constructor MkR
    {auto field : a}

r1 : R1 -- explicit fields access
r1 = MkR {field = "string"}

r2_shouldNotTypecheck1 : ?
r2_shouldNotTypecheck1 = MkR {field = the Nat 22} -- fail, impossible to disambiguate

interface Show a => (num : Num a) => MyIface a where -- Some interface with
  constructor MkIface
                                             -- constraints
  data MyData : a -> Type                    -- and a data declaration.
  someFunc : a -> a                          -- Constraints are now elaborated as auto implicits (as one would expect)
  giveBack : {x : a} -> MyData x -> a        -- (previously as explicit arguments of the interface
                                             -- constructor)

data MyDataImpl : a -> Type where            -- implementation of MyData
  MkMyData : (x : a) -> MyDataImpl x

-- implementation MyIface Int where
--    MyData = MyDataImpl
--    someFunc = id
--    giveBack (MkMyData x) = x

%hint
instanceMyIfaceInt : MyIface Integer -- this def, roughly speaking, is the 'same thing' as the above implementation
                                     -- Show Int, Num Int are auto implicits of MkIface
instanceMyIfaceInt = MkIface { MyData = MyDataImpl
                             , someFunc = id
                             , giveBack = \(MkMyData x) => x
                             , num = %search } -- auto implicit names are preserved

instanceOk : giveBack (MkMyData 22) = 22
instanceOk = Refl

interface Show' a where -- unlike Show, Show' reduces in types
   constructor MkShow'
   show' : a -> String

Show' String where
   show' = id

%hint
showMaybe' : Show' a => Show' (Maybe a)
showMaybe' = MkShow' { show' = fromMaybe "Nothing" . (("Just " ++ ) . show' <$>) }

showMaybe'Ok : show' (Just "nice") = "Just nice"
showMaybe'Ok = Refl

record AllFieldTypes a where
   constructor MkAllFieldTypes
   exp : a
   {imp : a}
   {auto aut : a}

testAllFieldTypesOk : MkAllFieldTypes { aut = "aut"
                                      , exp = "exp"
                                      , imp = "imp" }
                       = MkAllFieldTypes "exp" {imp = "imp"} @{"aut"}
testAllFieldTypesOk = Refl

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

------  The Update syntax --------

mapName : (String -> String) -> OrdinaryDog -> OrdinaryDog
mapName f = {name $= f}

setName : String -> OrdinaryDog -> OrdinaryDog
setName name' = {name := name'}

data Three : Type -> Type -> Type -> Type where
  MkThree : (x : a) -> (y : b) -> (z : c) -> Three a b c

mapSetMap : (a -> a') -> b' -> (c -> c') -> Three a b c -> Three a' b' c'
mapSetMap f y' g = {x $= f, y := y', z $= g}

setNameOld : String -> OrdinaryDog -> OrdinaryDog
setNameOld name' = {name := name'}

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

--------- Applications in presence of duplicate names ----------

-- Duplicate names are ok and treated sequentially
testDuplicateNames : {auto a : String} -> {auto a : String} -> String
testDuplicateNames @{a} @{a'} = show' a ++ ":" ++ show' a'

-- When binding arguments on LHS or listing arguments to
-- a function on RHS
-- unnamed arguments always take priority over named,
-- i.e they are bound/applied first,
-- regardless of their relative positions to named ones
testOrder1 : (a : String) -> (a : String) -> String
testOrder1 {a = a2} {-snd-} a1 {-fst-} = a1 ++ a2

testOrder1Ok : Main.testOrder1 "abc" "def" = "abcdef"
testOrder1Ok = Refl

-- unnamed explicit "1" is passed first, followed by named {a = "2"}
testAutoPriorityOk : Main.testOrder1 {a = "2"} "1" = "12"
testAutoPriorityOk = Refl

-- Two arguments with the same name can be successfully bound
-- if one of them is renamed in patterns.
-- As both arguments are requested by name
-- They are bound in the same order they are given
sameNamesOk : (a : String) -> (a : Nat) -> (String, Nat)
sameNamesOk {a {- = a-}, a = b} = (a, b)

-- All arguments are named and are of different `plicities`. Binds occur sequentially.
-- Arguments are renamed on LHS
eachArgType : (a : String) -> {a : String} -> {auto a : String} -> String
eachArgType {a = a1, a = a2, a = a3} = a1 ++ a2 ++ a3

eachArgTypeOk : eachArgType @{"3"} "1" {a = "2"} = "123"
eachArgTypeOk = Refl

-- Arguments with the same names are provided on RHS
-- which is ok, they are passed sequentially.
eachArgTypeOk2 : eachArgType {a = "1", a = "2", a = "3"} = "123"
eachArgTypeOk2 = Refl

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

--------- Bind-all-explicits pattern ----------

-- This pattern works like inexhaustible supply of
-- `_` (Match-any patterns).

-- Here to complete the definition we only
-- need to know the name of the OrdinaryDog.
-- We bind it with `{name, _}` also stating that
-- any extra explicits should be disregarded,
-- by inserting `_` into the braces.
onlyName : OrdinaryDog -> String
onlyName (MkDog {name, _ {-age, weight-} }) = name

dontCare : (x : String) -> (y : String) -> (z : String) -> String
dontCare {x, z, _} = x ++ z

dontCareOk : Main.dontCare "a" "b" "c" = "ac"
dontCareOk = Refl

-- If none of the explicit arguments are wanted
-- one `{}` can be used instead of writing an underscore for each.
dontCare2 : (x : Nat) -> Nat -> Nat -> Nat -> (y : Nat) -> x + y = y + x
dontCare2 {} = plusCommutative {}
-- dontCare2 _ _ _ _ _ = plusCommutative _ _

data Tree a = Leaf a | Node (Tree a) a (Tree a)

isNode : Tree a -> Bool
isNode (Node {}) = True
isNode _ = False

data IsNode : Tree a -> Type where
  Is : IsNode (Node {})

decIsNode : (x : Tree a) -> Dec (IsNode x)
decIsNode (Node {}) = Yes Is
decIsNode (Leaf {}) = No (\case Is impossible)
------------------------------------------------