Merge pull request #445 from pysmt/infix_euf

Added support for EUF infix notation

pysmt/fnode.py

@@ -650,27 +650,33 @@ def algebraic_approx_value(self, precision=10):
     # Infix Notation
     @assert_infix_enabled
     def _apply_infix(self, right, function, bv_function=None):
-        mgr = _mgr()
-        # BVs
         # Default bv_function to function
-        if bv_function is None: bv_function = function
+        if bv_function is None:
+            bv_function = function
+        right = self._infix_prepare_arg(right, self.get_type())
         if self.get_type().is_bv_type():
-            if is_python_integer(right):
-                right = mgr.BV(right, width=self.bv_width())
             return bv_function(self, right)
-        # Boolean, Integer and Arithmetic
-        if is_python_boolean(right):
-            right = mgr.Bool(right)
-        elif is_python_integer(right):
-            ty = self.get_type()
-            if ty.is_real_type():
-                right = mgr.Real(right)
-            else:
-                right = mgr.Int(right)
-        elif is_python_rational(right):
-            right = mgr.Real(right)
         return function(self, right)
 
+    @assert_infix_enabled
+    def _infix_prepare_arg(self, arg, expected_type):
+        mgr = _mgr()
+        if isinstance(arg, FNode):
+            return arg
+
+        # BVs
+        if expected_type.is_bv_type():
+            return mgr.BV(arg, width=expected_type.width)
+        # Boolean, Integer and Arithmetic
+        elif expected_type.is_bool_type():
+            return mgr.Bool(arg)
+        elif expected_type.is_int_type():
+            return mgr.Int(arg)
+        elif expected_type.is_real_type():
+            return mgr.Real(arg)
+        else:
+            raise PysmtValueError("Unsupported value '%s' in infix operator" % str(arg))
+
     def Implies(self, right):
         return self._apply_infix(right, _mgr().Implies)
 
@@ -854,6 +860,20 @@ def __rshift__(self, right):
 
     def __mod__(self, right):
         return self._apply_infix(right, None, bv_function=_mgr().BVURem)
+
+    @assert_infix_enabled
+    def __call__(self, *args):
+        if self.is_symbol() and self.symbol_type().is_function_type():
+            types = self.symbol_type().param_types
+            if (len(types) != len(args)):
+                raise PysmtValueError("Wrong number of parameters passed in "
+                                      "infix 'call' operator")
+            args = [self._infix_prepare_arg(x, t) for x,t in zip(args, types)]
+            return _mgr().Function(self, args)
+        else:
+            raise PysmtValueError("Call operator can be applied to symbol "
+                                  "types having function type only")
+
 # EOC FNode
 
 

pysmt/test/test_euf.py

@@ -16,10 +16,11 @@
 #   limitations under the License.
 #
 from pysmt.shortcuts import *
-from pysmt.typing import INT, REAL, FunctionType
+from pysmt.typing import INT, REAL, FunctionType, BV16
 from pysmt.logics import UFLRA, UFLIRA
 from pysmt.test import TestCase, main
-from pysmt.test import skipIfSolverNotAvailable, skipIfNoSolverForLogic
+from pysmt.test import skipIfNoSolverForLogic
+from pysmt.exceptions import PysmtModeError, PysmtValueError
 
 
 class TestEUF(TestCase):
@@ -37,6 +38,38 @@ def test_euf(self):
         self.assertSat(check, logic=UFLIRA,
                        msg="Formula was expected to be sat")
 
+
+    def test_infix(self):
+        ftype1 = FunctionType(REAL, [REAL])
+        ftype2 = FunctionType(REAL, [REAL, INT])
+        f = Symbol("f", ftype1)
+        g = Symbol("g", ftype2)
+
+        with self.assertRaises(PysmtModeError):
+            f(1.0)
+
+        get_env().enable_infix_notation = True
+
+        infix = Equals(f(1.0), g(2.0, 4))
+        explicit = Equals(Function(f, [Real(1.0)]), Function(g, [Real(2.0), Int(4)]))
+        self.assertEqual(infix, explicit)
+
+        ftype1 = FunctionType(REAL, [BV16])
+        ftype2 = FunctionType(BV16, [INT, BV16])
+        f = Symbol("bvf", ftype1)
+        g = Symbol("bvg", ftype2)
+        infix = Equals(f(g(2, 6)), Real(0))
+        explicit = Equals(Function(f, [Function(g, [Int(2), BV(6, 16)])]), Real(0))
+        self.assertEqual(infix, explicit)
+
+        with self.assertRaises(PysmtValueError):
+            f(BV(6, 16), BV(8, 16))
+
+        ftype3 = FunctionType(REAL, [])
+        h = Symbol("h", ftype3)
+        with self.assertRaises(PysmtValueError):
+            h()
+
     @skipIfNoSolverForLogic(UFLRA)
     def test_quantified_euf(self):
         ftype1 = FunctionType(REAL, [REAL, REAL])