From c9b7f76db88df1fb94c52b0228cfd416460c3899 Mon Sep 17 00:00:00 2001
From: Kevin Phoenix <kevin@kphoenix.us>
Date: Fri, 4 Oct 2024 17:26:14 -0700
Subject: [PATCH] Allow simplifications to handle annotations (#548)

---
 claripy/ast/base.py        |  7 +++++--
 claripy/operations.py      | 10 +++++-----
 claripy/simplifications.py | 34 ++++++++++++++++++++--------------
 3 files changed, 30 insertions(+), 21 deletions(-)

diff --git a/claripy/ast/base.py b/claripy/ast/base.py
index f48f530a2..fe862b854 100644
--- a/claripy/ast/base.py
+++ b/claripy/ast/base.py
@@ -246,7 +246,7 @@ def make_like(
         length: int | None = None,
     ) -> Self:
         # Try to simplify the expression again
-        simplified = claripy.simplifications.simplify(op, args) if simplify else None
+        simplified, annotated = claripy.simplifications.simplify(op, args) if simplify else (None, False)
         if simplified is not None:
             op = simplified.op
         if (  # fast path
@@ -294,7 +294,10 @@ def make_like(
         # special case: if self is one of the args, we do not copy annotations over from self since child
         # annotations will be re-processed during AST creation.
         if annotations is None:
-            annotations = self.annotations if not args or not any(self is arg for arg in args) else ()
+            if not annotated:
+                annotations = self.annotations if not args or not any(self is arg for arg in args) else ()
+            else:
+                annotations = simplified.annotations
         if variables is None and op in all_operations:
             variables = self.variables
         if uninitialized is None:
diff --git a/claripy/operations.py b/claripy/operations.py
index ed677f0e8..0ecc19383 100644
--- a/claripy/operations.py
+++ b/claripy/operations.py
@@ -54,9 +54,12 @@ def _op(*args):
                     raise ClaripyOperationError(msg)
 
         # pylint:disable=too-many-nested-blocks
-        simp = _handle_annotations(claripy.simplifications.simplify(name, fixed_args), args)
+        simp, annotated = claripy.simplifications.simplify(name, fixed_args)
         if simp is not None:
-            return simp
+            if not annotated:
+                simp = _handle_annotations(simp, fixed_args)
+            if simp is not None:
+                return simp
 
         kwargs = {}
         if calc_length is not None:
@@ -73,9 +76,6 @@ def _op(*args):
 
 
 def _handle_annotations(simp, args):
-    if simp is None:
-        return None
-
     # pylint:disable=isinstance-second-argument-not-valid-type
     ast_args = tuple(a for a in args if isinstance(a, claripy.ast.Base))
     preserved_relocatable = frozenset(simp._relocatable_annotations)
diff --git a/claripy/simplifications.py b/claripy/simplifications.py
index 4b2a6f948..b0125bdb0 100644
--- a/claripy/simplifications.py
+++ b/claripy/simplifications.py
@@ -5,10 +5,14 @@
 import itertools
 import operator
 from functools import reduce
+from typing import TYPE_CHECKING
 
 import claripy
 from claripy.fp import FSORT_DOUBLE, FSORT_FLOAT
 
+if TYPE_CHECKING:
+    from claripy.ast.base import Base
+
 SIMPLE_OPS = ("Concat", "SignExt", "ZeroExt")
 
 extract_distributable = {
@@ -829,7 +833,7 @@ def extract_simplifier(high, low, val):
 
             # avoid extracting if we need the full value
             if new_high == result.length - 1 and low_loc == 0:
-                return result
+                return result, True
 
             # else extract the part we need
             return result[new_high:low_loc]
@@ -1110,13 +1114,10 @@ def zeroext_comparing_against_simplifier(op, a, b):
     return None
 
 
-_simple_bool_simplifiers = {
+_all_simplifiers = {
     "And": boolean_and_simplifier,
     "Or": boolean_or_simplifier,
     "Not": boolean_not_simplifier,
-}
-
-_all_simplifiers = _simple_bool_simplifiers | {
     "Reverse": bv_reverse_simplifier,
     "Extract": extract_simplifier,
     "Concat": concat_simplifier,
@@ -1141,16 +1142,21 @@ def zeroext_comparing_against_simplifier(op, a, b):
     "__invert__": invert_simplifier,
 }
 
-# the actual functions
+# the actual function
 
 
-def simplify_bool_condition(op, args):
-    if op not in _simple_bool_simplifiers:
-        return None
-    return _simple_bool_simplifiers[op](*args)
-
+def simplify(op, args) -> tuple[Base | None, bool]:
+    """Simplifies the given operation with the given arguments. Returns the
+    simplified result if possible, along with a boolean representing whether
+    annotations were handled by the simplifier.
+    """
 
-def simplify(op, args):
     if op not in _all_simplifiers:
-        return None
-    return _all_simplifiers[op](*args)
+        return None, False
+
+    simplified = _all_simplifiers[op](*args)
+
+    if isinstance(simplified, tuple):
+        return simplified[0], simplified[1]
+
+    return simplified, False