{-
  An embeded DSL for prolog-style logic programming, based on the
  pure-prolog parser and interpreter originally written in Haskell

  Original author: Ken Friis Larsen      <kflarsen@diku.dk>
  Translated into Idris by Nathan Bedell <nbedell@tulane.edu>
-}

module Interp

import Ast
import Data.List

%default partial

----------------------------------------------------------------------
-- Interpreter
----------------------------------------------------------------------

public export
Unifier : Type
Unifier = List (Variable, Term)

-- A non-lazy version of maybe
public export
maybe' : b -> (a -> b) -> Maybe a -> b
maybe' n j Nothing  = n
maybe' n j (Just x) = j x

public export
subs : Unifier -> Term -> Term
subs u (Var x)   = maybe' (Var x) id (lookup x u)
subs u (Comp n ts) = Comp n $ map (subs u) ts

public export
compose : Unifier -> Unifier -> Unifier
compose u1 u2 = (map (\(v, t) => (v, subs u2 t)) u1) ++ u2

public export
occursIn : Variable -> Term -> Bool
occursIn v (Var x)     = v == x
occursIn v (Comp _ ms) = any (occursIn v) ms

mutual
  public export
  covering
  unify : Term -> Term -> Maybe Unifier
  unify (Var "_") t                         = pure []
  unify t (Var "_")                         = pure []
  unify (Var x) (Var y) with (x == y)
     unify (Var x) (Var y) | True           = pure []
     _ | False = pure [(x, Var y)]
  unify (Var x) t with (not (x `occursIn` t))
      unify (Var x) t   | True              = pure [(x, t)]
      _ | False = Nothing
  unify t v@(Var _)                         = unify v t
  unify (Comp m ms) (Comp n ns) with (m == n)
     unify (Comp m ms) (Comp n ns) | True   = unifyList ms ns
     _ | False = Nothing
  unify _ _ = Nothing

  public export
  unifyList : Terms -> Terms -> Maybe Unifier
  unifyList (t :: ts) (r :: rs) =
      do u1 <- unify t r
         u2 <- unifyList (map (subs u1) ts) (map (subs u1) rs)
  -- Note: I'm not sure if Idris supports backtick notation for operators
         pure $ u1 `compose` u2
  unifyList [] [] = Just []
  unifyList _ _   = Nothing

mutual
  public export
  vars : Term -> List Variable
  vars (Var "_") = []
  vars (Var x) = [x]
  vars (Comp _ ts) = varsList ts

  public export
  varsList : Terms -> List Variable
  varsList ts = [ v | t <- ts, v <- vars t]

public export
variables : Terms -> List Variable
variables ts = nub $ varsList ts

public export
freshen : List Variable -> (Term, Terms) -> (Term, Terms)
freshen bound (tc, tb) = (subs sub tc, map (subs sub) tb)
  where
    vars : List Variable
    vars = variables (tc :: tb)
    nextVar : Int -> Variable -> Variable
    nextVar i v =
      let v' = "_" ++ show i ++ "_" ++ v in
        if v' `elem` bound
          then nextVar (i+1) v
          else v'
    sub : Unifier
    sub = [ (v, Var (nextVar 0 v)) | v <- vars, v `elem` bound]

public export
data Solution : Type where
    MkSolution : List (Variable, Term) -> Solution

public export
implementation Eq Solution where
    (MkSolution x) == (MkSolution y) = assert_total $ x == y

public export
Show Solution where
  show (MkSolution [])            = "True"
  show (MkSolution xs ) = (concat $ intersperse "\n" $ map renderBindings xs)
    where renderBindings : (Variable, Term) -> String
          renderBindings (v, t) = v ++ " = " ++ show t

public export
isPlain : List Char -> Bool
isPlain (c :: cs) = isLower c && all (\c => c == '_' || isAlphaNum c) cs
isPlain [] = False

public export
data SearchTree : Type where
  Sol : Solution -> SearchTree
  Node : Goal -> List SearchTree -> SearchTree

public export
Eq SearchTree where
  (Sol x) == (Sol y) = assert_total $ x == y
  (Node x y) == (Node x' y') = assert_total $ (x == x') && (y == y')
  x == y = False

public export
isReportGoal : Term -> Bool
isReportGoal (Comp "_report" _) = True
isReportGoal _                  = False

public export
getSolution : Term -> Solution
getSolution (Comp "_report" args) = MkSolution sol
    where nontriv : (String, Term) -> Bool
          nontriv (x, (Var y)) with (x == y)
            nontriv (x, (Var y)) | True = False
            nontriv (x, (Var y)) | _ = True
          nontriv _ = True
          sol : ?
          sol = filter nontriv $ map (\ (Comp "=" [Comp v [], t]) => (v, t)) args
getSolution _ = MkSolution [("ERROR", Var "_Error")]

-- Uses the List monad for backtracking
public export
solve : Program -> Goal -> List SearchTree
solve prog [r] with (isReportGoal r)
    solve prog [r] | True  =  pure $ Sol $ getSolution r
    solve prog [r] | False = []
solve prog g@(t1 :: ts) = pure $ Node g trees
    where
      trees : List SearchTree
      trees = do
        c <- prog
        let (tc, tsc) = freshen (variables g) c
        case unify tc t1 of
         Just u => do
           let g' = map (subs u) $ tsc ++ ts
           solve prog g'
         Nothing => []

public export
makeReportGoal : Terms -> List Term
makeReportGoal goal = [Comp "_report" reportVars]
    where vars : List Variable
          vars = variables goal
          reportVars : Terms
          reportVars = map (\ v => Comp "=" [Comp v [], Var v]) vars

-- Use the trick of inserting an extra reporting goal
public export
makeReportTree : Program -> Goal -> SearchTree
makeReportTree prog goal = Node goal $ solve prog (goal ++ makeReportGoal goal)

----------------------------------------------------------------------
-- Traveral of Search Trees
----------------------------------------------------------------------

-- Depth first
public export
dfs : SearchTree -> List Solution
dfs (Sol sols) = [sols]
dfs (Node _ st) = [ s | t <- st, s <- dfs t]

-- Breath first
public export
bfs : SearchTree -> List Solution
bfs t = trav [t]
    where trav : List SearchTree -> List Solution
          trav [] = []
          trav ((Sol x) :: q) = x :: trav q
          trav ((Node _ st)  :: q) = trav (q ++ st)

-- Limited depth first
public export
limitedDfs : Int -> SearchTree -> List Solution
limitedDfs _ (Sol sols)  = [sols]
limitedDfs 0 _           = []
limitedDfs n (Node _ st) = [ s | t <- st, s <- limitedDfs (n-1) t]

-----------------------------------------------------------------------
--- Syntax for the DSL ---
-----------------------------------------------------------------------

-- Note: I might actually want to do something here like apply some
-- function to "h" and "b" here in order to supply some more syntatic
-- sugar.

-- syntax [h] ":-" [b] = (h, b)
infix 10 :-
export
(:-) : Term -> Terms -> Clause
head :- body = (head, body)

public export
empty : List Term
empty = []

-- syntax [p] "." = (p, empty)
-- suffix 10 .

export
data ProgramElement = Term | Clause

export
interface InterpProgramElement type (0 p : ProgramElement) | type where
  toClause : type -> Clause

export
InterpProgramElement Term Term where
  toClause p = (p, empty)

export
InterpProgramElement Clause Clause where
  toClause = id

public export
data DSLProgram : Type where
  Nil : DSLProgram
  (::) : InterpProgramElement type p => type -> DSLProgram -> DSLProgram

export
unDSL : DSLProgram -> Program
unDSL [] = []
unDSL (element :: elements) = toClause element :: unDSL elements

-- syntax "?-" [q] = nub $ bfs $ makeReportTree program [q]

prefix 10 ?-
export
(?-) : {auto program : Program} -> Term -> List Solution
(?-) {program} q = nub $ bfs $ makeReportTree program [q]


module Prelude.Either

||| Keep the payloads of all Left constructors in a list of Eithers
public export
lefts : List (Either a b) -> List a
lefts []      = []
lefts (x::xs) =
  case x of
    Left  l => l :: lefts xs
    Right r => lefts xs

||| Keep the payloads of all Right constructors in a list of Eithers
public export
rights : List (Either a b) -> List b
rights []      = []
rights (x::xs) =
  case x of
    Left  l => rights xs
    Right r => r :: rights xs

||| Split a list of Eithers into a list of the left elements and a list of the right elements
public export
partitionEithers : List (Either a b) -> (List a, List b)
partitionEithers l = (lefts l, rights l)


module Core.Hash

import Core.CaseTree
import Core.TT
import Core.CompileExpr

import Data.List
import Data.List1
import Libraries.Data.List.Lazy
import Data.String
import Libraries.Data.String.Iterator
import Data.Vect

%default covering

-- This is so that we can store a hash of the interface in ttc files, to avoid
-- the need for reloading modules when no interfaces have changed in imports.
-- As you can probably tell, I know almost nothing about writing good hash
-- functions, so better implementations will be very welcome...

public export
interface Hashable a where
  hash : a -> Int
  hashWithSalt : Int -> a -> Int

  hash = hashWithSalt 5381
  hashWithSalt h i = h * 33 + hash i

infixl 5 `hashWithSalt`

export
Hashable Int where
  hash = id

export
Hashable Integer where
  hash = fromInteger

export
Hashable Nat where
  hash = cast

export
Hashable Char where
  hash = cast

export
Hashable a => Hashable (Vect n a) where
  hashWithSalt h [] = abs h
  hashWithSalt h (x :: xs) = hashWithSalt (h * 33 + hash x) xs

export
Hashable a => Hashable (List a) where
  hashWithSalt h [] = abs h
  hashWithSalt h (x :: xs) = hashWithSalt (h * 33 + hash x) xs

export
Hashable a => Hashable (List1 a) where
  hashWithSalt h xxs = hashWithSalt (h * 33 + hash (head xxs)) (tail xxs)

export
Hashable a => Hashable (Maybe a) where
  hashWithSalt h Nothing = abs h
  hashWithSalt h (Just x) = hashWithSalt h x

export
Hashable a => Hashable b => Hashable (a, b) where
  hashWithSalt h (x, y) = h `hashWithSalt` x `hashWithSalt` y

export
Hashable String where
  hashWithSalt h = String.Iterator.foldl hashWithSalt h

export
Hashable Double where
  hash = hash . show

export
Hashable Namespace where
  hashWithSalt h ns = hashWithSalt h (unsafeUnfoldNamespace ns)

export
Hashable Name where
  hashWithSalt h (MN s _) = hashWithSalt h s
  hashWithSalt h (DN _ n) = hashWithSalt h n
  hashWithSalt h (NS ns n) = hashWithSalt (hashWithSalt h ns) n
  hashWithSalt h (Resolved i) = hashWithSalt h i
  hashWithSalt h n = hashWithSalt h (show n)

export
Hashable RigCount where
  hashWithSalt h = elimSemi
                     (hashWithSalt h 0)
                     (hashWithSalt h 1)
                     (const $ hashWithSalt h 2)

export
Hashable t => Hashable (PiInfo t) where
  hashWithSalt h Implicit = hashWithSalt h 0
  hashWithSalt h Explicit = hashWithSalt h 1
  hashWithSalt h AutoImplicit = hashWithSalt h 2
  hashWithSalt h (DefImplicit t) = h `hashWithSalt` 3 `hashWithSalt` t

export
Hashable ty => Hashable (Binder ty) where
  hashWithSalt h (Lam _ c p ty)
      = h `hashWithSalt` 0 `hashWithSalt` c `hashWithSalt` p `hashWithSalt` ty
  hashWithSalt h (Let _ c val ty)
      = h `hashWithSalt` 1 `hashWithSalt` c `hashWithSalt` val `hashWithSalt` ty
  hashWithSalt h (Pi _ c p ty)
      = h `hashWithSalt` 2 `hashWithSalt` c `hashWithSalt` p `hashWithSalt` ty
  hashWithSalt h (PVar _ c p ty)
      = h `hashWithSalt` 3 `hashWithSalt` c `hashWithSalt` p `hashWithSalt` ty
  hashWithSalt h (PLet _ c val ty)
      = h `hashWithSalt` 4 `hashWithSalt` c `hashWithSalt` val `hashWithSalt` ty
  hashWithSalt h (PVTy _ c ty)
      = h `hashWithSalt` 5 `hashWithSalt` c `hashWithSalt` ty

Hashable (Var vars) where
  hashWithSalt h (MkVar {i} _) = hashWithSalt h i

