From 3143fa7b0ac1544897a03dce5f08389507f90ca5 Mon Sep 17 00:00:00 2001 From: Steve Dunham Date: Tue, 23 Sep 2025 20:22:50 -0700 Subject: [PATCH] More inlining, fix issues in eval of case --- src/Lib/Compile.newt | 36 ++++++++------------- src/Lib/Elab.newt | 2 +- src/Lib/Eval.newt | 70 ++++++++++++++++++++++++++++------------ src/Lib/ProcessDecl.newt | 12 +++++-- src/Lib/Types.newt | 22 +++++++++++-- 5 files changed, 94 insertions(+), 48 deletions(-) diff --git a/src/Lib/Compile.newt b/src/Lib/Compile.newt index f30556b..b8584d1 100644 --- a/src/Lib/Compile.newt +++ b/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 tm Nil acc k = k (etaExpand env (cast etas) acc tm) - -- k (acc <>> Nil) - argsToJS tm (x :: xs) acc k = termToJS env x (\ x' => argsToJS tm xs (acc :< x') k) + argsToJS : ∀ e. JSEnv -> JSExp -> List CExp -> SnocList JSExp -> (JSExp -> JSStmt e) -> JSStmt e + argsToJS env tm Nil acc k = k (etaExpand env (cast etas) acc tm) + argsToJS env tm (x :: xs) acc k = termToJS env x (\ x' => argsToJS (incr env) 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{} - -- FIXME sc$ seemed to shadow something else, lets get this straightened out - -- we need freshName names that are not in env (i.e. do not play in debruijn) + -- TODO with inlining, we hit cases where the let gets pulled forward more than once + -- two cases as separate args, se we need actual unique names. For now, we're calling + -- 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" diff --git a/src/Lib/Elab.newt b/src/Lib/Elab.newt index 7338c7a..ac8e67f 100644 --- a/src/Lib/Elab.newt +++ b/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}" diff --git a/src/Lib/Eval.newt b/src/Lib/Eval.newt index d862796..bd4162b 100644 --- a/src/Lib/Eval.newt +++ b/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) - go env (arg :: args) (nm :: nms) = go (arg :: env) args nms - go env args Nil = do + pushArgs : Env -> List Val -> List String -> M (Maybe Val) + pushArgs env (arg :: args) (nm :: nms) = pushArgs (arg :: env) args nms + 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 - pure $ fromMaybe (VCase fc vsc alts) vcase + alts' <- traverse (evalAlt env) alts + 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) - (App fc t u) => do - u' <- zonk top l env u - zonkApp top l env t (u' :: Nil) - (Pi fc nm icit rig a b) => Pi fc nm icit rig <$> zonk top l env a <*> zonkBind top l env b - (Let fc nm t u) => Let fc nm <$> zonk top l env t <*> zonkBind top l env u - (LetRec fc nm ty t u) => LetRec fc nm <$> zonk top l env ty <*> zonkBind top l env t <*> zonkBind top l env u - (Case fc sc alts) => Case fc <$> zonk top l env sc <*> traverse (zonkAlt top l env) alts + (Lam fc nm icit rig u) => Lam fc nm icit rig <$> (zonk top (1 + l) env' u) + (App fc _ _) => zonkApp top l env t Nil + (Pi fc nm icit rig a b) => Pi fc nm icit rig <$> zonk top l env a <*> zonkBind top (l + 1) env' b + (Let fc nm t u) => Let fc nm <$> zonk top l env t <*> zonkBind top (l + 1) env' u + (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 + (Case fc sc alts) => do + sc' <- zonk top l env sc + 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 diff --git a/src/Lib/ProcessDecl.newt b/src/Lib/ProcessDecl.newt index 9664078..90f1d80 100644 --- a/src/Lib/ProcessDecl.newt +++ b/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 diff --git a/src/Lib/Types.newt b/src/Lib/Types.newt index 2829e47..031a9af 100644 --- a/src/Lib/Types.newt +++ b/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}"