Support higher-order function sorts and partial application

cvc5_pythonic_api/cvc5_pythonic.py

@@ -949,9 +949,74 @@ def __call__(self, *args):
         f(x, y)
         >>> f(x, x)
         f(x, ToReal(x))
+
+        Supplying fewer arguments than the arity builds a partial
+        application, whose sort is the rest of the function sort:
+
+        >>> f(x).sort()
+        (-> Real Bool)
+        >>> f(x)
+        f(x)
+
+        Applying that to the remaining arguments gives the same term as
+        applying `f` to all of them at once, and prints the same way:
+
+        >>> f(x)(y)
+        f(x, y)
+
+        Note that cvc5 only reasons about function terms under a
+        higher-order logic, one whose name carries the `HO_` prefix:
+
+        >>> s = SolverFor('HO_ALL')
+        >>> s.add(f(x)(y) != f(x, y))
+        >>> s.check()
+        unsat
         """
+        args = _get_args(args)
+        if 0 < len(args) and (len(args) < self.arity() or self.kind() == Kind.HO_APPLY):
+            return _partial_apply(self, args)
         return _higherorder_apply(self, args, Kind.APPLY_UF)
 
+    def __getitem__(self, arg):
+        """Shorthand for `self(arg)`.
+
+        Z3Py gives a lambda expression an array sort, so it spells the
+        application of one, and of anything defined by one, with `[]`. That
+        spelling is accepted here too, so such code carries over:
+
+        >>> x = Real('x')
+        >>> lo, hi = Reals('lo hi')
+        >>> body = Lambda([x], And(lo <= x, x <= hi))
+        >>> setof = Function('setof', RealSort(), RealSort(), body.sort())
+        >>> setof(lo, hi)[3]
+        setof(lo, hi, 3)
+
+        Several indices at once are applied in order:
+
+        >>> f = Function('f', IntSort(), IntSort(), IntSort())
+        >>> i = Int('i')
+        >>> f[i, i]
+        f(i, i)
+        """
+        if not isinstance(arg, tuple):
+            arg = (arg,)
+        return self(*arg)
+
+
+def _partial_apply(func, args):
+    """Apply `func` to `args`, one argument at a time.
+
+    Used when there are too few arguments to saturate `func`, and when `func`
+    is itself a partial application: neither can be expressed with APPLY_UF.
+    The result is a term of whatever is left of the function sort, which is
+    the range sort once the arguments run out.
+    """
+    t = func.ast
+    for i in range(len(args)):
+        arg = func.domain(i).cast(args[i])
+        t = func.ctx.tm.mkTerm(Kind.HO_APPLY, t, arg.as_ast())
+    return _to_expr_ref(t, func.ctx)
+
 
 def _higherorder_apply(func, args, kind):
     """Create an SMT application from a FuncDeclRef and a kind of application"""
@@ -982,23 +1047,68 @@ def is_func_decl(a):
     return isinstance(a, FuncDeclRef)
 
 
+def _to_function_sort(ctx, sig):
+    """Build the base function sort for the signature `sig`.
+
+    `sig` is a list of domain sorts followed by the range sort.
+
+    cvc5 normalizes higher-order sorts, so a function sort cannot appear as a
+    codomain. When the range is one, its own domains are spliced onto the
+    arguments instead: `Int -> (-> Real Bool)` is built as `(-> Int Real Bool)`.
+    Applying such a function to an argument for each of the leading domains
+    yields a term of the range sort again, so the distinction stays invisible.
+    """
+    arity = len(sig) - 1
+    rng = sig[arity]
+    doms = [sig[i].ast for i in range(arity)]
+    if isinstance(rng, FuncSortRef):
+        doms += [rng.domain_n(i).ast for i in range(rng.arity())]
+        cod = rng.range().ast
+    else:
+        cod = rng.ast
+    return ctx.tm.mkFunctionSort(doms, cod)
+
+
 def Function(name, *sig):
     """Create a new SMT uninterpreted function with the given sorts.
 
     >>> f = Function('f', IntSort(), IntSort())
     >>> f(f(0))
     f(f(0))
