More inlining, fix issues in eval of case

src/Lib/Compile.newt

@@ -63,12 +63,16 @@ Cont e = JSExp -> JSStmt e
 record JSEnv where
   constructor MkEnv
   jsenv : List JSExp
+  -- This is not depth, it is incremented as we go down the tree to get fresh names
   depth : Int
 
 -- this was like this, are we not using depth?
 push : JSEnv -> JSExp -> JSEnv
 push (MkEnv env depth) exp = MkEnv (exp :: env) depth
 
+incr : JSEnv → JSEnv
+incr env = MkEnv env.jsenv (1 + env.depth)
+
 emptyJSEnv : JSEnv
 emptyJSEnv = MkEnv Nil 0
 
@@ -78,7 +82,6 @@ litToJS (LChar c) = LitString $ pack (c :: Nil)
 litToJS (LInt i) = LitInt i
 
 -- Stuff nm.h1, nm.h2, ... into environment
--- TODO consider JSExp instead of nm, so we can have $foo.h1 instead of assigning a sc.
 mkEnv : JSExp -> Int -> JSEnv -> List String -> JSEnv
 mkEnv nm k env Nil = env
 mkEnv nm k env (x :: xs) = mkEnv nm (1 + k) (push env (Dot nm "h\{show k}")) xs
@@ -173,7 +176,7 @@ termToJS env (CConstr nm args) f = go args 0 (\ args => f $ LitObject (("tag", L
     go : ∀ e. List CExp -> Int -> (List (String × JSExp) -> JSStmt e) -> JSStmt e
     go Nil ix k = k Nil
     go (t :: ts) ix k = termToJS env t $ \ t' => go ts (ix + 1) $ \ args => k $ ("h\{show ix}", t') :: args
-termToJS env (CAppRef nm args etas) f = termToJS env (CRef nm) (\ t' => (argsToJS t' args Lin f))
+termToJS env (CAppRef nm args etas) f = termToJS env (CRef nm) (\ t' => (argsToJS env t' args Lin f))
   where
     etaExpand : JSEnv -> Nat -> SnocList JSExp -> JSExp -> JSExp
     etaExpand env Z args tm = Apply tm (args <>> Nil)
@@ -182,31 +185,25 @@ termToJS env (CAppRef nm args etas) f = termToJS env (CRef nm) (\ t' => (argsToJ
           env' = push env (Var nm')
        in JLam (nm' :: Nil) $ JReturn $ etaExpand (push env (Var nm')) etas (args :< Var nm') tm
 
-    argsToJS : ∀ e. JSExp -> List CExp -> SnocList JSExp -> (JSExp -> JSStmt e) -> JSStmt e
+    argsToJS : ∀ e. JSEnv -> JSExp -> List CExp -> SnocList JSExp -> (JSExp -> JSStmt e) -> JSStmt e
-    argsToJS tm Nil acc k = k (etaExpand env (cast etas) acc tm)
+    argsToJS env tm Nil acc k = k (etaExpand env (cast etas) acc tm)
-    -- k (acc <>> Nil)
+    argsToJS env tm (x :: xs) acc k = termToJS env x (\ x' => argsToJS (incr env) tm xs (acc :< x') k)
-    argsToJS tm (x :: xs) acc k = termToJS env x (\ x' => argsToJS tm xs (acc :< x') k)
 
 termToJS env (CApp t arg) f = termToJS env t (\ t' => termToJS env arg (\arg' => f (Apply t' (arg' :: Nil))))
 
 termToJS {e} env (CCase t alts) 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) => maybeCaseStmt env (Var nm) alts
     t' => do
-      -- TODO refactor nm to be a JSExp with Var{} or Dot{}
+      -- TODO with inlining, we hit cases where the let gets pulled forward more than once
-      -- FIXME sc$ seemed to shadow something else, lets get this straightened out
+      -- two cases as separate args, se we need actual unique names. For now, we're calling
-      -- we need freshName names that are not in env (i.e. do not play in debruijn)
+      -- incr when processing App, as a stopgap, we probably need a fresh names state monad
       let nm = "_sc$\{show env.depth}"
       -- increment the bit that goes into the name
