More inlining, fix issues in eval of case

This commit is contained in:
2025-09-23 20:22:50 -07:00
parent cc7d8b4968
commit 3143fa7b0a
5 changed files with 94 additions and 48 deletions

View File

@@ -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"

View File

@@ -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}"

View File

@@ -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

View File

@@ -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

View File

@@ -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}"