+
+    The range may itself be a function sort, as it is for a function defined
+    by a lambda expression:
+
+    >>> x = Real('x')
+    >>> lo, hi = Reals('lo hi')
+    >>> body = Lambda([x], And(lo <= x, x <= hi))
+    >>> setof = Function('setof', RealSort(), RealSort(), body.sort())
+    >>> setof.sort()
+    (-> Real Real Real Bool)
+    >>> setof(lo, hi).sort() == body.sort()
+    True
+
+    A partial application of `setof` can then be defined by `body`, which is
+    what a Z3Py definition of a function returning a lambda amounts to.
+    Reasoning about it needs a higher-order logic, and the quantified
+    definition needs `ho-elim` to be discharged rather than answered
+    `unknown`:
+
+    >>> s = SolverFor('HO_ALL')
+    >>> s.set('ho-elim', True)
+    >>> s.add(ForAll([lo, hi], setof(lo, hi) == body))
+    >>> s.add(Not(setof(0, 10)(3)))
+    >>> s.check()
+    unsat
     """
     sig = _get_args(sig)
     if debugging():
         _assert(len(sig) > 0, "At least two arguments expected")
-    arity = len(sig) - 1
-    rng = sig[arity]
+    rng = sig[len(sig) - 1]
     if debugging():
         _assert(is_sort(rng), "SMT sort expected")
     ctx = rng.ctx
-    sort = ctx.tm.mkFunctionSort([sig[i].ast for i in range(arity)], rng.ast)
-    e = ctx.get_var(name, _to_sort_ref(sort, ctx))
+    e = ctx.get_var(name, _to_sort_ref(_to_function_sort(ctx, sig), ctx))
     return FuncDeclRef(e, ctx)
 
 
@@ -1013,13 +1123,11 @@ def FreshFunction(*sig):
     sig = _get_args(sig)
     if debugging():
         _assert(len(sig) > 0, "At least two arguments expected")
-    arity = len(sig) - 1
-    rng = sig[arity]
+    rng = sig[len(sig) - 1]
     if debugging():
         _assert(is_sort(rng), "SMT sort expected")
     ctx = rng.ctx
-    sort = ctx.tm.mkFunctionSort([sig[i].ast for i in range(arity)], rng.ast)
-    name = ctx.next_fresh(sort, "freshfn")
+    name = ctx.next_fresh(_to_function_sort(ctx, sig), "freshfn")
     return Function(name, *sig)
 
 
@@ -5657,6 +5765,11 @@ def Store(a, i, v):
 def Select(a, i):
     """Return an SMT select array expression.
 
+    `Select` is an array operation. A lambda expression has a function sort
+    here, not an array sort, so it is applied with `[]` instead: Z3Py defines
+    `Select(a, i)` as `a[i]`, so rewriting a select of a lambda that way keeps
+    working under both.
+
     >>> a = Array('a', IntSort(), IntSort())
     >>> i = Int('i')
     >>> Select(a, i)
@@ -9170,6 +9283,38 @@ def sort(self):
             return _sort(self.ctx, self.as_ast())
         return BoolSort(self.ctx)
 