-      let env' = MkEnv env.jsenv (1 + env.depth)
+      let env' = incr env
       if simpleJSExp t'
         then (maybeCaseStmt env' t' alts)
         else JSnoc (JConst nm t') (maybeCaseStmt env' (Var nm) alts)
-
   where
     termToJSAlt : JSEnv -> JSExp -> CAlt -> JAlt
     termToJSAlt env nm (CConAlt name info args u) = JConAlt name (termToJS (mkEnv nm 0 env args) u f)
@@ -231,7 +228,7 @@ jsKeywords = (
   "implements" :: "class" :: "let" :: "package" :: "private" :: "protected" :: "public" ::
   "static" :: "yield" ::
   "null" :: "true" :: "false" ::
-  -- might not be a big issue with namespaces on names now.
+  -- might not occur now that we have namespaces on the names
   "String" :: "Number" :: "Array" :: "BigInt" :: Nil)
 
 
@@ -251,17 +248,13 @@ jsIdent id = if elem id jsKeywords then text ("$" ++ id) else text $ pack $ fix
       else
         '$' :: (toHex (cast x)) ++ fix xs
 
-
-
 stmtToDoc : ∀ e. JSStmt e -> Doc
 
-
 expToDoc : JSExp -> Doc
 expToDoc (LitArray xs) = fatalError "TODO - LitArray to doc"
 expToDoc (LitObject xs) = text "{" <+> folddoc (\ a e => a ++ text ", " <+/> e) (map entry xs) <+> text "}"
   where
     entry : (String × JSExp) -> Doc
-    -- TODO quote if needed
     entry (nm, exp) = jsIdent nm ++ text ":" <+> expToDoc exp
 
 expToDoc (LitString str) = text $ quoteString str
