Refactor code generation to prepare for optimization passes

port/Data/IORef.newt

@@ -5,11 +5,11 @@ import Prelude
 -- We should test this at some point
 
 ptype IORef : U → U
-pfunc primNewIORef uses (MkIORes MkUnit) : ∀ a. a → IO (IORef a) := `(_, a) => (w) => MkIORes(null, [a], w)`
-pfunc primReadIORef uses (MkIORes MkUnit) : ∀ a. IORef a → IO a := `(_, ref) => (w) => MkIORes(null, ref[0], w)`
+pfunc primNewIORef uses (MkIORes) : ∀ a. a → IO (IORef a) := `(_, a) => (w) => Prelude_MkIORes(null, [a], w)`
+pfunc primReadIORef uses (MkIORes) : ∀ a. IORef a → IO a := `(_, ref) => (w) => Prelude_MkIORes(null, ref[0], w)`
 pfunc primWriteIORef uses (MkIORes MkUnit) : ∀ a. IORef a → a → IO Unit := `(_, ref, a) => (w) => {
   ref[0] = a
-  return MkIORes(null,MkUnit,w)
+  return Prelude_MkIORes(null,Prelude_MkUnit,w)
 }`
 
 newIORef : ∀ io a. {{HasIO io}} → a → io (IORef a)

port/Lib/Compile.newt

@@ -278,86 +278,90 @@ maybeWrap : JSStmt Return -> JSExp
 maybeWrap (JReturn exp) = exp
 maybeWrap stmt = Apply (JLam Nil stmt) Nil
 
-entryToDoc : TopEntry -> M Doc
-entryToDoc (MkEntry _ name ty (Fn tm)) = do
+-- convert a Def to a Doc (compile to javascript)
+defToDoc : QName → Def → M Doc
+defToDoc name (Fn tm) = do
   debug $ \ _ => "compileFun \{render 90 $ pprint Nil tm}"
   ct <- compileFun tm
   let exp = maybeWrap $ termToJS emptyJSEnv ct JReturn
   pure $ text "const" <+> jsIdent (show name) <+> text "=" <+/> expToDoc exp ++ text ";"
-entryToDoc (MkEntry _ name type Axiom) = pure $ text ""
-entryToDoc (MkEntry _ name type (TCon strs)) = pure $ dcon name (piArity type)
-entryToDoc (MkEntry _ name type (DCon arity str)) = pure $ dcon name (cast arity)
-entryToDoc (MkEntry _ name type PrimTCon) = pure $ dcon name (piArity type)
-entryToDoc (MkEntry _ name _ (PrimFn src _)) = pure $ text "const" <+> jsIdent (show name) <+> text "=" <+> text src
+defToDoc name Axiom = pure $ text ""
+defToDoc name (DCon arity str) = pure $ dcon name (cast arity)
+defToDoc name (TCon arity strs) = pure $ dcon name (cast arity)
+defToDoc name (PrimTCon arity) = pure $ dcon name (cast arity)
+defToDoc name (PrimFn src _) = pure $ text "const" <+> jsIdent (show name) <+> text "=" <+> text src
 
