From f14d114d69e2228b54ba1236b16d696d3a806c6b Mon Sep 17 00:00:00 2001
From: Jonas Hoefer <hoferj@chalmers.se>
Date: Mon, 19 Aug 2024 11:23:21 +0200
Subject: [PATCH] Add option to display head unfolding information
 sequentially/summed

---
 src/Main.hs                      | 30 +++++++++++++++++++-----------
 src/PosTT/HeadLinearReduction.hs | 18 ++++++++----------
 2 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/src/Main.hs b/src/Main.hs
index 942fcf8..cb323bb 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -78,15 +78,17 @@ evalModule m p = do
         putStrLn $ "Head linear unfolding of " ++ pretty s
         putStrLn $ "Yields " ++ pretty s'
         putStrLn ""
-        putStrLn $ "Unfold counts: " ++ intercalate ", " [ show d ++ ": " ++ show c | (d, c) <- M.toList u]
-        putStrLn $ "Hence " ++ show (sum u) ++ " unfold steps"
+        print u
+        putStrLn $ "Hence " ++ show (length u) ++ " unfold steps"
     _ -> putStrLn "No definitions"
 
 
 --------------------------------------------------------------------------------
 ---- Repl
 
-data ReplCmd = Term String | Load Bool FilePath | Reload Bool | Quit | Unfold Int String | Infer String deriving (Show)
+data UnfoldStepPresentation = Sequential | Summed deriving (Show, Read)
+
+data ReplCmd = Term String | Load Bool FilePath | Reload Bool | Quit | Unfold Int UnfoldStepPresentation String | Infer String deriving Show
 
 inferCmd :: Mod CommandFields ReplCmd
 inferCmd = command ":infer"
@@ -113,6 +115,8 @@ quitCmd = command ":quit" (info (pure Quit) (progDesc "Quit repl"))
 unfoldCmd :: Mod CommandFields ReplCmd
 unfoldCmd = command ":unfold" $ info
   (Unfold <$> option auto (short 's' <> metavar "STEPS" <> showDefault <> value 1 <> help "Number of unfold steps")
+          <*> option auto (short 'd' <> metavar "DISPLAY-STYLE" <> showDefault <> value Sequential
+                             <> help "Presentation of unfolded steps. Either `Summed` or `Sequential`")
           <*> (unwords <$> many (argument str (metavar "DEF"))))
   (progDesc "Perform a single unfolding steps on given definition or, if not argument is given, on the last considered term.")
 
@@ -223,28 +227,32 @@ repl = do
                 Left err          -> outputStrLn $ show err
                 Right (_, val, _) -> outputStrLn $ bindStage terminalStage prettyVal val
           repl
-        Right (Unfold k "") -> do
+        Right (Unfold k s "") -> do
           lastUnfoldTarget >>= \case
             Nothing       ->
               outputStrLn "No current unfold target. Either unfold a definition using `:unfold DEF` or load a term using `:infer TERM`"
             Just (t, rep) -> do
-              unfoldTerm k t >>= rep
+              unfoldTerm k s t >>= rep
           repl
-        Right (Unfold k d) -> do
+        Right (Unfold k s d) -> do
           defUnfoldTarget d >>= \case
             Just (t, rep) -> 
-              unfoldTerm k t >>= rep
+              unfoldTerm k s t >>= rep
             Nothing       ->
               outputStrLn $ d ++ " is not a definition!"
           repl
 
-unfoldTerm :: Int -> Tm -> Repl Tm
-unfoldTerm k t = do
+unfoldTerm :: Int -> UnfoldStepPresentation -> Tm -> Repl Tm
+unfoldTerm k s t = do
   ρ <- gets environment
   let (u, t') = headUnfold ρ t (Just k)
-  outputStrLn $ pretty t'
+  outputStrLn (pretty t')
   outputStrLn ""
-  outputStrLn $ "Unfold counts: " ++ intercalate ", " [ show x ++ ": " ++ show c | (x, c) <- M.toList u]
+  case s of
+    Sequential -> outputStrLn (show u)
+    Summed     -> do
+      let u' = M.fromListWith (+) (map (,1::Int) u)
+      outputStrLn $ "Unfold counts: " ++ intercalate ", " [ show x ++ ": " ++ show c | (x, c) <- M.toList u']
   return t'
   
 -- | Returns the term associated to a definition, and a way to replace it, if this definition exists.
diff --git a/src/PosTT/HeadLinearReduction.hs b/src/PosTT/HeadLinearReduction.hs
index 935abb1..7af3f74 100644
--- a/src/PosTT/HeadLinearReduction.hs
+++ b/src/PosTT/HeadLinearReduction.hs
@@ -1,8 +1,6 @@
 module PosTT.HeadLinearReduction where
 
-import           Control.Applicative
-
-import qualified Data.Map as M
+import           Data.Tuple.Extra (first)
 
 import           PosTT.Common
 import           PosTT.Eval
@@ -12,8 +10,8 @@ import           PosTT.Terms
 import           PosTT.Values
 
 
-headUnfold :: Env -> Tm -> Maybe Int -> (M.Map Name Int, Tm)
-headUnfold δ = go M.empty
+headUnfold :: Env -> Tm -> Maybe Int -> ([Name], Tm)
+headUnfold δ = go
   where
     ρ = blockedEnv δ
     s = foldr sExtName terminalStage (blockedNames δ)
@@ -30,11 +28,11 @@ headUnfold δ = go M.empty
     unfold t = bindStage s $ (fmap . fmap) readBack
       $ singleReduction blockedLookup $ eval ρ t
 
-    go :: M.Map Name Int -> Tm -> Maybe Int -> (M.Map Name Int, Tm)
-    go u t (Just 0) = (u, t)
-    go u t steps    = case unfold t of
-      Nothing      -> (u, t)
-      Just (d, t') -> go (M.alter (\e -> fmap succ e <|> Just 1) d u) t' (pred <$> steps)
+    go :: Tm -> Maybe Int -> ([Name], Tm)
+    go t (Just 0) = ([], t)
+    go t steps    = case unfold t of
+      Nothing      -> ([], t)
+      Just (d, t') -> first (d:) $ go t' (pred <$> steps)
 
 
 --------------------------------------------------------------------------------