@@ -280,7 +273,6 @@ expToDoc (JPrimOp op t u) = parens 0 1 (expToDoc t) <+> text op <+> parens 0 1 (
 
 caseBody : ∀ e. JSStmt e -> Doc
 caseBody stmt@(JReturn x) = nest 2 (line ++ stmtToDoc stmt)
--- caseBody {e = Return} stmt@(JCase{}) = nest 2 (line ++ stmtToDoc stmt)
 caseBody {e} stmt@(JCase _ _) = nest 2 (line ++ stmtToDoc stmt </> text "break;")
 caseBody stmt = line ++ text "{" ++ nest 2 (line ++ stmtToDoc stmt </> text "break;") </> text "}"
 
@@ -472,5 +464,5 @@ compile = do
       let exec = stmtToDoc $ JPlain $ Apply (Var $ show name) Nil
       pure $ reverse (exec :: tmp)
     Nothing =>
-      -- TODO maybe dump everything if there is no main
+      -- TODO maybe emit everything if there is no main
       error emptyFC "No main function found"

src/Lib/Elab.newt

@@ -809,7 +809,7 @@ checkCase ctx prob scnm scty (dcName, arity, ty) = do
     _ => pure True
 
 -- ok, so this is a single constructor, CaseAlt
--- return Nothing if dcon doesn't unify with scrut
+-- return Nothing if dcon type doesn't unify with scrut
 buildCase : Context -> Problem -> String -> Val -> (QName × Int × Tm) -> M (Maybe CaseAlt)
 buildCase ctx prob scnm scty (dcName, arity, ty) = do
   debug $ \ _ => "CASE \{scnm} match \{show dcName} ty \{rpprint (names ctx) ty}"

src/Lib/Eval.newt

@@ -98,24 +98,26 @@ forceType env x = do
     | _ => pure x
   forceType env x'
 
+-- for cases applied to a value
+-- TODO this does not handle CaseLit
 evalCase : Env -> Val -> List CaseAlt -> M (Maybe Val)
 evalCase env sc@(VRef _ nm sp) (cc@(CaseCons name nms t) :: xs) = do
   top <- getTop
   if nm == name
     then do
       debug $ \ _ =>  "ECase \{show nm} \{show sp} \{show nms} \{showTm t}"
-      go env (sp <>> Nil) nms
+      pushArgs env (sp <>> Nil) nms
     else case lookup nm top of
       (Just (MkEntry _ str type (DCon _ k str1) _)) => evalCase env sc xs
       -- bail for a stuck function
       _ => pure Nothing
   where
-    go : Env -> List Val -> List String -> M (Maybe Val)
+    pushArgs : Env -> List Val -> List String -> M (Maybe Val)
-    go env (arg :: args) (nm :: nms) = go (arg :: env) args nms
+    pushArgs env (arg :: args) (nm :: nms) = pushArgs (arg :: env) args nms
-    go env args Nil = do
+    pushArgs env args Nil = do
       t' <- eval env t
       Just <$> vappSpine t' (Lin <>< args)
-    go env Nil rest = pure Nothing
+    pushArgs env Nil rest = pure Nothing
 -- REVIEW - this is handled in the caller already
 evalCase env sc@(VVar fc k sp) alts = case lookupVar env k of
     Just tt@(VVar fc' k' sp') => do
@@ -131,12 +133,19 @@ evalCase env sc@(VVar fc k sp) alts = case lookupVar env k of
     Nothing => do
         debug $ \ _ =>  "lookup \{show k} is Nothing in env \{show env}"
         pure Nothing
-evalCase env sc (CaseDefault u :: xs) = Just <$> eval (sc :: env) u
+evalCase env sc (CaseDefault u :: xs) = Just <$> eval env u
 evalCase env sc cc = do
   debug $ \ _ =>  "CASE BAIL sc \{show sc} vs " -- \{show cc}"
   debug $ \ _ =>  "env is \{show env}"
   pure Nothing
 
+-- neutral alts
+evalAlt : Env → CaseAlt → M VCaseAlt
+evalAlt env (CaseDefault tm) = VCaseDefault <$> eval env tm
+evalAlt env (CaseLit lit tm) = VCaseLit lit <$> eval env tm
+-- in the cons case, we're binding args
+evalAlt env (CaseCons nm args tm) = pure $ VCaseCons nm args env tm
+
 -- So smalltt says:
 --   Smalltt has the following approach:
 --   - Top-level and local definitions are lazy.
@@ -146,7 +155,6 @@ evalCase env sc cc = do
 
 -- TODO maybe add glueing
 
-
 eval env (Ref fc x) = pure $ VRef fc x Lin
 eval env (App _ t u) = do
   t' <- eval env t
@@ -165,7 +173,7 @@ eval env (Pi fc x icit rig a b) = do
   pure $ VPi fc x icit rig a' (MkClosure env b)
 eval env (Let fc nm t u) = do
   t' <- eval env t
-  u' <- eval (VVar fc (cast $ length env) Lin :: env) u
+  u' <- eval (VVar fc (length' env) Lin :: env) u
   pure $ VLet fc nm t' u'
 eval env (LetRec fc nm ty t u) = do
   ty' <- eval env ty
@@ -184,10 +192,14 @@ eval env tm@(Case fc sc alts) = do
   -- TODO we need to be able to tell eval to expand aggressively here.
   sc' <- eval env sc
   sc' <- unlet env sc' -- try to expand lets from pattern matching
+  -- possibly too aggressive?
   sc' <- forceType env sc'
+  Nothing <- evalCase env sc' alts
+    | Just v => pure v
+
   vsc <- eval env sc
-  vcase <- evalCase env sc' alts
+  alts' <- traverse (evalAlt env) alts
-  pure $ fromMaybe (VCase fc vsc alts) vcase
+  pure $ VCase fc vsc alts'
 
 quote : (lvl : Int) -> Val -> M Tm
 
@@ -198,6 +210,18 @@ quoteSp lvl t (xs :< x) = do
   x' <- quote lvl x
   pure $ App emptyFC t' x'
 
