First pass at a scheme backend

src/Lib/CompileScheme.newt

@@ -0,0 +1,338 @@
+module Lib.CompileScheme
+
+import Prelude
+import Lib.Common
+import Lib.Types
+import Lib.Eval
+import Lib.Prettier
+import Lib.CompileExp
+import Lib.TopContext
+import Lib.LiftWhere
+import Lib.LiftLambda
+import Lib.TCO
+import Lib.Ref2
+import Lib.Erasure
+import Data.String
+import Data.Int
+import Data.SortedMap
+import Data.IORef
+import Data.SnocList
+
+SCEnv : U
+SCEnv = List String
+
+getEnv : Int → SCEnv → String
+getEnv ix env = case getAt' ix env of
+  Just e => e
+  Nothing => fatalError "Bad bounds \{show ix}"
+
+scbind : String → SCEnv → (String × SCEnv)
+scbind nm env = let
+  x = length' env
+  nm' = "nm-\{show x}"
+  in (nm', nm' :: env)
+
+-- TODO
+scmKeywords : List String
+scmKeywords = "lambda" :: Nil
+
+scmIdent : String → String
+scmIdent id = if elem id scmKeywords then "$" ++ id else pack $ fix (unpack id)
+  where
+    fix : List Char -> List Char
+    fix Nil = Nil
+    fix (x :: xs) =
+      if isAlphaNum x || x == '_' || x == '.' then
+        x :: fix xs
+      -- make qualified names more readable
+      else if x == ',' then '_' :: fix xs
+      else if x == ' ' then '_' :: fix xs
+      else if x == '$' then
+        '$' :: '$' :: fix xs
+      else
+        '$' :: (toHex (cast x)) ++ fix xs
+
+
+scmLit : Literal → String
+scmLit (LString str) = quoteString str
+scmLit (LInt n) = show n
+scmLit (LChar c) = pack $ '#' :: '\\' :: c :: Nil -- FIXME
+scmLit (LBool b) = ite b "#t" "#f"
+
+scmName : QName → String
+scmName qn = scmIdent $ show qn
+
+cexpToScm : SCEnv → CExp → String
+
+withVar : SCEnv → CExp → (String → String) → String
+-- withVar env (CBnd _) f = ?
+-- withVar env (CLit _) f = ?
+withVar env t f = let nm = "wv$\{show $ length' env}"
+  in "(let ((\{nm} \{cexpToScm env t})) \{f nm})"
+
+cexpToScm env (CBnd n) = getEnv n env
+cexpToScm env (CLam nm t) = case scbind nm env of
+  (nm', env') => "(lambda (\{nm'}) \{cexpToScm env' t})"
+cexpToScm env (CFun args body) = case bindAll args Lin env  of
+    (nms,env') => "(lambda (\{joinBy " " nms}) \{cexpToScm env' body})"
+  where
+    bindAll : List (Quant × String) → SnocList String → SCEnv → List String × SCEnv
+    bindAll Nil acc env = (acc <>> Nil, env)
+    bindAll ((_,nm) :: rest) acc env = case scbind nm env of
+      (nm', env') => bindAll rest (acc :< nm') env'
+cexpToScm env (CApp t u) = "(\{cexpToScm env t} \{cexpToScm env u})"
+cexpToScm env (CAppRef nm args Nil) = go (scmName nm) $ map (cexpToScm env) args
+  where
+    go : String → List String → String
+    go acc Nil = acc
+    go acc (arg :: args) = go "(\{acc} \{arg})" args
+-- If we're short, we need to eta expand
+cexpToScm env (CAppRef nm args quants) =
+  go env "(\{scmName nm}" (map (cexpToScm env) args) quants
+  where
+    -- here `acc` is always missing a closing paren
+    go : SCEnv → String → List String → List Quant → String
+    go env acc Nil Nil = acc ++ ")"
+    -- extra args are applied
+    go env acc (arg :: args) Nil = go env "(\{acc}) \{arg}" args Nil
+    -- missing args get eta expanded
+    go env acc Nil (q :: qs) = case scbind "_" env of
+      (nm, env') => let acc = "\{acc} \{nm}" in "(lambda (\{nm}) \{go env' acc Nil qs})"
+    -- TODO / REVIEW Only for Many?
+    go env acc (arg :: args) (q :: qs) = go env "\{acc} \{arg}" args qs
+    -- go env acc (arg :: args) (q :: qs) = go env acc args qs
+-- so... we're not giving scrutinee a deBruijn index, but we may
+-- need to let it so we can pull data off for the CConAlt
+cexpToScm env (CCase sc alts) = do
+    -- assign sc, might need to do more to make the name unique, we only
+    -- get one of these per env at the moment
+    -- Not vector-ref for CLitAlt...
+    withVar env sc $ \ nm => doCase nm alts
+  where
+    -- add a bunch of lets then do body
+    conAlt : SCEnv → String → SnocList String → List String → CExp → String
+    conAlt env nm lets Nil body = "(let (\{joinBy " " (lets <>> Nil)}) \{cexpToScm env body})"
+    -- TODO let `vector-ref nm ..`` vs `#f` (erased) in env for erased fields
+    conAlt env nm lets (arg :: args) body = case scbind arg env of
+      (arg', env') => let ix = 1 + snoclen' lets
+        in conAlt env' nm (lets :< "(\{arg'} (vector-ref \{nm} \{show ix}))") args body
+
+    doAlt : String → CAlt → String
+    doAlt nm (CConAlt tag cname _ args body) = "((\{show tag}) \{conAlt env nm Lin args body})"
+    doAlt nm (CDefAlt body) = "(else \{cexpToScm env body})"
+    doAlt nm (CLitAlt lit body) = "((\{scmLit lit}) \{cexpToScm env body})"
+
+    doCase : String → List CAlt → String
+    doCase nm alts@(CLitAlt _ _ :: _) = "(case \{nm} \{joinBy " " $ map (doAlt nm) alts})"
+    doCase nm alts = "(case (vector-ref \{cexpToScm env sc} 0) \{joinBy " " $ map (doAlt nm) alts})"
+
+cexpToScm env (CRef nm) = scmName nm
+cexpToScm env (CMeta _) = fatalError "meta in code gen"
+cexpToScm env (CLit lit) = scmLit lit
+cexpToScm env (CLet nm t u) =
+  let (nm',env') = scbind nm env in
+  "(let ((\{nm'} \{cexpToScm env t})) \{cexpToScm env' u})"
+cexpToScm env (CLetRec nm t u) =
+  let (nm',env') = scbind nm env in
+  "(let ((\{nm'} \{cexpToScm env' t})) \{cexpToScm env' u})"
+cexpToScm env (CLetLoop _ _) = fatalError "CLetLoop in scheme codegen"
+cexpToScm env (CLoop _ _) = fatalError "CLoop in scheme codegen"
+cexpToScm env CErased = "#f"
+cexpToScm env (CConstr tag nm args quants) =
+  -- FIXME need to deal with quants
+  let args' = map (cexpToScm env) args in
+  "(vector \{show tag} \{joinBy " " args'})"
+-- TODO
+cexpToScm env (CRaw _ _) = "CRaw"
+-- TODO We still have a couple of these in CompileExp, for the magic Nat
+cexpToScm env (CPrimOp op a b) = "(\{op} \{cexpToScm env a} \{cexpToScm env b})"
+
+
+-- Collect the QNames used in a term
+getNames : Tm -> List QName -> List QName
+getNames (Ref x nm) acc = nm :: acc
+getNames (Lam x str _ _ t) acc = getNames t acc
+getNames (App x t u) acc = getNames u $ getNames t acc
+getNames (Pi x str icit y t u) acc = getNames u $ getNames t $ QN primNS "PiType" :: acc
+getNames (Let x str t u) acc = getNames u $ getNames t acc
+getNames (LetRec x str _ t u) acc = getNames u $  getNames t acc
+getNames (Case x t alts) acc = getNames t $ foldl getAltNames acc alts
+  where
+    getAltNames : List QName -> CaseAlt -> List QName
+    getAltNames acc (CaseDefault t) = getNames t acc
+    getAltNames acc (CaseCons name args t) = name :: getNames t acc
+    getAltNames acc (CaseLit lit t) = getNames t acc
+getNames _ acc = acc
+
+-- returns a QName -> Def of in-use entries
+-- This will be what we work on for optimization passes
+getEntries : SortedMap QName Def → QName → M (SortedMap QName Def)
+getEntries acc name = do
+  top <- getTop
+  case lookup name top of
+    Nothing => do
+      putStrLn "bad name \{show name}"
+      pure acc
+    Just (MkEntry _ name type def@(Fn exp) _) => case lookupMap' name acc of
+      Just _ => pure acc
+      Nothing => do
+        top <- getTop
+        exp <- zonk top 0 Nil exp
+        let acc = updateMap name (Fn exp) acc
+        foldlM getEntries acc $ getNames exp Nil
+    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
+
+/-
+
+## Sort names by dependencies
+
+This code is _way too subtle. The problem goes away if I wrap `() =>` around 0-ary top level functions. But I'm
+stubborn, so I'm giving it a try. Changing codegen may repeatedly break this and require switching to `() =>`.
+
+The idea here is to get a list of names to emit in order of dependencies. But top level 0-ary functions can reference
+and call names at startup. They can't reference something that hasn't been declared yet and can't call something that
+hasn't been defined.
+
+As an example, a recursive show instance:
+- References the `show` function
+- `show` function references the instance under a lambda (not inlining this yet)
+- We hit the bare function first, it depends on the instance (because of recursion), which depends on the
+  function, but loop prevention cuts.
+
+We have main at the top, it is called so we visit it deep. We do a depth-first traversal, but will distinguish whether
+we're visiting shallow or deep. We're trying to avoid hitting issues with indirect circular references.
+
+- Anything we visit deep, we ensure is visited shallow first
+- Shallow doesn't go into function bodies, but we do go into lambdas
+- Anything invoked with arguments while shallow is visited deep, anything referenced or partially applied is still shallow.
+- We keep track of both shallow and deep visits in our accumuulator
+- Shallow represents the declaration, so we filter to those at the end
+
+TODO this could be made faster by keeping a map of the done information
+REVIEW could I avoid most of this by using `function` instead of arrow functions?
+
+-/
+
+sortedNames : SortedMap QName CExp → List QName → List QName
+sortedNames defs names =
+  map snd $ filter (not ∘ fst) $ foldl (go Nil) Nil $ map (True,) names
+  where
+    getBody : CAlt → CExp
+    getBody (CConAlt _ _ _ _ t) = t
+    getBody (CLitAlt _ t) = t
+    getBody (CDefAlt t) = t
+
+    -- deep if this qn is being applied to something
+    getNames : (deep : Bool) → List (Bool × QName) → CExp →  List (Bool × QName)
+    -- liftIO calls a lambda statically
+    getNames deep acc (CLam _ t) = getNames deep acc t
+    getNames deep acc (CLetLoop _ t) = getNames deep acc t
+    -- top level 0-ary function, doesn't happen
+    getNames deep acc (CFun _ t) = if deep then getNames deep acc t else acc
+    -- REVIEW - True or deep?
+    getNames deep acc (CLoop args qs) = foldl (getNames True) acc args
+    getNames deep acc (CAppRef nm args qs) =
+      if length' args == length' qs
+        then case args of
+          Nil => (True, nm) :: acc
+          ts => foldl (getNames True) ((True, nm) :: acc) ts
+        else
+          foldl (getNames deep) ((deep, nm) :: acc) args
+   -- TODO look at which cases generate CApp
+    getNames deep acc (CApp t u) = getNames True (getNames deep acc u) t
+    getNames deep acc (CCase t alts) = foldl (getNames deep) acc $ t :: map getBody alts
+    -- we're not calling it
+    getNames deep acc (CRef qn) = (deep, qn) :: acc
+    getNames deep acc (CLet _ t u) = getNames deep (getNames deep acc t) u
+    getNames deep acc (CLetRec _ t u) = getNames deep (getNames deep acc t) u
+    getNames deep acc (CConstr _ _ ts _) = foldl (getNames deep) acc ts
+    -- if the CRaw is called, then the deps are called
+    getNames deep acc (CRaw _ deps) = acc -- map (_,_ deep) deps ++ acc
+    -- wrote these out so I get an error when I add a new constructor
+    getNames deep acc (CLit _) = acc
+    getNames deep acc (CMeta _) = acc
+    getNames deep acc (CBnd _) = acc
+    getNames deep acc CErased = acc
+    getNames deep acc (CPrimOp op t u) = getNames deep (getNames deep acc t) u
+
+    -- recurse on all dependencies, pushing onto acc
+    go : List (Bool × QName) → List (Bool × QName) → (Bool × QName) → List (Bool × QName)
+      -- Need to force shallow if we're doing deep and haven't done shallow.
+    go loop acc this@(deep, qn) =
+      -- there is a subtle issue here with an existing (False, qn) vs (True, qn)
+      let acc = if deep && not (elem (False, qn) acc) && not (elem (False, qn) loop)
+                  then go loop acc (False, qn)
+                  else acc
+      in if elem this loop
+        then acc
+        else if elem this acc then acc
+        else
+          case lookupMap' qn defs of
+            Nothing => acc   -- only `bouncer`
+            Just exp => this :: foldl (go $ this :: loop) acc (getNames deep Nil exp)
+
+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 : List QName → M (List String)
+process names = do
+  top <- getTop
+  entries <- foldlM getEntries emptyMap names
+
+  -- Maybe move this dance into liftWhere
+  ref <- newIORef entries
+  let foo = MkRef ref -- for the autos below
+  eraseEntries
+  liftWhere
+  entries <- readIORef ref
+  -- Now working with defs
+  exprs <- mapM defToCExp $ toList entries
+  let cexpMap = foldMap const emptyMap exprs
+  -- not needed for scheme
+  -- cexpMap <- tailCallOpt cexpMap
+  -- Not needed for scheme
+  -- cexpMap <- liftLambda cexpMap
+  let names = sortedNames cexpMap names
+  pure $ mapMaybe (go cexpMap) names
+  where
+    go : ExpMap → QName → Maybe String
+    go cexpMap name = do
+      (qn, cexp) <- lookupMap name cexpMap
+      case cexp of
+        CRaw _ _ => Nothing
+        _ => Just $ "(define \{scmName qn} \{cexpToScm Nil cexp})"
+
+compileScheme : M (List String)
+compileScheme = do
+  top <- getTop
+  let exports = getExports Nil $ listValues top.currentMod.modDefs
+  let mainName = (QN top.currentMod.modName "main")
+  let main = lookup mainName top
+  let todo = case main of
+    Nothing => exports
+    Just _ => mainName :: exports
+  defs <- process todo
+  case lookup mainName top of
+    Just _ => -- tack on call to main function
+      let exec = "(\{show mainName} 'world)"
+      in pure $ "(import (chezscheme))" :: "(include \"prim.ss\")" :: reverse (exec :: defs)
+    Nothing => pure $ reverse defs
+  where
+    getExports : List QName → List TopEntry → List QName
+    getExports acc Nil = acc
+    getExports acc ((MkEntry fc name@(QN ns nm) type def eflags) :: rest) =
+      let acc = if elem Export eflags then name :: acc else acc
+      in getExports acc rest
+