From d39c9aa9b2f532b6d74fe61401ea9fe8e5e1bbcf Mon Sep 17 00:00:00 2001
From: Steve Dunham <dunhamsteve@gmail.com>
Date: Sun, 11 Aug 2024 22:22:09 -0700
Subject: [PATCH] Add CExp

---
 src/Lib/Compile.idr    |  97 +++++++++++++++++++-------------
 src/Lib/CompileExp.idr | 124 +++++++++++++++++++++++++++++++++++++++++
 src/Main.idr           |   2 +-
 3 files changed, 183 insertions(+), 40 deletions(-)
 create mode 100644 src/Lib/CompileExp.idr

diff --git a/src/Lib/Compile.idr b/src/Lib/Compile.idr
index 370fc84..d5ba251 100644
--- a/src/Lib/Compile.idr
+++ b/src/Lib/Compile.idr
@@ -2,6 +2,7 @@ module Lib.Compile
 
 import Lib.Types
 import Lib.Prettier
+import Lib.CompileExp
 import Data.String
 
 data Kind = Plain | Return | Assign String
@@ -33,42 +34,56 @@ data JSStmt : Kind -> Type where
 
 Cont e = JSExp -> JSStmt e
 
+-- FIXME - add names to env so we can guarantee fresh names in the generated javascript.
+JSEnv = List JSExp
+
+-- Stuff nm.h1, nm.h2, ... into environment
+mkEnv : String -> Nat -> List JSExp -> List String -> List JSExp
+mkEnv nm k env [] = env
+mkEnv nm k env (x :: xs) = mkEnv nm (S k) (Dot (Var nm) "h\{show k}" :: env) xs
+
+
 -- This is inspired by A-normalization, look into the continuation monad
 -- There is an index on JSStmt, adopted from Stefan Hoeck's code.
 --
 -- Here we turn a Term into a statement (which may be a sequence of statements), there
 -- is a continuation, which turns the final JSExpr into a JSStmt, and the function returns
 -- a JSStmt, wrapping recursive calls in JSnoc if necessary.
-termToJS : List JSExp -> Tm -> Cont e -> JSStmt e
-termToJS env (Bnd _ k) f = case getAt k env of
+termToJS : List JSExp -> CExp -> Cont e -> JSStmt e
+termToJS env (CBnd k) f = case getAt k env of
   (Just e) => f e
   Nothing => ?bad_bounds
-termToJS env (Ref _ nm y) f = f $ Var nm
-termToJS env (Meta _ k) f = f $ LitString "META \{show k}"
-termToJS env (Lam _ nm t) f =
-  let nm' = "nm$\{show $ length env}"
+termToJS env (CLam nm t) f =
+  let nm' = "\{nm}$\{show $ length env}"
       env' = (Var nm' :: env)
   in f $ JLam [nm'] (termToJS env' t JReturn)
-termToJS env (App _ t u) f = termToJS env t (\ t' => termToJS env u (\ u' => f (Apply t' [u'])))
-termToJS env (U _) f = f $ LitString "U"
-termToJS env (Pi _ nm icit t u) f = f $ LitString "Pi \{nm}"
-termToJS env (Case _ t alts) f =
-  -- need to assign the scrutinee to a variable
-  -- and add (Bnd -> JSExpr map)
-  termToJS env t (\ t' =>
-    let (l,c) = getFC t in
-    let nm = "sc$\{show l}$\{show c}" in
-    JSnoc (JConst nm t')
-      (JCase (Dot (Var nm) "tag") (map (termToJSAlt nm) alts) Nothing))
+termToJS env (CFun nms t) f =
+  let nms' = map (\nm => "\{nm}$\{show $ length env}") nms
+      env' = foldl (\ e, nm => Var nm :: e) env nms'
+  in f $ JLam nms' (termToJS env' t JReturn)
+termToJS env (CRef nm) f = f $ Var nm
+termToJS env (CMeta k) f = f $ LitString "META \{show k}"
+termToJS env (CApp t args) f = termToJS env t (\ t' => argsToJS args [<] (\ args' => f (Apply t' args')))
   where
-    -- Stuff nm.h1, nm.h2, ... into environment
-    mkEnv : String -> Nat -> List JSExp -> List String -> List JSExp
-    mkEnv nm k env [] = env
-    mkEnv nm k env (x :: xs) = mkEnv nm (S k) (Dot (Var nm) "h\{show k}" :: env) xs
+    argsToJS : List CExp -> SnocList JSExp -> (List JSExp -> JSStmt e) -> JSStmt e
+    argsToJS [] acc k = k (acc <>> [])
+    argsToJS (x :: xs) acc k = termToJS env x (\ x' => argsToJS xs (acc :< x') k)
 
