first pass at liftWhere

src/Lib/Compile.newt

@@ -7,10 +7,13 @@ import Lib.Types
 import Lib.Prettier
 import Lib.CompileExp
 import Lib.TopContext
+import Lib.LiftWhere
+import Lib.Ref2
 import Lib.Erasure
 import Data.String
 import Data.Int
 import Data.SortedMap
+import Data.IORef
 
 data StKind = Plain | Return | Assign String
 
@@ -280,18 +283,18 @@ maybeWrap (JReturn exp) = exp
 maybeWrap stmt = Apply (JLam Nil stmt) Nil
 
 -- convert a Def to a Doc (compile to javascript)
-defToDoc : QName → Def → M Doc
+defToDoc : {{Ref2 Defs St}} → QName → Def → M Doc
 defToDoc name (Fn tm) = do
   debug $ \ _ => "compileFun \{render 90 $ pprint Nil tm}"
-  tm' <- erase Nil tm Nil
-  ct <- compileFun tm'
+  -- tm' <- erase Nil tm Nil
+  ct <- compileFun tm
   let exp = maybeWrap $ termToJS emptyJSEnv ct JReturn
   pure $ text "const" <+> jsIdent (show name) <+> text "=" <+/> expToDoc exp ++ text ";"
 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
+defToDoc name (PrimFn src _ _) = pure $ text "const" <+> jsIdent (show name) <+> text "=" <+> text src
 
 -- Collect the QNames used in a term
 getNames : Tm -> List QName -> List QName
@@ -323,7 +326,7 @@ getEntries acc name = do
       Nothing =>
         let acc = updateMap name def acc in
         foldlM getEntries acc $ getNames exp Nil
-    Just (MkEntry _ name type def@(PrimFn _ used)) =>
+    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
@@ -341,14 +344,32 @@ sortedNames defs qn = go Nil Nil qn
       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 (PrimFn src _ used) => qn :: foldl (go $ qn :: loop) acc used
         Just def => qn :: acc
 
+eraseEntries : {{Ref2 Defs St}} → M Unit
+eraseEntries = do
+  defs <- getRef Defs
+  ignore $ traverse go $ toList defs
+  where
+    go : {{Ref2 Defs St}} → (QName × Def) → M Unit
+    go (qn, Fn tm) = do
+      tm' <- erase Nil tm Nil
+      modifyRef Defs $ updateMap qn (Fn tm')
+    go _ = pure MkUnit
+
 -- 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
+
+  -- Maybe move this dance into liftWhere
+  ref <- newIORef entries
+  let foo = MkRef ref -- for the autos below
+  eraseEntries
+  liftWhere
+  entries <- readIORef ref
   let names = sortedNames entries name
   for names $ \ nm => case lookupMap nm entries of
       Nothing => error emptyFC "MISS \{show nm}"

src/Lib/CompileExp.newt

@@ -13,6 +13,8 @@ import Lib.Types -- Name / Tm
 import Lib.TopContext
 import Lib.Prettier
 import Lib.Util
+import Lib.Ref2
+import Data.SortedMap
 
 CExp : U
 
@@ -27,11 +29,7 @@ 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
@@ -52,35 +50,28 @@ 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 : FC -> QName -> M Nat
+arityForName : {{Ref2 Defs St}} → FC -> QName -> M Nat
 arityForName fc nm = do
-  top <- getTop
-  case lookup nm top of
-  -- let the magic hole through for now (will generate bad JS)
+  defs <- getRef Defs
+  case lookupMap' nm defs of
     Nothing => error fc "Name \{show nm} not in scope"
-    (Just (MkEntry _ name type Axiom)) => pure Z
-    (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 arity))) => pure $ cast arity
-    -- Assuming a primitive can't return a function
-    (Just (MkEntry _ name type (PrimFn t used))) => pure $ piArity type
+    (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 : Tm -> M CExp
+compileTerm : {{Ref2 Defs St}} → 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
+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) ty = pure $ CApp t (acc <>> Nil) (1 + cast k)
--- FIXME - this should be handled by Erasure.newt
--- We somehow hit the error below, with a Pi?
-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, got \{showTm ty}. Overapplied function that escaped type checking?"
+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 ty = go (CApp t (acc <>> Nil) 0) ts
+apply t ts acc Z = go (CApp t (acc <>> Nil) 0) ts
   where
     go : CExp -> List CExp -> M CExp
     -- drop zero arg call
