Repository: leonardt/hwtypes

Issue:
SMT Int class [Enhancement]
For building up constraints for synthesis problems it is useful to have access to an Int class that plays nicely with SMTBit.

Current implementation in branch: smt_int

Semantics defined in: https://smtlib.cs.uiowa.edu/theories-Ints.shtml

Would be nice to include methods to translate between SMTInt and SMTBitVector.


Comments:

Patch:
diff --git a/hwtypes/__init__.py b/hwtypes/__init__.py
index 6858443..ea48e32 100644
--- a/hwtypes/__init__.py
+++ b/hwtypes/__init__.py
@@ -7,3 +7,4 @@ from .fp_vector_abc import *
 from .fp_vector import *
 from .smt_fp_vector import *
 from .modifiers import *
+from .smt_int import *
diff --git a/hwtypes/smt_bit_vector.py b/hwtypes/smt_bit_vector.py
index 75d515c..0aa8c4a 100644
--- a/hwtypes/smt_bit_vector.py
+++ b/hwtypes/smt_bit_vector.py
@@ -643,11 +643,15 @@ class SMTBitVector(AbstractBitVector):
 #    def __int__(self):
 #        return self.as_uint()
 #
-#    def as_uint(self):
-#        return self._value.bv_unsigned_value()
-#
-#    def as_sint(self):
-#        return self._value.bv_signed_value()
+    def as_uint(self):
+        #Avoids circular import
+        from . import SMTInt
+        return SMTInt(self)
+
+    def as_sint(self):
+        #Avoids circular import
+        from . import SMTInt
+        return SMTInt(self[:-1]) - SMTInt(self & (1<<(self.size-1)))
 #
 #    @classmethod
 #    def random(cls, width):
diff --git a/hwtypes/smt_int.py b/hwtypes/smt_int.py
new file mode 100644
index 0000000..e4eeb16
--- /dev/null
+++ b/hwtypes/smt_int.py
@@ -0,0 +1,148 @@
+import itertools as it
+import functools as ft
+from .smt_bit_vector import SMTBit, SMTBitVector
+
+import pysmt
+import pysmt.shortcuts as smt
+from pysmt.typing import INT
+
+from collections import defaultdict
+import re
+import warnings
+import weakref
+
+__ALL__ = ['SMTInt']
+
+_var_counters = defaultdict(it.count)
+_name_table = weakref.WeakValueDictionary()
+
+def _gen_name(prefix='V'):
+    name = f'{prefix}_{next(_var_counters[prefix])}'
+    while name in _name_table:
+        name = f'{prefix}_{next(_var_counters[prefix])}'
+    return name
+
+_name_re = re.compile(r'V_\d+')
+
+class _SMYBOLIC:
+    def __repr__(self):
+        return 'SYMBOLIC'
+
+class _AUTOMATIC:
+    def __repr__(self):
+        return 'AUTOMATIC'
+
+SMYBOLIC = _SMYBOLIC()
+AUTOMATIC = _AUTOMATIC()
+
+def int_cast(fn):
+    @ft.wraps(fn)
+    def wrapped(self, other):
+        if isinstance(other, SMTInt):
+            return fn(self, other)
+        else:
+            try:
+                other = SMTInt(other)
+            except TypeError:
+                return NotImplemented
+            return fn(self, other)
+    return wrapped
+
+class SMTInt:
+    def __init__(self, value=SMYBOLIC, *, name=AUTOMATIC, prefix=AUTOMATIC):
+        if (name is not AUTOMATIC or prefix is not AUTOMATIC) and value is not SMYBOLIC:
+            raise TypeError('Can only name symbolic variables')
+        elif name is not AUTOMATIC and prefix is not AUTOMATIC:
+            raise ValueError('Can only set either name or prefix not both')
+        elif name is not AUTOMATIC:
+            if not isinstance(name, str):
+                raise TypeError('Name must be string')
+            elif name in _name_table:
+                raise ValueError(f'Name {name} already in use')
+            elif _name_re.fullmatch(name):
+                warnings.warn('Name looks like an auto generated name, this might break things')
+            _name_table[name] = self
+        elif prefix is not AUTOMATIC:
+            name = _gen_name(prefix)
+            _name_table[name] = self
+        elif name is AUTOMATIC and value is SMYBOLIC:
+            name = _gen_name()
+            _name_table[name] = self
+
+        if value is SMYBOLIC:
+            self._value = smt.Symbol(name, INT)
+        elif isinstance(value, pysmt.fnode.FNode):
+            if value.get_type().is_int_type():
+                self._value = value
+            elif value.get_type().is_bv_type():
+                self._value = smt.BVToNatural(value)
+            else:
+                raise TypeError(f'Expected int type not {value.get_type()}')
+        elif isinstance(value, SMTInt):
+            self._value = value._value
+        elif isinstance(value, SMTBitVector):
+            self._value = smt.BVToNatural(value.value)
+        elif isinstance(value, bool):
+            self._value = smt.Int(int(value))
+        elif isinstance(value, int):
+            self._value = smt.Int(value)
+        elif hasattr(value, '__int__'):
+            self._value = smt.Int(int(value))
+        else:
+            raise TypeError("Can't coerce {} to Int".format(type(value)))
+
+        self._name = name
+        self._value = smt.simplify(self._value)
+
+    def __repr__(self):
+        if self._name is not AUTOMATIC:
+            return f'{type(self)}({self._name})'
+        else:
+            return f'{type(self)}({self._value})'
+
+    @property
+    def value(self):
+        return self._value
+
+    def __neg__(self):
+        return SMTInt(0) - self
+
+    @int_cast
+    def __sub__(self, other: 'SMTInt') -> 'SMTInt':
+        return SMTInt(self.value - other.value)
+
+    @int_cast
+    def __add__(self, other: 'SMTInt') -> 'SMTInt':
+        return SMTInt(self.value + other.value)
+
+    @int_cast
+    def __mul__(self, other: 'SMTInt') -> 'SMTInt':
+        return SMTInt(self.value * other.value)
+
+    @int_cast
+    def __floordiv__(self, other: 'SMTInt') -> 'SMTInt':
+        return SMTInt(smt.Div(self.value, other.value))
+
+    @int_cast
+    def __ge__(self, other: 'SMTInt') -> SMTBit:
+        return SMTBit(self.value >= other.value)
+
+    @int_cast
+    def __gt__(self, other: 'SMTInt') -> SMTBit:
+        return SMTBit(self.value > other.value)
+
+    @int_cast
+    def __le__(self, other: 'SMTInt') -> SMTBit:
+        return SMTBit(self.value <= other.value)
+
+    @int_cast
+    def __lt__(self, other: 'SMTInt') -> SMTBit:
+        return SMTBit(self.value < other.value)
+
+    @int_cast
+    def __eq__(self, other: 'SMTInt') -> SMTBit:
+        return SMTBit(smt.Equals(self.value, other.value))
+
+    @int_cast
+    def __ne__(self, other: 'SMTInt') -> SMTBit:
+        return SMTBit(smt.NotEquals(self.value, other.value))