-    termToJSAlt : String -> CaseAlt -> (String, JSStmt e)
-    termToJSAlt nm (CaseDefault u) = ?handle_default_case
-    termToJSAlt nm (CaseCons name args u) =
+
+termToJS env (CCase t alts def) f =
+  -- need to assign the scrutinee to a variable (unless it is a var already?)
+  -- and add (Bnd -> JSExpr map)
+  -- TODO default case, let's drop the extra field.
+
+  termToJS env t $ \case
+     (Var nm) =>       (JCase (Dot (Var nm) "tag") (map (termToJSAlt nm) alts) Nothing)
+     t' =>
+      let nm = "sc$\{show $ length env}" in
+        JSnoc (JConst nm t')
+          (JCase (Dot (Var nm) "tag") (map (termToJSAlt nm) alts) Nothing)
+  where
+    termToJSAlt : String -> CAlt -> (String, JSStmt e)
+    termToJSAlt nm (CConAlt name args u) =
       let env' = mkEnv nm 0 env args in
       (name, termToJS env' u f)
 
@@ -110,30 +125,34 @@ stmtToDoc (JCase sc alts y) =
 
 -- FIXME - if the result is JSnoc, we get extra top level code
 -- If we make top level 0-arity values lazy, this won't happen
-function : String -> Tm -> Doc
-function nm tm = stmtToDoc $ termToJS [] tm (JConst nm)
+-- function : String -> Tm -> Doc
+-- function nm tm = stmtToDoc $ termToJS [] tm (JConst nm)
+
+
+mkArgs : Nat -> List String -> List String
+mkArgs Z acc = acc
+mkArgs (S k) acc = mkArgs k ("h\{show k}" :: acc)
 
 dcon : String -> Nat -> Doc
 dcon nm arity =
   let args := mkArgs arity []
       obj := ("tag", LitString nm) :: map (\x => (x, Var x)) args
   in stmtToDoc $ JConst nm (JLam args (JReturn (LitObject obj)))
-  where
-    -- FIXME arity wrong
-    mkArgs : Nat -> List String -> List String
-    mkArgs Z acc = acc
-    mkArgs (S k) acc = mkArgs k ("h\{show k}" :: acc)
 
-entryToDoc : TopEntry -> Maybe Doc
-entryToDoc (MkEntry name type (Fn tm)) =
-  let body = stmtToDoc $ termToJS [] tm JPlain in
-  Just (text "const" <+> text name <+> "=" <+/> body)
-entryToDoc (MkEntry name type Axiom) = Nothing
-entryToDoc (MkEntry name type (TCon strs)) = Nothing
-entryToDoc (MkEntry name type (DCon arity str)) = Just $ dcon name arity
+
+entryToDoc : TopEntry -> M Doc
+entryToDoc (MkEntry name ty (Fn tm)) = do
+  -- so this has a bunch of lambdas on it now, which we want to consolidate
+  -- and we might need betas? It seems like a mirror of what happens in CExp
+  ct <- compileFun tm
+  let body = stmtToDoc $ termToJS [] ct JPlain
+  pure (text "const" <+> text name <+> text "=" <+/> body)
+entryToDoc (MkEntry name type Axiom) = pure ""
+entryToDoc (MkEntry name type (TCon strs)) = pure ""
+entryToDoc (MkEntry name type (DCon arity str)) = pure $ dcon name arity
 
 export
 compile : M Doc
 compile = do
   top <- get
-  pure $ stack $ mapMaybe entryToDoc top.defs
+  pure $ stack $ !(traverse entryToDoc top.defs)
diff --git a/src/Lib/CompileExp.idr b/src/Lib/CompileExp.idr
new file mode 100644
index 0000000..4605467
--- /dev/null
+++ b/src/Lib/CompileExp.idr
@@ -0,0 +1,124 @@
+||| First pass of compilation
+||| - work out arities and fully apply functions / constructors
+||| - expand metas
+||| 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.TT -- lookupMeta
+
+public export
+data CExp : Type
+
+public export
+data CAlt : Type where
+  CConAlt : String -> List String -> CExp -> CAlt
+  -- literal
+
+data CExp : Type where
+  CBnd : Nat -> CExp
+  CLam : Name -> CExp -> CExp
+  CFun : List Name -> CExp -> CExp
+  CApp : CExp -> List CExp -> CExp
+  -- TODO make DCon/TCon app separate so we can specialize
+  -- U / Pi are compiled to type constructors
+  CCase : CExp -> List CAlt -> Maybe CExp -> CExp
+  CRef : Name -> CExp
+  CMeta : Nat -> CExp
+
+funArgs : Tm -> (Tm, List Tm)
+funArgs tm = go tm []
+  where
+    go : Tm -> List Tm -> (Tm, List Tm)
+    go (App _ t u) args = go t (u :: args)
+    go t args = (t, args)
+
+||| I'm counting Lam in the term for arity.  This matches what I need in
+||| code gen.
+export
+getArity : Tm -> Nat
+getArity (Lam _ _ t) = S (getArity t)
+getArity _ = Z
+
+arityForName : FC -> Name -> M Nat
+arityForName fc nm = case lookup nm !get of
+  Nothing => error fc "Name \{show nm} not in scope"
+  (Just (MkEntry name type Axiom)) => pure 0
+  (Just (MkEntry name type (TCon strs))) => pure 0 -- FIXME
+  (Just (MkEntry name type (DCon k str))) => pure k
+  (Just (MkEntry name type (Fn t))) => pure $ getArity t
+
+export
+compileTerm : Tm -> M CExp
+
+-- need to eta out extra args, fill in the rest of the apps
+apply : CExp -> List CExp -> SnocList CExp -> Nat -> M CExp
+-- out of args, make one up
+apply t [] acc (S k) = pure $
+  CLam "eta\{show k}" !(apply t [] (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 <>> [])) ts
+  where
+    go : CExp -> List CExp -> M CExp
+    go t [] = pure t
+    go t (arg :: args) = go (CApp t [arg]) args
+
+compileTerm (Bnd _ k) = pure $ CBnd k
+-- need to eta expand to arity
+compileTerm t@(Ref fc nm _) = apply (CRef nm) [] [<] !(arityForName fc nm)
+
+-- need to zonk
+compileTerm (Meta _ k) = pure $ CRef "meta$\{show k}" -- FIXME
+compileTerm (Lam _ nm t) = pure $ CLam nm !(compileTerm t)
+compileTerm tm@(App _ _ _) with (funArgs tm)
+  _ | (Meta _ k, args) = do
+        -- FIXME get arity or zonk?
+        -- Maybe we should be storing the Term without the lambdas...
+        -- we don't have a lot here, but JS has an "environment" with names and
+        -- we can assume fully applied.
+        meta <- lookupMeta k
+        args' <- traverse compileTerm args
+        -- apply (CRef "Meta\{show k}") args' [<] 0
+        arity <- case meta of
+                -- maybe throw
+                (Unsolved x j xs) => pure 0
+                (Solved j tm) => pure $ getArity !(quote 0 tm)
+        apply (CRef "Meta\{show k}") args' [<] arity
+  _ | (t@(Ref fc nm _), args) = do
+        t' <- compileTerm t
+        args' <- traverse compileTerm args
+        apply t' args' [<] !(arityForName fc nm)
+  _ | (t, args) = do
+        debug "apply \{pprint [] t}"
+        t' <- compileTerm t
+        args' <- traverse compileTerm args
+        apply t' args' [<] 0
+compileTerm (U _) = pure $ CRef "U"
+compileTerm (Pi _ nm icit t u) = pure $ CApp (CRef "PiType") [ !(compileTerm t), CLam nm !(compileTerm u)]
+compileTerm (Case _ t alts) = do
+  t' <- compileTerm t
+  alts' <- catMaybes <$> traverse (\case
+    CaseDefault tm => pure Nothing
+    CaseCons nm args tm => pure $ Just $ CConAlt nm args !(compileTerm tm)) alts
+  def <- getDefault alts
+  pure $ CCase t' alts' def
+  where
+    getDefault : List CaseAlt -> M (Maybe CExp)
+    getDefault [] = pure Nothing
+    getDefault (CaseDefault u :: _) = Just <$> compileTerm u
+    getDefault (_ :: xs) = getDefault xs
+
+
+export
+compileFun : Tm -> M CExp
+compileFun tm = go tm []
+  where
+    go : Tm -> List String -> M CExp
+    go (Lam _ nm t) acc = go t (nm :: acc)
+    go tm args = pure $ CFun (reverse args) !(compileTerm tm)
+
+
+
diff --git a/src/Main.idr b/src/Main.idr
index 62a2deb..da5291e 100644
--- a/src/Main.idr
+++ b/src/Main.idr
@@ -71,5 +71,5 @@ main = do
   -- we'll need to reset for each file, etc.
   ctx <- empty
   Right _  <- runEitherT $ runStateT ctx $ main'
-    | Left (E (c, r) str) => putStrLn "Error: \{show c} \{show r} \{show str}"
+    | Left (E (c, r) str) => putStrLn "ERROR at (\{show c}, \{show r}): \{show str}"
   putStrLn "done"