mutual
  export
  Hashable (Term vars) where
    hashWithSalt h (Local fc x idx y)
        = h `hashWithSalt` 0 `hashWithSalt` idx
    hashWithSalt h (Ref fc x name)
        = h `hashWithSalt` 1 `hashWithSalt` name
    hashWithSalt h (Meta fc x y xs)
        = h `hashWithSalt` 2 `hashWithSalt` y `hashWithSalt` xs
    hashWithSalt h (Bind fc x b scope)
        = h `hashWithSalt` 3 `hashWithSalt` b `hashWithSalt` scope
    hashWithSalt h (App fc fn arg)
        = h `hashWithSalt` 4 `hashWithSalt` fn `hashWithSalt` arg
    hashWithSalt h (As fc _ nm pat)
        = h `hashWithSalt` 5 `hashWithSalt` nm `hashWithSalt` pat
    hashWithSalt h (TDelayed fc x y)
        = h `hashWithSalt` 6 `hashWithSalt` y
    hashWithSalt h (TDelay fc x t y)
        = h `hashWithSalt` 7 `hashWithSalt` t `hashWithSalt` y
    hashWithSalt h (TForce fc r x)
        = h `hashWithSalt` 8 `hashWithSalt` x
    hashWithSalt h (PrimVal fc c)
        = h `hashWithSalt` 9 `hashWithSalt` (show c)
    hashWithSalt h (Erased fc _)
        = hashWithSalt h 10
    hashWithSalt h (TType fc)
        = hashWithSalt h 11

  export
  Hashable Pat where
    hashWithSalt h (PAs fc nm pat)
        = h `hashWithSalt` 0 `hashWithSalt` nm `hashWithSalt` pat
    hashWithSalt h (PCon fc x tag arity xs)
        = h `hashWithSalt` 1 `hashWithSalt` x `hashWithSalt` xs
    hashWithSalt h (PTyCon fc x arity xs)
        = h `hashWithSalt` 2 `hashWithSalt` x `hashWithSalt` xs
    hashWithSalt h (PConst fc c)
        = h `hashWithSalt` 3 `hashWithSalt` (show c)
    hashWithSalt h (PArrow fc x s t)
        = h `hashWithSalt` 4 `hashWithSalt` s `hashWithSalt` t
    hashWithSalt h (PDelay fc r t p)
        = h `hashWithSalt` 5 `hashWithSalt` t `hashWithSalt` p
    hashWithSalt h (PLoc fc x)
        = h `hashWithSalt` 6 `hashWithSalt` x
    hashWithSalt h (PUnmatchable fc x)
        = h `hashWithSalt` 7 `hashWithSalt` x

  export
  Hashable (CaseTree vars) where
    hashWithSalt h (Case idx x scTy xs)
        = h `hashWithSalt` 0 `hashWithSalt` idx `hashWithSalt` xs
    hashWithSalt h (STerm _ x)
        = h `hashWithSalt` 1 `hashWithSalt` x
    hashWithSalt h (Unmatched msg)
        = h `hashWithSalt` 2
    hashWithSalt h Impossible
        = h `hashWithSalt` 3

  export
  Hashable (CaseAlt vars) where
    hashWithSalt h (ConCase x tag args y)
        = h `hashWithSalt` 0 `hashWithSalt` x `hashWithSalt` args
            `hashWithSalt` y
    hashWithSalt h (DelayCase t x y)
        = h `hashWithSalt` 2 `hashWithSalt` (show t)
            `hashWithSalt` (show x) `hashWithSalt` y
    hashWithSalt h (ConstCase x y)
        = h `hashWithSalt` 3 `hashWithSalt` (show x) `hashWithSalt` y
    hashWithSalt h (DefaultCase x)
        = h `hashWithSalt` 4 `hashWithSalt` x

export
Hashable CFType where
  hashWithSalt h = \case
    CFUnit =>
      h `hashWithSalt` 0
    CFInt =>
      h `hashWithSalt` 1
    CFUnsigned8 =>
      h `hashWithSalt` 2
    CFUnsigned16 =>
      h `hashWithSalt` 3
    CFUnsigned32 =>
      h `hashWithSalt` 4
    CFUnsigned64 =>
      h `hashWithSalt` 5
    CFString =>
      h `hashWithSalt` 6
    CFDouble =>
      h `hashWithSalt` 7
    CFChar =>
      h `hashWithSalt` 8
    CFPtr =>
      h `hashWithSalt` 9
    CFGCPtr =>
      h `hashWithSalt` 10
    CFBuffer =>
      h `hashWithSalt` 11
    CFWorld =>
      h `hashWithSalt` 12
    CFFun a b =>
      h `hashWithSalt` 13 `hashWithSalt` a `hashWithSalt` b
    CFIORes a =>
      h `hashWithSalt` 14 `hashWithSalt` a
    CFStruct n fs =>
      h `hashWithSalt` 15 `hashWithSalt` n `hashWithSalt` fs
    CFUser n xs =>
      h `hashWithSalt` 16 `hashWithSalt` n `hashWithSalt` xs
    CFInt8 =>
      h `hashWithSalt` 17
    CFInt16 =>
      h `hashWithSalt` 18
    CFInt32 =>
      h `hashWithSalt` 19
    CFInt64 =>
      h `hashWithSalt` 20
    CFForeignObj =>
      h `hashWithSalt` 21
    CFInteger =>
      h `hashWithSalt` 22

export
Hashable Constant where
  hashWithSalt h = \case
    I i =>
      h `hashWithSalt` 0 `hashWithSalt` i
    BI x =>
      h `hashWithSalt` 1 `hashWithSalt` x
    B8 x =>
      h `hashWithSalt` 2 `hashWithSalt` x
    B16 x =>
      h `hashWithSalt` 3 `hashWithSalt` x
    B32 x =>
      h `hashWithSalt` 4 `hashWithSalt` x
    B64 x =>
      h `hashWithSalt` 5 `hashWithSalt` x
    Str x =>
      h `hashWithSalt` 6 `hashWithSalt` x
    Ch x =>
      h `hashWithSalt` 7 `hashWithSalt` x
    Db x =>
      h `hashWithSalt` 8 `hashWithSalt` x

    WorldVal =>
      h `hashWithSalt` 9

    IntType =>
      h `hashWithSalt` 10
    IntegerType =>
      h `hashWithSalt` 11
    Bits8Type =>
      h `hashWithSalt` 12
    Bits16Type =>
      h `hashWithSalt` 13
    Bits32Type =>
      h `hashWithSalt` 14
    Bits64Type =>
      h `hashWithSalt` 15
    StringType =>
      h `hashWithSalt` 16
    CharType =>
      h `hashWithSalt` 17
    DoubleType =>
      h `hashWithSalt` 18
    WorldType =>
      h `hashWithSalt` 19

    I8 x => h `hashWithSalt` 20 `hashWithSalt` x
    I16 x => h `hashWithSalt` 21 `hashWithSalt` x
    I32 x => h `hashWithSalt` 22 `hashWithSalt` x
    I64 x => h `hashWithSalt` 23 `hashWithSalt` x

    Int8Type => h `hashWithSalt` 24
    Int16Type => h `hashWithSalt` 25
    Int32Type => h `hashWithSalt` 26
    Int64Type => h `hashWithSalt` 27

export
Hashable LazyReason where
  hashWithSalt h = \case
    LInf => h `hashWithSalt` 0
    LLazy => h `hashWithSalt` 1
    LUnknown => h `hashWithSalt` 2

export
Hashable (PrimFn arity) where
  hashWithSalt h = \case
    Add ty =>
      h `hashWithSalt` 0 `hashWithSalt` ty
    Sub ty =>
      h `hashWithSalt` 1 `hashWithSalt` ty
    Mul ty =>
      h `hashWithSalt` 2 `hashWithSalt` ty
    Div ty =>
      h `hashWithSalt` 3 `hashWithSalt` ty
    Mod ty =>
      h `hashWithSalt` 4 `hashWithSalt` ty
    Neg ty =>
      h `hashWithSalt` 5 `hashWithSalt` ty
    ShiftL ty =>
      h `hashWithSalt` 6 `hashWithSalt` ty
    ShiftR ty =>
      h `hashWithSalt` 7 `hashWithSalt` ty

    BAnd ty =>
      h `hashWithSalt` 8 `hashWithSalt` ty
    BOr ty =>
      h `hashWithSalt` 9 `hashWithSalt` ty
    BXOr ty =>
      h `hashWithSalt` 10 `hashWithSalt` ty

    LT ty =>
      h `hashWithSalt` 11 `hashWithSalt` ty
    LTE ty =>
      h `hashWithSalt` 12 `hashWithSalt` ty
    EQ ty =>
      h `hashWithSalt` 13 `hashWithSalt` ty
    GTE ty =>
      h `hashWithSalt` 14 `hashWithSalt` ty
    GT ty =>
      h `hashWithSalt` 15 `hashWithSalt` ty

    StrLength =>
      h `hashWithSalt` 16
    StrHead =>
      h `hashWithSalt` 17
    StrTail =>
      h `hashWithSalt` 18
    StrIndex =>
      h `hashWithSalt` 19
    StrCons =>
      h `hashWithSalt` 20
    StrAppend =>
      h `hashWithSalt` 21
    StrReverse =>
      h `hashWithSalt` 22
    StrSubstr =>
      h `hashWithSalt` 23

    DoubleExp =>
      h `hashWithSalt` 24
    DoubleLog =>
      h `hashWithSalt` 25
    DoubleSin =>
      h `hashWithSalt` 26
    DoubleCos =>
      h `hashWithSalt` 27
    DoubleTan =>
      h `hashWithSalt` 28
    DoubleASin =>
      h `hashWithSalt` 29
    DoubleACos =>
      h `hashWithSalt` 30
    DoubleATan =>
      h `hashWithSalt` 31
    DoubleSqrt =>
      h `hashWithSalt` 32
    DoubleFloor =>
      h `hashWithSalt` 33
    DoubleCeiling =>
      h `hashWithSalt` 34

    Cast f t =>
      h `hashWithSalt` 35 `hashWithSalt` f `hashWithSalt` t
    BelieveMe =>
      h `hashWithSalt` 36
    Crash =>
      h `hashWithSalt` 37

    DoublePow =>
      h `hashWithSalt` 38

export
Hashable ConInfo where
  hashWithSalt h = \case
    DATACON => h `hashWithSalt` 0
    TYCON => h `hashWithSalt` 1
    NIL => h `hashWithSalt` 2
    CONS => h `hashWithSalt` 3
    ENUM => h `hashWithSalt` 4
    NOTHING => h `hashWithSalt` 5
    JUST => h `hashWithSalt` 6
    RECORD => h `hashWithSalt` 7
    ZERO => h `hashWithSalt` 8
    SUCC => h `hashWithSalt` 9

mutual
  export
  Hashable NamedCExp where
    hashWithSalt h = \case
      NmLocal fc n =>
        h `hashWithSalt` 0 `hashWithSalt` n
      NmRef fc n =>
        h `hashWithSalt` 1 `hashWithSalt` n
      NmLam fc x rhs =>
        h `hashWithSalt` 2 `hashWithSalt` x `hashWithSalt` rhs
      NmLet fc x val rhs =>
        h `hashWithSalt` 3 `hashWithSalt` x `hashWithSalt` val `hashWithSalt` rhs
      NmApp fc f xs =>
        h `hashWithSalt` 4 `hashWithSalt` f `hashWithSalt` xs
      NmCon fc cn ci t args =>
        h `hashWithSalt` 5 `hashWithSalt` cn `hashWithSalt` ci `hashWithSalt` t `hashWithSalt` args
      NmOp fc fn args =>
        h `hashWithSalt` 6 `hashWithSalt` fn `hashWithSalt` args
      NmExtPrim fc p args =>
        h `hashWithSalt` 7 `hashWithSalt` p `hashWithSalt` args
      NmForce fc r x =>
        h `hashWithSalt` 8 `hashWithSalt` r `hashWithSalt` x
      NmDelay fc r x =>
        h `hashWithSalt` 9 `hashWithSalt` r `hashWithSalt` x
      NmConCase fc sc alts df =>
        h `hashWithSalt` 10 `hashWithSalt` sc `hashWithSalt` alts `hashWithSalt` df
      NmConstCase fc sc alts df =>
        h `hashWithSalt` 11 `hashWithSalt` sc `hashWithSalt` alts `hashWithSalt` df
      NmPrimVal fc c =>
        h `hashWithSalt` 12 `hashWithSalt` c
      NmErased fc =>
        h `hashWithSalt` 13
      NmCrash fc msg =>
        h `hashWithSalt` 14 `hashWithSalt` msg

  export
  Hashable NamedConAlt where
    hashWithSalt h (MkNConAlt n ci tag args rhs) =
      h `hashWithSalt` n `hashWithSalt` ci `hashWithSalt` tag `hashWithSalt` args `hashWithSalt` rhs

  export
  Hashable NamedConstAlt where
    hashWithSalt h (MkNConstAlt c rhs) =
      h `hashWithSalt` c `hashWithSalt` rhs

export
Hashable NamedDef where
  hashWithSalt h = \case
    MkNmFun args ce =>
      h `hashWithSalt` 0 `hashWithSalt` args `hashWithSalt` ce
    MkNmCon tag arity nt =>
      h `hashWithSalt` 1 `hashWithSalt` tag `hashWithSalt` arity `hashWithSalt` nt
    MkNmForeign ccs fargs cft =>
      h `hashWithSalt` 2 `hashWithSalt` ccs `hashWithSalt` fargs `hashWithSalt` cft
    MkNmError e =>
      h `hashWithSalt` 3 `hashWithSalt` e


module Utils.Either

%default total

export
mapError : (a -> c) -> Either a b -> Either c b
mapError f e = either (Left . f) (Right . id) e

export
partitionEithers : List (Either a b) -> (List a, List b)
partitionEithers [] = ([], [])
partitionEithers (Left x :: es)
        = let (ls, rs) = partitionEithers es in
              (x :: ls, rs)
partitionEithers (Right y :: es)
        = let (ls, rs) = partitionEithers es in
              (ls, y :: rs)


||| A small library introducing string positions
||| This can be used as an alternative to unpacking a string into a list of
||| characters
module Data.String.Position

import Data.String

%default total

||| @ValidPosition points to an existing character into
||| @str  its parameter string
||| Do NOT publicly export so that users cannot manufacture arbitrary positions
export
record ValidPosition (str : String) where
  constructor MkValidPosition
  ||| @currentIndex is the valid position into str
  currentIndex    : Int
  ||| @parameterLength is the length of the parameter str
  parameterLength : Int
-- TODO: add invariants?
--  0 isLength : length str === parameterLength
--  0 isIndex  : (0 `LTE` currentIndex, 0 `LT` parameterLength)

||| A position is either a ValidPosition or the end of the string
public export
Position : String -> Type
Position str = Maybe (ValidPosition str)

||| Manufacture a position by checking it is strictly inside the string
export
mkPosition : (str : String) -> Int -> Position str
mkPosition str idx
  = do let len : Int = cast (length str)
       guard (0 <= idx && idx < len)
       pure (MkValidPosition idx len)

||| Move a position (note that we do not need access to the string here)
export
mvPosition : ValidPosition str -> Int -> Position str
mvPosition (MkValidPosition pos len) idx
  = do guard (0 <= idx && idx < len)
       pure (MkValidPosition idx len)

||| Find the initial position inside the input string
export
init : (str : String) -> Position str
init str = mkPosition str 0

||| Move from a valid position to the next position in the string
export
next : ValidPosition str -> Position str
next (MkValidPosition pos len)
  = do let idx = pos + 1
       guard (idx < len)
       pure (MkValidPosition idx len)

||| We can safely access the character at a valid position
export
index : {str : _} -> ValidPosition str -> Char
index pos = assert_total (strIndex str pos.currentIndex)