+quoteAlt : Int → VCaseAlt → M CaseAlt
+quoteAlt l (VCaseDefault val) = CaseDefault <$> quote l val
+quoteAlt l (VCaseLit lit val) = CaseLit lit <$> quote l val
+quoteAlt l (VCaseCons nm args env tm) = do
+  val <- eval (mkenv l env args) tm
+  tm <- quote (length' args + l) val
+  pure $ CaseCons nm args tm
+  where
+    mkenv : Int → Env → List String → Env
+    mkenv l env Nil = env
+    mkenv l env (n :: ns) = mkenv (l + 1) (VVar emptyFC l Lin :: env) ns
+
 quote l (VVar fc k sp) = if k < l
   then quoteSp l (Bnd fc (lvl2ix l k )) sp -- level to index
   else error fc "Bad index in quote \{show k} depth \{show l}"
@@ -226,8 +250,9 @@ quote l (VLetRec fc nm ty t u) = do
   pure $ LetRec fc nm ty' t' u'
 quote l (VU fc) = pure (UU fc)
 quote l (VRef fc n sp) = quoteSp l (Ref fc n) sp
-quote l (VCase fc sc alts) = do
+quote l (VCase fc sc valts) = do
   sc' <- quote l sc
+  alts <- traverse (quoteAlt l) valts
   pure $ Case fc sc' alts
 quote l (VLit fc lit) = pure $ Lit fc lit
 quote l (VErased fc) = pure $ Erased fc
@@ -312,6 +337,7 @@ zonkApp top l env t sp = do
     zonk top l env t')
     (\_ => pure $ appSpine t' sp)
   where
+    -- lookup name and return Def if flagged inline
     inlineDef : Tm → Maybe Tm
     inlineDef (Ref _ nm) = case lookup nm top of
       Just (MkEntry _ _ ty (Fn tm) flags) => if elem Inline flags then Just tm else Nothing
@@ -327,16 +353,20 @@ zonkAlt top l env (CaseCons name args t) = CaseCons name args <$> go l env args
     go l env Nil tm = zonk top l env t
     go l env (x :: xs) tm = go (1 + l) (VVar (getFC tm) l Lin :: env) xs tm
 
-zonk top l env t = case t of
+zonk top l env t =
+  let env' = VVar emptyFC l Lin :: env in
+  case t of
   (Meta fc k) => zonkApp top l env t Nil
-  (Lam fc nm icit rig u) => Lam fc nm icit rig <$> (zonk top (1 + l) (VVar fc l Lin :: env) u)
+  (Lam fc nm icit rig u) => Lam fc nm icit rig <$> (zonk top (1 + l) env' u)
-  (App fc t u) => do
+  (App fc _ _) => zonkApp top l env t Nil
-    u' <- zonk top l env u
+  (Pi fc nm icit rig a b) => Pi fc nm icit rig <$> zonk top l env a <*> zonkBind top (l + 1) env' b
-    zonkApp top l env t (u' :: Nil)
+  (Let fc nm t u) => Let fc nm <$> zonk top l env t <*> zonkBind top (l + 1) env' u
-  (Pi fc nm icit rig a b) => Pi fc nm icit rig <$> zonk top l env a <*> zonkBind top l env b
+  (LetRec fc nm ty t u) => LetRec fc nm <$> zonk top l env ty <*> zonkBind top (l + 1) env' t <*> zonkBind top (l + 1) env' u
-  (Let fc nm t u) => Let fc nm <$> zonk top l env t <*> zonkBind top l env u
+  (Case fc sc alts) => do
-  (LetRec fc nm ty t u) => LetRec fc nm <$> zonk top l env ty <*> zonkBind top l env t <*> zonkBind top l env u
+      sc' <- zonk top l env sc
-  (Case fc sc alts) => Case fc <$> zonk top l env sc <*> traverse (zonkAlt top l env) alts
+      alts' <- traverse (zonkAlt top l env) alts
+      pure $ Case fc sc' alts'
+
   UU fc => pure $ UU fc
   Lit fc lit => pure $ Lit fc lit
   Bnd fc ix => pure $ Bnd fc ix

src/Lib/ProcessDecl.newt

@@ -143,6 +143,8 @@ complexity (Ref _ _) = 1
 complexity (Lam _ _ _ _ sc) = 1 + complexity sc
 complexity (App _ t u) = complexity t + complexity u
 complexity (Bnd _ _) = 1
+-- These turn into a projection
+complexity (Case _ sc (CaseCons _ _ t :: Nil)) = 1 + complexity sc + complexity t
 complexity _ = 100
 
 processDef : List String → FC → String → List (Raw × Raw) → M Unit
@@ -168,12 +170,18 @@ processDef ns fc nm clauses = do
   -- TODO - make nf that expands all metas and drop zonk
   -- Idris2 doesn't expand metas for performance - a lot of these are dropped during erasure.
   -- Day1.newt is a test case
-  -- NOW - might not need this if we do it at compile time
+  -- This inlines metas and functions flagged Inline.
   tm' <- zonk top 0 Nil tm
   debug $ \ _ => "NF\n\{render 80 $ pprint Nil tm'}"
   debug $ \ _ => "Add def \{nm} \{render 90 $ pprint Nil tm'} : \{render 90 $ pprint Nil ty}"
   updateDef (QN ns nm) fc ty (Fn tm')
-  if complexity tm' < 10
+  -- putStrLn "complexity \{show (QN ns nm)} \{show $ complexity tm'}"
+  -- putStrLn $ show tm'
+  -- TODO we need some protection against inlining a function calling itself.
+  -- 14 gets us to 6.21s, higher than 11 breaks Zoo4eg.newt with a unification error (probably need to inline at the end instead)
+  -- But we need better heuristics, maybe fuel and deciding while inlining.
+  -- bind is explicit here because the complexity has a 100 in it.
+  if complexity tm' < 11 || show (QN ns nm) == "Prelude.Prelude.Monad Prelude.IO,bind"
     then setFlag (QN ns nm) fc Inline
     else pure MkUnit
 

src/Lib/Types.newt

@@ -213,13 +213,24 @@ Val : U
 -- Yaffle is vars -> vars here
 
 Closure : U
+Env : U
+
+data VCaseAlt : U where
+  -- Like a Closure, but with a lot of args
+  VCaseCons : (name : QName) -> (args : List String) -> Env -> Tm ->  VCaseAlt
+  VCaseLit : Literal -> Val -> VCaseAlt
+  VCaseDefault : Val -> VCaseAlt
+  -- VCaseCons : (name : QName) -> (args : List String) -> Tm ->  VCaseAlt
+  -- VCaseLit : Literal -> Tm -> VCaseAlt
+  -- VCaseDefault : Tm -> VCaseAlt
+
 
 data Val : U where
   -- This will be local / flex with spine.
   VVar : FC -> (k : Int) -> (sp : SnocList Val) -> Val
   VRef : FC -> (nm : QName) -> (sp : SnocList Val) -> Val
   -- neutral case
-  VCase : FC -> (sc : Val) -> List CaseAlt -> Val
+  VCase : FC -> (sc : Val) -> List VCaseAlt -> Val
   -- we'll need to look this up in ctx with IO
   VMeta : FC -> QName -> (sp : SnocList Val) -> Val
   VLam : FC -> Name -> Icit -> Quant -> Closure -> Val
@@ -230,7 +241,7 @@ data Val : U where
   VErased : FC -> Val
   VLit : FC -> Literal -> Val
 
-Env : U
+
 Env = List Val
 
 data Closure = MkClosure Env Tm
@@ -262,7 +273,12 @@ instance Show Val where
   show (VPi fc str Implicit rig x y) = "(%pi {\{show rig} \{str} : \{show  x}}. \{showClosure  y})"
   show (VPi fc str Explicit rig x y) = "(%pi (\{show rig} \{str} : \{show  x}). \{showClosure  y})"
   show (VPi fc str Auto rig x y) = "(%pi {{\{show rig} \{str} : \{show  x}}}. \{showClosure  y})"
-  show (VCase fc sc alts) = "(%case \{show sc} ...)"
+  show (VCase fc sc alts) = "(%case \{show sc} \{unwords $ map showAlt alts})"
+    where
+      showAlt : VCaseAlt → String
+      showAlt (VCaseDefault v) = "(%cdef \{show v})"
+      showAlt (VCaseLit lit v) = "(%clit \{show v})"
+      showAlt (VCaseCons nm args env v) = "(%ccon \{show nm} \{unwords $ map show args} [\{show $ length env} env] \{show v}"
   show (VU _) = "U"
   show (VLit _ lit) = show lit
   show (VLet _ nm a b) = "(%let \{show nm} = \{show a} in \{show b}"