Recording mutated closures

asterius/src/Asterius/Builtins.hs

@@ -165,6 +165,7 @@ rtsAsteriusModule opts =
     <> dirtyMVarFunction opts
     <> dirtyStackFunction opts
     <> recordClosureMutatedFunction opts
+    <> recordMutableCapFunction opts
     <> tryWakeupThreadFunction opts
     <> raiseExceptionHelperFunction opts
     <> barfFunction opts
@@ -900,15 +901,19 @@ dirtyTSO _ tso =
   if'
     []
     (eqZInt32 $ loadI32 tso offset_StgTSO_dirty)
-    (storeI32 tso offset_StgTSO_dirty $ constI32 1)
+    (do
+      storeI32 tso offset_StgTSO_dirty $ constI32 1
+      recordMutable tso)
     mempty
 
 dirtySTACK :: Expression -> Expression -> EDSL ()
 dirtySTACK _ stack =
   if'
     []
     (eqZInt32 $ loadI32 stack offset_StgStack_dirty)
-    (storeI32 stack offset_StgStack_dirty $ constI32 1)
+    (do
+      storeI32 stack offset_StgStack_dirty $ constI32 1
+      recordMutable stack)
     mempty
 
 -- `_scheduleTSO(tso,func)` executes the given tso starting at the given
@@ -1057,6 +1062,7 @@ newCAFFunction _ = runEDSL "newCAF" $ do
     loadI64 reg offset_StgRegTable_rCurrentTSO
   storeI64 caf offset_StgIndStatic_indirectee bh
   storeI64 caf 0 $ symbol "stg_IND_STATIC_info"
+  recordMutable caf
   emit bh
 
 -- Repeatedly calls the function pointed by ``__asterius_pc`` until this
@@ -1300,19 +1306,29 @@ threadPausedFunction _ = runEDSL "threadPaused" $ do
   _ <- params [I64, I64]
   pure ()
 
+-- | Write barrier for generational GC
+
 dirtyMutVarFunction :: BuiltinsOptions -> AsteriusModule
 dirtyMutVarFunction _ = runEDSL "dirty_MUT_VAR" $ do
   [_, p] <- params [I64, I64]
   if'
     []
     (loadI64 p 0 `eqInt64` symbol "stg_MUT_VAR_CLEAN_info")
-    (storeI64 p 0 $ symbol "stg_MUT_VAR_DIRTY_info")
+    (do
+      storeI64 p 0 $ symbol "stg_MUT_VAR_DIRTY_info"
+      recordMutable p)
     mempty
 
 dirtyMVarFunction :: BuiltinsOptions -> AsteriusModule
 dirtyMVarFunction _ = runEDSL "dirty_MVAR" $ do
-  [_basereg, _mvar] <- params [I64, I64]
-  mempty
+  [_, p] <- params [I64, I64]
+  if'
+    []
+    (loadI64 p 0 `eqInt64` symbol "stg_MVAR_CLEAN_info")
+    (do
+      storeI64 p 0 $ symbol "stg_MVAR_DIRTY_info"
+      recordMutable p)
+    mempty
 
 dirtyStackFunction :: BuiltinsOptions -> AsteriusModule
 dirtyStackFunction _ = runEDSL "dirty_STACK" $ do
@@ -1324,6 +1340,19 @@ recordClosureMutatedFunction _ = runEDSL "recordClosureMutated" $ do
   [_cap, _closure] <- params [I64, I64]
   mempty
 
+recordMutable :: Expression -> EDSL ()
+recordMutable _ =
+  pure () -- STUB
+
+recordMutableCap :: Expression -> Expression -> EDSL ()
+recordMutableCap _ _ =
+  pure () -- STUB
+
+recordMutableCapFunction :: BuiltinsOptions -> AsteriusModule
+recordMutableCapFunction _ = runEDSL "recordMutableCap" $ do
+  [cap, p] <- params [I64, I64]
+  recordMutableCap cap p
+
 tryWakeupThreadFunction :: BuiltinsOptions -> AsteriusModule
 tryWakeupThreadFunction _ = runEDSL "tryWakeupThread" $ do
   [_cap, tso] <- params [I64, I64]

ghc-toolkit/boot-libs/rts/PrimOps.cmm

@@ -1723,6 +1723,12 @@ stg_putMVarzh ( P_ mvar, /* :: MVar a */
         jump stg_block_putmvar(mvar,val);
     }
 
+    // We are going to mutate the closure, make sure its current pointers
+    // are marked.
+    if (info == stg_MVAR_CLEAN_info) {
+        ccall dirty_MVAR(BaseReg "ptr", mvar "ptr");
+    }
+
     q = StgMVar_head(mvar);
 loop:
     if (q == stg_END_TSO_QUEUE_closure) {

ghc-toolkit/ghc-libdir/include/Cmm.h

@@ -815,36 +815,17 @@
 #define END_TSO_QUEUE             stg_END_TSO_QUEUE_closure
 #define STM_AWOKEN                stg_STM_AWOKEN_closure
 
-#define recordMutableCap(p, gen)
-#define recordMutable(p)
-
-/*
-#define recordMutableCap(p, gen)                                        \
-  W_ __bd;                                                              \
-  W_ mut_list;                                                          \
-  mut_list = Capability_mut_lists(MyCapability()) + WDS(gen);           \
- __bd = W_[mut_list];                                                   \
-  if (bdescr_free(__bd) >= bdescr_start(__bd) + BLOCK_SIZE) {           \
-      W_ __new_bd;                                                      \
-      ("ptr" __new_bd) = foreign "C" allocBlock_lock();                 \
-      bdescr_link(__new_bd) = __bd;                                     \
-      __bd = __new_bd;                                                  \
-      W_[mut_list] = __bd;                                              \
-  }                                                                     \
-  W_ free;                                                              \
-  free = bdescr_free(__bd);                                             \
-  W_[free] = p;                                                         \
-  bdescr_free(__bd) = free + WDS(1);
-
-#define recordMutable(p)                                        \
-      P_ __p;                                                   \
-      W_ __bd;                                                  \
-      W_ __gen;                                                 \
-      __p = p;                                                  \
-      __bd = Bdescr(__p);                                       \
-      __gen = TO_W_(bdescr_gen_no(__bd));                       \
-      if (__gen > 0) { recordMutableCap(__p, __gen); }
-*/
+#define recordMutableCap(p, gen) \
+    ccall recordMutableCap(p "ptr", gen)
+
+#define recordMutable(p)                                      \
+    P_ __p;                                                   \
+    W_ __bd;                                                  \
+    W_ __gen;                                                 \
+    __p = p;                                                  \
+    __bd = Bdescr(__p);                                       \
+    __gen = TO_W_(bdescr_gen_no(__bd));                       \
+    if (__gen > 0) { recordMutableCap(__p, __gen); }
 
 /* -----------------------------------------------------------------------------
    Arrays