||| We can successfully uncons the substring starting at a `ValidPosition`.
||| Note that we can use `map uncons pos` to uncons the substring starting
||| a `Position`.
export
uncons : {str : _} -> ValidPosition str -> (Char, Position str)
uncons pos = (index pos, next pos)

||| @span keeps munching characters of the substring starting at a given
||| position as long as the predicate is true of them. It returns the position
||| after the last successful munch.
||| Using `subString` to recover the string selected by `span` should yield
||| the same result as Data.String's `span`. That is to say we should have:
||| ```
||| subString pos (Position.span p pos) = String.span p (subString pos Nothing)
||| ```
export
span : {str : _} -> (Char -> Bool) -> Position str -> Position str
span p pos = do (c, str) <- map uncons pos
                if p c then assert_total (span p str) else pos

||| @count computes the length of the substring one would obtain if one were
||| to filter out characters based on the predicate passed as an argument.
||| It replaces the `length (filter p (fastUnpack str))` pattern.
export
count : {str : _} -> (Char -> Bool) -> Position str -> Nat
count p Nothing = 0
count p (Just pos) =
  if p (index pos)
    then S (assert_total (count p (next pos)))
    else assert_total (count p (next pos))

||| Distance between a starting position and an end one.
||| We assume that the end position is *after* the starting one, otherwise the
||| output may be negative.
export
distance : (start : Position str) ->
           (end : Position str) -> Int
distance Nothing _ = 0
distance (Just pos) pos' =
  let start = pos.currentIndex
      end = maybe pos.parameterLength currentIndex pos'
  in end - start

||| the substring of length `distance start end` that is contained between
||| start (inclusive) and end (exclusive).
export
subString : {str : _} ->
            (start : Position str) ->
            (end : Position str) -> String
subString Nothing pos' = ""
subString (Just pos) pos' =
  let start = pos.currentIndex
      end = maybe pos.parameterLength currentIndex pos'
  in assert_total (strSubstr start (end - start) str)


module Main

StringOrInt : Bool -> Type
StringOrInt x = case x of
  True  => Int
  False => String

getStringOrInt : (x : Bool) -> StringOrInt x
getStringOrInt x = case x of
  True  => 94
  False => "Ninety four"

valToString : (x : Bool) -> StringOrInt x -> String
valToString x val = case x of
  True  => cast val
  False => val

main : IO ()
main = do
  putStrLn (valToString True (getStringOrInt True))
  putStrLn (valToString False (getStringOrInt False))


module Circuits.Idealised.Interp

import Decidable.Equality

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

import Toolkit.Decidable.Informative
import Toolkit.Data.Graph.EdgeBounded
import Toolkit.Data.Graph.EdgeBounded.HasExactDegree.All

import Toolkit.Data.Whole

import Circuits.Idealised.Types
import Circuits.Idealised.Terms

%default total

public export
InterpTy : Ty -> Type
InterpTy TyUnit
  = Unit

InterpTy (TyPort x)
  = Vertex String

InterpTy TyGate
  = Graph String

namespace Environments
  public export
  data Env : List (Ty, Usage) -> Type where
    Empty  : Env Nil
    Extend : InterpTy type -> Env rest -> Env ((type,usage) :: rest)

  export
  lookup : (envO  : Env old )
        -> (prf   : Elem (type,FREE) old)
        -> (which : Use old prf new)
        -> (Env new, InterpTy type)
  lookup (Extend x y) Here H = (Extend x y, x)
  lookup (Extend y z) (There later) (T x)
    = let (newEnv, res) = lookup z later x
      in (Extend y newEnv, res)

namespace Result

  public export
  data Result : (context : List (Ty, Usage))
             -> (type    : Ty)
                        -> Type
    where
      R : (counter : Nat)
       -> (env     : Env new)
       -> (graph   : Graph String)
       -> (result  : InterpTy type)
                  -> Result new type
  public export
  getResult : Result ctxt type -> InterpTy type
  getResult (R _ _ _ res) = res

  public export
  data GetGraph : (res : Result ctxt type) -> Graph String -> Type
    where
      G : (g : Graph String) -> GetGraph (R c e g r) g

  public export
  getGraph : (r : Result ctxt type) -> DPair (Graph String) (GetGraph r)
  getGraph (R counter env graph result) = MkDPair graph (G graph)

interp : (env     : Env old)
      -> (counter : Nat)
      -> (graph   : Graph String)
      -> (term    : Term old type new)
                 -> Result new type

-- [ NOTE ]
--
-- As normal
interp env counter graph (Var prf x)
  = let (newEnv, res) = lookup env prf x
    in R counter newEnv graph res

