From 4e34924f7ca9c372c753299cc527c0f21caf4aab Mon Sep 17 00:00:00 2001
From: Georgy Lukyanov <mail@geo2a.info>
Date: Fri, 18 Oct 2024 15:15:21 +0200
Subject: [PATCH] Reuse substitution internalisation code

Handle with care
---
 booster/library/Booster/Pattern/ApplyEquations.hs  | 3 ++-
 booster/library/Booster/Pattern/Rewrite.hs         | 5 +++--
 booster/library/Booster/Syntax/Json/Externalise.hs | 3 ++-
 booster/library/Booster/Syntax/Json/Internalise.hs | 6 ++++++
 4 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/booster/library/Booster/Pattern/ApplyEquations.hs b/booster/library/Booster/Pattern/ApplyEquations.hs
index 166b18aaea..bcf778558f 100644
--- a/booster/library/Booster/Pattern/ApplyEquations.hs
+++ b/booster/library/Booster/Pattern/ApplyEquations.hs
@@ -71,6 +71,7 @@ import Booster.Pattern.Util
 import Booster.Prettyprinter (renderOneLineText)
 import Booster.SMT.Interface qualified as SMT
 import Booster.Syntax.Json.Externalise (externaliseTerm)
+import Booster.Syntax.Json.Internalise (extractSubsitution)
 import Booster.Util (Bound (..))
 import Kore.JsonRpc.Types.ContextLog (CLContext (CLWithId), IdContext (CtxCached))
 import Kore.Util (showHashHex)
@@ -483,7 +484,7 @@ evaluatePattern' pat@Pattern{term, ceilConditions} = withPatternContext pat $ do
     -- evaluatedConstraints. To avoid duplicating constraints (i.e. having equivalent entities
     -- in pat.predicate and pat.substitution), we discard the old substitution here
     -- and extract a possible simplified one from evaluatedConstraints.
-    let (simplifiedSubsitution, simplifiedConstraints) = partitionPredicates (Set.toList evaluatedConstraints)
+    let (simplifiedSubsitution, simplifiedConstraints) = extractSubsitution (Set.toList evaluatedConstraints)
 
     pure
         Pattern
diff --git a/booster/library/Booster/Pattern/Rewrite.hs b/booster/library/Booster/Pattern/Rewrite.hs
index bd946290de..da49ff2a0f 100644
--- a/booster/library/Booster/Pattern/Rewrite.hs
+++ b/booster/library/Booster/Pattern/Rewrite.hs
@@ -71,6 +71,7 @@ import Booster.Pattern.Util
 import Booster.Prettyprinter
 import Booster.SMT.Interface qualified as SMT
 import Booster.Syntax.Json.Externalise (externaliseTerm)
+import Booster.Syntax.Json.Internalise (extractSubsitution)
 import Booster.Util (Flag (..))
 
 newtype RewriteT io a = RewriteT
@@ -367,7 +368,7 @@ applyRule pat@Pattern{ceilConditions} rule =
                     let (newSubsitution, newConstraints) = partitionPredicates ensuredConditions
 
                     -- compose the existing substitution pattern and the newly acquired one
-                    let modifiedPatternSubst = newSubsitution `compose` pat.substitution
+                    let (modifiedPatternSubst, leftoverConstraints) = extractSubsitution . asEquations $ newSubsitution `compose` pat.substitution
 
                     let rewrittenTerm = substituteInTerm (modifiedPatternSubst `compose` ruleSubstitution) rule.rhs
                         substitutedNewConstraints =
@@ -379,7 +380,7 @@ applyRule pat@Pattern{ceilConditions} rule =
                             Pattern
                                 rewrittenTerm
                                 -- adding new constraints that have not been trivially `Top`, substituting the Ex# variables
-                                (pat.constraints <> substitutedNewConstraints)
+                                (pat.constraints <> substitutedNewConstraints <> Set.fromList leftoverConstraints)
                                 modifiedPatternSubst -- ruleSubstitution is not needed, do not attach it to the result
                                 ceilConditions
                     withContext CtxSuccess $
diff --git a/booster/library/Booster/Syntax/Json/Externalise.hs b/booster/library/Booster/Syntax/Json/Externalise.hs
index f130a801d7..7c75e7aaec 100644
--- a/booster/library/Booster/Syntax/Json/Externalise.hs
+++ b/booster/library/Booster/Syntax/Json/Externalise.hs
@@ -19,6 +19,7 @@ import Data.Text.Encoding qualified as Text
 import Booster.Pattern.Base (externaliseKmapUnsafe)
 import Booster.Pattern.Base qualified as Internal
 import Booster.Pattern.Bool qualified as Internal
+import Booster.Pattern.Substitution qualified as Substitution
 import Booster.Pattern.Util (sortOfTerm)
 import Data.Map (Map)
 import Data.Map qualified as Map
@@ -35,7 +36,7 @@ externalisePattern ::
 externalisePattern Internal.Pattern{term = term, constraints, ceilConditions, substitution = ensuredSubstitution} inputSubstitution =
     -- need a sort for the predicates in external format
     let sort = externaliseSort $ sortOfTerm term
-        substitutions = ensuredSubstitution <> inputSubstitution -- TODO ensuredSubstitution takes priority. Do we even need inputSubstitution?
+        substitutions = inputSubstitution <> (ensuredSubstitution `Substitution.compose` inputSubstitution)
         externalisedSubstitution =
             if null substitutions
                 then Nothing
diff --git a/booster/library/Booster/Syntax/Json/Internalise.hs b/booster/library/Booster/Syntax/Json/Internalise.hs
index c90550ee30..7b0b96d604 100644
--- a/booster/library/Booster/Syntax/Json/Internalise.hs
+++ b/booster/library/Booster/Syntax/Json/Internalise.hs
@@ -35,6 +35,8 @@ module Booster.Syntax.Json.Internalise (
     pattern CheckSubsorts,
     pattern IgnoreSubsorts,
     logPatternError,
+    -- substitution mining
+    extractSubsitution,
     -- for test only
     InternalisedPredicate (..),
 ) where
@@ -508,6 +510,10 @@ mbSubstitution = \case
   where
     boolPred = BoolPred . Internal.Predicate
 
+extractSubsitution ::
+    [Internal.Predicate] -> (Map Internal.Variable Internal.Term, [Internal.Predicate])
+extractSubsitution = mkSubstitution . map mbSubstitution . coerce
+
 internalisePred ::
     Flag "alias" ->
     Flag "subsorts" ->