+    def __getitem__(self, arg):
+        """Apply the lambda expression `self` to `arg`.
+
+        Z3Py gives a lambda an array sort and applies it with `[]`; the same
+        spelling works here, even though the sort is a function sort. Note
+        that `Select` stays an array operation, so Z3Py code written as
+        `Select(L, i)` should be rewritten as `L[i]`, which Z3Py accepts too.
+
+        >>> x, y = Ints('x y')
+        >>> i = Int('i')
+        >>> Lambda([x], x + 1)[i]
+        Lambda(x, x + 1)(i)
+        >>> simplify(Lambda([x], x + 1)[3])
+        4
+        >>> simplify(Lambda([x, y], x + y)[3, 4])
+        7
+        """
+        if debugging():
+            _assert(self.is_lambda(), "Only lambda expressions can be applied")
+        if not isinstance(arg, tuple):
+            arg = (arg,)
+        # Applied one argument at a time rather than through a FuncDeclRef
+        # view of `self`: the wrapper would claim a type the term does not
+        # have, and a lambda is not printable as a declaration.
+        sort = self.sort()
+        t = self.as_ast()
+        for i in range(len(arg)):
+            t = self.ctx.tm.mkTerm(
+                Kind.HO_APPLY, t, sort.domain_n(i).cast(arg[i]).as_ast()  # type: ignore
+            )
+        return _to_expr_ref(t, self.ctx)
+
     def is_forall(self):
         """Return `True` if `self` is a universal quantifier.
 

cvc5_pythonic_api/cvc5_pythonic_printer.py

@@ -1208,6 +1208,8 @@ def pp_app(self, a, d, xs):
                 return self.pp_unary(a, d, xs)
             elif k == Kind.APPLY_UF:
                 return self.pp_uf_apply(a, d, xs)
+            elif k == Kind.HO_APPLY:
+                return self.pp_ho_apply(a, d, xs)
             elif k in [Kind.APPLY_CONSTRUCTOR, Kind.APPLY_SELECTOR, Kind.APPLY_TESTER]:
                 return self.pp_dt_apply(a, d, xs)
             elif k == Kind.SEXPR:
@@ -1227,6 +1229,38 @@ def pp_uf_apply(self, a, d, xs):
                 break
         return seq1(self.pp_name(first), r)
 
+    def pp_ho_apply(self, a, d, xs):
+        # A partial application is a spine of binary HO_APPLY nodes. Collapse
+        # it into a single application, so that a term built as f(x)(y) reads
+        # as f(x, y) -- the same way a saturated application of f is printed,
+        # and the two are equal terms.
+        args = []
+        head = a
+        while head.kind() == Kind.HO_APPLY:
+            children = head.children()
+            head = children[0]
+            args.append(children[1])
+        args.reverse()
+        r = []
+        sz = 0
+        for child in args:
+            r.append(self.pp_expr(child, d + 1, xs))
+            sz = sz + 1
+            if sz > self.max_args:
+                r.append(self.pp_ellipses())
+                break
+        # Not seq1: the head of the spine need not be a plain name -- it can
+        # be any term of function sort, such as an If over two functions --
+        # and seq1 measures its header, which only works for a name.
+        return group(
+            compose(
+                self.pp_expr(head, d + 1, xs),
+                to_format("("),
+                indent(1, seq(r)),
+                to_format(")"),
+            )
+        )
+
     def pp_dt_apply(self, a, d, xs):
         r = []
         sz = 0

test/pgm_outputs/higher_order.py.out

@@ -0,0 +1,19 @@
+(-> Interval Real Bool)
+(-> Real Bool)
+True
+setof(i)
+setof(i, x)
+ForAll(i,
+       setof(i) == Lambda(x, And(lo(i) <= x, x <= hi(i))))
+unsat
+unsat
+False
+unsat
+If(c, f, g)(y)
+If(c, f, g)(y, y)
+setof(i, x)
+True
+First argument must be an SMT array expression
+Lambda(x, And(lo(i) <= x, x <= hi(i)))(3)
+4
+unsat

test/pgms/higher_order.py

@@ -0,0 +1,68 @@
+from cvc5_pythonic_api import *
+
+# A function whose range is the sort of a lambda: the sort is flattened, so
+# saturating the leading domains yields the range sort back.
+Interval = Datatype('Interval')
+Interval.declare('mk', ('lo', RealSort()), ('hi', RealSort()))
+Interval = Interval.create()
+
+i = Const('i', Interval)
+x = Real('x')
+body = Lambda([x], And(Interval.lo(i) <= x, x <= Interval.hi(i)))
+
+setof = Function('setof', Interval, body.sort())
+print(setof.sort())
+print(setof(i).sort())
+print(setof(i).sort() == body.sort())
+
+# Partial applications print as ordinary applications.
+print(setof(i))
+print(setof(i)(x))
+
+# The definition a Z3Py `define` would build now typechecks.
+defn = ForAll([i], setof(i) == body)
+print(defn)
+
+# ... and is usable, under a higher-order logic.
+s = SolverFor('HO_ALL')
+s.set('ho-elim', True)
+s.add(defn)
+s.add(Not(setof(Interval.mk(0, 10))(3)))
+print(s.check())
+
+s = SolverFor('HO_ALL')
+s.set('ho-elim', True)
+s.add(defn)
+s.add(setof(Interval.mk(0, 10))(42))
+print(s.check())
+
+# Currying an ordinary function agrees with applying it outright.
+f = Function('f', IntSort(), IntSort(), IntSort())
+y = Int('y')
+print(f(y)(y).eq(f(y, y)))
+s = SolverFor('HO_ALL')
+s.add(f(y)(y) != f(y, y))
+print(s.check())
+
+# The head of an application need not be a name.
+g = Function('g', IntSort(), IntSort(), IntSort())
+print(If(Bool('c'), f, g)(y))
+print(If(Bool('c'), f, g)(y)(y))
+
+# Z3Py spells the application of a lambda, and of anything defined by one,
+# with []. Both spellings work here, and build the same term. Select stays
+# an array operation, and does not accept a function.
+print(setof(i)[x])
+print(setof(i)[x].eq(setof(i)(x)))
+try:
+    Select(setof(i), x)
+except SMTException as e:
+    print(e)
+print(body[3])
+print(simplify(Lambda([y], y + 1)[3]))
+
+s = SolverFor('HO_ALL')
+s.set('ho-elim', True)
+s.add(defn)
+s.add(Not(setof(Interval.mk(0, 10))[3]))
+print(s.check())