rts: retainer: Make visit callback easier to implement

Currently it is necessary for user code to expend at least one extra bit in
the closure header just to know whether visit() should return true or
false, to indicate if children should be traversed.

The generic traversal code already has this information in the visited bit
so simply pass it to the visit callback.

2 files changed, 25 insertions, 17 deletions

diff --git a/rts/RetainerProfile.c b/rts/RetainerProfile.c
index 2bb5a503f3..1d2807f468 100644
--- a/rts/RetainerProfile.c
+++ b/rts/RetainerProfile.c
@@ -961,16 +961,18 @@ endRetainerProfiling( void )
 
 /**
  * Make sure a closure's profiling data is initialized to zero if it does not
- * conform to the current value of the flip bit.
+ * conform to the current value of the flip bit, returns true in this case.
  *
  * See Note [Profiling heap traversal visited bit].
  */
-void
+bool
 traverseMaybeInitClosureData(StgClosure *c)
 {
     if (!isTravDataValid(c)) {
         setTravDataToZero(c);
+        return true;
     }
+    return false;
 }
 
 /* -----------------------------------------------------------------------------
@@ -1333,7 +1335,7 @@ traversePAP (traverseState *ts,
 }
 
 static bool
-retainVisitClosure( const StgClosure *c, const StgClosure *cp, const stackData data, stackData *out_data )
+retainVisitClosure( const StgClosure *c, const StgClosure *cp, const stackData data, const bool first_visit, stackData *out_data )
 {
     retainer r = data.c_child_r;
     RetainerSet *s, *retainerSetOfc;
@@ -1379,7 +1381,7 @@ retainVisitClosure( const StgClosure *c, const StgClosure *cp, const stackData d
     } else {
         // This is not the first visit to *c.
         if (isMember(r, retainerSetOfc))
-            return 1;          // no need to process child
+            return 0;          // no need to process child
 
         if (s == NULL)
             associate(c, addElement(r, retainerSetOfc));
@@ -1398,7 +1400,7 @@ retainVisitClosure( const StgClosure *c, const StgClosure *cp, const stackData d
         }
 
         if (isRetainer(c))
-            return 1;          // no need to process child
+            return 0;          // no need to process child
 
         // compute c_child_r
         out_data->c_child_r = r;
@@ -1407,7 +1409,7 @@ retainVisitClosure( const StgClosure *c, const StgClosure *cp, const stackData d
     // now, RSET() of all of *c, *cp, and *r is valid.
     // (c, c_child_r) are available.
 
-    return 0;
+    return 1;
 }
 
 static void
@@ -1438,8 +1440,10 @@ resetMutableObjects(void)
 }
 
 /**
- * Traverse all closures on the traversal work-stack, calling 'visit_cb'
- * on each closure. See 'visitClosure_cb' for details.
+ * Traverse all closures on the traversal work-stack, calling 'visit_cb' on each
+ * closure. See 'visitClosure_cb' for details. This function flips the 'flip'
+ * bit and hence every closure's profiling data will be reset to zero upon
+ * visiting. See Note [Profiling heap traversal visited bit].
  */
 void
 traverseWorkStack(traverseState *ts, visitClosure_cb visit_cb)
@@ -1536,9 +1540,10 @@ inner_loop:
     }
 
     // If this is the first visit to c, initialize its data.
-    traverseMaybeInitClosureData(c);
-
-    if(visit_cb(c, cp, data, (stackData*)&child_data))
+    bool first_visit = traverseMaybeInitClosureData(c);
+    bool traverse_children
+        = visit_cb(c, cp, data, first_visit, (stackData*)&child_data);
+    if(!traverse_children)
         goto loop;
 
     // process child
diff --git a/rts/TraverseHeap.h b/rts/TraverseHeap.h
index 95c6cfbcdf..a82bf0ec6d 100644
--- a/rts/TraverseHeap.h
+++ b/rts/TraverseHeap.h
@@ -85,22 +85,25 @@ typedef struct traverseState_ {
 /**
  * Callback called when heap traversal visits a closure.
  *
- * Before this callback is called the profiling header of the visited closure
- * 'c' is zero'd with 'setTravDataToZero' if this closure hasn't been visited in
- * this run yet. See Note [Profiling heap traversal visited bit].
+ * The callback can assume that the closure's profiling data has been
+ * initialized to zero if this is the first visit during a pass.
  *
- * Return 'true' when this is not the first visit to this element. The generic
- * traversal code will then skip traversing the children.
+ * See Note [Profiling heap traversal visited bit].
+ *
+ * Returning 'false' will instruct the heap traversal code to skip processing
+ * this closure's children. If you don't need to traverse any closure more than
+ * once you can simply return 'first_visit'.
  */
 typedef bool (*visitClosure_cb) (
     const StgClosure *c,
     const StgClosure *cp,
     const stackData data,
+    const bool first_visit,
     stackData *child_data);
 
 void traverseWorkStack(traverseState *ts, visitClosure_cb visit_cb);
 void traversePushClosure(traverseState *ts, StgClosure *c, StgClosure *cp, stackData data);
-void traverseMaybeInitClosureData(StgClosure *c);
+bool traverseMaybeInitClosureData(StgClosure *c);
 
 void initializeTraverseStack(traverseState *ts);
 void closeTraverseStack(traverseState *ts);