Move SI information from BVS args to annotation (#520)

* Move SI information from BVS args to annotation

* Add docstrings

* Refactor BackendVSA convert and abstract

Author: twizmwazin

claripy/annotation.py

@@ -1,8 +1,12 @@
 from __future__ import annotations
 
-import claripy
+from typing import TYPE_CHECKING
+
 from claripy.errors import ClaripyOperationError
 
+if TYPE_CHECKING:
+    import claripy
+
 
 class Annotation:
     """
@@ -64,6 +68,8 @@ def relocate(self, src: claripy.ast.Base, dst: claripy.ast.Base):  # pylint:disa
 
 
 class SimplificationAvoidanceAnnotation(Annotation):
+    """SimplificationAvoidanceAnnotation is an annotation that prevents simplification of an AST."""
+
     @property
     def eliminatable(self):
         return False
@@ -73,6 +79,33 @@ def relocatable(self):
         return False
 
 
+class StridedIntervalAnnotation(SimplificationAvoidanceAnnotation):
+    """StridedIntervalAnnotation allows annotating a BVS to represent a strided interval."""
+
+    stride: int
+    lower_bound: int
+    upper_bound: int
+
+    def __init__(self, stride: int, lower_bound: int, upper_bound: int):
+        self.stride = stride
+        self.lower_bound = lower_bound
+        self.upper_bound = upper_bound
+
+    def __hash__(self):
+        return hash((self.stride, self.lower_bound, self.upper_bound))
+
+    def __eq__(self, other):
+        return (
+            isinstance(other, StridedIntervalAnnotation)
+            and self.stride == other.stride
+            and self.lower_bound == other.lower_bound
+            and self.upper_bound == other
+        )
+
+    def __repr__(self):
+        return f"<StridedIntervalAnnotation {self.stride}:{self.lower_bound} - {self.upper_bound}>"
+
+
 class RegionAnnotation(SimplificationAvoidanceAnnotation):
     """
     Use RegionAnnotation to annotate ASTs. Normally, an AST annotated by
@@ -84,12 +117,6 @@ def __init__(self, region_id, region_base_addr, offset):
         self.region_base_addr = region_base_addr
         self.offset = offset
 
-        # Do necessary conversion here
-        if isinstance(self.region_base_addr, claripy.ast.Base):
-            self.region_base_addr = claripy.backends.vsa.convert(self.region_base_addr)
-        if isinstance(self.offset, claripy.ast.Base):
-            self.offset = claripy.backends.vsa.convert(self.offset)
-
     #
     # Overriding base methods
     #

claripy/ast/base.py

@@ -35,6 +35,7 @@
 ArgType = Union["Base", bool, int, float, str, FSort, tuple["ArgType"], None]
 
 T = TypeVar("T", bound="Base")
+A = TypeVar("A", bound="Annotation")
 
 
 class ASTCacheKey(Generic[T]):
@@ -578,7 +579,7 @@ def has_annotation_type(self, annotation_type: type[Annotation]) -> bool:
         """
         return any(isinstance(a, annotation_type) for a in self.annotations)
 
-    def get_annotations_by_type(self, annotation_type: type[Annotation]) -> tuple[Annotation, ...]:
+    def get_annotations_by_type(self, annotation_type: type[A]) -> tuple[A, ...]:
         """
         Get all annotations of a given type.
 
@@ -587,6 +588,18 @@ def get_annotations_by_type(self, annotation_type: type[Annotation]) -> tuple[An
         """
         return tuple(a for a in self.annotations if isinstance(a, annotation_type))
 
+    def get_annotation(self, annotation_type: type[A]) -> A | None:
+        """
+        Get the first annotation of a given type.
+
+        :param annotation_type:     The type of the annotation.
+        :return:                    The annotation of the given type, or None if not found.
+        """
+        for a in self.annotations:
+            if isinstance(a, annotation_type):
+                return a
+        return None
+
     def append_annotation(self, a: Annotation) -> Self:
         """
         Appends an annotation to this AST.
@@ -766,17 +779,7 @@ def _op_repr(
     ):
         if details < ReprLevel.FULL_REPR:
             if op == "BVS":
-                extras = []
-                if args[1] is not None:
-                    fmt = "%#x" if isinstance(args[1], int) else "%s"
-                    extras.append("min=%s" % (fmt % args[1]))
-                if args[2] is not None:
-                    fmt = "%#x" if isinstance(args[2], int) else "%s"
-                    extras.append("max=%s" % (fmt % args[2]))
-                if args[3] is not None:
-                    fmt = "%#x" if isinstance(args[3], int) else "%s"
-                    extras.append("stride=%s" % (fmt % args[3]))
-                return "{}{}".format(args[0], "{{{}}}".format(", ".join(extras)) if extras else "")
+                return f"{args[0]}"
 
             if op == "BoolV":
                 return str(args[0])

claripy/ast/bv.py

@@ -199,9 +199,6 @@ def identical(self, other: Self, strict=False) -> bool:
 def BVS(  # pylint:disable=redefined-builtin
     name,
     size,
-    min=None,
-    max=None,
-    stride=None,
     explicit_name=None,
     **kwargs,
 ) -> BV:
@@ -213,17 +210,11 @@ def BVS(  # pylint:disable=redefined-builtin
 
     :param name:            The name of the symbol.
     :param size:            The size (in bits) of the bit-vector.
-    :param min:             The minimum value of the symbol, used only for value-set analysis
-    :param max:             The maximum value of the symbol, used only for value-set analysis
-    :param stride:          The stride of the symbol, used only for value-set analysis
     :param bool explicit_name:   If False, an identifier is appended to the name to ensure uniqueness.
 
     :returns:               a BV object representing this symbol.
     """
 
-    if stride == 0 and max != min:
-        raise ClaripyValueError("BVSes of stride 0 should have max == min")
-
     if isinstance(name, bytes):
         name = name.decode()
     if not isinstance(name, str):
@@ -234,7 +225,7 @@ def BVS(  # pylint:disable=redefined-builtin
 
     return BV(
         "BVS",
-        (n, min, max, stride),
+        (n,),
         variables=frozenset((n,)),
         length=size,
         symbolic=True,
@@ -301,16 +292,11 @@ def SI(
         si = claripy.backends.backend_vsa.CreateStridedInterval(
             name=name, bits=bits, lower_bound=lower_bound, upper_bound=upper_bound, stride=stride, to_conv=to_conv
         )
-        return BVS(
-            name, si._bits, min=si._lower_bound, max=si._upper_bound, stride=si._stride, explicit_name=explicit_name
+        return BVS(name, si._bits, explicit_name=explicit_name).annotate(
+            claripy.annotation.StridedIntervalAnnotation(si._stride, si._lower_bound, si._upper_bound)
         )
-    return BVS(
-        name,
-        bits,
-        min=lower_bound,
-        max=upper_bound,
-        stride=stride,
-        explicit_name=explicit_name,
+    return BVS(name, bits, explicit_name=explicit_name).annotate(
+        claripy.annotation.StridedIntervalAnnotation(stride, lower_bound, upper_bound)
     )
 
 
@@ -331,6 +317,8 @@ def ValueSet(bits, region=None, region_base_addr=None, value=None, name=None, va
         region_base_addr = 0
 
     v = region_base_addr + value
+    if isinstance(v, claripy.ast.Base):
+        v = claripy.simplify(v)
 
     # Backward compatibility
     if isinstance(v, numbers.Number):
@@ -340,19 +328,25 @@ def ValueSet(bits, region=None, region_base_addr=None, value=None, name=None, va
         min_v, max_v = v.lower_bound, v.upper_bound
         stride = v.stride
     elif isinstance(v, claripy.ast.Base):
-        sv = claripy.simplify(v)
-        if sv.op == "BVS":
-            min_v = sv.args[1]
-            max_v = sv.args[2]
-            stride = sv.args[3]
+        si_anno = v.get_annotation(claripy.annotation.StridedIntervalAnnotation)
+        if si_anno is not None:
+            min_v = si_anno.lower_bound
+            max_v = si_anno.upper_bound
+            stride = si_anno.stride
+        elif v.op == "BVV":
+            min_v = v.args[0]
+            max_v = v.args[0]
+            stride = 0
         else:
-            raise ClaripyValueError(f"ValueSet() does not take `value` ast with op {sv.op}")
+            raise ClaripyValueError(f"ValueSet() does not take `value` ast with op {v.op}")
     else:
         raise ClaripyValueError(f"ValueSet() does not take `value` of type {type(value)}")
 
     if name is None:
         name = "ValueSet"
-    bvs = BVS(name, bits, min=region_base_addr + min_v, max=region_base_addr + max_v, stride=stride)
+    bvs = BVS(name, bits).annotate(
+        claripy.annotation.StridedIntervalAnnotation(stride, region_base_addr + min_v, region_base_addr + max_v)
+    )
 
     # Annotate the bvs and return the new AST
     return bvs.annotate(claripy.annotation.RegionAnnotation(region, region_base_addr, value))

claripy/backends/backend_vsa/backend_vsa.py

@@ -6,10 +6,11 @@
 import operator
 from functools import reduce
 
-from claripy.annotation import RegionAnnotation
+from claripy.annotation import RegionAnnotation, StridedIntervalAnnotation
 from claripy.ast.base import Base
-from claripy.ast.bv import ESI, SI, TSI
+from claripy.ast.bv import BV, BVV, ESI, SI, TSI, VS
 from claripy.backends.backend import Backend
+from claripy.backends.backend_vsa.errors import ClaripyVSAError
 from claripy.balancer import Balancer
 from claripy.errors import BackendError
 from claripy.operations import backend_operations_vsa_compliant, expression_set_operations
@@ -128,13 +129,28 @@ def _abstract(self, e):
                 return TSI(e.bits, explicit_name=e.name)
             if e.is_bottom:
                 return ESI(e.bits)
+            if e.stride in {0, 1} and e.lower_bound == e.upper_bound:
+                return BVV(e.lower_bound, e.bits)
             return SI(
                 name=e.name,
                 bits=e.bits,
                 lower_bound=e.lower_bound,
                 upper_bound=e.upper_bound,
                 stride=e.stride,
             )
+        if isinstance(e, ValueSet):
+            if len(e.regions) == 0:
+                return VS(bits=e.bits, name=e.name)
+            if len(e.regions) == 1:
+                region = next(iter(e.regions))
+                return VS(
+                    bits=e.bits,
+                    region=region,
+                    region_base_addr=e._region_base_addrs[region].eval(1)[0] if e._region_base_addrs else 0,
+                    value=e.regions[region].eval(1)[0],
+                    name=e.name,
+                )
+            raise ClaripyVSAError("Cannot abstract ValueSet with multiple regions")
         raise BackendError(f"Don't know how to abstract {type(e)}")
 
     def _eval(self, expr, n, extra_constraints=(), solver=None, model_callback=None):
@@ -230,18 +246,43 @@ def apply_annotation(self, o, a):
         :rtype: BackendObject
         """
 
-        # Currently we only support RegionAnnotation
-
-        if not isinstance(a, RegionAnnotation):
-            return o
-
-        if not isinstance(o, ValueSet):
-            # Convert it to a ValueSet first
-            # Note that the original value is not kept at all. If you want to convert a StridedInterval to a ValueSet,
-            # you gotta do the conversion by calling AST.annotate() from outside.
-            o = ValueSet.empty(o.bits)
+        if isinstance(o, StridedInterval):
+            if isinstance(a, StridedIntervalAnnotation):
+                return CreateStridedInterval(
+                    bits=o.bits,
+                    stride=a.stride,
+                    lower_bound=a.lower_bound,
+                    upper_bound=a.upper_bound,
+                    name=o.name,
+                )
+
+            if isinstance(a, RegionAnnotation):
+                offset = self.convert(a.offset)
+                if isinstance(offset, numbers.Number):
+                    offset = StridedInterval(bits=o.bits, stride=0, lower_bound=offset, upper_bound=offset)
+                vs = ValueSet.empty(o.bits)
+                if isinstance(offset, StridedInterval):
+                    vs._merge_si(a.region_id, a.region_base_addr, offset)
+                elif isinstance(offset, ValueSet):
+                    for si in offset.regions.values():
+                        vs._merge_si(a.region_id, a.region_base_addr, si)
+                else:
+                    raise ClaripyVSAError(f"Unsupported offset type {type(offset)}")
+                return vs
+
+        if isinstance(o, ValueSet) and isinstance(a, StridedIntervalAnnotation):
+            si = CreateStridedInterval(
+                bits=o.bits,
+                stride=a.stride,
+                lower_bound=a.lower_bound,
+                upper_bound=a.upper_bound,
+                name=o.name,
+            )
+            vs = o.copy()
+            vs._merge_si(a.region_id, a.region_base_addr, si)
+            return vs
 
-        return o.apply_annotation(a)
+        raise ValueError(f"Unsupported annotation type {type(a)} for object {type(o)}")
 
     @staticmethod
     def BVV(ast):
@@ -302,16 +343,8 @@ def SGE(a, b):
         return a.SGE(b)
 
     @staticmethod
-    def BVS(ast: Base):
-        size = ast.size()
-        name, mn, mx, stride = ast.args
-        return CreateStridedInterval(
-            name=name,
-            bits=size,
-            lower_bound=mn,
-            upper_bound=mx,
-            stride=stride,
-        )
+    def BVS(ast: BV):
+        return CreateStridedInterval(name=ast.args[0], bits=ast.size())
 
     def If(self, cond, t, f):
         if not self.has_true(cond):

claripy/backends/backend_vsa/valueset.py

@@ -209,19 +209,6 @@ def get_si(self, region):
     def stridedinterval(self):
         return self._si
 
-    def apply_annotation(self, annotation):
-        """
-        Apply a new annotation onto self, and return a new ValueSet object.
-
-        :param RegionAnnotation annotation: The annotation to apply.
-        :return: A new ValueSet object
-        :rtype: ValueSet
-        """
-
-        vs = self.copy()
-        vs._merge_si(annotation.region_id, annotation.region_base_addr, annotation.offset)
-        return vs
-
     def __repr__(self):
         s = ""
         for region, si in self._regions.items():

claripy/balancer.py

@@ -52,7 +52,7 @@ def _replacements_iter(self):
             min_int = 0
             mn = self._lower_bounds.get(k, min_int)
             mx = self._upper_bounds.get(k, max_int)
-            bound_si = BVS("bound", len(k.ast), min=mn, max=mx)
+            bound_si = BVS("bound", len(k.ast)).annotate(claripy.annotation.StridedIntervalAnnotation(1, mn, mx))
             l.debug("Yielding bound %s for %s.", bound_si, k.ast)
             if k.ast.op == "Reverse":
                 yield (k.ast.args[0], k.ast.intersection(bound_si).reversed)

tests/test_expression.py

@@ -180,7 +180,7 @@ def test_expression(self):
 
     def test_cardinality(self):
         x = claripy.BVS("x", 32)
-        y = claripy.BVS("y", 32, min=100, max=120)
+        y = claripy.BVS("y", 32).annotate(claripy.annotation.StridedIntervalAnnotation(1, 100, 120))
         n = claripy.BVV(10, 32)
         m = claripy.BVV(20, 32)
 

tests/test_simplify.py

@@ -4,6 +4,7 @@
 import unittest
 
 import claripy
+import claripy.annotation
 
 
 class TestSimplify(unittest.TestCase):
@@ -56,7 +57,7 @@ def assert_correct(a, b):
         assert_correct(x % y, claripy.backends.z3.simplify(x % y))
 
     def test_rotate_shift_mask_simplification(self):
-        a = claripy.BVS("N", 32, max=0xC, min=0x1)
+        a = claripy.BVS("N", 32).annotate(claripy.annotation.StridedIntervalAnnotation(1, 0x1, 0xC))
         extend_ = claripy.BVS("extend", 32, uninitialized=True)
         a_ext = extend_.concat(a)
         expr = ((a_ext << 3) | (claripy.LShR(a_ext, 61))) & 0x7FFFFFFF8