3 files changed, 63 insertions(+), 30 deletions(-)

M src/Codegen.hs
M src/Extern.hs
M src/Gen.hs

@@ 1,4 1,4 @@
-{-# LANGUAGE OverloadedStrings, LambdaCase, TupleSections, FlexibleContexts #-}
+{-# LANGUAGE OverloadedStrings, LambdaCase, TupleSections, FlexibleContexts, RankNTypes #-}
 
 -- | Generation of LLVM IR code from our monomorphic AST.
 module Codegen (codegen) where


@@ 14,6 14,7 @@ import qualified LLVM.AST.Type as LLType
 import qualified LLVM.AST.Constant as LLConst
 import Data.String
 import System.FilePath
+import Control.Monad.Reader
 import Control.Monad.Writer
 import Control.Monad.Except
 import qualified Data.Map as Map


@@ 26,9 27,8 @@ import Data.List
 import Data.Function
 import Data.Functor
 import Data.Functor.Identity
-import Data.Bifunctor
 import Control.Applicative
-import Lens.Micro.Platform (use, assign, view)
+import Lens.Micro.Platform (use, assign, Lens')
 
 import Misc
 import SrcPos


@@ 51,7 51,8 @@ codegen layout moduleFilePath (Program (Topo defs) tdefs externs) = runExcept $ 
             $ augment dataTypes tdefs''
             $ withBuiltins
             $ withExternSigs externs
-            $ withGlobDefSigs (map (second unpos) defs')
+            $ withGlobFunDefSigs funDefs
+            $ withGlobVarDefSigs varDefs
             $ do
                   es <- genExterns externs
                   funDefs' <- mapM genGlobFunDef funDefs


@@ 78,8 79,33 @@ codegen layout moduleFilePath (Program (Topo defs) tdefs externs) = runExcept $ 
                                   ]
         }
   where
-    withGlobDefSigs sigs ga = do
-        sigs' <- forM sigs $ \(v@(TypedVar x t), (us, _)) -> do
+    withGlobFunDefSigs = withGlobDefSigs globalEnv
+    -- TODO: This is a workaround for global vars not being registered in the GC when
+    --       running in JIT.
+    --
+    --       The plan is to keep global defs in globalEnv, and not capture these in
+    --       closures, as they can always be reached at the top, regardless of the scope
+    --       etc. The bug is that when running for example "sieve.carth" with the JIT
+    --       interpreter, it freezes/crashes after just a few houndred iterations (not
+    --       when compiling though, which is wierd). After some trial an error, my current
+    --       hypothesis is that global non-function variables can require heap allocations
+    --       when being initialized in `init`, but the static locations that the created
+    --       values are stored in are not registered as roots in the Boehm GC, so some
+    --       values that are actually in use are garbage collected after a little while.
+    --
+    --       By keeping the global non-function vars in the `localEnv`, all values in use
+    --       are guaranteed to be in captured envs and reachable when starting at the
+    --       stack / the registers.
+    withGlobVarDefSigs = withGlobDefSigs localEnv
+
+    withGlobDefSigs
+        :: MonadReader Env m
+        => Lens' Env (Map TypedVar Operand)
+        -> [(TypedVar, WithPos ([M.Type], e))]
+        -> m x
+        -> m x
+    withGlobDefSigs env sigs ga = do
+        sigs' <- forM sigs $ \(v@(TypedVar x t), WithPos _ (us, _)) -> do
             t' <- genType' t
             pure
                 ( v


@@ 258,7 284,7 @@ genExpr (Expr pos expr) = locally srcPos (pos <|>) $ do
 
 genExprLambda :: TypedVar -> (Expr, M.Type) -> Gen Val
 genExprLambda p (b, bt) = do
-    let fvXs = Set.toList (Set.delete p (freeVars b))
+    fvXs <- lambdaBodyFreeVars p b
     bt' <- genRetType bt
     genLambda fvXs p (genTailExpr b, bt')
 


@@ 283,22 309,21 @@ genApp fe' ae' = genBetaReduceApp genExpr pure NoTail (fe', [ae'])
 -- | Beta-reduction and closure application
 genBetaReduceApp
     :: (Expr -> Gen a) -> (Val -> Gen a) -> TailCallKind -> (Expr, [Expr]) -> Gen a
-genBetaReduceApp genExpr' returnMethod tail' applic = view env >>= \env' ->
-    case applic of
-        (Expr _ (Fun (p, (b, _))), ae : aes) -> do
-            a <- genExpr ae
-            withVal p a (genBetaReduceApp genExpr' returnMethod tail' (b, aes))
-        (Expr _ (App fe ae _), aes) ->
-            genBetaReduceApp genExpr' returnMethod tail' (fe, ae : aes)
-        (fe, []) -> genExpr' fe
-        (Expr _ (Var x), aes) | not (Map.member x env') ->
-            returnMethod =<< genAppBuiltinVirtual x (map genExpr aes)
-        (fe, aes) -> do
-            f <- genExpr fe
-            as <- mapM genExpr (init aes)
-            closure <- foldlM (app (Just NoTail)) f as
-            arg <- genExpr (last aes)
-            returnMethod =<< app (Just tail') closure arg
+genBetaReduceApp genExpr' returnMethod tail' applic = ask >>= \env -> case applic of
+    (Expr _ (Fun (p, (b, _))), ae : aes) -> do
+        a <- genExpr ae
+        withVal p a (genBetaReduceApp genExpr' returnMethod tail' (b, aes))
+    (Expr _ (App fe ae _), aes) ->
+        genBetaReduceApp genExpr' returnMethod tail' (fe, ae : aes)
+    (fe, []) -> genExpr' fe
+    (Expr _ (Var x), aes) | isNothing (lookupVar' x env) ->
+        returnMethod =<< genAppBuiltinVirtual x (map genExpr aes)
+    (fe, aes) -> do
+        f <- genExpr fe
+        as <- mapM genExpr (init aes)
+        closure <- foldlM (app (Just NoTail)) f as
+        arg <- genExpr (last aes)
+        returnMethod =<< app (Just tail') closure arg
 
 app :: Maybe TailCallKind -> Val -> Val -> Gen Val
 app tailkind closure a = do


@@ 360,7 385,7 @@ genLet' def genBody = case def of
     RecDefs ds -> do
         (binds, cs) <- fmap unzip $ forM ds $ \case
             (lhs, WithPos _ (_, (p, (fb, fbt)))) -> do
-                let fvXs = Set.toList (Set.delete p (freeVars fb))
+                fvXs <- lambdaBodyFreeVars p fb
                 tcaptures <- fmap typeStruct (mapM (\(TypedVar _ t) -> genType t) fvXs)
                 captures <- genHeapAllocGeneric tcaptures
                 fbt' <- genRetType fbt


@@ 371,6 396,12 @@ genLet' def genBody = case def of
             forM_ cs (uncurry populateCaptures)
             genBody
 
+lambdaBodyFreeVars :: MonadReader Env m => TypedVar -> Expr -> m [TypedVar]
+lambdaBodyFreeVars param body = ask <&> \env ->
+    let globsToNotCapture = Set.difference (Map.keysSet (_globalEnv env))
+                                           (Map.keysSet (_localEnv env))
+    in  Set.toList $ Set.difference (Set.delete param (freeVars body)) globsToNotCapture
+
 genTailMatch :: Expr -> DecisionTree -> Type -> Gen ()
 genTailMatch m dt tbody = do
     m' <- getLocal =<< genExpr m

@@ 49,7 49,7 @@ withExternSigs es ga = do
             , ConstantOperand
                 $ LLConst.GlobalReference (LLType.ptr t') (mkName ("_wrapper_" ++ name))
             )
-    augment env (Map.fromList es') ga
+    augment globalEnv (Map.fromList es') ga
 
 genExterns :: [(String, M.Type, SrcPos)] -> Gen' [Definition]
 genExterns = fmap join . mapM genExtern

@@ 65,7 65,8 @@ data FunInstr = WithRetType Instr Type
 data Env = Env
     -- TODO: Could operands in env be Val instead? I.e., either stack-allocated
     --       or local?
-    { _env :: Map TypedVar Operand -- ^ Environment of stack allocated variables
+    { _localEnv :: Map TypedVar Operand -- ^ Environment of stack allocated variables
+    , _globalEnv :: Map TypedVar Operand
     , _enumTypes :: Map Name Word32
     , _dataTypes :: Map Name [Type]
     , _builtins :: Map String ([Parameter], Type)


@@ 304,7 305,8 @@ compileUnitId = MetadataNodeID 0
 runGen' :: Monad m => StateT St (ReaderT Env m) a -> m a
 runGen' g = runReaderT (evalStateT g initSt) initEnv
   where
-    initEnv = Env { _env = Map.empty
+    initEnv = Env { _localEnv = Map.empty
+                  , _globalEnv = Map.empty
                   , _enumTypes = Map.empty
                   , _dataTypes = Map.empty
                   , _builtins = Map.empty


@@ 423,7 425,7 @@ withVars = withXs withVar
 -- | Takes a local, stack allocated value, and runs a generator in the
 --   environment with the variable
 withVar :: TypedVar -> Operand -> Gen a -> Gen a
-withVar x v = locally env (Map.insert x v)
+withVar x v = locally localEnv (Map.insert x v)
 
 withVals :: [(TypedVar, Val)] -> Gen a -> Gen a
 withVals = withXs withVal


@@ 461,8 463,8 @@ lookupVar x = lookupVar' x >>= \case
     Nothing -> genAppBuiltinVirtual x []
 
 lookupVar' :: MonadReader Env m => TypedVar -> m (Maybe Val)
-lookupVar' x = do
-    view (env . to (Map.lookup x)) >>= pure . fmap VVar
+lookupVar' x =
+    ask <&> \e -> fmap VVar (Map.lookup x (_localEnv e) <|> Map.lookup x (_globalEnv e))
 
 genAppBuiltinVirtual :: TypedVar -> [Gen Val] -> Gen Val
 genAppBuiltinVirtual (TypedVar g t) aes = do