From 711637d2a193f195e51af1529d457eb432fee140 Mon Sep 17 00:00:00 2001
From: Georgy Lukyanov <mail@geo2a.info>
Date: Fri, 25 Oct 2024 14:51:24 +0200
Subject: [PATCH] Introduce RewriteBranchNextState

---
 booster/library/Booster/JsonRpc.hs         |  3 +-
 booster/library/Booster/Pattern/Rewrite.hs | 38 ++++++++++++++++++----
 2 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/booster/library/Booster/JsonRpc.hs b/booster/library/Booster/JsonRpc.hs
index a6f1841aa4..f1dbab6537 100644
--- a/booster/library/Booster/JsonRpc.hs
+++ b/booster/library/Booster/JsonRpc.hs
@@ -49,6 +49,7 @@ import Booster.Pattern.Base qualified as Pattern
 import Booster.Pattern.Implies (runImplies)
 import Booster.Pattern.Pretty
 import Booster.Pattern.Rewrite (
+    RewriteBranchNextState (..),
     RewriteConfig (..),
     RewriteFailed (..),
     RewriteResult (..),
@@ -485,7 +486,7 @@ execResponse req (d, traces, rr) unsupported = case rr of
                     , nextStates =
                         Just
                             $ map
-                                ( \(_, muid, p', mrulePred, ruleSubst) -> toExecState p' unsupported (Just muid) mrulePred (Just ruleSubst)
+                                ( \(RewriteBranchNextState{ruleUniqueId, rewrittenPat, mRulePredicate, ruleSubstitution}) -> toExecState rewrittenPat unsupported (Just ruleUniqueId) mRulePredicate (Just ruleSubstitution)
                                 )
                             $ toList nexts
                     , rule = Nothing
diff --git a/booster/library/Booster/Pattern/Rewrite.hs b/booster/library/Booster/Pattern/Rewrite.hs
index b30da58ed5..97252938a9 100644
--- a/booster/library/Booster/Pattern/Rewrite.hs
+++ b/booster/library/Booster/Pattern/Rewrite.hs
@@ -13,6 +13,7 @@ module Booster.Pattern.Rewrite (
     RewriteConfig (..),
     RewriteFailed (..),
     RewriteResult (..),
+    RewriteBranchNextState (..),
     RewriteTrace (..),
     pattern CollectRewriteTraces,
     pattern NoCollectRewriteTraces,
@@ -247,7 +248,14 @@ rewriteStep cutLabels terminalLabels pat = do
         RewriteBranch base $
             NE.fromList $
                 map
-                    ( \(rule, RewriteRuleAppliedData{rewritten, rulePredicate, ruleSubstitution}) -> (ruleLabelOrLocT rule, uniqueId rule, rewritten, rulePredicate, ruleSubstitution)
+                    ( \(rule, RewriteRuleAppliedData{rewritten, rulePredicate, ruleSubstitution}) ->
+                        RewriteBranchNextState
+                            { ruleLabel = ruleLabelOrLocT rule
+                            , ruleUniqueId = uniqueId rule
+                            , rewrittenPat = rewritten
+                            , mRulePredicate = rulePredicate
+                            , ruleSubstitution
+                            }
                     )
                     leafs
 
@@ -789,10 +797,20 @@ ruleLabelOrLoc rule =
     fromMaybe "unknown rule" $
         fmap pretty rule.attributes.ruleLabel <|> fmap pretty rule.attributes.location
 
+data RewriteBranchNextState pat = RewriteBranchNextState
+    { ruleLabel :: Text
+    , ruleUniqueId :: UniqueId
+    , rewrittenPat :: pat
+    , mRulePredicate :: Maybe Predicate
+    , ruleSubstitution :: Substitution
+    }
+    deriving stock (Eq, Show)
+    deriving (Functor, Foldable, Traversable)
+
 -- | Different rewrite results (returned from RPC execute endpoint)
 data RewriteResult pat
     = -- | branch point
-      RewriteBranch pat (NonEmpty (Text, UniqueId, pat, Maybe Predicate, Substitution))
+      RewriteBranch pat (NonEmpty (RewriteBranchNextState pat))
     | -- | no rules could be applied, config is stuck
       RewriteStuck pat
     | -- | cut point rule, return current (lhs) and single next state
@@ -1021,14 +1039,18 @@ performRewrite rewriteConfig pat = do
                 Nothing -> pure $ RewriteTrivial orig
                 Just p' -> do
                     -- simplify the 3rd component, i.e. the pattern
-                    let simplifyP3rd (a, b, c, e, f) =
-                            fmap (a,b,,e,f) <$> simplifyP c
-                    nexts' <- catMaybes <$> mapM simplifyP3rd (toList nexts)
+                    let simplifyRewritten pattr@RewriteBranchNextState{rewrittenPat} = do
+                            ( fmap @Maybe
+                                    ( \rewrittenSimplified -> (pattr{rewrittenPat = rewrittenSimplified})
+                                    )
+                                )
+                                <$> simplifyP rewrittenPat
+                    nexts' <- catMaybes <$> mapM simplifyRewritten (toList nexts)
                     pure $ case nexts' of
                         -- The `[]` case should be `Stuck` not `Trivial`, because `RewriteTrivial p'`
                         -- means the pattern `p'` is bottom, but we know that is not the case here.
                         [] -> RewriteStuck p'
-                        [(lbl, uId, n, _rp, _rs)] -> RewriteFinished (Just lbl) (Just uId) n
+                        [RewriteBranchNextState{ruleLabel, ruleUniqueId, rewrittenPat}] -> RewriteFinished (Just ruleLabel) (Just ruleUniqueId) rewrittenPat
                         ns -> RewriteBranch p' $ NE.fromList ns
         r@RewriteStuck{} -> pure r
         r@RewriteTrivial{} -> pure r
@@ -1098,7 +1120,9 @@ performRewrite rewriteConfig pat = do
                                     incrementCounter
                                     doSteps False single
                                 RewriteBranch pat'' branches -> withPatternContext pat' $ do
-                                    emitRewriteTrace $ RewriteBranchingStep pat'' $ fmap (\(lbl, uid, _, _, _) -> (lbl, uid)) branches
+                                    emitRewriteTrace $
+                                        RewriteBranchingStep pat'' $
+                                            fmap (\RewriteBranchNextState{ruleLabel, ruleUniqueId} -> (ruleLabel, ruleUniqueId)) branches
                                     pure simplified
                                 _other -> withPatternContext pat' $ error "simplifyResult: Unexpected return value"
                         Right (cutPoint@(RewriteCutPoint lbl _ _ _), _) -> withPatternContext pat' $ do