EXPERIMENTS: shorten cache lookup chains, invalidate cache on new ensures (#3998)

### Fixes #3982

To avoid situations where the cache contains successive mappings `e ->
e1 -> e2 -> e3 ... -> e_n` (a "lookup chain"), the following
modifications are made:

1. When inserting `e -> e1`, first check whether `e1 -> e2` is in the
cache, and insert `e -> e2` instead in this case.
2. When looking up `e` and a binding `e -> e1` is found, check whether a
binding `e1 -> e2` is also in the cache.
In this case, update the binding for `e` to `e -> e2` and return `e2`.

This will eventually lead to removing all lookup chains that may exist
in a cache (for "interesting" cached values that are being looked up),
without the risk of a loop.

Assume the following:
* the cache was initially empty;
* insertions and lookups are performed as described above;
* the cache contains a lookup chain `e -> e1 -> e2`.

If `e1 -> e2` was inserted first, insertion of `e -> e1` would have
resulted in storing `e -> e2`. Therefore, `e -> e1` must have been
inserted first.

A subsequent lookup of `e` will update the cache to contain `e -> e2`.

Insertion of `e -> e1`, then `e1 -> e2`, then `e2 -> e3`, will create
the lookup chain `e -> e1 -> e2 -> e3`.

- A lookup of `e` would leave the cache containing two shorter chains `e
-> e2 -> e3` and `e1 -> e2 -> e3`.
- A lookup of `e1` would also leave the cache containing one shorter
chain `e -> e1 -> e3`, and `e2 -> e3`.

Every lookup potentially shortens a lookup chain, possibly to several
chains.

### Fixes #3993 

When a new path condition is added (from an `ensures` clause), the
`Equations` cache is purged. While some entries may still be valid,
* all entries `t -> t` (no simplification possible) may be invalid
because new equations might be applicable with the new path condition
* even entries `t -> t'` with `t /= t'` may be invalid because while
they already memoise some applicable simplifications, they may store
sub-terms that could now be evaluated or simplified further (new
equations may apply)

---------

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

Author: jberthold

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