Test patch:
diff --git a/tests/test_smt_int.py b/tests/test_smt_int.py
new file mode 100644
index 0000000..554f83d
--- /dev/null
+++ b/tests/test_smt_int.py
@@ -0,0 +1,89 @@
+from hwtypes import SMTInt, SMTBitVector
+import pysmt.shortcuts as smt
+import pytest
+
+def test_sat_unsat():
+    x = SMTInt(prefix='x')
+    f1 = (x > 0) & (x < 0)
+    logic = None
+    #Test unsat
+    with smt.Solver(logic=logic, name='z3') as solver:
+        solver.add_assertion(f1.value)
+        res = solver.solve()
+        assert not res
+
+    f2 = (x >= 0) & (x*x+x == 30)
+    #test sat
+    with smt.Solver(logic=logic, name='z3') as solver:
+        solver.add_assertion(f2.value)
+        res = solver.solve()
+        assert res
+        x_val = solver.get_value(x.value)
+        assert x_val.constant_value() == 5
+
+bin_ops = dict(
+    sub = lambda x,y: x-y,
+    add = lambda x,y: x+y,
+    mul = lambda x,y: x*y,
+    div = lambda x,y: x//y,
+    lte = lambda x,y: x<=y,
+    lt = lambda x,y: x<y,
+    gte = lambda x,y: x>=y,
+    gt = lambda x,y: x>y,
+    eq = lambda x,y: x != y,
+    neq = lambda x,y: x == y,
+)
+
+unary_ops = dict(
+    neg = lambda x: -x,
+)
+
+import random
+
+@pytest.mark.parametrize("name, fun", bin_ops.items())
+def test_bin_ops(name, fun):
+    x = SMTInt(prefix='x')
+    y = SMTInt(prefix='y')
+    x_val = random.randint(-100, 100)
+    y_val = random.choice(list(set(x for x in range(-100,100))-set((0,))))
+    f_val = fun(x_val,y_val)
+    f = (fun(x, y) == f_val) & (y != 0)
+    with smt.Solver(name='z3') as solver:
+        solver.add_assertion(f.value)
+        res = solver.solve()
+        assert res
+        x_solved = solver.get_value(x.value).constant_value()
+        y_solved = solver.get_value(y.value).constant_value()
+        assert fun(x_solved, y_solved) == f_val
+
+
+@pytest.mark.parametrize("name, fun", unary_ops.items())
+def test_unary_ops(name, fun):
+    x = SMTInt(prefix='x')
+    x_val = random.randint(-100, 100)
+    f_val = fun(x_val)
+    f = (fun(x) == f_val)
+    with smt.Solver(name='z3') as solver:
+        solver.add_assertion(f.value)
+        res = solver.solve()
+        assert res
+        x_solved = solver.get_value(x.value).constant_value()
+        assert fun(x_solved) == f_val
+
+def test_init_bv():
+    SMTInt(SMTBitVector[5]())
+    SMTInt(SMTBitVector[5](10))
+    SMTInt(SMTBitVector[5](10).value)
+
+
+def test_as_sint():
+    x = SMTBitVector[5](-5)
+    x_int = x.as_sint()
+    assert isinstance(x_int, SMTInt)
+    assert x_int.value.constant_value() == -5
+
+def test_as_uint():
+    x = SMTBitVector[5](5)
+    x_int = x.as_uint()
+    assert isinstance(x_int, SMTInt)
+    assert x.as_uint().value.constant_value() == 5
\ No newline at end of file