HOTFIX avoid caching extra predicates, only simplify term during rewrite (#4103)

The change in #3953 introduced additional arguments to evaluation
functions which allow callers to supply some "known true" predicates for
the simplification and evaluation. However, doing so means that the
cache will get populated with associations that might only be true if
this known truth does not change.

The case in point was a predicate being cached as "true" because of an
earlier evaluation, and then _removed_ from the path condition (see
added integration test).

This change 
* refactors `checkConstraint` functions to fix the pattern constraints
* skips the cache update if any additional "known truth" arguments were
supplied to `checkConstraint`, and
* runs internal computations that manipulate the predicates with a fresh
empty cache.

Furthermore, to avoid a significant performance penalty, we
* simplify the pattern constraints once before starting to rewrite
* only simplify the pattern _term_ during rewriting
   - the pattern constraints are already evaluated and stay the same
- if a new constraint is added (by way of an `ensures`), this constraint
is evaluated before adding it
* At the end of the rewrite, we still simplify everything again
(including the entire path condition).

We are not using the SMT solver to detect a vacuous state at the start,
though, could add that separately if desired. At the moment, the
behaviour is largely the same as before (integration test expectations
identical except for some evaluation in vacuous states).

---------

Co-authored-by: github-actions <[email protected]>

Author: jberthold

booster/library/Booster/Pattern/ApplyEquations.hs

@@ -333,6 +333,7 @@ runEquationT definition llvmApi smtSolver sCache known (EquationT m) = do
                         , logger
                         , prettyModifiers
                         }
+    -- NB the returned cache assumes the known predicates
     pure (res, endState.cache)
 
 iterateEquations ::
@@ -423,19 +424,24 @@ llvmSimplify term = do
 ----------------------------------------
 -- Interface functions
 
--- | Evaluate and simplify a term.
+{- | Evaluate and simplify a term.
+
+  The returned cache should only be reused with the same known predicates.
+-}
 evaluateTerm ::
     LoggerMIO io =>
     Direction ->
     KoreDefinition ->
     Maybe LLVM.API ->
     SMT.SMTContext ->
+    SimplifierCache ->
     Set Predicate ->
     Term ->
     io (Either EquationFailure Term, SimplifierCache)
-evaluateTerm direction def llvmApi smtSolver knownPredicates =
-    runEquationT def llvmApi smtSolver mempty knownPredicates
-        . evaluateTerm' direction
+evaluateTerm direction def llvmApi smtSolver cache knownPredicates term =
+    runEquationT def llvmApi smtSolver cache knownPredicates $
+        withTermContext term $
+            evaluateTerm' direction term
 
 -- version for internal nested evaluation
 evaluateTerm' ::
@@ -447,6 +453,9 @@ evaluateTerm' direction = iterateEquations direction PreferFunctions
 
 {- | Simplify a Pattern, processing its constraints independently.
      Returns either the first failure or the new pattern if no failure was encountered
+
+   The returned cache may only be reused if pat.constraints are known
+   to remain true in the next usage context.
 -}
 evaluatePattern ::
     LoggerMIO io =>
@@ -510,35 +519,28 @@ evaluatePattern' pat@Pattern{term, ceilConditions} = withPatternContext pat $ do
         err -> throw err
 
 -- evaluate the given predicate assuming all others
+-- This manipulates the known predicates so it should run without cache
 simplifyAssumedPredicate :: LoggerMIO io => Predicate -> EquationT io ()
 simplifyAssumedPredicate p = do
-    allPs <- predicates <$> getState
-    let otherPs = Set.delete p allPs
-    eqState $ modify $ \s -> s{predicates = otherPs}
+    prior <- getState
+    let otherPs = Set.delete p (prior.predicates)
+    eqState $ modify $ \s -> s{predicates = otherPs, cache = mempty}
     newP <- simplifyConstraint' True $ coerce p
-    pushConstraints $ Set.singleton $ coerce newP
+    eqState $ modify $ \s -> s{cache = prior.cache, predicates = otherPs <> Set.singleton (coerce newP)}
 
 evaluateConstraints ::
     LoggerMIO io =>
     KoreDefinition ->
     Maybe LLVM.API ->
     SMT.SMTContext ->
-    SimplifierCache ->
     Set Predicate ->