-- [ NOTE ]
--
-- As normal
interp env counter graph (Seq x y)
  = let R c1 e1 g1 x' = interp env counter graph x
    in interp e1 c1 ((merge g1 x')) y

-- [ NOTE ]
--
-- Ports are leaf nodes in the graph.
interp env counter graph (NewSignal flow dtype body)

  = let n = vertexFromFlow flow (S counter) 1

    in let env' = Extend n env

    in let g'   = insertNode n graph

    in interp env' (S counter) g' body

  where
    vertexFromFlow : Direction -> Nat -> Nat -> Vertex String
    vertexFromFlow INPUT  = driver  ("INPUT : "  <+> show dtype )
    vertexFromFlow OUTPUT = catcher ("OUTPUT : " <+> show dtype )

-- [ NOTE ]
--
-- Wires are internal connected nodes in the graph
interp env counter graph (Wire dtype body)

  = let c1 = S counter
    in let c2 = S c1

    in let chanIN  = node ("CHAN_IN : "  <+> show dtype) c1 1 1
    in let chanOUT = node ("CHAN_OUT : " <+> show dtype) c2 1 1

    in interp (Extend chanIN (Extend chanOUT env))
              c2
              (updateWith graph
                          [chanIN, chanOUT]
                          [(MkPair (ident chanIN) (ident chanOUT))])
              body

-- [ NOTE ]
--
-- Multiplexers are internal nodes connected to other nodes.
interp env counter graph (Mux o c l r)

  = let c1 = S counter

    in let mux = node "MUX" c1 3 1

    in let R c2 e1 g1 vo = interp env c1 graph  o
    in let R c3 e2 g2 vc = interp e1  c2 g1     c
    in let R c4 e3 g3 vl = interp e2  c2 g2     l
    in let R c5 e4 g4 vr = interp e3  c3 g3     r

    in let mg = fromLists [mux]
                          [ (ident vl,  ident mux)
                          , (ident vr,  ident mux)
                          , (ident vc,  ident mux)
                          , (ident mux, ident vo)
                          ]
    in R c5 e4 g4 mg

-- [ NOTE ]
--
-- Multiplexers are internal nodes connected to other nodes.
interp env counter graph (Dup a b i)
  = let c1 = S counter

    in let dup = node "DUP" c1 1 2

    in let R c2 e1 g1 va = interp env c1 graph a
    in let R c3 e2 g2 vb = interp e1  c2 g1    b
    in let R c4 e3 g3 vi = interp e2  c3 g2    i

    in let dg = fromLists [dup]
                          [ MkPair (ident vi)   (ident dup)
                          , MkPair (ident dup) (ident va)
                          , MkPair (ident dup) (ident vb)
                          ]

    in R c4 e3 g3 dg

-- [ NOTE ]
--
-- internal nodes connected to other nodes.
interp env counter graph (Not output input)
  = let c1 = S counter

    in let vnot = node "NOT" c1 1 1

    in let R c2 e1 g1 vo = interp env counter graph output
    in let R c3 e2 g2 vi = interp e1  c2      g1    input

    in let ng = fromLists [vnot]
                          [ MkPair (ident vnot) (ident vo)
                          , MkPair (ident vi)   (ident vnot)
                          ]
    in R c3 e2 g2 ng

-- [ NOTE ]
--
-- internal nodes connected to other nodes.
interp env counter graph (Gate k output inputA inputB)
   = let c1 = S counter

    in let bgate = node (show k) c1 2 1

    in let R c2 e1 g1 vo = interp env c1 graph output
    in let R c3 e2 g2 vl = interp e1  c2 g1    inputA
    in let R c4 e3 g3 vr = interp e2  c3 g2    inputB

    in let bg = fromLists [bgate]
                          [ MkPair (ident bgate) (ident vo)
                          , MkPair (ident vl) (ident bgate)
                          , MkPair (ident vr) (ident bgate)
                          ]

    in R c4 e3 g3 bg

-- [ NOTE ]
--
-- internal nodes connected to other nodes.
interp env counter graph (IndexSingleton o i)
  = let c1 = S counter

    in let vnot = node "IDX_SINGLETON" c1 1 1

    in let R c2 e1 g1 vo = interp env c1  graph o
    in let R c3 e2 g2 vi = interp e1  c2  g1    i

    in let ng = fromLists [vnot]
                          [ MkPair (ident vnot) (ident vo)
                          , MkPair (ident vi)   (ident vnot)
                          ]
    in R c3 e2 g2 ng

-- [ NOTE ]
--
-- internal nodes connected to other nodes.
interp env counter graph (IndexEdge p idx oused ofree i)

  = let c1 = S counter

    in let dup = node "IDX_EDGE" c1 1 2

    in let R c2 e1 g1 va = interp env c1 graph oused
    in let R c3 e2 g2 vb = interp e1  c2 g1    ofree
    in let R c4 e3 g3 vi = interp e2  c3 g2    i

    in let dg = fromLists [dup]
                          [ MkPair (ident vi)  (ident dup)
                          , MkPair (ident dup) (ident va)
                          , MkPair (ident dup) (ident vb)
                          ]

    in R c4 e3 g3 dg

-- [ NOTE ]
--
-- internal nodes connected to other nodes.
interp env counter graph (IndexSplit p idx used freeA freeB i)

  = let c1 = S counter

    in let dup = node "IDX_SPLIT" c1 1 3

    in let R c2 e1 g1 vu = interp env c1 graph used
    in let R c3 e2 g2 va = interp e1  c2 g1    freeA
    in let R c4 e3 g3 vb = interp e2  c3 g2    freeB
    in let R c5 e4 g4 vi = interp e3  c4 g3    i

    in let dg = fromLists [dup]
                          [ MkPair (ident vi)  (ident dup)
                          , MkPair (ident dup) (ident vu)
                          , MkPair (ident dup) (ident va)
                          , MkPair (ident dup) (ident vb)
                          ]

    in R c5 e4 g4 dg

-- [ NOTE ]
--
-- internal nodes connected to other nodes.
interp env counter graph (Merge2L2V output inputA inputB)

  = let c1 = S counter

    in let bgate = node "MERGE_2L2V" c1 2 1

    in let R c2 e1 g1 vo = interp env c1 graph output
    in let R c3 e2 g2 vl = interp e1  c2 g1    inputA
    in let R c4 e3 g3 vr = interp e2  c3 g2    inputB

    in let bg = fromLists [bgate]
                          [ MkPair (ident bgate) (ident vo)
                          , MkPair (ident vl) (ident bgate)
                          , MkPair (ident vr) (ident bgate)
                          ]

    in R c4 e3 g3 bg

-- [ NOTE ]
--
-- internal nodes connected to other nodes.
interp env counter graph (Merge2V2V prf output inputA inputB)

  = let c1 = S counter

    in let bgate = node "MERGE_2V2V" c1 2 1

    in let R c2 e1 g1 vo = interp env c1 graph output
    in let R c3 e2 g2 vl = interp e1  c2 g1    inputA
    in let R c4 e3 g3 vr = interp e2  c3 g2    inputB

    in let bg = fromLists [bgate]
                          [ MkPair (ident bgate) (ident vo)
                          , MkPair (ident vl) (ident bgate)
                          , MkPair (ident vr) (ident bgate)
                          ]

    in R c4 e3 g3 bg

-- [ NOTE ]
--
-- internal nodes connected to other nodes.
interp env counter graph (MergeSingleton output input)
  = let c1 = S counter

    in let vnot = node "MERGE_SINGLETON" c1 1 1

    in let R c2 e1 g1 vo = interp env counter graph output
    in let R c3 e2 g2 vi = interp e1  c2      g1    input

    in let ng = fromLists [vnot]
                          [ MkPair (ident vnot) (ident vo)
                          , MkPair (ident vi)   (ident vnot)
                          ]
    in R c3 e2 g2 ng

interp env counter graph (Stop x)
  = R counter Empty graph MkUnit

public export
data ValidGraph : Graph type -> Type where
  IsValid : HasExactDegrees vs es
         -> ValidGraph (MkGraph vs es)

export
validGraph : {type : Type}
          -> (g    : Graph type)
                  -> DecInfo (HasExactDegree.Error type)
                             (ValidGraph g)
validGraph (MkGraph nodes edges) with (hasExactDegrees nodes edges)
  validGraph (MkGraph nodes edges) | (Yes prf)
    = Yes (IsValid prf)
  validGraph (MkGraph nodes edges) | (No msg contra)
    = No msg (\(IsValid prf) => contra prf)

public export
data Valid : (res  : Result ctxt TyUnit)
                  -> Type
  where
    D : {res : Result ctxt TyUnit}
     -> (g   : Graph String)
            -> GetGraph res g
            -> ValidGraph g
            -> Valid res

isValid : (r : Result ctxt TyUnit)
            -> DecInfo (Graph String, HasExactDegree.Error String)
                       (Valid r)
isValid r with (getGraph r)
  isValid (R c e g r) | (MkDPair g (G g)) with (validGraph g)
    isValid (R c e g r) | (MkDPair g (G g)) | (Yes prf)
      = Yes (D g (G g) prf)
    isValid (R c e g r) | (MkDPair g (G g)) | (No msg contra)
      = No (g,msg) (\(D graph (G graph) prf) => contra prf)

public export
data Run : (term : Term Nil TyUnit Nil) -> Type where
  R : (prf  : Valid (interp Empty Z (MkGraph Nil Nil) term))
           -> Run term

public export
getGraph : Run term -> DPair (Graph String) ValidGraph
getGraph (R (D g x y)) = MkDPair g y

export
run : (term : Term Nil TyUnit Nil)
           -> DecInfo (Graph String, HasExactDegree.Error String) (Run term)
run term with (isValid (interp Empty Z empty term))
  run term | (Yes prf) = Yes (R prf)
  run term | (No msg contra) = No msg (\(R prf) => contra prf)

export
runIO : (term : Term Nil TyUnit Nil)
             -> IO (Either (Graph String, HasExactDegree.Error String)
                           (Run term))
runIO term = pure $ (asEither (run term))

-- [ EOF ]


{-
  An embeded DSL for prolog-style logic programming, based on the
  pure-prolog parser and interpreter originally written in Haskell

  Original author: Ken Friis Larsen      <kflarsen@diku.dk>
  Translated into Idris by Nathan Bedell <nbedell@tulane.edu>
-}

module Ast

%access public export

----------------------------------------------------------------------
-- Abstract Syntax Tree
----------------------------------------------------------------------

Atom : Type
Atom = String

Variable : Type
Variable = String

data Term : Type where
 Var : Variable -> Term
 Comp : Atom -> List Term -> Term

and' : List Bool -> Bool
and' [] = True
and' xs = foldr1 p xs
  where p True True = True
        p x y = False

implementation Eq Term where
  (Var x) == (Var y) = assert_total $ x == y
  (Comp x y) == (Comp x' y') = assert_total $ (x == x') && (and' $ Prelude.List.zipWith (==) y y')
  x == y = assert_total $ False

implementation Show Term where
  show (Var s) = s
  show (Comp s []) = s
  show (Comp s xs) = s ++ "(" ++ (concat $ intersperse "," (map show xs)) ++ ")"

Terms : Type
Terms = List Term

Clause : Type
Clause = (Term, Terms) -- head and body

Clauses : Type
Clauses = List Clause

Goal : Type
Goal = List Term

Program : Type
Program = Clauses


{- Ctrl-Alt-A: Add definition
   Ctrl-Alt-C: Case split
   Ctrl-Alt-D: Documentation
   Ctrl-Alt-L: Lift hole: Lifts a hole to the top level as a new function definition
   Ctrl-Alt-M: Match: Replaces a hole with a case expression that matches on an intermediate result
   Ctrl-Alt-R: Reloads and typechecks the current buffer
   Ctrl-Alt-S: Search
   Ctrl-Alt-T: Type-check: Displays the type under cursor -}
module Main

import Data.SortedMap
import Data.SortedSet
import Control.Monad.State
import Control.Monad.Reader
import Prelude.Traversable
import Text.PrettyPrint.WL

interface Pretty p where
  doc : p -> Doc

(Pretty p1, Pretty p2) => Pretty (Pair p1 p2) where
  doc (p1, p2) = brackets (doc p1 |++| text "," |++| doc p2)

Pretty p => Pretty (List p) where
  doc ps = vsep (map doc ps)

Pretty Nat where
  doc = text . show

-- Arity
data Arity = MkArity Nat
Eq Arity where
  (MkArity a1) == (MkArity a2) = a1 == a2
Show Arity where
  show (MkArity a) = show a
Pretty Arity where
  doc (MkArity a) = text $ show a

-- Var
data Var = MkVar String
Eq Var where
  (MkVar v1) == (MkVar v2) = v1 == v2
Ord Var where
  compare (MkVar v1) (MkVar v2) = compare v1 v2
Pretty Var where
  doc (MkVar v) = text v

--  Constructor
data Cnst
  = MkCnst String
  | PCnst String
Eq Cnst where
  (MkCnst v1) == (MkCnst v2) = v1 == v2
  (PCnst c1) == (PCnst c2) = c1 == c2
  _ == _ = False
Show Cnst where
  show (MkCnst v) = show v
  show (PCnst v) = show v
Ord Cnst where
  compare (MkCnst c1) (MkCnst c2) = compare c1 c2
  compare (PCnst c1) (PCnst c2) = compare c1 c2
  compare (MkCnst c1) (PCnst c2) = LT
  compare (PCnst c1) (MkCnst c2) = GT
Pretty Cnst where
  doc (MkCnst v) = text v
  doc (PCnst c) = text c

-- Expr
data Expr
  = Call Cnst (List Var)
  | Pap Cnst (List Var)
  | App Var Var
  | Ctor Int (List Var)
  | Proj Int Var
  | Reset Var
  | Reuse Var Int (List Var)
Eq Expr where
  (Call c1 vs1)     == (Call c2 vs2)      = c1 == c2 && vs1 == vs2
  (Pap c1 vs1)      == (Pap c2 vs2)       = c1 == c2 && vs1 == vs2
  (App v11 v12)     == (App v21 v22)      = v11 == v21 && v21 == v22
  (Ctor i1 v1)      == (Ctor i2 v2)       = i1 == i2 && v1 == v2
  (Proj i1 v1)      == (Proj i2 v2)       = i1 == i2 && v1 == v2
  (Reset v1)        == (Reset v2)         = v1 == v2
  (Reuse v1 i1 vs1) == (Reuse v2 i2 vs2)  = v1 == v2 && i1 == i2 && vs1 == vs2
  _ == _ = False
Pretty Expr where
  doc e = case e of
    Call c vs     => hsep (doc c :: map doc vs)
    Pap c vs      => hsep (text "pap" :: doc c :: map doc vs)
    App x y       => hsep [doc x, doc y]
    Ctor i vs     => hsep (text "ctor" :: text (show i) :: map doc vs)
    Proj i v      => hsep [text "proj", text (show i), doc v]
    Reset v       => hsep [text "reset", doc v]
    Reuse v i vs  => hsep (text "reuse" :: doc v :: text "in" :: text "ctor" :: text (show i) :: map doc vs)

data Phase
  = Pure
  | RC -- Reset/Reuse
  | OC -- Owned Const

data FnBody : Phase -> Type where
  Ret  : Var                            -> FnBody p
  Let  : Var -> Expr -> FnBody p        -> FnBody p
  Case : Var -> List (Arity, FnBody p)  -> FnBody p
  Inc  : Var -> FnBody p                -> FnBody p
  Dec  : Var -> FnBody p                -> FnBody p
  LetC : Cnst -> Cnst -> FnBody p       -> FnBody p -- let cO := c
Pretty (FnBody p) where
  doc e = case e of
    Ret  v      => text "ret" |++| doc v
    Let  v e b  => text "let" |++| doc v |++| text "=" |++| doc e |+| text ";" |$$| doc b
    Case v alts => text "case" |++| doc v |$$| indent 2 (vsep (map (\(a,b) => parens (indent 1 (doc b) |+| text " ")) alts))
    Inc  v b    => text "inc" |++| doc v |++| text ";" |$$| doc b
    Dec  v b    => text "dec" |++| doc v |++| text ";" |$$| doc b
    LetC o c b  => text "let" |++| doc o |++| text ":=" |$$| doc b

eqFnBody : FnBody p -> FnBody p -> Bool
eqFnBody b1 b2 = case (b1, b2) of
  (Ret v1, Ret v2)              => v1 == v2
  (Let v1 e1 b1, Let v2 e2 b2)  => v1 == v2 && e1 == e2 && eqFnBody b1 b2
  (Case v1 as1, Case v2 as2)    => v1 == v2 && length as1 == length as2 && and (zipWith (\(a1,c1) , (a2,c2) => a1 == a2 && eqFnBody c1 c2) as1 as2)
  (Inc v1 b1, Inc v2 b2)        => v1 == v2 && eqFnBody b1 b2
  (Dec v1 b1, Dec v2 b2)        => v1 == v2 && eqFnBody b1 b2
  (LetC o1 c1 b1, LetC o2 c2 b2) => o1 == o2 && c1 == c2 && eqFnBody b1 b2
  (_, _)                        => False

data Fun : Phase -> Type where
  MkFun : List Var -> FnBody p -> Fun p
Pretty (Fun p) where
  doc (MkFun ps b) = hsep (map doc ps) |++| text "=" |$$| indent 2 (doc b)

data Program : Phase -> Type where
  MkProgram : List (Cnst, Fun p) -> Program p
Pretty (Program p) where
  doc (MkProgram fs) = vsep $ map (\(n,f) => doc n |++| doc f) fs

swapExample : Program Pure
swapExample = MkProgram
  [ ( swap
    , MkFun [xs]
      $ Case xs
        [ (MkArity 0, Ret xs)
        , (MkArity 2, Let t1 (Proj 2 xs)
                    $ Case t1
                      [ (MkArity 0, Ret xs)
                      , (MkArity 2, Let h1 (Proj 1 xs)
                                  $ Let h2 (Proj 2 t1)
                                  $ Let t2 (Proj 2 t1)
                                  $ Let r1 (Ctor 2 [h1, t2])
                                  $ Let r2 (Ctor 2 [h2, r1])
                                  $ Ret r2
  )])])]
  where
    swap = MkCnst "swap"
    xs = MkVar "xs"
    t1 = MkVar "t1"
    t2 = MkVar "t2"
    h1 = MkVar "h1"
    h2 = MkVar "h2"
    r1 = MkVar "r1"
    r2 = MkVar "r2"

mapExample : Program Pure
mapExample = MkProgram
  [ ( map
    , MkFun [f, xs]
        $ Case xs
          [ (MkArity 0, Ret xs)
          , (MkArity 2, Let x  (Proj 1 xs)
                      $ Let s  (Proj 2 xs)
                      $ Let y  (App f x)
                      $ Let ys (Call map [f, s])
                      $ Let r  (Ctor 2 [y, ys])
                      $ the (FnBody Pure) (Ret r))
          ]
    )
  ]
  where
    map = MkCnst "map"
    f = MkVar "f"
    x = MkVar "x"
    s = MkVar "s"
    r = MkVar "r"
    y = MkVar "y"
    ys = MkVar "ys"
    xs = MkVar "xs"

-- Originally S
insertReuse : Var -> Arity -> FnBody RC -> FnBody RC
insertReuse w n f = case f of
  Ret x     => Ret x
  Case x fs => Case x (map (\(a, f) => (a, insertReuse w n f)) fs)
  Let x e g => case e of
    Ctor i ys => if MkArity (List.length ys) == n
                    then Let x (Reuse w i ys) g
                    else Let x (Ctor i ys) (insertReuse w n g)
    e' => Let x e (insertReuse w n g)
  Dec x g => Dec x g
  Inc x g => Inc x g
  LetC c1 c2 g => LetC c1 c2 g

-- NewName
data NewNameT : (f : Type -> Type) -> (a : Type) -> Type where
  NewName : StateT Nat m a -> NewNameT m a

unNewName : NewNameT m a -> StateT Nat m a
unNewName (NewName m) = m

runNewNameT : Monad m => NewNameT m a -> m a
runNewNameT (NewName s) = fst <$> runStateT s 0

runNewName : NewNameT Identity a -> a
runNewName = runIdentity . runNewNameT

Functor f => Functor (NewNameT f) where
  map f (NewName s) = NewName (map f s)

(Functor (NewNameT f), Applicative (StateT Nat f)) => Applicative (NewNameT f) where
  pure x                      = NewName (pure x)
  (NewName f) <*> (NewName x) = NewName (f <*> x)

(Monad (StateT Nat f), Applicative (NewNameT f)) => Monad (NewNameT f) where
   (NewName m) >>= k = NewName $ do
     x <- m
     unNewName (k x)

freshVar : Monad m => NewNameT m Var
freshVar = NewName $ do
  n <- get
  put (S n) -- TODO: How to increment nat??? :)
  pure $ MkVar $ "w" ++ show n

freshCnst : Monad m => NewNameT m Cnst
freshCnst = NewName $ do
  n <- get
  put (S n)
  pure $ PCnst ("c" ++ show n)

total
inExpr : Var -> Expr -> Bool
inExpr v e = case e of
  (Call c vs)     => v `elem` vs
  (Pap c vs)      => v `elem` vs
  (App x y)       => x == v || y == v
  (Ctor i vs)     => v `elem` vs
  (Proj i x)      => v == x
  (Reuse x i vs)  => (v `elem` vs) || (v == x)
  (Reset x)       => v == x

-- total : TODO: List is not strong enough
inFBody : Var -> FnBody p -> Bool
inFBody v f = case f of
  Ret  x      => v == x
  Let  x e b  => v == x || (inExpr v e) || (inFBody v b)
  Case x alts => v == x || (foldl (\x,y => x && y) True (map (inFBody v . snd) alts)) -- TODO: How to use and here?
  Inc  x b    => v == x || (inFBody v b)
  Dec  x b    => v == x || (inFBody v b)
  LetC _ _ b  => inFBody v b

insertResetIns : Monad m => Var -> Arity -> FnBody RC -> NewNameT m (FnBody RC)
insertResetIns z n f = do
  w <- freshVar
  let g = insertReuse w n f
  pure $ if not $ eqFnBody g f
            then Let w (Reset z) g
            else f

-- Originally D
insertReset : Monad m => Var -> Arity -> FnBody RC -> NewNameT m (FnBody RC)
insertReset z n f = case f of
  Dec x g   => pure $ Dec x g
  Inc x g   => pure $ Inc x g
  LetC o c g => LetC o c <$> insertReset z n g
  Ret x     => pure $ Ret x
  Case x gs => Case x <$> traverse (\(a, g) => (\v => (a,v)) <$> insertReset z n g) gs
  Let x e g =>
    case ((inExpr z e) || (inFBody z g)) of
      True  => Let x e <$> insertReset z n g
      False => insertResetIns z n f

-- Originally R
resetReuse : Monad m => FnBody Pure -> NewNameT m (FnBody RC)
resetReuse b = case b of
  Ret x     => pure $ Ret x
  Let x e f => Let x e <$> resetReuse f
  LetC o c f => LetC o c <$> resetReuse f
  Case x bs => Case x <$> traverse (\(arity, body) => do
                                        body' <- resetReuse body
                                        (MkPair arity) <$> (insertReset x arity body'))
                                   bs

resetReuseFun : Monad m => Fun Pure -> NewNameT m (Fun RC)
resetReuseFun (MkFun ps body) = MkFun ps <$> resetReuse body

resetReuseProgram : Monad m => Program Pure -> NewNameT m (Program RC)
resetReuseProgram (MkProgram defs) = MkProgram <$> traverse (\(n,f) => (MkPair n <$> resetReuseFun f)) defs

||| Inserts local constant
insertLetCFnBody : Monad m => FnBody p -> NewNameT m (FnBody p) -- TODO: Right phase parameter
insertLetCFnBody b = case b of
  Ret  v => pure $ Ret v
  Let  v e g => case e of
                  (Pap c xs) => do
                    c2 <- freshCnst
                    (LetC c2 c . Let v (Pap c2 xs)) <$> insertLetCFnBody g
                  e0 => Let v e0 <$> insertLetCFnBody g
  Case v alts => Case v <$> traverse (\(a,g) => MkPair a <$> insertLetCFnBody g) alts
  Inc  v g => Inc v <$> insertLetCFnBody g
  Dec  v g => Dec v <$> insertLetCFnBody g
  LetC c1 c2 g => LetC c1 c2 <$> insertLetCFnBody g

insertLetCFun : Monad m => Fun p -> NewNameT m (Fun p)
insertLetCFun (MkFun ps body) = MkFun ps <$> insertLetCFnBody body

insertLetCProgram : Monad m => Program p -> NewNameT m (Program p)
insertLetCProgram (MkProgram defs) = MkProgram <$> traverse (\(n,f) => (MkPair n <$> insertLetCFun f)) defs

{-
Inferring borrowing signatures:
Every function parameter is annotated with Owned or Borrowed

Our heruristing pcollects which parameters and variables should be owned.
We say a parameter should be owned
  - if x or one of its projections is used in a reset
  - is passed to a function that takes an owned reference.
    (This is heuristic and not required for correctness)
-}

data Borrowing
  = Owned
  | Borrowed
Eq Borrowing where
  Owned     == Owned    = True
  Borrowed  == Borrowed = True
  _         == _        = False
Show Borrowing where
  show Owned = "Owned"
  show Borrowing = "Borrowing"
Pretty Borrowing where
  doc Owned     = text "Owned"
  doc Borrowed  = text "Borrowed"

BorrowingMap : Type
BorrowingMap = SortedMap (Cnst, Nat) Borrowing

borrowed : BorrowingMap -> (Cnst, Nat) -> Maybe Borrowing
borrowed b (PCnst _, _) = Just Owned
borrowed b (c,n)        = Data.SortedMap.lookup (c,n) b

constBorrowingParams : BorrowingMap -> Cnst -> List Borrowing
constBorrowingParams bm c
  = map snd
  $ sortBy (\((_, e1), _) , ((_, e2), _) => compare e1 e2)
  $ filter ((c ==) . fst . fst) $ Data.SortedMap.toList bm

number : List a -> List (Nat, a)
number = snd . foldl (\(n, ns), a => (S n, (n,a) :: ns)) (Z, [])

||| Collects the owned variables in a function body
collectOwnedVars : BorrowingMap -> FnBody rc -> SortedSet Var
collectOwnedVars beta b = case b of -- TODO: Unicode Beta
  Let z (Ctor i xs)     f => collectOwnedVars beta f
  Let z (Reuse x i xs)  f => collectOwnedVars beta f
  Let z (Call c xs)     f
    => union (collectOwnedVars beta f)
             (fromList
              (concatMap (\(n,x) => maybe [x]
                                          (\b => if b == Owned then [x] else [])
                                          (borrowed beta (c,n)))
                         (number xs)))
  Let z (Reset x)       f => insert x (collectOwnedVars beta f)
  Let z (App x y)       f => union (collectOwnedVars beta f) (fromList [x, y])
  Let z (Pap cO xs)     f => union (collectOwnedVars beta f) (fromList xs)
  -- TODO: Define a wrapper constant transformation for partial application
  -- as we assume that partially applied functions always have Owned parameters.
  Let z (Proj i x)      f => let c = collectOwnedVars beta f
                             in if contains z c
                                  then insert x c
                                  else c
  Ret x => empty
  Case z alts => foldl1 union $ map (collectOwnedVars beta . snd) alts
  Inc z f => collectOwnedVars beta f
  Dec z f => collectOwnedVars beta f
  LetC o c f => collectOwnedVars beta f

ownedFunctionVar : BorrowingMap -> Fun p -> List (Nat, Borrowing)
ownedFunctionVar beta (MkFun ps body) = snd $ foldl borrowed (0, []) ps
  where
    c : SortedSet Var
    c = collectOwnedVars beta body

    borrowed : (Nat, List (Nat, Borrowing)) -> Var -> (Nat, List (Nat, Borrowing))
    borrowed (n, vs) v = (n + 1, (n, if contains v c then Owned else Borrowed) :: vs)

ownedParameters : BorrowingMap -> Program p -> BorrowingMap
ownedParameters beta (MkProgram funs) = foldl insertFun empty funs
  where
    insertFun : BorrowingMap -> (Cnst, Fun p) -> BorrowingMap
    insertFun s (name, fun) = foldl (\s1, (n, b) => insert (name, n) b s1) s (ownedFunctionVar beta fun)

startBorrowingMap : Program p -> BorrowingMap
startBorrowingMap (MkProgram funs) = foldl insertFunSignature empty funs
  where
    insertFunSignature : BorrowingMap -> (Cnst, Fun p) -> BorrowingMap
    insertFunSignature beta (cnst, MkFun ps _) = fromList
      [ ((cnst, n), Borrowed) | (n,_) <- number ps ]

Eq BorrowingMap where
  m1 == m2 = toList m1 == toList m2

calculateOwnedParameters : Program p -> BorrowingMap
calculateOwnedParameters program = go empty (startBorrowingMap program)
  where
    go : BorrowingMap -> BorrowingMap -> BorrowingMap
    go previous actual = if previous == actual
                            then actual
                            else ownedParameters previous program

BorrowVarMap : Type
BorrowVarMap = Var -> Borrowing

varsInExp : Expr -> SortedSet Var
varsInExp e = case e of
  Call c  vs    => SortedSet.fromList vs
  Pap  c  vs    => SortedSet.fromList vs
  App  v1 v2    => SortedSet.fromList [v2] -- ???
  Ctor i  vs    => SortedSet.fromList vs
  Proj i  v     => SortedSet.fromList [v]
  Reset v       => SortedSet.fromList [v]
  Reuse v i vs  => SortedSet.fromList $ v :: vs

freeVars : FnBody p -> SortedSet Var
freeVars f = case f of
  Ret  x        => fromList [x]
  Let  x e b    => delete x $ union (varsInExp e) (freeVars b)
  Case x alts   => insert x $ foldl1 union $ map (freeVars . snd) alts
  -- Special cases
  Inc  x b      => insert x $ freeVars b
  Dec  x b      => insert x $ freeVars b
  LetC c1 c2 b  => freeVars b

ownPlus : Var -> SortedSet Var -> BorrowVarMap -> FnBody p -> FnBody p
ownPlus x vars beta f = if (not (contains x vars)) && (beta x == Owned)
  then f
  else Inc x f

ownMinus_ : Var -> SortedSet Var -> BorrowVarMap -> FnBody p -> FnBody p
ownMinus_ x freeVars beta f = if (not (contains x freeVars)) && (beta x == Owned)
  then Dec x f
  else f

ownMinuses_ : List Var -> SortedSet Var -> BorrowVarMap -> FnBody p -> FnBody p
ownMinuses_ []        freeVars beta f = f
ownMinuses_ (x :: xs) freeVars beta f = ownMinuses_ xs freeVars beta (ownMinus_ x freeVars beta f)

ownMinus : Var -> BorrowVarMap -> FnBody p -> FnBody p
ownMinus x beta f = ownMinus_ x (freeVars f) beta f

ownMinuses : List Var -> BorrowVarMap -> FnBody p -> FnBody p
ownMinuses xs beta f = ownMinuses_ xs (freeVars f) beta f

assignBorrowMap : Var -> Borrowing -> BorrowVarMap -> BorrowVarMap
assignBorrowMap x b f v = if x == v then b else f v

insertIncDecApp : Monad m => List Var -> List Borrowing -> BorrowVarMap -> FnBody p -> m (FnBody p)
insertIncDecApp (y :: ys) (Owned    :: bs) beta (Let z e f)
  = ownPlus y (union (freeVars f) (fromList ys)) beta <$> insertIncDecApp ys bs beta (Let z e f)
insertIncDecApp (y :: ys) (Borrowed :: bs) beta (Let z e f)
  = insertIncDecApp ys bs beta (Let z e (ownMinus y beta f))
insertIncDecApp ys os beta f = pure f

--insertIncDec1 : (Monad m, MonadReader BorrowingMap m) => BorrowVarMap -> FnBody p -> m (FnBody p)
--insertIncDec1 beta (Ret v)        = pure (ownPlus v empty beta (Ret v))
--insertIncDec1 beta (Inc v b)      = Inc v <$> insertIncDec1 beta b
--insertIncDec1 beta (Dec v b)      = Dec v <$> insertIncDec1 beta b
--insertIncDec1 beta (LetC c1 c2 b) = LetC c1 c2 <$> insertIncDec1 beta b
--insertIncDec1 beta f@(Case v alts)
--  = Case v <$> traverse
--      (\(arity, b0) => map (MkPair arity . ownMinuses (SortedSet.toList (freeVars f)) beta) $ insertIncDec1 beta b0)
--      alts
--insertIncDec1 beta (Let w e b) with (e)
--  insertIncDec1 beta (Let w e b) | Proj i x     = insertIncDec1 beta b
--  insertIncDec1 beta (Let w e b) | Reset x      = insertIncDec1 beta b
--  insertIncDec1 beta (Let w e b) | Call c ys    = insertIncDec1 beta b
--  insertIncDec1 beta (Let w e b) | Pap c ys     = insertIncDec1 beta b
--  insertIncDec1 beta (Let w e b) | App x y      = insertIncDec1 beta b
--  insertIncDec1 beta (Let w e b) | Ctor i ys    = insertIncDec1 beta b
--  insertIncDec1 beta (Let w e b) | Reuse x i ys = insertIncDec1 beta b

insertIncDec : (Monad m, MonadReader BorrowingMap m) => BorrowVarMap -> FnBody p -> m (FnBody p)
insertIncDec beta f = case f of
  Ret  v     => pure $ ownPlus v empty beta (Ret v)
  Let  w e b => case e of
                  (Proj i x)     => case beta x of
                                      Owned  => Let w (Proj i x) . Inc w . ownMinus x beta <$> insertIncDec beta b
                                      Borrowed => Let w (Proj i x) <$> insertIncDec (assignBorrowMap w Borrowed beta) b
                  (Reset x)      => Let w (Reset x) <$> insertIncDec beta b
                  (Call c ys)    => do
                    bm <- ask
                    b1 <- insertIncDec beta b
                    insertIncDecApp ys (constBorrowingParams bm c) beta $ Let w (Call c ys) b1
                  (Pap c ys)     => do
                    bm <- ask
                    b1 <- insertIncDec beta b
                    insertIncDecApp ys (constBorrowingParams bm c) beta $ Let w (Pap c ys) b1
                  (App x y)      => do
                    b1 <- insertIncDec beta b
                    insertIncDecApp [x,y] [Owned, Owned] beta $ Let w (App x y) b1
                  (Ctor i ys)    => do
                    b1 <- insertIncDec beta b
                    insertIncDecApp ys (const Owned <$> ys) beta $ Let w (Ctor i ys) b1
                  (Reuse x i ys) => do
                    b1 <- insertIncDec beta b
                    insertIncDecApp ys (const Owned <$> ys) beta $ Let w (Reuse x i ys) b1
  Case v alts => let fv = SortedSet.toList (freeVars f)
                 in Case v <$> traverse
                                (\(arity, b0) => do
                                    b1 <- insertIncDec beta b0
                                    let b2 = ownMinuses fv beta b1
                                    pure (arity, b2))
                                alts
  Inc  v     b  => Inc  v      <$> insertIncDec beta b
  Dec  v     b  => Dec  v      <$> insertIncDec beta b
  LetC c1 c2 b  => LetC c1 c2  <$> insertIncDec beta b

funInsertIncDec : (Monad m, MonadReader BorrowingMap m) => Cnst -> Fun p -> m (Fun p)
funInsertIncDec funName (MkFun args body) = do
  bm <- ask
  let argNum = \x => findIndex (x==) args
  let beta = \x => case argNum x of
                    Nothing => Owned
                    Just n => case Data.SortedMap.lookup (funName, n) bm of
                      Nothing => Owned
                      Just r  => r
  (MkFun args . ownMinuses args beta) <$> insertIncDec beta body

programInsertIncDec : (Monad m, MonadReader BorrowingMap m) => Program p -> m (Program p)
programInsertIncDec (MkProgram funDefs) = MkProgram <$> traverse (\(c, f) => MkPair c <$> funInsertIncDec c f) funDefs

prettyPutStrLn : (Pretty p) => p -> IO ()
prettyPutStrLn = putStrLn . toString 1.0 80 . doc

main : IO ()
main = do
  let mapExampleRC = runNewName $ do
        p1 <- resetReuseProgram mapExample
        p2 <- insertLetCProgram p1
        pure p2
  let swapExampleRC = runNewName $ do
        p1 <- resetReuseProgram swapExample
        p2 <- insertLetCProgram p1
        pure p2
  putStrLn ""
  prettyPutStrLn mapExample
  prettyPutStrLn mapExampleRC
  putStrLn ""
  prettyPutStrLn swapExample
  prettyPutStrLn swapExampleRC
  putStrLn "Owned parameters: Map example"
  printLn $ ownedParameters empty mapExampleRC
  prettyPutStrLn $ toList $ ownedParameters empty mapExampleRC
  putStrLn "Owned parameters: Swap example"
  printLn $ ownedParameters empty swapExampleRC
  prettyPutStrLn $ toList $ ownedParameters empty swapExampleRC
  printLn "Start Borrowing Maps"
  let mco = calculateOwnedParameters mapExampleRC
  let sco = calculateOwnedParameters swapExampleRC
  printLn mco
  printLn sco
  let mapExampleRC2 = runIdentity $ runReaderT (programInsertIncDec mapExampleRC) sco
  prettyPutStrLn (the (Program RC) mapExampleRC2)
  let swapExampleRC2 = runIdentity $ runReaderT (programInsertIncDec swapExampleRC) sco
  prettyPutStrLn (the (Program RC) swapExampleRC2)

{-
Implementation of the Counting Immutable Beans (CIB) for the GRIN compiler.

CIB uses instrumentation of the original program. There are four new instructions in the syntax
that are inserted via instrumentation. This can be categorized into two;
- reference counter instructions:
  - inc
  - dec
- heap location reuse:
  - reset
  - reuse

In the CIB approach every heap location has a reference counter associated with it.
Inc increments the counter for the location, and also increments all the referred locations transitively.
Dec decrements the counter for the location, and also decrements all the referred locations transitively.

Reset, reuse:
From the CIB paper:
let y = reset x.
If x is a shared value, than y is set to a special pointer value BOX, otherwise to the heap location associated with x.
If x is not shared than reset decrements the reference counters of the components of x, and y is set to x.

let z = reuse y in ctor_i w.
If y is BOX reuse allocates a new heap for the constructor.
If y is not Box the runtime reuses the heap location for storing the constructor.

Application of the same idea for GRIN:
Differences: meanwhile Lean's IR put every variable on the heap, GRIN uses variables as were registers
and only a subset of the registers are associated with heap locations. A register is associated
with heap location if its type is Loc. This means the GRIN implementation of the CIB approach needs
a type environment which tells which variables can be affected by the CIB operations.

In GRIN:
 * The CIB instrumentation should happen after the optimization steps.
 * Special interpreter should be implemented which handles the CIB instructions.
 * Probably it should have its own LLVM code generator and LLVM implemented runtime, preferably a plugin
   for the existing one.

We need to add 4 new instructions:
 * `x <- reset y;`   where y is a heap location, x can be a special heap location, which can be BOX too.
 * `z <- reuse x y;` where x is a special heap location created by reset, and y is a Node value.
 * `z <- inc x;`     where x is a heap location, it transitively increments the reference counters in the locations.
                     cycle detection should happen. The increment operation computes unit as its return value.
 * `z <- dec x;`     where x is a heap location, it transitively decrements the reference counters in the locations.
                     cycle detection should happen. When the counter reaches zero, the runtime must deallocate
                     the location. The decrement operation computes unit as its return value.

The GRIN nodes store primitive values, but the runtime makes the difference between location values and
primitive values, thus it is able to create the transitive closure of the reachability relation of a location
and manipulate its reference counters.

Every of the four instructions needs to be implemented in the GRIN runtime/interpreter.

In the original paper reuse of the constructors could happen only of the arity of the constructors
are the same. But in GRIN as the runtime needs to allocate heaps based on the type of the heap location.
This means every heap location can have its own arity, and reuse if the heap location is possible
only if the new node does not exceeds the arity of the heap node. Otherwise a new node needs to
be allocated, with the maximum arity.

The most important change is the reuse construction. It changes certain instances of the
store operation to the reuse operation.

```
Before:
x <- store y;

After:
z <- reset w;
...
x <- reuse z y;
```

In this case we need to decide to reuse the heap location associated with w only if w can accommodate
all the possible values of x. This means the max-arity(w) >= max-arity(x). Meanwhile Lean's
approach uses the arity of the constructors in the alternatives, we can use the abstract information
of all the possible runs.

Implementation steps:
[ ] Import abstracting the definitional interpreters from the mini-grin repo
[ ] Change the implementation to use base functors instead of Expr datatype
[ ] Implement reference statistics with the new interpreter, as a warm-up exercise
[ ] Implement CIB program instrumentation for GRIN producing ExprF :+: CibF AST
[ ] Implement interpreter for CIB extended GRIN program
[ ] Extra: Implement LLVM codegen plugin for CIB instructions
-}


module Cond

%default total
%access abstract

{-
Usage example:

foo : Int -> String
foo x = cond $ do
  x == 5 ?: "five"
  x == 17 ?: "bigger"
  otherwise "fnord"

• An otherwise at the end is required.
• Note that (?:) is a low-precedence operator and will act as such: parentheses
  may be required in some cases. Note also that otherwise is a function, so ($)
  or parentheses are required for non-simple default expressions.
• The expression for a non-matched clause is not evaluated, so this construct
  can be used to check preconditions on partial functions (but please encode
  these in types instead to make the functions total, if feasible).
-}

infix 0 ?:

data Alt a = Unmatched | Matched (Lazy a)

data Res a = It (Lazy a)

(>>=) : Alt a -> (() -> Res a) -> Res a
(>>=) Unmatched f = f ()
(>>=) (Matched x) f = It x

(?:) : Bool -> Lazy a -> Alt a
(?:) False x = Unmatched
(?:) True x = Matched x

otherwise : Lazy a -> Res a
otherwise x = It x

cond : Res a -> a
cond (It x) = x


module Expr

%access public export

data TypeT : Type where
  UnitT : TypeT
  IntT : TypeT
  DoubleT : TypeT
  BoolT : TypeT
  ArrayT : TypeT -> TypeT
  VarT : TypeT -> TypeT
  ArrowT : TypeT -> TypeT -> TypeT

Show TypeT where
  show UnitT = "unit"
  show IntT = "int"
  show DoubleT = "double"
  show BoolT = "bool"
  show (ArrayT typ) = show typ ++ " []"
  show (VarT typ) = show typ ++ " ref"
  show (ArrowT arg res) = "(" ++ show arg ++ " -> " ++ show res ++ ")"

interface Symantics (rep : TypeT -> Type) where
  defaultof : rep a
  int : Int -> rep IntT
  double : Double -> rep DoubleT
  bool : Bool -> rep BoolT
  (==) : rep IntT -> rep IntT -> rep BoolT
  (>) : rep IntT -> rep IntT -> rep BoolT
  (<) : rep IntT -> rep IntT -> rep BoolT
  (&&) : rep BoolT -> rep BoolT -> rep BoolT
  (||) : rep BoolT -> rep BoolT -> rep BoolT
  not : rep BoolT -> rep BoolT
  (+) : rep IntT -> rep IntT -> rep IntT
  (<+>) : rep DoubleT -> rep DoubleT -> rep DoubleT
  (*) : rep IntT -> rep IntT -> rep IntT
  (<*>) : rep DoubleT -> rep DoubleT -> rep DoubleT
  ite : rep BoolT -> rep a -> rep a -> rep a
  it : rep BoolT -> rep UnitT -> rep UnitT
  deref : rep (VarT a) -> rep a
  assign : rep a -> rep (VarT a) -> rep UnitT
  newVar : rep a -> (rep (VarT a) -> rep b) -> rep b
  letVal : rep a -> (rep a -> rep b) -> rep b
  index : rep IntT -> rep (ArrayT a) -> rep a
  length : rep (ArrayT a) -> rep IntT
  while : rep BoolT -> rep UnitT -> rep UnitT
  seq : rep a -> rep b -> rep b
  seqs : List (rep UnitT) -> rep UnitT
  lam : (rep a -> rep b) -> rep (ArrowT a b)
  app : rep (ArrowT a b) -> rep a -> rep b

data Code : TypeT -> Type where
  C : (Int -> String) -> Code a

Symantics Code where
  defaultof {a = UnitT} = C (\_ => "()")
  defaultof {a = IntT} = C (\_ => "0")
  defaultof {a = DoubleT} = C (\_ => "0")
  defaultof {a = BoolT} = C (\_ => "false")
  defaultof {a = (ArrayT x)} = C (\_ => "null")
  defaultof {a = (VarT x)} = C (\_ => "null")
  defaultof {a = (ArrowT x y)} = C (\_ => "null")
  int x = C (\_ => show x)
  double x = C (\_ => show x)
  bool x = C (\_ => if x then "true" else "false")
  (==) (C l) (C r) = C (\v => "( " ++ l v ++ " = " ++ r v ++ " )")
  (>) (C l) (C r) = C (\v => "( " ++ l v ++ " > " ++ r v ++ " )")
  (<) (C l) (C r) = C (\v => "( " ++ l v ++ " < " ++ r v ++ " )")
  (&&) (C l) (C r) = C (\v => "( " ++ l v ++ " && " ++ r v ++ " )")
  (||) (C l) (C r) = C (\v => "( " ++ l v ++ " || " ++ r v ++ " )")
  not (C c) = C (\v => "not (" ++ c v ++ " )")
  (+) (C l) (C r) = C (\v => "( " ++ l v ++ " + " ++ r v ++ " )")
  (<+>) (C l) (C r) = C (\v => "( " ++ l v ++ " + " ++ r v ++ " )")
  (*) (C l) (C r) = C (\v => "( " ++ l v ++ " * " ++ r v ++ " )")
  (<*>) (C l) (C r) = C (\v => "( " ++ l v ++ " * " ++ r v ++ " )")
  ite (C b) (C l) (C r) = C (\v => "(if " ++ b v ++ " then " ++ l v ++ " else " ++ r v ++ ")")
  it (C b) (C t) = C (\v => "(if " ++ b v ++ " then " ++ t v ++ ")")
  deref (C x) = C (\v => "!" ++ x v)
  assign (C val) (C var) = C (\v => var v ++ " := " ++ val v)
  newVar {a} (C s) f = C (\v => let x = "x" ++ show v in
                                let (C c) = f $ the (Code (VarT a)) (C (\_ => x)) in
                                "let " ++ x ++ " = ref " ++ s v ++ " in " ++ c (v + 1))
  letVal (C s) f = C (\v => let x = "x" ++ show v in
                            let (C c) = f $ C (\_ => x) in
                            "let " ++ x ++ " = " ++ s v ++ " in " ++ c (v + 1))
  index (C i) (C arr) = C (\v => arr v ++ ".[" ++ i v ++ "]")
  length (C arr) = C (\v => arr v ++ ".Length")
  while (C p) (C b) = C (\v => "(while " ++ p v ++ " do " ++ b v ++ ")")
  seq (C fs) (C sn) = C (\v => "(" ++ fs v ++ "; " ++ sn v ++ ")")
  seqs steps = foldr seq defaultof steps
  lam {a} f = C (\v => let x = "x" ++ show v in
                       let (C g) = f (C (\_ => x)) in
                       "(fun (" ++ x ++ " : " ++ show a ++ ") -> " ++ g (v + 1) ++ ")")
  app (C f) (C g) = C (\v => "(" ++ f v ++ " " ++ g v ++ ")")

compile : Code a -> String
compile code = let (C c) = code in c 0


module Effekt.AvoidEta

import Effekt.CpsUnstaged
import Effekt.CpsStaged
import Effekt.IteratedUnstaged

mutual
  public export
  data CTX : List Type -> Type -> Type -> Type where
    Static :   (Exp a -> STM rs b)   -> CTX rs a b
    Dynamic :  Exp (a -> Stm rs b)   -> CTX rs a b

  public export
  STM : List Type -> Type -> Type
  STM [] a = Exp a
  STM (r :: rs) a = CTX rs a r -> STM rs r

mutual
  export
  reify : STM rs a -> Exp (Stm rs a)
  reify {rs = []} m = m
  reify {rs = (q :: qs)} m =
    (Lam $ \ k =>  reify (m (Dynamic k)))

  export
  reflect : Exp (Stm rs a) -> STM rs a
  reflect {rs = []} m = m
  reflect {rs = (q :: qs)} m =
    \k => reflect (App (m)  ((reifyContext k)))

  export
  reifyContext : CTX rs a r -> Exp (a -> Stm rs r)
  reifyContext (Static k) = (Lam $ \ a =>  reify (k a))
  reifyContext (Dynamic k) = k

export
resume : CTX rs a r -> (Exp a -> STM rs r)
resume (Static k) = k
resume (Dynamic k) = \a => reflect (App (k)  (a))

export
pure : Exp a -> STM rs a
pure{  rs= []}         a = a
pure{  rs= r :: rs}  a = \k => resume k a

push : CTX (r :: rs) a b -> CTX rs b r -> CTX rs a r
push f k = Static (\a => resume f a k)

bind : STM rs a -> CTX rs a b -> STM rs b
bind{  rs= []}         m f = resume f m
bind{  rs= r :: rs}  m f = \k => m (push f k)

export
shift0 : (CTX rs a r -> STM rs r) -> STM (r :: rs) a
shift0 = id

export
runIn0 : STM (r :: rs) a -> CTX rs a r -> STM rs r
runIn0 = id

export
reset0 : STM (a :: rs) a -> STM rs a
reset0 m = runIn0 m (Static pure)

export
lift : STM rs a -> STM (r :: rs) a
lift = bind

export
(>>=) : STM rs a -> (Exp a -> STM rs b) -> STM rs b
m >>= f = bind m (Static f)

add : STM rs Int -> STM rs Int -> STM rs Int
add mx my = do
  x <- mx
  y <- my
  pure ((Add x y))

display : Exp String -> STM (IO () :: rs) ()
display {rs} s = shift0 (\c => do
  a <- resume c Uni
  pure {rs} ((Seq ((Dis s)) a)))

export
displayTwice : STM [] (IO ())
displayTwice = reset0 {rs=[]} (
  (>>=) {rs=[IO ()]} (display {rs=[]} (Str "hello ")) (\u =>
  (>>=) {rs=[IO ()]} (display {rs=[]} (Str "world")) (\o =>
  pure {rs=[IO ()]} Ski)))

export
ifthenelse : Exp Bool -> STM rs a -> STM rs a -> STM rs a
ifthenelse {rs} Tru t _ = t
ifthenelse {rs} Fls _ e = e
ifthenelse {rs=[]} b t e = Ite b t e
ifthenelse {rs=q::qs} b t e = \k => ifthenelse b (t k) (e k)

export
letrec : ((Exp a -> STM rs b) -> Exp a -> STM rs b) -> Exp a -> STM rs b
letrec body a = reflect (App (Rec (\f => \x =>
  reify (body (\y => reflect (App (f)  (y))) x)))  (a))

-- Iteration happens to be recursion.
export
loop : (Exp a -> STM (r :: rs) a) -> Exp a -> STM rs r
loop m = letrec (\k => \a => m a (Static k))

export
break : Exp r -> STM (r :: rs) a
break r = shift0 (\_ => pure r)


module HoareAssertAssume

import Assn
import Expr
import Logic
import Maps

%access public export

%default total

data Com : Type where
  CSkip : Com
  CAss : Id -> AExp -> Com
  CSeq : Com -> Com -> Com
  CIf : BExp -> Com -> Com -> Com
  CWhile : BExp -> Com -> Com
  CAssert : BExp -> Com
  CAssume : BExp -> Com

infix 5 ::=

SKIP : Com
SKIP = CSkip

(::=) : Id -> AExp -> Com
(::=) = CAss

(>>=) : Com -> (() -> Com) -> Com
(>>=) c f = CSeq c (f ())

syntax IFB [b] THEN [c1] ELSE [c2] FI = CIf b c1 c2

WHILE : BExp -> Com -> Com
WHILE = CWhile

ASSERT : BExp -> Com
ASSERT = CAssert

ASSUME : BExp -> Com
ASSUME = CAssume

data Result : Type where
  RNormal : State -> Result
  RError : Result

rNormalInjective : RNormal st1 = RNormal st2 -> st1 = st2
rNormalInjective Refl = Refl

Uninhabited (RError = RNormal st) where
  uninhabited Refl impossible

data CEval : Com -> State -> Result -> Type where
  E_Skip : CEval CSkip st (RNormal st)
  E_Ass : aeval st a1 = n -> CEval (CAss x a1) st (RNormal (t_update x n st))
  E_SeqNormal : CEval c1 st1 (RNormal st2) -> CEval c2 st2 r ->
                CEval (CSeq c1 c2) st1 r
  E_SeqError : CEval c1 st RError -> CEval (CSeq c1 c2) st RError
  E_IfTrue : beval st b = True -> CEval c1 st r -> CEval (CIf b c1 c2) st r
  E_IfFalse : beval st b = False -> CEval c2 st r -> CEval (CIf b c1 c2) st r
  E_WhileEnd : beval st b = False -> CEval (CWhile b c) st (RNormal st)
  E_WhileLoopNormal : beval st b = True ->
                      CEval c st (RNormal st1) -> CEval (CWhile b c) st1 r ->
                      CEval (CWhile b c) st r
  E_WhileLoopError : beval st b = True -> CEval c st RError ->
                     CEval (CWhile b c) st RError
  E_AssertTrue : beval st b = True -> CEval (CAssert b) st (RNormal st)
  E_AssertFalse : beval st b = False -> CEval (CAssert b) st RError
  E_Assume : beval st b = True -> CEval (CAssume b) st (RNormal st)

syntax [c1] "/" [st1] "\\\\" [st2] = CEval c1 st1 st2

HoareTriple : (p : Assertion) -> (c : Com) -> (q : Assertion) -> Type
HoareTriple p c q = (st : State) -> (r : Result) -> (c / st \\ r) -> p st ->
                    (st' : State ** (r = RNormal st', q st'))

assert_assume_differ : DPair (Assertion, BExp, Assertion)
                             (\(p, b, q) => ( HoareTriple p (ASSUME b) q
                                            , Not (HoareTriple p (ASSERT b) q) ))
assert_assume_differ = ((const (), 0 == 1, const ()) ** (assume, not_assert))
where assume : HoareTriple (const ()) (ASSUME (0 == 1)) (const ())
      assume st (RNormal _) (E_Assume _) _ = (st ** (Refl, ()))
      not_assert : Not (HoareTriple (const ()) (ASSERT (0 == 1)) (const ()))
      not_assert ht =
        let (_ ** (prf, _)) = ht empty_state RError (E_AssertFalse Refl) ()
        in absurd prf

assert_implies_assume : (p : Assertion) -> (b : BExp) -> (q : Assertion) ->
                        HoareTriple p (ASSERT b) q -> HoareTriple p (ASSUME b) q
assert_implies_assume p b q ht st (RNormal st) (E_Assume prf) p_st =
  ht st (RNormal st) (E_AssertTrue prf) p_st

hoare_assign : (q : Assertion) -> (x : Id) -> (a : AExp) ->
               HoareTriple (AssignSub x a q) (x ::= a) q
hoare_assign q x a st _ (E_Ass {n} prf) p_st =
  let q_st = replace {P=\m => q (t_update x m st)} prf p_st
  in (t_update x n st ** (Refl, q_st))

hoare_consequence_pre : (p, p', q : Assertion) -> (c : Com) ->
                        HoareTriple p' c q -> p ->> p' -> HoareTriple p c q
hoare_consequence_pre p p' q c htc imp st r rel p_st = htc st r rel (imp st p_st)

hoare_consequence_post : (p, q, q' : Assertion) -> (c : Com) ->
                         HoareTriple p c q' -> q' ->> q -> HoareTriple p c q
hoare_consequence_post p q q' c htc imp st r rel p_st =
  let (st' ** (prf, q'_st')) = htc st r rel p_st
  in (st' ** (prf, imp st' q'_st'))

hoare_seq : (p, q, r : Assertion) -> (c1, c2 : Com) ->
            HoareTriple q c2 r -> HoareTriple p c1 q ->
            HoareTriple p (do c1; c2) r
hoare_seq p q r c1 c2 htc2 htc1 st res (E_SeqNormal {st2} cc1 cc2) p_st =
  let (st2 ** (pf, q_st2)) = htc1 st (RNormal st2) cc1 p_st
      cc2' = replace {P=\x => CEval c2 x res} (rNormalInjective pf) cc2
  in htc2 st2 res cc2' q_st2
hoare_seq p q r c1 c2 htc2 htc1 st RError (E_SeqError cc1) p_st =
  let (_ ** (prf, _)) = htc1 st RError cc1 p_st
  in absurd prf

hoare_skip : (p : Assertion) -> HoareTriple p SKIP p
hoare_skip p st (RNormal st) E_Skip p_st = (st ** (Refl, p_st))

hoare_if : (p, q : Assertion) -> (b : BExp) -> (ct, cf : Com) ->
           HoareTriple (\st => (p st, BAssn b st)) ct q ->
           HoareTriple (\st => (p st, Not (BAssn b st))) cf q ->
           HoareTriple p (CIf b ct cf) q
hoare_if p q b ct cf htct htcf st r (E_IfTrue prf cct) p_st =
  htct st r cct (p_st, prf)
hoare_if p q b ct cf htct htcf st r (E_IfFalse prf ccf) p_st =
  htcf st r ccf (p_st, bexp_eval_false prf)

hoare_while : (p : Assertion) -> (b : BExp) -> (c : Com) ->
              HoareTriple (\st => (p st, BAssn b st)) c p ->
              HoareTriple p (CWhile b c) (\st => (p st, Not (BAssn b st)))
hoare_while p b c htc st (RNormal st) (E_WhileEnd prf) p_st =
  (st ** (Refl, p_st, bexp_eval_false prf))
hoare_while p b c htc st r w@(E_WhileLoopNormal {st1} prf cc cnext) p_st =
  let (st2 ** (pf, p_st2)) = htc st (RNormal st1) cc (p_st, prf)
      cnext' = assert_smaller w
                              (replace {P=\x => CEval (CWhile b c) x r}
                                       (rNormalInjective pf) cnext)
  in hoare_while p b c htc st2 r cnext' p_st2
hoare_while p b c htc st RError (E_WhileLoopError prf cc) p_st =
  let (_ ** (pf, _)) = htc st RError cc (p_st, prf)
  in absurd pf

hoare_assert : (p : Assertion) -> (b : BExp) ->
               p ->> (\st => BAssn b st) -> HoareTriple p (ASSERT b) p
hoare_assert p b _ st (RNormal st) (E_AssertTrue prf) p_st = (st ** (Refl, p_st))
hoare_assert p b imp st RError (E_AssertFalse prf) p_st =
  absurd $ trans (sym prf) (imp st p_st)

hoare_assume : (p : Assertion) -> (b : BExp) ->
               HoareTriple p (ASSUME b) (\st => (p st, BAssn b st))
hoare_assume p b st (RNormal st) (E_Assume prf) p_st = (st ** (Refl, p_st, prf))

assert_assume_example : HoareTriple (const ())
                                    (do ASSUME (X == 1)
                                        X ::= X + 1
                                        ASSERT (X == 2))
                                    (const ())
assert_assume_example =
  let ht_assm = hoare_consequence_post
                  (const ())
                  (\st => st X = 1)
                  (\st => ((), BAssn (X == 1) st))
                  (ASSUME (X == 1))
                  (hoare_assume (const ()) (X == 1))
                  (\st, (_, prf) => fst (nat_beq_iff (st X) 1) prf)
      ht_asgn = hoare_consequence_pre
                  (\st => st X = 1)
                  (\st => st X + 1 = 2)
                  (\st => st X = 2)
                  (X ::= X + 1)
                  (hoare_assign (\st => st X = 2) X (X + 1))
                  (\st, p_st => plusConstantRight (st X) 1 1 p_st)
      ht_asrt = hoare_consequence_post
                  (\st => st X = 2)
                  (const ())
                  (\st => st X = 2)
                  (ASSERT (X == 2))
                  (hoare_assert (\st => st X = 2)
                                (X == 2)
                                (\st, prf => snd (nat_beq_iff (st X) 2) prf))
                  (\_, _ => ())
  in hoare_seq (const ())
               (\st => st X = 1)
               (const ())
               (ASSUME (X == 1))
               (do X ::= X + 1
                   ASSERT (X == 2))
               (hoare_seq (\st => st X = 1)
                          (\st => st X = 2)
                          (const ())
                          (X ::= X + 1)
                          (ASSERT (X == 2))
                          ht_asrt
                          ht_asgn)
               ht_assm


||| Derive Eq and DecEq
|||
||| You'll get cryptic errors if you try to use it without importing
||| `Language.Reflection` and `Extra.Prelude`.
module Extra.Language.Derive

import Extra.Prelude
import Language.Reflection
import Decidable.Equality

%language ElabReflection

argTys : (ty : TTImp) -> List TTImp
argTys (IPi _ _ ExplicitArg _ argTy retTy) = argTy :: argTys retTy
argTys (IPi _ _ _ _ _ retTy) = argTys retTy
argTys _ = []

countArgs : (ty : TTImp) -> Nat
countArgs = length . argTys

fullyBind : FC -> TTImp -> String -> Int -> (ty : TTImp) -> TTImp
fullyBind pos f argName i (IPi _ _ ExplicitArg _ _ retTy)
  = fullyBind pos (IApp pos f (IBindVar pos (argName ++ show i))) argName (i+1) retTy
fullyBind pos f argName i (IPi _ _ _ _ _ retTy)
  = fullyBind pos f argName i retTy
fullyBind pos f argName i _ = f

fullyApply : FC -> TTImp -> String -> Int -> (ty : TTImp) -> TTImp
fullyApply pos f argName i (IPi _ _ ExplicitArg _ _ retTy)
  = fullyApply pos (IApp pos f (IVar pos $ UN $ argName ++ show i)) argName (i+1) retTy
fullyApply pos f argName i (IPi _ _ _ _ _ retTy)
  = fullyApply pos f argName i retTy
fullyApply pos f argName i _ = f

lookupConstr : Name -> Elab (Name, TTImp)
lookupConstr constr = do
  [(n, ty)] <- getType constr
    | [] => fail $ "constr " ++ show constr ++ " does not exist? hmmm"
    | _ => fail $ "ambiguous name for constr " ++ show constr
  pure (n, ty)

lookupType : Name -> Elab (Name, TTImp)
lookupType typeName = do
  [(n, ty)] <- getType typeName
    | [] => fail $ "name " ++ show typeName ++ " not in scope"
    | _ => fail $ "ambiguous name " ++ show typeName
  pure (n, ty)

||| Derive `Eq` for a type.
|||
||| Use:
||| mutual
|||   myEq : Eq a => MyType a -> MyType a -> Bool
|||   %runElab deriveEq `{{ myEq }} `{{ MyType }}
|||
|||   Eq a => Eq (MyType a) where
|||     (==) = myEq
export
deriveEq : Name -> Name -> Elab ()
deriveEq funcName typeName = do
  let pos = MkFC ("deriveEq for " ++ show typeName) (0,0) (0,0)
  (n, _) <- lookupType typeName
  constrs <- getCons n
  let and : TTImp -> TTImp -> TTImp
      and x y = `(~(x) && ~(y))
      compareEq : String -> String -> TTImp
      compareEq x y = `((~(IVar pos $ UN x)) == (~(IVar pos $ UN y)))
      makeClause : Name -> Elab Clause
      makeClause constr = do
        (_, ty) <- lookupConstr constr
        let nArgs = countArgs ty
        let xs = map (\i => "x_" ++ show i) $ take nArgs [1..]
        let ys = map (\i => "y_" ++ show i) $ take nArgs [1..]
        let px = foldl (IApp pos) (IVar pos constr) $ map (IBindVar pos) xs
        let py = foldl (IApp pos) (IVar pos constr) $ map (IBindVar pos) ys
        pure $ PatClause pos `(MkPair ~(px) ~(py))
             $ foldl and `(True) $ zipWith compareEq xs ys
      finalClause : Clause
      finalClause = PatClause pos `(_) `(False)
  clauses <- traverse makeClause constrs
  let allClauses = clauses ++ [finalClause]
      caseExpr = ICase pos `(MkPair x y) (Implicit pos True) allClauses
      result = `(\x, y => ~(caseExpr))
      clauses = [PatClause pos (IVar pos funcName) result]
  declare [IDef pos funcName clauses]

-- Wow! what an opportunity to derive equality
eqNamespace : Namespace -> Namespace -> Bool
%runElab deriveEq `{{ eqNamespace }} `{{ Namespace }}

export
Eq Namespace where
  (==) = eqNamespace

mutual
  eqName : Name -> Name -> Bool
  %runElab deriveEq `{{ eqName }} `{{ Name }}

  export
  Eq Name where
    (==) = eqName

||| Derive `DecEq` for a type.
|||
||| Use:
||| mutual
|||   myDecEq : DecEq a => (x : MyType a) -> (y : MyType a) -> Dec (x = y)
|||   %runElab deriveDecEq `{{ myDecEq }} `{{ MyType }}
|||
|||   DecEq a => DecEq (MyType a) where
|||     decEq = myDecEq
export
deriveDecEq : Name -> Name -> Elab ()
deriveDecEq funcName typeName = do
  let pos = MkFC ("deriveDecEq for " ++ show typeName) (0,0) (0,0)
      thisFunc = IVar pos funcName
  (n, _) <- lookupType typeName
  constrs <- getCons n
  let makeClause : Int -> (Int -> TTImp) -> TTImp -> List TTImp -> TTImp -> Clause
      makeClause i mkIndTy lhsSoFar withBinds ty = case ty of
        IPi _ _ ExplicitArg _ argTy retTy =>
          let nArgs = 1 + countArgs retTy
              px = fullyBind pos lhsSoFar "x__" i ty
              py = fullyBind pos lhsSoFar "y__" i ty
              thisX = IVar pos $ UN $ "x__" ++ show i
              thisY = IVar pos $ UN $ "y__" ++ show i
              decEqXY = `(decEq ~(thisX) ~(thisY))
              lhs = foldl (IApp pos) `(~(thisFunc) ~(px) ~(py)) withBinds
          in WithClause pos lhs decEqXY [
            -- No clause
            PatClause pos `(~(lhs) (No ~(IBindVar pos "bad")))
              $ ILocal pos `[ ind : ~(mkIndTy i)
                              ind Refl = Refl ]
                           `(No (bad . ind)),
            -- Yes clause
            makeClause (i+1)
                       mkIndTy
                       (IApp pos lhsSoFar thisX)
                       (withBinds ++ [`(Yes Refl)])
                       retTy
            ]
        IPi _ _ _ _ _ retTy =>
          makeClause i mkIndTy lhsSoFar withBinds retTy -- Ignore implicit arguments
        _ => let startingLhs = `(~(thisFunc) ~(lhsSoFar) ~(lhsSoFar))
                 lhs = foldl (IApp pos) startingLhs withBinds
             in PatClause pos lhs `(Yes Refl)
      -- bindImplicits "x" Constr (typeof Constr) thing = forall xi. thing
      bindImplicits : String -> Name -> TTImp -> TTImp -> TTImp
      bindImplicits varName constr ty thing =
        let args = the (List (Int, _)) $ zip [1..] $ argTys ty in
        let implicitPi : (String,TTImp) -> TTImp -> TTImp
            implicitPi (n,arg) ty = IPi pos M0 ImplicitArg (Just $ UN n) arg ty
        in foldr implicitPi thing $ map (\(i,t) => (varName ++ show i, t)) args
      makeSameClause : Name -> Elab Clause
      makeSameClause constr = do
        (name, ty) <- lookupConstr constr
        let boundA = fullyApply pos (IVar pos name) "a__" 1 ty
            boundB = fullyApply pos (IVar pos name) "b__" 1 ty
        let makeIndTy : Int -> TTImp
            makeIndTy i = let a = IVar pos $ UN $ "a__" ++ show i
                              b = IVar pos $ UN $ "b__" ++ show i
                          in bindImplicits "a__" name ty
                            $ bindImplicits "b__" name ty
                            $ `(~(boundA) === ~(boundB) -> ~(a) === ~(b))
        pure $ makeClause 1 makeIndTy (IVar pos name) [] ty
      makeDiffClause : Name -> Name -> Elab Clause
      makeDiffClause c1 c2 = do
        (n1, ty1) <- lookupConstr c1
        (n2, ty2) <- lookupConstr c2
        let thing1 = fullyApply pos (IVar pos n1) "a__" 1 ty1
            thing2 = fullyApply pos (IVar pos n2) "b__" 1 ty2
            badTy = bindImplicits "a__" n1 ty1
                  $ bindImplicits "b__" n2 ty2
                  $ `(~(thing1) = ~(thing2) -> Void)
        let pat1 : TTImp
            pat1 = foldl (\t, _ => IApp pos t (Implicit pos True)) (IVar pos n1) $ argTys ty1
            pat2 : TTImp
            pat2 = foldl (\t, _ => IApp pos t (Implicit pos True)) (IVar pos n2) $ argTys ty2
            lhs : TTImp
            lhs = `(~(thisFunc) ~(pat1) ~(pat2))
            bad : List Decl
            bad = [IClaim pos MW Private [] $ MkTy pos (UN "bad") badTy]
                ++ `[ bad Refl impossible ]
        pure $ PatClause pos lhs
             $ ILocal pos bad `(No bad)
  clauses <- sequence $ do
    a <- constrs
    b <- constrs
    pure $ if a == b then makeSameClause a
                     else makeDiffClause a b
  declare [IDef pos funcName clauses]


module Main

record Point where
  constructor MkPoint
  x : Double
  y : Double

-- a record creates two projections:
--
-- .x : Point -> Double
-- .y : Point -> Double
--
-- because %prefix_record_projections are on by default, we also get:
--
-- x : Point -> Double
-- y : Point -> Double

-- to prevent cluttering the ordinary name space with short identifiers
%prefix_record_projections off

record Rect where
  constructor MkRect
  topLeft : Point
  bottomRight : Point

-- .topLeft : Rect -> Point
-- .bottomRight : Rect -> Point
--
-- For Rect, we don't get the prefix projections:
--
-- Main> :t topLeft
-- (interactive):1:4--1:11:Undefined name topLeft
-- Main> :t .topLeft
-- \{rec:0} => .topLeft rec : ?_ -> Point

pt : Point
pt = MkPoint 4.2 6.6

rect : Rect
rect =
  MkRect
    (MkPoint 1.1 2.5)
    (MkPoint 4.3 6.3)

-- New lexical structure:
--
-- Foo.bar.baz with uppercase F is one lexeme: NSIdent ["baz", "bar", "Foo"]
-- foo.bar.baz with lowercase f is three lexemes: foo, .bar, .baz
-- .foo.bar.baz is three lexemes: .foo, .bar, .baz
--
-- If you want Constructor.field, you have to write (Constructor).field.
--
-- New syntax of simpleExpr:
--
-- Expressions binding tighter than application (simpleExpr),
-- such as variables or parenthesised expressions,
-- have been renamed to simplerExpr,
-- and an extra layer of syntax has been inserted.
--
--   simpleExpr ::= (.field)+               -- parses as PPostfixAppPartial
--                | simplerExpr (.field)+   -- parses as PPostfixApp
--                | simplerExpr             -- (parses as whatever it used to)
--
-- (.foo) is a name, so you can use it to e.g. define a function called .foo (see .squared below)
-- (.foo.bar) is a parenthesised expression
--
-- New desugaring rules
--
-- (.field1 .field2 .field3) desugars to (\x => .field3 (.field2 (.field1 x)))
-- (simpleExpr .field1 .field2 .field3) desugars to ((.field .field2 .field3) simpleExpr).
--
-- There are more details such as namespacing; see below.

-- user-defined projections work, too (should they?)
(.squared) : Double -> Double
(.squared) x = x * x

-- #649: Forced patterns remain unaffected by postfix projections.
test649 : (x, y : Bool) -> x = y -> Bool
test649 .(x) x Refl = x

-- #441: Locals do not shadow record projections.
test441 : (Point -> Double) -> (Point -> Double) -> Point -> Double
test441 x y pt = pt.x + pt.y
  -- no ambiguity errors: the locals "x" and "y" are not considered in pt.x and pt.y


module Compiler.Erlang.IR.AbstractFormat

import public Data.Vect

import public Compiler.Erlang.Utils.PrimTerm

%default total

-- This file contain representations of Erlang's Abstract Format,
-- as described here: https://erlang.org/doc/apps/erts/absform.html

public export
Line : Type
Line = Int

public export
data Literal : Type where
  ALAtom : Line -> String -> Literal
  ALChar : Line -> Char -> Literal
  ALFloat : Line -> Double -> Literal
  ALInteger : Line -> Integer -> Literal
  ALCharlist : Line -> String -> Literal

public export
data MapFieldAssoc : Type -> Type where
  MkAssoc : Line -> (key : a) -> (value : a) -> MapFieldAssoc a

public export
data MapFieldExact : Type -> Type where
  MkExact : Line -> (key : a) -> (value : a) -> MapFieldExact a

public export
data BitType = ABInteger | ABFloat | ABBinary | ABBytes | ABBitstring | ABBits | ABUtf8 | ABUtf16 | ABUtf32

public export
data BitSignedness = ABUnsigned | ABSigned

public export
data BitEndianness = ABBig | ABLittle | ABNative

-- Unit is a value from 1-256 where `FZ` is representing the value 1
public export
data BitUnit = MkBitUnit (Fin 256)

public export
data BitSize : Type where
  ABSDefault : BitSize
  ABSInteger : Line -> Integer -> BitSize
  ABSVar : Line -> String -> BitSize

public export
record TypeSpecifierList where
  constructor MkTSL
  signedness : Maybe BitSignedness
  endianness : Maybe BitEndianness
  type : Maybe BitType
  unit : Maybe BitUnit

public export
data BitSegment : Type -> Type where
  MkBitSegment : Line -> a -> (size : BitSize) -> TypeSpecifierList -> BitSegment a

public export
data BitPattern : Type where
  ABPInteger : Line -> Integer -> BitPattern
  ABPFloat : Line -> Double -> BitPattern
  ABPCharlist : Line -> String -> BitPattern
  ABPUniversal : Line -> BitPattern -- Wildcard matching
  ABPVar : Line -> (name : String) -> BitPattern

public export
data Pattern : Type where
  APLiteral : Literal -> Pattern
  APBitstring : Line -> List (BitSegment BitPattern) -> Pattern
  -- Compound pattern
  APCons : Line -> Pattern -> Pattern -> Pattern
  APMap : Line -> List (MapFieldExact Pattern) -> Pattern
  APNil : Line -> Pattern
  -- Operator pattern
  -- Paranthesized pattern
  -- Record field index pattern
  -- Record pattern
  APTuple : Line -> List Pattern -> Pattern
  APUniversal : Line -> Pattern -- Wildcard matching
  APVar : Line -> (name : String) -> Pattern

public export
data Guard : Type where
  AGLiteral : Literal -> Guard
  -- Bitstring constructor
  AGCons : Line -> Guard -> Guard -> Guard
  AGFunCall : Line -> (fnName : String) -> List Guard -> Guard -- NOTE: Function calls are restricted to the 'erlang' module
  -- Map creation
  -- Map update
  AGNil : Line -> Guard
  AGOp : Line -> (op : String) -> Guard -> Guard -> Guard
  -- Paranthesized guard
  -- Record creation
  -- Record field access
  -- Record field index
  AGTuple : Line -> List Guard -> Guard
  AGVar : Line -> (name : String) -> Guard

-- A guard sequence `Gd1, Gd2; Gd3, Gd4` has the same meaning as:
-- `(Gd1 andalso Gd2) orelse (Gd3 andalso Gd4)`
--
-- A guard sequence is represented by the type `List GuardAlt`:
-- * You can have zero or more `GuardAlt` expressions.
-- * Each `GuardAlt` needs to have at least one `Guard` expression.
public export
data GuardAlt : Type where
  MkGuardAlt : (guards : Vect (S k) Guard) -> GuardAlt

mutual
  public export
  data Expr : Type where
    AELiteral : Literal -> Expr
    AEBitstring : Line -> List (BitSegment Expr) -> Expr
    -- Bitstring comprehension
    AEBlock : Line -> (body : Vect (S k) Expr) -> Expr
    AECase : Line -> Expr -> (clauses : Vect (S k) CaseClause) -> Expr
    -- Catch expression
    AECons : Line -> Expr -> Expr -> Expr
    -- Function reference
    AEFun : Line -> (arity : Nat) -> (clauses : Vect (S k) (FunClause arity)) -> Expr
    AEFunCall : Line -> (fn : Expr) -> (args : List Expr) -> Expr
    AERemoteRef : Line -> (modName : Expr) -> (fnName : Expr) -> Expr
    -- If expression
    -- List comprehension
    AEMapNew : Line -> List (MapFieldAssoc Expr) -> Expr
    -- Map update
    AEMatch : Line -> (lhs : Pattern) -> (rhs : Expr) -> Expr
    AENil : Line -> Expr
    AEOp : Line -> (op : String) -> (lhs : Expr) -> (rhs : Expr) -> Expr
    -- Paranthesized expression
    AEReceive : Line -> List CaseClause -> ReceiveTimeout -> Expr
    -- Record creation expression
    -- Record update expression
    AETuple : Line -> List Expr -> Expr
    AETry : Line -> (sc : Vect (S k) Expr) -> List CaseClause -> List CatchClause -> (afterBody : List Expr) -> Expr
    AEVar : Line -> (name : String) -> Expr

  public export
  data FunClause : Nat -> Type where
    MkFunClause : Line -> Vect arity Pattern -> List GuardAlt -> (body : Vect (S k) Expr) -> FunClause arity

  public export
  data CaseClause : Type where
    MkCaseClause : Line -> Pattern -> List GuardAlt -> (body : Vect (S k) Expr) -> CaseClause

  public export
  data CatchClause : Type where
    MkCatchClause : Line -> (class : Pattern) -> (reason : Pattern) -> (stacktrace : Pattern) -> List GuardAlt -> (body : Vect (S k) Expr) -> CatchClause

  public export
  data ReceiveTimeout : Type where
    NoTimeout : ReceiveTimeout
    TimeoutAfter : (ms : Expr) -> (body : Vect (S k) Expr) -> ReceiveTimeout

public export
data Decl : Type where
  ADExport : Line -> List (String, Nat) -> Decl
  -- ADImport
  ADModule : Line -> (name : String) -> Decl
  -- ADFile
  ADFunDef : Line -> (name : String) -> (arity : Nat) -> (clauses : Vect (S k) (FunClause arity)) -> Decl
  -- ADFunSpec
  -- ADRecord
  -- ADType
  ADAttribute : Line -> (attr : String) -> (value : PrimTerm) -> Decl

-- LINE HELPERS

export
getLiteralLine : Literal -> Line
getLiteralLine (ALAtom l _) = l
getLiteralLine (ALChar l _) = l
getLiteralLine (ALFloat l _) = l
getLiteralLine (ALInteger l _) = l
getLiteralLine (ALCharlist l _) = l

export
getGuardLine : Guard -> Line
getGuardLine (AGLiteral lit) = getLiteralLine lit
getGuardLine (AGCons l _ _) = l
getGuardLine (AGFunCall l _ _) = l
getGuardLine (AGNil l) = l
getGuardLine (AGOp l _ _ _) = l
getGuardLine (AGTuple l _) = l
getGuardLine (AGVar l _) = l