Reuse substitution internalisation code

Handle with care

Author: geo2a

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,11 +484,11 @@ 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
-            { constraints = Set.fromList simplifiedConstraints
+            { constraints = simplifiedConstraints
             , term = newTerm
             , ceilConditions
             , substitution = simplifiedSubsitution

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
@@ -364,22 +365,21 @@ applyRule pat@Pattern{ceilConditions} rule =
                         lift . RewriteT . lift . modify $ \s -> s{equations = mempty}
 
                     -- partition ensured constrains into substitution and predicates
-                    let (newSubsitution, newConstraints) = partitionPredicates ensuredConditions
+                    let (newSubsitution, newConstraints) = extractSubsitution 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 =
-                            Set.fromList $
-                                map
-                                    (coerce . substituteInTerm (modifiedPatternSubst `compose` ruleSubstitution) . coerce)
-                                    newConstraints
+                            Set.map
+                                (coerce . substituteInTerm (modifiedPatternSubst `compose` ruleSubstitution) . coerce)
+                                newConstraints
                     let rewritten =
                             Pattern
                                 rewrittenTerm
                                 -- adding new constraints that have not been trivially `Top`, substituting the Ex# variables
-                                (pat.constraints <> substitutedNewConstraints)
+                                (pat.constraints <> substitutedNewConstraints <> leftoverConstraints)
                                 modifiedPatternSubst -- ruleSubstitution is not needed, do not attach it to the result
                                 ceilConditions
                     withContext CtxSuccess $

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

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
@@ -52,7 +54,7 @@ import Data.Coerce (coerce)
 import Data.Foldable ()
 import Data.Generics (extQ)
 import Data.Graph (SCC (..), stronglyConnComp)
-import Data.List (foldl1', nub)
+import Data.List (foldl', foldl1', nub)
 import Data.Map (Map)
 import Data.Map qualified as Map
 import Data.Set (Set)
@@ -508,6 +510,25 @@ mbSubstitution = \case
   where
     boolPred = BoolPred . Internal.Predicate
 
+extractSubsitution ::
+    [Internal.Predicate] -> (Map Internal.Variable Internal.Term, Set Internal.Predicate)
+extractSubsitution ps =
+    let (potentialSubstituion, otherPreds) = partitionSubstitutionPreds . map mbSubstitution . coerce $ ps
+        (newSubstitution, leftoverPreds) = mkSubstitution potentialSubstituion
+     in (newSubstitution, Set.fromList $ leftoverPreds <> map unsafeFromBoolPred otherPreds)
+  where
+    partitionSubstitutionPreds = foldl' accumulate ([], [])
+      where
+        accumulate (accSubst, accOther) = \case
+            pSubst@SubstitutionPred{} -> (pSubst : accSubst, accOther)
+            pOther -> (accSubst, pOther : accOther)
+
+    -- this conversion is safe withing this function since we pass Internal.Predicate as input
+    unsafeFromBoolPred :: InternalisedPredicate -> Internal.Predicate
+    unsafeFromBoolPred = \case
+        BoolPred p -> coerce p
+        _ -> error "extractSubsitution.unsafeFromBoolPred: impossible case"
+
 internalisePred ::
     Flag "alias" ->
     Flag "subsorts" ->