-    io (Either EquationFailure (Set Predicate), SimplifierCache)
-evaluateConstraints def mLlvmLibrary smtSolver cache =
-    runEquationT def mLlvmLibrary smtSolver cache mempty . evaluateConstraints'
-
-evaluateConstraints' ::
-    LoggerMIO io =>
-    Set Predicate ->
-    EquationT io (Set Predicate)
-evaluateConstraints' constraints = do
-    pushConstraints constraints
-    -- evaluate all existing constraints, once
-    traverse_ simplifyAssumedPredicate . predicates =<< getState
-    -- this may yield additional new constraints, left unevaluated
-    predicates <$> getState
+    io (Either EquationFailure (Set Predicate))
+evaluateConstraints def mLlvmLibrary smtSolver constraints =
+    fmap fst $ runEquationT def mLlvmLibrary smtSolver mempty constraints $ do
+        -- evaluate all existing constraints, once
+        traverse_ simplifyAssumedPredicate . predicates =<< getState
+        -- this may yield additional new constraints, left unevaluated
+        predicates <$> getState
 
 ----------------------------------------
 
@@ -1075,9 +1077,8 @@ applyEquation term rule =
     (to decide whether or not a rule can apply, not to retain the
     ensured conditions).
 
-    If and as soon as this function is used inside equation
-    application, it needs to run within the same 'EquationT' context
-    so we can detect simplification loops and avoid monad nesting.
+    Consistency of the returned SimplifierCache must be tracked by the
+    caller, the cache incorporates the given knownPredicates.
 -}
 simplifyConstraint ::
     LoggerMIO io =>
@@ -1088,8 +1089,11 @@ simplifyConstraint ::
     Set Predicate ->
     Predicate ->
     io (Either EquationFailure Predicate, SimplifierCache)
-simplifyConstraint def mbApi smt cache knownPredicates (Predicate p) = do
-    runEquationT def mbApi smt cache knownPredicates $ (coerce <$>) . simplifyConstraint' True $ p
+simplifyConstraint def mbApi smt cache knownPredicates =
+    runEquationT def mbApi smt cache knownPredicates
+        . (coerce <$>)
+        . simplifyConstraint' True
+        . coerce
 
 simplifyConstraints ::
     LoggerMIO io =>

booster/library/Booster/Pattern/Implies.hs

@@ -150,8 +150,8 @@ runImplies def mLlvmLibrary mSMTOptions antecedent consequent =
                                         subst
                                 else -- FIXME This is incomplete because patL.constraints are not assumed in the check.
 
-                                    ApplyEquations.evaluateConstraints def mLlvmLibrary solver mempty filteredConsequentPreds >>= \case
-                                        (Right newPreds, _) ->
+                                    ApplyEquations.evaluateConstraints def mLlvmLibrary solver filteredConsequentPreds >>= \case
+                                        Right newPreds ->
                                             if all (== Predicate TrueBool) newPreds
                                                 then
                                                     implies
@@ -161,7 +161,7 @@ runImplies def mLlvmLibrary mSMTOptions antecedent consequent =
                                                         subst
                                                 else -- here we conservatively abort (incomplete)
                                                     pure . Left . RpcError.backendError $ RpcError.Aborted "unknown constraints"
-                                        (Left other, _) ->
+                                        Left other ->
                                             pure . Left . RpcError.backendError $ RpcError.Aborted (Text.pack . constructorName $ other)
 
             case (internalised antecedent.term, internalised consequent.term) of

booster/library/Booster/Pattern/Rewrite.hs

@@ -38,7 +38,7 @@ import Data.List.NonEmpty (NonEmpty (..), toList)
 import Data.List.NonEmpty qualified as NE
 import Data.Map qualified as Map
 import Data.Maybe (catMaybes, fromMaybe)
-import Data.Sequence (Seq, (|>))
+import Data.Sequence as Seq (Seq, fromList, (|>))
 import Data.Set qualified as Set
 import Data.Text as Text (Text, pack)
 import Numeric.Natural
