Merge remote-tracking branch 'origin/master' into release

Author: rv-auditor

booster/library/Booster/Pattern/ApplyEquations.hs

@@ -22,7 +22,8 @@ module Booster.Pattern.ApplyEquations (
     handleSimplificationEquation,
     simplifyConstraint,
     simplifyConstraints,
-    SimplifierCache,
+    SimplifierCache (..),
+    CacheTag (..),
     evaluateConstraints,
 ) where
 
@@ -234,19 +235,46 @@ popRecursion = do
             throw $ InternalError "Trying to pop an empty recursion stack"
         else eqState $ put s{recursionStack = tail s.recursionStack}
 
-toCache :: Monad io => CacheTag -> Term -> Term -> EquationT io ()
-toCache tag orig result = eqState . modify $ \s -> s{cache = updateCache tag s.cache}
-  where
-    insertInto = Map.insert orig result
-    updateCache LLVM cache = cache{llvm = insertInto cache.llvm}
-    updateCache Equations cache = cache{equations = insertInto cache.equations}
-
-fromCache :: Monad io => CacheTag -> Term -> EquationT io (Maybe Term)
-fromCache tag t = eqState $ Map.lookup t <$> gets (select tag . (.cache))
-  where
-    select :: CacheTag -> SimplifierCache -> Map Term Term
-    select LLVM = (.llvm)
-    select Equations = (.equations)
+toCache :: LoggerMIO io => CacheTag -> Term -> Term -> EquationT io ()
+toCache LLVM orig result = eqState . modify $
+    \s -> s{cache = s.cache{llvm = Map.insert orig result s.cache.llvm}}
+toCache Equations orig result = eqState $ do
+    s <- get
+    -- Check before inserting a new result to avoid creating a
+    -- lookup chain e -> result -> olderResult.
+    newEqCache <- case Map.lookup result s.cache.equations of
+        Nothing ->
+            pure $ Map.insert orig result s.cache.equations
+        Just furtherResult -> do
+            when (result /= furtherResult) $ do
+                withContextFor Equations . logMessage $
+                    "toCache shortening a chain "
+                        <> showHashHex (getAttributes orig).hash
+                        <> "->"
+                        <> showHashHex (getAttributes furtherResult).hash
+            pure $ Map.insert orig furtherResult s.cache.equations
+    put s{cache = s.cache{equations = newEqCache}}
+
+fromCache :: LoggerMIO io => CacheTag -> Term -> EquationT io (Maybe Term)
+fromCache tag t = eqState $ do
+    s <- get
+    case tag of
+        LLVM -> pure $ Map.lookup t s.cache.llvm
+        Equations -> do
+            case Map.lookup t s.cache.equations of
+                Nothing -> pure Nothing
+                Just t' -> case Map.lookup t' s.cache.equations of
+                    Nothing -> pure $ Just t'
+                    Just t'' -> do
+                        when (t'' /= t') $ do
+                            withContextFor Equations . logMessage $
+                                "fromCache shortening a chain "
+                                    <> showHashHex (getAttributes t).hash
+                                    <> "->"
+                                    <> showHashHex (getAttributes t'').hash
+                            let newEqCache = Map.insert t t'' s.cache.equations
+                            put s{cache = s.cache{equations = newEqCache}}
+                        pure $ Just t''
 
 logWarn :: LoggerMIO m => Text -> m ()
 logWarn msg =
@@ -899,6 +927,12 @@ applyEquation term rule =
                             Left other ->
                                 liftIO $ Exception.throw other
                         lift $ pushConstraints $ Set.fromList ensuredConditions
+                        -- when a new path condition is added, invalidate the equation cache
+                        unless (null ensuredConditions) $ do
+                            withContextFor Equations . logMessage $
+                                ("New ensured condition from evaluation, invalidating cache" :: Text)
+                            lift . eqState . modify $
+                                \s -> s{cache = s.cache{equations = mempty}}
                         pure $ substituteInTerm subst rule.rhs
   where
     filterOutKnownConstraints :: Set Predicate -> [Predicate] -> EquationT io [Predicate]

booster/library/Booster/Pattern/Rewrite.hs

@@ -50,8 +50,9 @@ import Booster.Definition.Base
 import Booster.LLVM as LLVM (API)
 import Booster.Log
 import Booster.Pattern.ApplyEquations (
+    CacheTag (Equations),
     EquationFailure (..),
-    SimplifierCache,
+    SimplifierCache (..),
     evaluatePattern,
     simplifyConstraint,
  )
@@ -401,6 +402,12 @@ applyRule pat@Pattern{ceilConditions} rule =
                         Left other ->
                             liftIO $ Exception.throw other
 
+                    -- if a new constraint is going to be added, the equation cache is invalid
+                    unless (null newConstraints) $ do
+                        withContextFor Equations . logMessage $
+                            ("New path condition ensured, invalidating cache" :: Text)
+                        lift . RewriteT . lift . modify $ \s -> s{equations = mempty}
+
                     -- existential variables may be present in rule.rhs and rule.ensures,
                     -- need to strip prefixes and freshen their names with respect to variables already
                     -- present in the input pattern and in the unification substitution