-
-process : (List QName × List Doc) -> QName -> M (List QName × List Doc)
-walkTm : Tm -> (List QName × List Doc) -> M (List QName × List Doc)
-walkAlt : (List QName × List Doc) -> CaseAlt -> M (List QName × List Doc)
-
-walkAlt acc (CaseDefault t) = walkTm t acc
-walkAlt acc (CaseCons name args t) = walkTm t acc
-walkAlt acc (CaseLit lit t) = walkTm t acc
-
-
-walkTm (Ref x nm) acc = process acc nm
-walkTm (Lam x str _ _ t) acc = walkTm t acc
-walkTm (App x t u) acc = walkTm u acc >>= walkTm t
-walkTm (Pi x str icit y t u) acc = walkTm u acc >>= walkTm t
-walkTm (Let x str t u) acc = walkTm u acc >>= walkTm t
-walkTm (LetRec x str _ t u) acc = walkTm u acc >>=  walkTm t
-walkTm (Case x t alts) acc = foldlM walkAlt acc alts
-walkTm _ acc = pure acc
-
--- This version (call `reverse ∘ snd <$> process "main" (Nil × Nil)`) will do dead
--- code elimination, but the Prelude js primitives are reaching for
--- stuff like True, False, MkUnit, fs which get eliminated
-process (done,docs) nm = do
-  let (False) = elem nm done | _ => pure (done,docs)
-  top <- get
-  case lookup nm top of
-    Nothing => error emptyFC "\{show nm} not in scope"
-    Just entry@(MkEntry _ name ty (PrimFn src used)) => do
-      (done,docs) <- foldlM assign (nm :: done, docs) used
-      edoc <- entryToDoc entry
-      pure (done, edoc :: docs)
-    Just (MkEntry _ name ty (Fn tm)) => do
-                debug $ \ _ => "compileFun \{render 90 $ pprint Nil tm}"
-                ct <- compileFun tm
-                -- If ct has zero arity and is a compount expression, this fails..
-                let exp = maybeWrap $ termToJS emptyJSEnv ct JReturn
-                let doc = text "const" <+> jsIdent (show name) <+> text "=" <+/> expToDoc exp ++ text ";"
-                (done,docs) <- walkTm tm (nm :: done, docs)
-                pure (done, doc :: docs)
-    Just entry => do
-      edoc <- entryToDoc entry
-      pure (nm :: done, edoc :: docs)
+-- Collect the QNames used in a term
+getNames : Tm -> List QName -> List QName
+getNames (Ref x nm) acc = nm :: acc
+getNames (Lam x str _ _ t) acc = getNames t acc
+getNames (App x t u) acc = getNames u $ getNames t acc
+getNames (Pi x str icit y t u) acc = getNames u $ getNames t acc
+getNames (Let x str t u) acc = getNames u $ getNames t acc
+getNames (LetRec x str _ t u) acc = getNames u $  getNames t acc
+getNames (Case x t alts) acc = foldl getAltNames acc alts
   where
-    assign : (List QName × List Doc) -> String -> M (List QName × List Doc)
-    assign (done, docs) nm = do
-      top <- get
-      case lookupRaw nm top of
-        Nothing => pure (done, docs)
-        (Just (MkEntry fc name type def)) => do
-          let tag = QN Nil nm
-          let (False) = elem tag done | _ => pure (done,docs)
-          (done,docs) <- process (done, docs) name
-          let doc = text "const" <+> jsIdent nm <+> text "=" <+> jsIdent (show name) ++ text ";"
-          pure (tag :: done, doc :: docs)
+    getAltNames : List QName -> CaseAlt -> List QName
+    getAltNames acc (CaseDefault t) = getNames t acc
+    getAltNames acc (CaseCons name args t) = getNames t acc
+    getAltNames acc (CaseLit lit t) = getNames t acc
+getNames _ acc = acc
 
+-- returns a QName -> Def of in-use entries
+-- This will be what we work on for optimization passes
+getEntries : SortedMap QName Def → QName → M (SortedMap QName Def)
+getEntries acc name = do
+  top <- get
+  case lookup name top of
+    Nothing => do
+      putStrLn "bad name \{show name}"
+      pure acc
+    Just (MkEntry _ name type def@(Fn exp)) => case lookupMap' name acc of
+      Just _ => pure acc
+      Nothing =>
+        let acc = updateMap name def acc in
+        foldlM getEntries acc $ getNames exp Nil
+    Just (MkEntry _ name type def@(PrimFn _ used)) =>
+      let acc = updateMap name def acc in
+      foldlM getEntries acc used
+    Just entry => pure $ updateMap name entry.def acc
 
+-- sort names by dependencies
+-- In JS this is only really needed for references that don't fall
+-- under a lambda.
+sortedNames : SortedMap QName Def → QName → List QName
+sortedNames defs qn = go Nil Nil qn
+  where
+    go : List QName → List QName → QName → List QName
+    go loop acc qn =
+      -- O(n^2) it would be more efficient to drop qn from the map
+      if elem qn loop || elem qn acc then acc else
+      case lookupMap' qn defs of
+        Nothing => acc
+        Just (Fn tm) => qn :: foldl (go $ qn :: loop) acc (getNames tm Nil)
+        Just (PrimFn src used) => qn :: foldl (go $ qn :: loop) acc used
+        Just def => qn :: acc
 