@@ -91,14 +82,11 @@ apply t ts acc Z ty = go (CApp t (acc <>> Nil) 0) ts
 compileTerm (Bnd _ k) = pure $ CBnd k
 -- need to eta expand to arity
 compileTerm t@(Ref fc nm) = do
-  top <- getTop
-  let (Just (MkEntry _ _ type _)) = lookup nm top
-          | Nothing => error fc "Undefined name \{show nm}"
   arity <- arityForName fc nm
   case arity of
     -- we don't need to curry functions that take one argument
     (S Z) => pure $ CRef (show nm)
-    _ => apply (CRef (show nm)) Nil Lin arity type
+    _ => apply (CRef (show nm)) Nil Lin arity
 
 compileTerm (Meta _ k) = pure $ CRef "meta$\{show k}" -- FIXME
 compileTerm (Lam _ nm _ _ t) = CLam nm <$> compileTerm t
@@ -111,14 +99,14 @@ compileTerm tm@(App _ _ _) = case funArgs tm of
         args' <- traverse compileTerm args
         arity <- arityForName fc nm
         top <- getTop
-        let (Just (MkEntry _ _ type _)) = lookup nm top
-          | Nothing => error fc "Undefined name \{show nm}"
-        apply (CRef (show nm)) args' Lin arity type
+        -- let (Just (MkEntry _ _ type _)) = lookup nm top
+        --   | Nothing => error fc "Undefined name \{show nm}"
+        apply (CRef (show 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 (UU emptyFC)
+        apply t' args' Lin Z
         -- error (getFC t) "Don't know how to apply \{showTm t}"
 compileTerm (UU _) = pure $ CRef "U"
 compileTerm (Pi _ nm icit rig t u) = do
@@ -144,7 +132,7 @@ compileTerm (LetRec _ nm _ t u) = do
   pure $ CLetRec nm t' u'
 compileTerm (Erased _) = pure CErased
 
-compileFun : Tm -> M CExp
+compileFun : {{Ref2 Defs St}} → Tm -> M CExp
 compileFun tm = go tm Lin
   where
     go : Tm -> SnocList String -> M CExp

src/Lib/LiftWhere.newt

@@ -0,0 +1,93 @@
+module Lib.LiftWhere
+
+import Prelude
+import Lib.Common
+import Lib.Types
+import Lib.TopContext
+import Lib.Ref2
+import Data.SortedMap
+import Monad.State
+import Data.IORef
+
+-- track depth and whether we need to replace Bnd with a top level call
+LiftEnv : U
+LiftEnv = List (Maybe (QName × Nat))
+
+liftWhereTm : {{Ref2 Defs St}} → QName → LiftEnv → Tm → M Tm
+liftWhereTm name env (Lam fc nm icit quant t) =
+  Lam fc nm icit quant <$> liftWhereTm name (Nothing :: env) t
+liftWhereTm name env (App fc t u) =
+  App fc <$> liftWhereTm name env t <*> liftWhereTm name env u
+liftWhereTm name env (Pi fc nm icit quant t u) = do
+  t <- liftWhereTm name env t
+  u <- liftWhereTm name (Nothing :: env) u
+  pure $ Pi fc nm icit quant t u
+liftWhereTm name env (Let fc nm v sc) = do
+  v <- liftWhereTm name env v
+  sc <- liftWhereTm name (Nothing :: env) sc
+  pure $ Let fc nm v sc
+liftWhereTm name env tm@(Case fc t alts) = do
+  t <- liftWhereTm name env t
+  alts' <- traverse liftWhereAlt alts
+  pure $ Case fc t alts'
+  where
+    -- This is where I wish I had put indexes on things
+    liftWhereAlt : CaseAlt → M CaseAlt
+    liftWhereAlt (CaseDefault tm) = CaseDefault <$> liftWhereTm name env tm
+    liftWhereAlt (CaseLit l tm) = CaseLit l <$> liftWhereTm name env tm
+    liftWhereAlt (CaseCons qn args tm) =
+      CaseCons qn args <$> liftWhereTm name (map (const Nothing) args ++ env) tm
+-- This is where the magic happens
+liftWhereTm name env (LetRec fc nm ty t u) = do
+  let l = length env
+  -- FIXME we need a namespace and a name, these collide everywhere.
+  qn <- getName name nm
+  let env' = (Just (qn, S l) :: env)
+  -- environment should subst this function (see next case)
+  t' <- liftWhereTm qn env' t
+  -- TODO we could have subst in this var and dropped the extra argument
+  modifyRef Defs (updateMap qn (Fn $ wrapLam (S l) t'))
+
+  -- The rest
+  u' <- liftWhereTm qn env' u
+  pure $ LetRec fc nm (Erased fc) (Erased fc) u'
+  where
+    getName : QName → String → M QName
+    getName qn@(QN ns nm) ext = do
+      let qn' = QN ns (nm ++ "." ++ ext)
+      top <- getRef Defs
+      let (Just _) = lookupMap qn' top | _ => pure qn'
+      getName qn (ext ++ "'")
+
+    -- Hacky - CompileExp expects a pi type that matches arity
+    wrapPi : Nat → Tm → Tm
+    wrapPi Z tm = tm
+    wrapPi (S k) tm = Pi fc "_" Explicit Many (Erased fc) $ wrapPi k tm
+
+    wrapLam : Nat → Tm → Tm
+    wrapLam Z tm = tm
+    -- REVIEW We've already erased, hopefully we don't need quantity
+    wrapLam (S k) tm = Lam fc "_" Explicit Many $ wrapLam k tm
+
+-- And where it lands
+liftWhereTm name env tm@(Bnd fc k) = case getAt (cast k) env of
+  Just (Just (qn, v)) => pure $ apply (length' env) (cast v) (Ref fc qn)
+  _ => pure tm
+  where
+    apply : Int → Int → Tm → Tm
+    apply l 0 tm = tm
+    -- (l - k) is like lvl2ix, but witih k one bigger
+    apply l k tm = App fc (apply l (k - 1) tm) (Bnd fc (l - k))
+liftWhereTm name env tm = pure tm
+
+liftWhereFn : {{Ref2 Defs St}} → QName × Def → M Unit
+liftWhereFn (name, Fn tm) = do
+  tm' <- liftWhereTm name Nil tm
+  modifyRef Defs $ updateMap name (Fn tm')
+  -- updateDef name fc type (Fn tm')
+liftWhereFn _ = pure MkUnit
+
+liftWhere :  {{Ref2 Defs St}} → M Unit
+liftWhere = do
+  defs <- getRef Defs
+  ignore $ traverse liftWhereFn $ toList defs

src/Lib/ProcessDecl.newt

@@ -132,7 +132,8 @@ processDecl ns  (PFunc fc nm used ty src) = do
   used' <- for used $ \ name => case lookupRaw name top of
     Nothing => error fc "\{name} not in scope"
     Just entry => pure entry.name
-  setDef (QN ns nm) fc ty' (PrimFn src used')
+  let arity = piArity ty'
+  setDef (QN ns nm) fc ty' (PrimFn src arity used')
 
 processDecl ns (Def fc nm clauses) = do
   log 1 $ \ _ => "-----"

src/Lib/Ref2.newt

@@ -0,0 +1,27 @@
+module Lib.Ref2
+
+import Prelude
+import Lib.Common
+import Lib.Types
+import Data.IORef
+import Data.SortedMap
+
+data Defs : U where
+
+-- St holds our code while we're optimizing
+St : U
+St = SortedMap QName Def
+
+-- This is inspired by Idris.
+-- Mainly to get an extra state variable into M
+-- I tried parameterizing M, but inference was having trouble
+-- in the existing code.
+data Ref2 : (l : U) → U → U where
+  MkRef : ∀ a . {0 x : U} → IORef a → Ref2 x a
+
+getRef : ∀ io a. {{HasIO io}} → (l : U) → {{Ref2 l a}} → io a
+getRef l {{MkRef a}} = readIORef a
+
+modifyRef : ∀ io a. {{HasIO io}} → (l : U) → {{Ref2 l a}} → (a → a) → io Unit
+-- TODO inference needs liftIO here
+modifyRef l {{MkRef a}} f = liftIO $ modifyIORef a f

src/Lib/Types.newt

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