From 001cba26ee3a3b7059818154a7fb9d55352b12cd Mon Sep 17 00:00:00 2001 From: Steve Dunham Date: Sat, 15 Feb 2025 21:07:44 -0800 Subject: [PATCH] first pass at liftWhere --- src/Lib/Compile.newt | 33 +++++++++++--- src/Lib/CompileExp.newt | 56 ++++++++++-------------- src/Lib/LiftWhere.newt | 93 ++++++++++++++++++++++++++++++++++++++++ src/Lib/ProcessDecl.newt | 3 +- src/Lib/Ref2.newt | 27 ++++++++++++ src/Lib/Types.newt | 4 +- 6 files changed, 173 insertions(+), 43 deletions(-) create mode 100644 src/Lib/LiftWhere.newt create mode 100644 src/Lib/Ref2.newt diff --git a/src/Lib/Compile.newt b/src/Lib/Compile.newt index ae4018b..90f533b 100644 --- a/src/Lib/Compile.newt +++ b/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}" diff --git a/src/Lib/CompileExp.newt b/src/Lib/CompileExp.newt index fc04a4c..473719d 100644 --- a/src/Lib/CompileExp.newt +++ b/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 diff --git a/src/Lib/LiftWhere.newt b/src/Lib/LiftWhere.newt new file mode 100644 index 0000000..3a6adcc --- /dev/null +++ b/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 diff --git a/src/Lib/ProcessDecl.newt b/src/Lib/ProcessDecl.newt index d1f8804..80de1c5 100644 --- a/src/Lib/ProcessDecl.newt +++ b/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 $ \ _ => "-----" diff --git a/src/Lib/Ref2.newt b/src/Lib/Ref2.newt new file mode 100644 index 0000000..e1ba876 --- /dev/null +++ b/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 diff --git a/src/Lib/Types.newt b/src/Lib/Types.newt index 9492ec5..d4a7266 100644 --- a/src/Lib/Types.newt +++ b/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