Newt in Newt compiles (but does not run)

done/Lib/CompileExp.newt

@@ -0,0 +1,170 @@
+-- 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 Data.List
+
+import Lib.Types -- Name / Tm
+import Lib.TopContext
+import Lib.Prettier
+import Lib.Util
+
+
+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
+  CLam : Name -> CExp -> CExp
+  CFun : List Name -> CExp -> CExp
+  -- REVIEW This feels like a hack, but if we put CLam here, the
+  -- deBruijn gets messed up in code gen
+  CApp : CExp -> List CExp -> Int -> CExp
+  -- TODO make DCon/TCon app separate so we can specialize
+  -- U / Pi are compiled to type constructors
+  CCase : CExp -> List CAlt -> CExp
+  CRef : Name -> CExp
+  CMeta : Int -> CExp
+  CLit : Literal -> CExp
+  CLet : Name -> CExp -> CExp -> CExp
+  CLetRec : Name -> CExp -> CExp -> CExp
+  CErased : 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
+
+
+piArity : Tm -> Nat
+piArity (Pi _ _ _ quant _ b) = S (piArity b)
+piArity _ = 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 : FC -> QName -> M Nat
+arityForName fc nm = do
+  top <- get
+  case lookup nm top of
+  -- let the magic hole through for now (will generate bad JS)
+    Nothing => error fc "Name \{show nm} not in scope"
+    (Just (MkEntry _ name type Axiom)) => pure Z
+    (Just (MkEntry _ name type (TCon strs))) => pure $ piArity type
+    (Just (MkEntry _ name type (DCon k str))) => pure $ cast k
+    (Just (MkEntry _ name type (Fn t))) => pure $ lamArity t
+    (Just (MkEntry _ name type (PrimTCon))) => pure $ piArity type
+    -- Assuming a primitive can't return a function
+    (Just (MkEntry _ name type (PrimFn t used))) => pure $ piArity type
+
+
+compileTerm : Tm -> M CExp
+
+-- need to eta out extra args, fill in the rest of the apps
+apply : CExp -> List CExp -> SnocList CExp -> Nat -> Tm -> M CExp
+-- out of args, make one up (fix that last arg)
+apply t Nil acc (S k) ty = pure $ CApp t (acc <>> Nil) (1 + cast k)
+  -- inserting Clam, index wrong?
+  -- CLam "eta\{show k}" !(apply t Nil (acc :< CBnd k) k ty)
+apply t (x :: xs) acc (S k) (Pi y str icit Zero a b) =  apply t xs (acc :< CErased) k b
+apply t (x :: xs) acc (S k) (Pi y str icit Many a b) =  apply t xs (acc :< x) k b
+-- see if there is anything we have to handle here
+apply t (x :: xs) acc (S k) ty = error (getFC ty) "Expected pi \{showTm ty}. Overapplied function that escaped type checking?"
+-- once we hit zero, we fold the rest
+apply t ts acc Z ty = 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
+
+-- apply : CExp -> List CExp -> SnocList CExp -> Int -> M CExp
+-- -- out of args, make one up
+-- apply t Nil acc (S k) = pure $
+--   CLam "eta\{show k}" !(apply t Nil (acc :< CBnd k) k)
+-- apply t (x :: xs) acc (S k) = apply t xs (acc :< x) k
+-- apply t ts acc 0 = go (CApp t (acc <>> Nil)) ts
+--   where
+--     go : CExp -> List CExp -> M CExp
+--     -- drop zero arg call
+--     go (CApp t Nil) args = go t args
+--     go t Nil = pure t
+--     go t (arg :: args) = go (CApp t (arg :: Nil)) args
+
+compileTerm (Bnd _ k) = pure $ CBnd k
+-- need to eta expand to arity
+compileTerm t@(Ref fc nm _) = do
+  top <- get
+  let (Just (MkEntry _ _ type _)) = lookup nm top
+          | Nothing => error fc "Undefined name \{show nm}"
+  arity <- arityForName fc nm
+  apply (CRef (show nm)) Nil Lin arity type
+
+compileTerm (Meta _ k) = pure $ CRef "meta$\{show k}" -- FIXME
+compileTerm (Lam _ nm _ _ t) = CLam nm <$> compileTerm t
+compileTerm tm@(App _ _ _) = case funArgs tm of
+  (Meta _ k, args) => do
+        -- this will be undefined, should only happen for use metas
+        pure $ CApp (CRef "Meta\{show k}") Nil 0
+  (t@(Ref fc nm _), args) => do
+        args' <- traverse compileTerm args
+        arity <- arityForName fc nm
+        top <- get
+        let (Just (MkEntry _ _ type _)) = lookup nm top
+          | Nothing => error fc "Undefined name \{show nm}"
+        apply (CRef (show nm)) args' Lin arity type
+  (t, args) => do
+        debug $ \ _ => "apply other \{render 90 $ pprint Nil t}"
+        t' <- compileTerm t
+        args' <- traverse compileTerm args
+        apply t' args' Lin Z (UU emptyFC)
+        -- error (getFC t) "Don't know how to apply \{showTm t}"
+compileTerm (UU _) = pure $ CRef "U"
+compileTerm (Pi _ nm icit rig t u) = do
+  t' <- compileTerm t
+  u' <- compileTerm u
+  pure $ CApp (CRef "PiType") (t' :: 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 : 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
+
+