+-- given a initial function, return a dependency-ordered list of javascript source
+process : QName → M (List Doc)
+process name = do
+  let wat = QN ("Prelude" :: Nil) "arrayToList"
+  entries <- getEntries EmptyMap name
+  let names = sortedNames entries name
+  for names $ \ nm => case lookupMap nm entries of
+      Nothing => error emptyFC "MISS \{show nm}"
+      Just _ => pure MkUnit
+  mapM (uncurry defToDoc) $ mapMaybe (\x => lookupMap x entries) names
 
 compile : M (List Doc)
 compile = do
   top <- get
   case lookupRaw "main" top of
     Just (MkEntry fc name type def) => do
-      tmp <- snd <$> process (Nil, Nil) name
+      tmp <- process name
+      -- tack on call to main function
       let exec = stmtToDoc $ JPlain $ Apply (Var $ show name) Nil
       pure $ reverse (exec :: tmp)
-    -- If there is no main, compile everything for the benefit of the playground
-    Nothing => do
-      top <- get
-      traverse entryToDoc $ map snd $ toList top.defs
-
+    Nothing =>
+      -- TODO maybe dump everything if there is no main
+      error emptyFC "No main function found"

port/Lib/CompileExp.newt

@@ -48,9 +48,6 @@ 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
@@ -62,10 +59,10 @@ arityForName fc nm = do
   -- 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 (TCon arity strs))) => pure $ cast arity
     (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
+    (Just (MkEntry _ name type (PrimTCon arity))) => pure $ cast arity
     -- Assuming a primitive can't return a function
     (Just (MkEntry _ name type (PrimFn t used))) => pure $ piArity type
 
@@ -76,8 +73,7 @@ compileTerm : Tm -> M CExp
 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)
+-- FIXME - this should be handled by Erasure.newt (wdiff of esbuild output says this is still used)
 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
@@ -91,19 +87,6 @@ apply t ts acc Z ty = go (CApp t (acc <>> Nil) 0) ts
     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
@@ -159,7 +142,6 @@ compileTerm (LetRec _ nm _ t u) = do
   pure $ CLetRec nm t' u'
 compileTerm (Erased _) = pure CErased
 
-
 compileFun : Tm -> M CExp
 compileFun tm = go tm Lin
   where

port/Lib/Elab.newt

@@ -620,7 +620,7 @@ primType : FC -> QName -> M Val
 primType fc nm = do
   top <- get
   case lookup nm top of
-      Just (MkEntry _ name ty PrimTCon) => pure $ VRef fc name Lin
+      Just (MkEntry _ name ty (PrimTCon _)) => pure $ VRef fc name Lin
       _ => error fc "Primitive type \{show nm} not in scope"
 
 infer : Context -> Raw  -> M (Tm × Val)
@@ -688,7 +688,7 @@ getConstructors ctx scfc (VRef fc nm _) = do
     lookupTCon str = do
       top <- get
       case lookup nm top of
-        (Just (MkEntry _ name type (TCon names))) => pure names
+        (Just (MkEntry _ name type (TCon _ names))) => pure names
         _ => error scfc "Not a type constructor \{show nm}"
     lookupDCon : QName -> M (QName × Int × Tm)
     lookupDCon nm = do
@@ -1287,6 +1287,8 @@ undo prev ((DoExpr fc tm) :: Nil) = pure tm
 -- TODO decide if we want to use >> or just >>=
 undo prev ((DoExpr fc tm) :: xs) = do
   xs' <- undo fc xs
