Fix selector thunk optimization

ChangeLog.md

@@ -1,5 +1,11 @@
 # Revision history for visualize-cbn
 
+## 0.2.1 -- 2024-01-10
+
+* Fixes to the selector thunk optimization: also apply it at the top-level,
+  and correctly apply `--hide-selector-thunk-opt` (previously `--hide-gc`
+  was hiding selector thunk optimization steps by mistake).
+
 ## 0.2.0 -- 2023-12-20
 
 * Support multiple (mutually recursive) bindings in `let`

examples/selthunkopt2.hs

@@ -0,0 +1,20 @@
+break = (\xs ->
+      case xs of {
+          Nil -> Pair Nil Nil
+        ; Cons x xs' ->
+            if eq x 0
+              then Pair Nil xs'
+              else let b = @break xs'
+                   in Pair (Cons x (fst b)) (snd b)
+        }
+    )
+
+last = (\def -> \xs ->
+      case xs of {
+          Nil -> def
+        ; Cons x' xs' -> @last x' xs'
+        }
+    )
+
+main = let broken = @break (Cons 1 (Cons 2 (Cons 3 (Cons 4 (Cons 0 (Cons 5 (Cons 6 (Cons 7 (Cons 8 Nil)))))))))
+       in eq (@last 0 (fst broken)) (@last 0 (snd broken))

src/CBN/SelThunkOpt.hs

@@ -11,6 +11,7 @@
     Specifically section 2.5.7, "Selector optimization"
 
   - "Three runtime optimizations done by GHC's GC", Ömer Sinan Ağacan
+    <https://osa1.net/posts/2018-03-16-gc-optimizations.html>
     Specifically section 3, "Selector thunk evaluation"
 
   - "GHC Commentary: The Layout of Heap Objects", section "Selector thunks"
@@ -32,8 +33,18 @@ import CBN.Heap
 import CBN.Language
 
 -- | Apply selector thunk optimization
-selThunkOpt :: Heap Term -> (Heap Term, Set Ptr)
-selThunkOpt = findAll Set.empty
+selThunkOpt :: Heap Term -> Term -> (Heap Term, Term, Bool, Set Ptr)
+selThunkOpt hp0 e0 =
+    let (hp1, e1, atToplevel) = case applyInTerm hp0 e0 of
+                                  Nothing        -> (hp0, e0, False)
+                                  Just (hp', e') -> (hp', e', True)
+        (hp2, ptrs)           = applyInHeap hp1
+
+    in (hp2, e1, atToplevel, ptrs)
+
+-- | Apply selector thunk optimization
+applyInHeap :: Heap Term -> (Heap Term, Set Ptr)
+applyInHeap = findAll Set.empty
   where
     findAll :: Set Ptr -> Heap Term -> (Heap Term, Set Ptr)
     findAll acc hp =
@@ -89,7 +100,7 @@ applyInTerm = \hp term -> do
     -- This code is a bit simpler than the corresponding code in evaluation,
     -- because we /only/ deal with selectors, not general case statements. This
     -- means we don't need to care about substitution, but can literally just
