From 95b2fd6855834f000ae97ede8abb24ff7af68e67 Mon Sep 17 00:00:00 2001
From: Edsko de Vries <edsko@well-typed.com>
Date: Wed, 20 Dec 2023 19:53:35 +0100
Subject: [PATCH] Add option to hide prelude from certain step

---
 examples/circular_hos.hs |  2 +-
 examples/repmin.hs       |  2 +-
 src/CBN/Options.hs       |  5 +++--
 src/CBN/Trace.hs         | 12 +++++++-----
 4 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/examples/circular_hos.hs b/examples/circular_hos.hs
index e5895f8..6305b43 100644
--- a/examples/circular_hos.hs
+++ b/examples/circular_hos.hs
@@ -8,7 +8,7 @@
 --
 -- > cabal run visualize-cbn -- \
 -- >   --show-trace \
--- >   --hide-prelude \
+-- >   --hide-prelude=1 \
 -- >   --gc \
 -- >   --selector-thunk-opt \
 -- >   --inline-heap \
diff --git a/examples/repmin.hs b/examples/repmin.hs
index 6937d21..f832352 100644
--- a/examples/repmin.hs
+++ b/examples/repmin.hs
@@ -5,7 +5,7 @@
 --
 -- > cabal run visualize-cbn -- \
 -- >   --show-trace \
--- >   --hide-prelude \
+-- >   --hide-prelude=1 \
 -- >   --gc \
 -- >   --selector-thunk-opt \
 -- >   --inline-heap \
diff --git a/src/CBN/Options.hs b/src/CBN/Options.hs
index bf8dbc9..e50e0bf 100644
--- a/src/CBN/Options.hs
+++ b/src/CBN/Options.hs
@@ -88,9 +88,10 @@ parseSummarizeOptions = SummarizeOptions
            , value 1000
            , metavar "N"
            ])
-    <*> (switch $ mconcat [
+    <*> (optional $ option auto $ mconcat [
              long "hide-prelude"
-           , help "Hide the prelude from the help"
+           , metavar "STEP"
+           , help "Hide the prelude from the help from the given step"
            ])
     <*> (many $ option str $ mconcat [
              long "hide-term"
diff --git a/src/CBN/Trace.hs b/src/CBN/Trace.hs
index 65c2dfd..cc97811 100644
--- a/src/CBN/Trace.hs
+++ b/src/CBN/Trace.hs
@@ -102,7 +102,7 @@ traceTerm shouldGC shouldInline enableSelThunkOpt = go
 data SummarizeOptions = SummarizeOptions {
       summarizeCollapseBeta :: Bool
     , summarizeMaxNumSteps  :: Int
-    , summarizeHidePrelude  :: Bool
+    , summarizeHidePrelude  :: Maybe Int
     , summarizeHideTerms    :: [String]
     , summarizeHideGC       :: Bool
     , summarizeHideSelThunk :: Bool
@@ -167,7 +167,7 @@ summarize SummarizeOptions{..} = go 0
 
       where
         showSrc :: TraceCont -> Trace
-        showSrc = Trace (goHeap hp, e)
+        showSrc = Trace (goHeap n hp, e)
 
     -- | We already saw one beta reduction; skip any subsequent ones
     goBeta :: Int -> Trace -> TraceCont
@@ -186,13 +186,15 @@ summarize SummarizeOptions{..} = go 0
           _otherwise  -> False
 
     -- | Cleanup the heap
-    goHeap :: Heap Term -> Heap Term
-    goHeap (Heap next heap) =
+    goHeap :: Int -> Heap Term -> Heap Term
+    goHeap n (Heap next heap) =
         Heap next $ Map.filterWithKey shouldShow heap
       where
         shouldShow :: Ptr -> Term -> Bool
         shouldShow (Ptr Nothing (Just name)) _ = and [
-              not summarizeHidePrelude
+              case summarizeHidePrelude of
+                Nothing -> True
+                Just n' -> n < n'
             , not (name `elem` summarizeHideTerms)
             ]
         shouldShow (Ptr _ _) _ = True