@@ -55,9 +55,12 @@ import Booster.LLVM as LLVM (API)
 import Booster.Log
 import Booster.Pattern.ApplyEquations (
     CacheTag (Equations),
+    Direction (..),
     EquationFailure (..),
     SimplifierCache (..),
+    evaluateConstraints,
     evaluatePattern,
+    evaluateTerm,
     simplifyConstraint,
  )
 import Booster.Pattern.Base
@@ -516,9 +519,16 @@ applyRule pat@Pattern{ceilConditions} rule =
                                                 , rulePredicate = Just rulePredicate
                                                 }
   where
-    filterOutKnownConstraints :: Set.Set Predicate -> [Predicate] -> RewriteT io [Predicate]
-    filterOutKnownConstraints priorKnowledge constraitns = do
-        let (knownTrue, toCheck) = partition (`Set.member` priorKnowledge) constraitns
+    -- These predicates are known (and do not change) during the
+    -- entire rewrite step. The simplifier cache cannot be retained
+    -- when additional predicates are used (see 'checkConstraint').
+    knownPatternPredicates =
+        pat.constraints <> (Set.fromList . asEquations $ pat.substitution)
+
+    filterOutKnownConstraints :: [Predicate] -> RewriteT io [Predicate]
+    filterOutKnownConstraints constraints = do
+        let (knownTrue, toCheck) =
+                partition (`Set.member` knownPatternPredicates) constraints
         unless (null knownTrue) $
             getPrettyModifiers >>= \case
                 ModifiersRep (_ :: FromModifiersT mods => Proxy mods) ->
@@ -537,14 +547,16 @@ applyRule pat@Pattern{ceilConditions} rule =
         Set.Set Predicate ->
         Predicate ->
         RewriteRuleAppT (RewriteT io) (Maybe a)
-    checkConstraint onUnclear onBottom knownPredicates p = do
+    checkConstraint onUnclear onBottom extraPredicates p = do
         RewriteConfig{definition, llvmApi, smtSolver} <- lift $ RewriteT ask
-        RewriteState{cache = oldCache} <- lift . RewriteT . lift $ get
-        (simplified, cache) <-
+        RewriteState{cache} <- lift . RewriteT . lift $ get
+        let knownPredicates = knownPatternPredicates <> extraPredicates
+        (simplified, newCache) <-
             withContext CtxConstraint $
-                simplifyConstraint definition llvmApi smtSolver oldCache knownPredicates p
-        -- update cache
-        lift $ updateRewriterCache cache
+                simplifyConstraint definition llvmApi smtSolver cache knownPredicates p
+        -- Important: only retain new cache if no extraPredicates were supplied!
+        when (Set.null extraPredicates) $
+            lift (updateRewriterCache newCache)
         case simplified of
             Right (Predicate FalseBool) -> onBottom
             Right (Predicate TrueBool) -> pure Nothing
@@ -559,14 +571,9 @@ applyRule pat@Pattern{ceilConditions} rule =
         -- apply substitution to rule requires
         let ruleRequires =
                 concatMap (splitBoolPredicates . substituteInPredicate matchingSubst) rule.requires
-            knownConstraints = pat.constraints <> (Set.fromList . asEquations $ pat.substitution)
 
         -- filter out any predicates known to be _syntactically_ present in the known prior
-        toCheck <-
-            lift $
-                filterOutKnownConstraints
-                    knownConstraints
-                    ruleRequires
+        toCheck <- lift $ filterOutKnownConstraints ruleRequires
 
         -- simplify the constraints (one by one in isolation). Stop if false, abort rewrite if indeterminate.
         unclearRequires <-
@@ -575,17 +582,13 @@ applyRule pat@Pattern{ceilConditions} rule =
                     ( checkConstraint
                         id
                         returnNotApplied
-                        knownConstraints
+                        mempty -- checkConstraint already considers knownConstraints
                     )
                     toCheck
 
         -- unclear conditions may have been simplified and
         -- could now be syntactically present in the path constraints, filter again
-        stillUnclear <-
-            lift $
-                filterOutKnownConstraints
-                    knownConstraints
-                    unclearRequires
+        stillUnclear <- lift $ filterOutKnownConstraints unclearRequires
 
         -- check unclear requires-clauses in the context of known constraints (priorKnowledge)
         solver <- lift $ RewriteT $ (.smtSolver) <$> ask