-    -- select the right argument (using
+    -- select the right argument.
 
     go term@(TCase e (Selector s)) = do
         (hp, _) <- get

src/CBN/Trace.hs

@@ -42,7 +42,9 @@ data TraceCont =
   | TraceGC (Set Ptr) Trace
 
     -- | The selector thunk optimization was applied
-  | TraceSelThunk (Set Ptr) Trace
+    --
+    -- We separately record if the selector thunk was applied at the top-level.
+  | TraceSelThunk Bool (Set Ptr) Trace
 
     -- | We simplified the heap by inlining some definitions
   | TraceInline (Set Ptr) Trace
@@ -58,11 +60,11 @@ traceTerm shouldGC shouldInline enableSelThunkOpt = go
         Step d (hp1, e1) ->
           let (traceSelThunkOpt, hp2, e2)
                 | enableSelThunkOpt
-                = let (hp', optimized) = selThunkOpt hp1
-                  in if Set.null optimized then
+                = let (hp', e', atToplevel, optimized) = selThunkOpt hp1 e1
+                  in if not atToplevel && Set.null optimized then
                        (id, hp1, e1)
                      else
-                       (Trace (hp1, e1) . TraceSelThunk optimized, hp', e1)
+                       (Trace (hp1, e1) . TraceSelThunk atToplevel optimized, hp', e')
                 | otherwise
                 = (id, hp1, e1) in
 
@@ -145,10 +147,10 @@ summarize SummarizeOptions{..} = go 0
             if summarizeHideGC
               then go (n + 1) t'
               else showSrc $ TraceGC ps $ go (n + 1) t'
-          TraceSelThunk ps t' ->
-            if summarizeHideGC
+          TraceSelThunk atToplevel ps t' ->
+            if summarizeHideSelThunk
               then go (n + 1) t'
-              else showSrc $ TraceSelThunk ps $ go (n + 1) t'
+              else showSrc $ TraceSelThunk atToplevel ps $ go (n + 1) t'
           TraceInline ps t' ->
             if summarizeHideInlining
               then go (n + 1) t'

src/CBN/Trace/Graph.hs

@@ -27,13 +27,13 @@ render tr =
     go :: Int -> Trace -> String
     go index (Trace (hp, t) cont) =
       case cont of
-        TraceWHNF _          -> mkFrame Set.empty Nothing "whnf"
-        TraceStuck err       -> mkFrame Set.empty Nothing (mkErr err)
-        TraceStopped         -> mkFrame Set.empty Nothing "stopped"
-        TraceStep d tr'      -> mkFrame Set.empty (mkFocus d) (mkDesc d) ++ go (index + 1) tr'
-        TraceGC ps tr'       -> mkFrame ps Nothing "gc"       ++ go (index + 1) tr'
-        TraceSelThunk ps tr' -> mkFrame ps Nothing "selector" ++ go (index + 1) tr'
-        TraceInline ps tr'   -> mkFrame ps Nothing "inline"   ++ go (index + 1) tr'
+        TraceWHNF _            -> mkFrame Set.empty Nothing "whnf"
+        TraceStuck err         -> mkFrame Set.empty Nothing (mkErr err)
+        TraceStopped           -> mkFrame Set.empty Nothing "stopped"
+        TraceStep d tr'        -> mkFrame Set.empty (mkFocus d) (mkDesc d) ++ go (index + 1) tr'
+        TraceGC ps tr'         -> mkFrame ps Nothing "gc"       ++ go (index + 1) tr'
+        TraceSelThunk _ ps tr' -> mkFrame ps Nothing "selector" ++ go (index + 1) tr'
+        TraceInline ps tr'     -> mkFrame ps Nothing "inline"   ++ go (index + 1) tr'
       where
         mkFrame :: Set Ptr -> Maybe Ptr -> T.Text -> String
         mkFrame garbage focus status =

src/CBN/Trace/JavaScript.hs

@@ -42,13 +42,13 @@ render name graph = \tr ->
     go :: Int -> Trace -> String
     go n (Trace (hp, e) c) =
         case c of
-          TraceWHNF _          -> mkFrame Set.empty Nothing "whnf"
-          TraceStuck err       -> mkFrame Set.empty Nothing (mkErr err)
-          TraceStopped         -> mkFrame Set.empty Nothing "stopped"
-          TraceStep d tr'      -> mkFrame Set.empty (mkFocus d) (mkDesc d) ++ go (n + 1) tr'
-          TraceGC ps tr'       -> mkFrame ps Nothing "gc"       ++ go (n + 1) tr'
-          TraceSelThunk ps tr' -> mkFrame ps Nothing "selector" ++ go (n + 1) tr'
-          TraceInline ps tr'   -> mkFrame ps Nothing "inline"   ++ go (n + 1) tr'
+          TraceWHNF _            -> mkFrame Set.empty Nothing "whnf"
+          TraceStuck err         -> mkFrame Set.empty Nothing (mkErr err)
+          TraceStopped           -> mkFrame Set.empty Nothing "stopped"
+          TraceStep d tr'        -> mkFrame Set.empty (mkFocus d) (mkDesc d) ++ go (n + 1) tr'
+          TraceGC ps tr'         -> mkFrame ps Nothing "gc"       ++ go (n + 1) tr'
+          TraceSelThunk _ ps tr' -> mkFrame ps Nothing "selector" ++ go (n + 1) tr'
+          TraceInline ps tr'     -> mkFrame ps Nothing "inline"   ++ go (n + 1) tr'
       where
         mkFrame :: Set Ptr -> Maybe Ptr -> String -> String
         mkFrame garbage focus status =

src/CBN/Trace/Textual.hs

@@ -22,13 +22,13 @@ renderIO disableAnsi = go 0
     go :: Int -> Trace -> IO ()
     go n (Trace (hp, e) c) = do
       case c of
-        TraceWHNF _          -> mkFrame Set.empty Nothing (putStr $ "whnf")
-        TraceStuck err       -> mkFrame Set.empty Nothing (putStr $ "stuck: " ++ err)
-        TraceStopped         -> mkFrame Set.empty Nothing (putStr $ "stopped")
-        TraceStep d tr'      -> mkFrame Set.empty (mkFocus d) (pretty d)  >> go (n + 1) tr'
-        TraceGC ps tr'       -> mkFrame ps Nothing (ptrs "collecting" ps) >> go (n + 1) tr'
-        TraceSelThunk ps tr' -> mkFrame ps Nothing (ptrs "apply selectors" ps) >> go (n + 1) tr'
-        TraceInline ps tr'   -> mkFrame ps Nothing (ptrs "inlining" ps) >> go (n + 1) tr'
+        TraceWHNF _              -> mkFrame Set.empty Nothing (putStr $ "whnf")
+        TraceStuck err           -> mkFrame Set.empty Nothing (putStr $ "stuck: " ++ err)
+        TraceStopped             -> mkFrame Set.empty Nothing (putStr $ "stopped")
+        TraceStep d tr'          -> mkFrame Set.empty (mkFocus d) (pretty d)  >> go (n + 1) tr'
+        TraceGC ps tr'           -> mkFrame ps Nothing (ptrs False "collecting" ps) >> go (n + 1) tr'
+        TraceSelThunk top ps tr' -> mkFrame ps Nothing (ptrs top   "apply selectors" ps) >> go (n + 1) tr'
+        TraceInline ps tr'       -> mkFrame ps Nothing (ptrs False "inlining" ps) >> go (n + 1) tr'
       where
         mkFrame :: Set Ptr -> Maybe Ptr -> IO () -> IO ()
         mkFrame garbage focus msg = do
@@ -38,10 +38,13 @@ renderIO disableAnsi = go 0
           putChar '\n'
           putStr "(" ; msg ; putStrLn ")\n"
 
-    ptrs :: String -> Set Ptr -> IO ()
-    ptrs label ps = do
+    ptrs :: Bool -> String -> Set Ptr -> IO ()
+    ptrs atToplevel label ps = do
       putStr (label ++ " ")
-      sequence_ . intersperse (putStr ", ") . map pretty $ Set.toList ps
+      sequence_ . intersperse (putStr ", ") $ concat [
+          [putStr "top-level" | atToplevel]
+        , map pretty $ Set.toList ps
+        ]
 
     pretty :: ToDoc a => a -> IO ()
     pretty =

visualize-cbn.cabal

@@ -1,6 +1,6 @@
 cabal-version:       >=1.10
 name:                visualize-cbn
-version:             0.2.0
+version:             0.2.1
 synopsis:            Visualize CBN reduction
 description:         CBN interpretation and visualization tool.
                      Exports in text format, coloured text (ANSI) or HTML/JavaScript.