+  -- output is bigger, not sure if it helps inference or not
+  -- pure $ RApp fc (RApp fc (RVar fc "_>>_") tm Explicit) xs' Explicit
   pure $ RApp fc (RApp fc (RVar fc "_>>=_") tm Explicit) (RLam fc (BI fc "_" Explicit Many) xs') Explicit
 undo prev ((DoLet fc nm tm) :: xs) = RLet fc nm (RImplicit fc) tm <$> undo fc xs
 undo prev ((DoArrow fc left@(RVar fc' nm) right Nil) :: xs) = do

port/Lib/Eval.newt

@@ -53,7 +53,7 @@ lookupVar env k = let l = cast $ length env in
       Nothing => Nothing
 
 -- hoping to apply what we got via pattern matching
-
+-- TODO see if we can drop this after updating pattern matching
 unlet : Env -> Val -> M Val
 unlet env t@(VVar fc k sp) = case lookupVar env k of
       Just tt@(VVar fc' k' sp') => do
@@ -65,7 +65,7 @@ unlet env t@(VVar fc k sp) = case lookupVar env k of
         pure t
 unlet env x = pure x
 
-
+-- Try applying VRef to spine, back out if it is stuck
 tryEval : Env -> Val -> M (Maybe Val)
 tryEval env (VRef fc k sp) = do
   top <- get

port/Lib/ProcessDecl.newt

@@ -120,18 +120,19 @@ processDecl ns  (TypeSig fc names tm) = do
 processDecl ns (PType fc nm ty) = do
   top <- get
   ty' <- check (mkCtx fc) (maybe (RU fc) id ty) (VU fc)
-  setDef (QN ns nm) fc ty' PrimTCon
+  let arity = cast $ piArity ty'
+  setDef (QN ns nm) fc ty' (PrimTCon arity)
 
 processDecl ns  (PFunc fc nm used ty src) = do
   top <- get
   ty <- check (mkCtx fc) ty (VU fc)
   ty' <- nf Nil ty
   log 1 $ \ _ => "pfunc \{nm} : \{render 90 $ pprint Nil ty'} = \{show src}"
-  -- TODO wire through fc?
-  for used $ \ name => case lookupRaw name top of
+  -- TODO wire through fc for not in scope error
+  used' <- for used $ \ name => case lookupRaw name top of
     Nothing => error fc "\{name} not in scope"
-    _ => pure MkUnit
-  setDef (QN ns nm) fc ty' (PrimFn src used)
+    Just entry => pure entry.name
+  setDef (QN ns nm) fc ty' (PrimFn src used')
 
 processDecl ns (Def fc nm clauses) = do
   log 1 $ \ _ => "-----"
@@ -265,7 +266,7 @@ processDecl ns (Instance instfc ty decls) = do
 
   let (Ref _ tconName, args) = funArgs codomain
     | (tm, _) => error tyFC "\{render 90 $ pprint Nil codomain} doesn't appear to be a TCon application"
-  let (Just (MkEntry _ name type (TCon cons))) = lookup tconName top
+  let (Just (MkEntry _ name type (TCon _ cons))) = lookup tconName top
     | _ => error tyFC "\{show tconName} is not a type constructor"
   let (con :: Nil) = cons
     | _ => error tyFC "\{show tconName} has multiple constructors \{show cons}"
@@ -404,7 +405,8 @@ processDecl ns (Data fc nm ty cons) = do
           pure $ map (QN ns) names
       decl => throwError $ E (getFC decl) "expected constructor declaration"
   log 1 $ \ _ => "setDef \{nm}  TCon \{show $ join cnames}"
-  updateDef (QN ns nm) fc tyty (TCon (join cnames))
+  let arity = cast $ piArity tyty
+  updateDef (QN ns nm) fc tyty (TCon arity (join cnames))
   where
     binderName : Binder → Name
     binderName (MkBinder _ nm _ _ _) = nm

port/Lib/Types.newt

@@ -307,15 +307,15 @@ record MetaContext where
   next : Int
   mcmode : MetaMode
 
-data Def = Axiom | TCon (List QName) | DCon Int QName | Fn Tm | PrimTCon
-         | PrimFn String (List String)
+data Def = Axiom | TCon Int (List QName) | DCon Int QName | Fn Tm | PrimTCon Int
+         | PrimFn String (List QName)
 
 instance Show Def where
   show Axiom = "axiom"
-  show (TCon strs) = "TCon \{show strs}"
+  show (TCon _ strs) = "TCon \{show strs}"
   show (DCon k tyname) = "DCon \{show k} \{show tyname}"
   show (Fn t) = "Fn \{show t}"
-  show (PrimTCon) = "PrimTCon"
+  show (PrimTCon _) = "PrimTCon"
   show (PrimFn src used) = "PrimFn \{show src} \{show used}"
 
 -- entry in the top level context

port/Lib/Util.newt

@@ -5,6 +5,10 @@ import Lib.Common
 import Lib.Types
 import Data.List1
 
+-- pi arity is primitive functions
+piArity : Tm -> Nat
+piArity (Pi _ _ _ quant _ b) = S (piArity b)
+piArity _ = Z
 
 funArgs : Tm -> (Tm × List Tm)
 funArgs tm = go tm Nil

port/Node.newt

@@ -2,35 +2,36 @@ module Node
 
 import Prelude
 
-pfunc fs : JSObject := `require('fs')`
-pfunc getArgs : List String := `arrayToList(String, process.argv.slice(1))`
-pfunc readFile uses (fs MkIORes Left Right) : (fn : String) -> IO (Either String String) := `(fn) => (w) => {
+pfunc getArgs uses (arrayToList) : List String := `Prelude_arrayToList(null, process.argv.slice(1))`
+pfunc readFile uses (MkIORes Left Right) : (fn : String) -> IO (Either String String) := `(fn) => (w) => {
+  let fs = require('fs')
   let result
   try {
     let content = fs.readFileSync(fn, 'utf8')
-    result = Right(null, null, content)
+    result = Prelude_Right(null, null, content)
   } catch (e) {
     let err = ""+e
-    result = Left(null, null, e)
+    result = Prelude_Left(null, null, e)
   }
-  return MkIORes(null, result, w)
+  return Prelude_MkIORes(null, result, w)
 }`
 
 -- I wonder if I should automatically `uses` the constructors in the types
-pfunc writeFile uses (fs MkIORes MkUnit) : String → String → IO (Either String Unit) := `(fn, content) => (w) => {
+pfunc writeFile uses (MkIORes MkUnit) : String → String → IO (Either String Unit) := `(fn, content) => (w) => {
+  let fs = require('fs')
   let result
   try {
     fs.writeFileSync(fn, content, 'utf8')
-    result = Right(null, null, MkUnit)
+    result = Prelude_Right(null, null, Prelude_MkUnit)
   } catch (e) {
       let err = ""+e
-      result = Left(null, null, e)
+      result = Prelude_Left(null, null, e)
   }
-  return MkIORes(null, result, w)
+  return Prelude_MkIORes(null, result, w)
 }`
 
 -- maybe System.exit or something, like the original putStrLn msg >> exitFailure
-pfunc exitFailure : ∀ a. String → a := `(_, msg) => { 
+pfunc exitFailure : ∀ a. String → a := `(_, msg) => {
   console.log(msg);
   process.exit(1);
 }`

port/Serialize.newt

@@ -9,13 +9,14 @@ import Data.SortedMap
 -- this was an experiment, prepping for dumping module information
 -- it ends up with out of memory dumping defs of some of the files.
 -- Prelude is 114MB pretty-printed... gzip to 1M
-pfunc dumpObject uses (MkIORes MkUnit fs): ∀ a. String → a → IO Unit := `(_,fn,a) => (w) => {
+pfunc dumpObject uses (MkIORes MkUnit): ∀ a. String → a → IO Unit := `(_,fn,a) => (w) => {
+  let fs = require('fs')
   try {
     let {EncFile} = require('./serializer')
     let enc = EncFile.encode(a)
     fs.writeFileSync(fn, enc)
   } catch (e) {}
-  return MkIORes(null, MkUnit, w)
+  return Prelude_MkIORes(null, Prelude_MkUnit, w)
 }`
 
 -- for now, include src and use that to see if something changed