-- First pass of compilation
-- - work out arities and fully apply functions / constructors (currying)
--     currying is problemmatic because we need to insert lambdas (η-expand) and
--     it breaks all of the de Bruijn indices
-- - expand metas (this is happening earlier)
-- - erase stuff (there is another copy that essentially does the same thing)
-- I could make names unique (e.q. on lambdas), but I might want that to vary per backend?
module Lib.CompileExp

import Prelude
import Lib.Common
import Lib.Types -- Name / Tm
import Lib.TopContext
import Lib.Prettier
import Lib.Util
import Lib.Ref2
import Data.SortedMap

CExp : U

data CAlt : U where
  CConAlt : String -> List String -> CExp -> CAlt
  -- REVIEW keep var name?
  CDefAlt : CExp -> CAlt
  -- literal
  CLitAlt : Literal -> CExp -> CAlt

data CExp : U where
  CBnd : Int -> CExp
  -- How is CLam different from CFun with one arg?
  CLam : Name -> CExp -> CExp
  CFun : List Name -> CExp -> CExp
  CApp : CExp -> List CExp -> Int -> CExp
  CCase : CExp -> List CAlt -> CExp
  CRef : QName -> CExp
  CMeta : Int -> CExp
  CLit : Literal -> CExp
  CLet : Name -> CExp -> CExp -> CExp
  CLetRec : Name -> CExp -> CExp -> CExp
  CErased : CExp
  -- Data / type constructor
  CConstr : Name -> List CExp -> CExp
  -- Raw javascript for `pfunc`
  CRaw : String -> List QName -> CExp

-- I'm counting Lam in the term for arity.  This matches what I need in
-- code gen.

lamArity : Tm -> Nat
lamArity (Lam _ _ _ _ t) = S (lamArity t)
lamArity _ = Z



-- This is how much we want to curry at top level
-- leading lambda Arity is used for function defs and metas
-- TODO - figure out how this will work with erasure
arityForName : {{Ref2 Defs St}} → FC -> QName -> M Nat
arityForName fc nm = do
  defs <- getRef Defs
  case lookupMap' nm defs of
    Nothing => error fc "Name \{show nm} not in scope"
    (Just Axiom) => pure Z
    (Just (TCon arity strs)) => pure $ cast arity
    (Just (DCon k str)) => pure $ cast k
    (Just (Fn t)) => pure $ lamArity t
    (Just (PrimTCon arity)) => pure $ cast arity
    (Just (PrimFn t arity used)) => pure arity


compileTerm : {{Ref2 Defs St}} → Tm -> M CExp

-- need to eta out extra args, fill in the rest of the apps
-- NOW - maybe eta here instead of Compile.newt, drop number on CApp
-- The problem would be deBruijn. We have to put the app under CLam
-- which would mess up all of the deBruijn (unless we push it out)

apply : CExp -> List CExp -> SnocList CExp -> Nat -> M CExp
-- out of args, make one up (fix that last arg)
apply t Nil acc (S k) =
  pure $ CApp t (acc <>> Nil) (1 + cast k)
apply t (x :: xs) acc (S k)  =  apply t xs (acc :< x) k
-- once we hit zero, we fold the rest
apply t ts acc Z = go (CApp t (acc <>> Nil) 0) ts
  where
    go : CExp -> List CExp -> M CExp
    -- drop zero arg call
    go (CApp t Nil 0) args = go t args
    go t Nil = pure t
    go t (arg :: args) = go (CApp t (arg :: Nil) 0) args

compileTerm (Bnd _ k) = pure $ CBnd k
-- need to eta expand to arity
compileTerm t@(Ref fc nm) = do
  arity <- arityForName fc nm
  case arity of
    -- we don't need to curry functions that take one argument
    (S Z) => pure $ CRef nm
    _ => apply (CRef nm) Nil Lin arity

compileTerm (Meta _ k) = pure $ CRef (QN Nil "meta$\{show k}") -- FIXME should be exception
compileTerm (Lam _ nm _ _ t) = CLam nm <$> compileTerm t
compileTerm tm@(App _ _ _) = case funArgs tm of
  (Meta _ k, args) => do
        error (getFC tm) "Compiling an unsolved meta \{show tm}"
        -- info (getFC tm) "Compiling an unsolved meta \{show tm}"
        -- pure $ CApp (CRef "Meta\{show k}") Nil 0
  (t@(Ref fc nm), args) => do
        args' <- traverse compileTerm args
        arity <- arityForName fc nm
        apply (CRef nm) args' Lin arity
  (t, args) => do
        debug $ \ _ => "apply other \{render 90 $ pprint Nil t}"
        t' <- compileTerm t
        args' <- traverse compileTerm args
        apply t' args' Lin Z
        -- error (getFC t) "Don't know how to apply \{showTm t}"
compileTerm (UU _) = pure $ CRef (QN Nil "U")
compileTerm (Pi _ nm icit rig t u) = do
  t' <- compileTerm t
  u' <- compileTerm u
  pure $ CApp (CRef (QN Nil "PiType")) (t' :: CLam nm u' :: Nil) 0
compileTerm (Case _ t alts) = do
  t' <- compileTerm t
  alts' <- for alts $ \case
    CaseDefault tm => CDefAlt <$> compileTerm tm
    -- we use the base name for the tag, some primitives assume this
    CaseCons (QN ns nm) args tm => CConAlt nm args <$> compileTerm tm
    CaseLit lit tm => CLitAlt lit <$> compileTerm tm
  pure $ CCase t' alts'
compileTerm (Lit _ lit) = pure $ CLit lit
compileTerm (Let _ nm t u) = do
  t' <- compileTerm t
  u' <- compileTerm u
  pure $ CLet nm t' u'
compileTerm (LetRec _ nm _ t u) = do
  t' <- compileTerm t
  u' <- compileTerm u
  pure $ CLetRec nm t' u'
compileTerm (Erased _) = pure CErased

compileFun : {{Ref2 Defs St}} → Tm -> M CExp
compileFun tm = go tm Lin
  where
    go : Tm -> SnocList String -> M CExp
    go (Lam _ nm _ _ t) acc = go t (acc :< nm)
    go tm Lin = compileTerm tm
    go tm args = CFun (args <>> Nil) <$> compileTerm tm

-- What are the Defs used for above? (Arity for name)
compileDCon : QName → Int → CExp
compileDCon (QN _ nm) 0 = CConstr nm Nil
compileDCon (QN _ nm) arity =
  let args = map (\k => "h\{show k}") (range 0 arity) in
  CFun args $ CConstr nm $ map (\k => CBnd $ arity - k - 1) (range 0 arity)

-- probably want to drop the Ref2 when we can
defToCExp : {{Ref2 Defs St}} → (QName × Def) -> M (QName × CExp)
defToCExp (qn, Axiom)          = pure $ (qn, CErased)
defToCExp (qn, DCon arity _)   = pure $ (qn, compileDCon qn arity)
defToCExp (qn, TCon arity _)   = pure $ (qn, compileDCon qn arity)
defToCExp (qn, PrimTCon arity) = pure $ (qn, compileDCon qn arity)
defToCExp (qn, PrimFn src _ deps) = pure $ (qn, CRaw src deps)
defToCExp (qn, Fn tm)          = (_,_ qn) <$> compileFun tm