@@ -614,17 +617,14 @@ applyRule pat@Pattern{ceilConditions} rule =
         -- apply substitution to rule ensures
         let ruleEnsures =
                 concatMap (splitBoolPredicates . coerce . substituteInTerm matchingSubst . coerce) rule.ensures
-            knownConstraints =
-                pat.constraints
-                    <> (Set.fromList . asEquations $ pat.substitution)
-                    <> Set.fromList unclearRequiresAfterSmt
         newConstraints <-
             catMaybes
                 <$> mapM
                     ( checkConstraint
                         id
                         returnTrivial
-                        knownConstraints
+                        -- supply required path conditions as extra constraints
+                        (Set.fromList unclearRequiresAfterSmt)
                     )
                     ruleEnsures
 
@@ -672,7 +672,7 @@ applyRule pat@Pattern{ceilConditions} rule =
         let ruleRequires =
                 concatMap (splitBoolPredicates . coerce . substituteInTerm matchingSubst . coerce) rule.requires
         collapseAndBools . catMaybes
-            <$> mapM (checkConstraint id returnNotApplied pat.constraints) ruleRequires
+            <$> mapM (checkConstraint id returnNotApplied mempty) ruleRequires
 
 ruleGroupPriority :: [RewriteRule a] -> Maybe Priority
 ruleGroupPriority = \case
@@ -1001,9 +1001,16 @@ performRewrite ::
     Pattern ->
     io (Natural, Seq (RewriteTrace ()), RewriteResult Pattern)
 performRewrite rewriteConfig pat = do
-    (rr, RewriteStepsState{counter, traces}) <-
-        flip runStateT rewriteStart $ doSteps False pat
-    pure (counter, traces, rr)
+    simplifiedConstraints <-
+        withContext CtxSimplify $ evaluateConstraints definition llvmApi smtSolver pat.constraints
+    case simplifiedConstraints of
+        Right constraints ->
+            (flip runStateT rewriteStart $ doSteps False pat{constraints})
+                >>= \(rr, RewriteStepsState{counter, traces}) -> pure (counter, traces, rr)
+        Left r@(SideConditionFalse{}) ->
+            pure (0, fromList [RewriteSimplified (Just r)], error "Just return #Bottom here")
+        Left err ->
+            error (show err)
   where
     RewriteConfig
         { definition
@@ -1034,6 +1041,27 @@ performRewrite rewriteConfig pat = do
 
     updateCache simplifierCache = modify $ \rss -> (rss :: RewriteStepsState){simplifierCache}
 
+    -- only simplifies the _term_ of the pattern
+    simplifyT :: Pattern -> StateT RewriteStepsState io (Maybe Pattern)
+    simplifyT p = withContext CtxSimplify $ do
+        cache <- simplifierCache <$> get
+        evaluateTerm BottomUp definition llvmApi smtSolver cache p.constraints p.term >>= \(res, newCache) -> do
+            updateCache newCache
+            case res of
+                Right newTerm -> do
+                    emitRewriteTrace $ RewriteSimplified Nothing
+                    pure $ Just p{term = newTerm}
+                Left r@SideConditionFalse{} -> do
+                    emitRewriteTrace $ RewriteSimplified (Just r)
+                    pure Nothing
+                Left r@UndefinedTerm{} -> do
+                    emitRewriteTrace $ RewriteSimplified (Just r)
+                    pure Nothing
+                Left other -> do
+                    emitRewriteTrace $ RewriteSimplified (Just other)
+                    pure $ Just p
+
+    -- simplifies term and constraints of the pattern
     simplifyP :: Pattern -> StateT RewriteStepsState io (Maybe Pattern)
     simplifyP p = withContext CtxSimplify $ do
         st <- get
@@ -1228,7 +1256,7 @@ performRewrite rewriteConfig pat = do
                                 else withSimplified pat' msg (pure . RewriteAborted failure)
       where
         withSimplified p msg cont = do
-            (withPatternContext p $ simplifyP p) >>= \case
+            (withPatternContext p $ simplifyT p) >>= \case
                 Nothing -> do
                     logMessage ("Rewrite stuck after simplification." :: Text)
                     pure $ RewriteStuck p

booster/test/rpc-integration/resources/3934-smt.kompile

@@ -0,0 +1 @@
+kompile-from-double-definition.sh