Implement enumerator for functions. (#1243)

* Implement enumerator for functions.

* Address review.

* Minor

* Format

* Improve comment.

* Format

5 files changed, 159 insertions, 26 deletions

diff --git a/src/theory/builtin/kinds b/src/theory/builtin/kinds
index c1edd81cb..a04c12903 100644
--- a/src/theory/builtin/kinds
+++ b/src/theory/builtin/kinds
@@ -321,6 +321,9 @@ cardinality FUNCTION_TYPE \
     "::CVC4::theory::builtin::FunctionProperties::computeCardinality(%TYPE%)" \
     "theory/builtin/theory_builtin_type_rules.h"
 well-founded FUNCTION_TYPE false
+enumerator FUNCTION_TYPE \
+    ::CVC4::theory::builtin::FunctionEnumerator \
+    "theory/builtin/type_enumerator.h"
 operator SEXPR_TYPE 0: "the type of a symbolic expression"
 cardinality SEXPR_TYPE \
     "::CVC4::theory::builtin::SExprProperties::computeCardinality(%TYPE%)" \
diff --git a/src/theory/builtin/theory_builtin_rewriter.cpp b/src/theory/builtin/theory_builtin_rewriter.cpp
index 4948216e5..edc934d0d 100644
--- a/src/theory/builtin/theory_builtin_rewriter.cpp
+++ b/src/theory/builtin/theory_builtin_rewriter.cpp
@@ -99,6 +99,34 @@ RewriteResponse TheoryBuiltinRewriter::postRewrite(TNode node) {
   }
 }
 
+TypeNode TheoryBuiltinRewriter::getFunctionTypeForArrayType(TypeNode atn,
+                                                            Node bvl)
+{
+  std::vector<TypeNode> children;
+  for (unsigned i = 0; i < bvl.getNumChildren(); i++)
+  {
+    Assert(atn.isArray());
+    Assert(bvl[i].getType() == atn.getArrayIndexType());
+    children.push_back(atn.getArrayIndexType());
+    atn = atn.getArrayConstituentType();
+  }
+  children.push_back(atn);
+  return NodeManager::currentNM()->mkFunctionType(children);
+}
+
+TypeNode TheoryBuiltinRewriter::getArrayTypeForFunctionType(TypeNode ftn)
+{
+  Assert(ftn.isFunction());
+  // construct the curried array type
+  unsigned nchildren = ftn.getNumChildren();
+  TypeNode ret = ftn[nchildren - 1];
+  for (int i = (static_cast<int>(nchildren) - 2); i >= 0; i--)
+  {
+    ret = NodeManager::currentNM()->mkArrayType(ftn[i], ret);
+  }
+  return ret;
+}
+
 Node TheoryBuiltinRewriter::getLambdaForArrayRepresentationRec( TNode a, TNode bvl, unsigned bvlIndex, 
                                                                 std::unordered_map< TNode, Node, TNodeHashFunction >& visited ){
   std::unordered_map< TNode, Node, TNodeHashFunction >::iterator it = visited.find( a );
diff --git a/src/theory/builtin/theory_builtin_rewriter.h b/src/theory/builtin/theory_builtin_rewriter.h
index 3667cf159..d729e5687 100644
--- a/src/theory/builtin/theory_builtin_rewriter.h
+++ b/src/theory/builtin/theory_builtin_rewriter.h
@@ -62,31 +62,58 @@ private:
   /** recursive helper for getArrayRepresentationForLambda */
   static Node getArrayRepresentationForLambdaRec( TNode n, bool reqConst, TypeNode retType );
 public:
-  /** 
-   * Given an array constant a, returns a lambda expression that it corresponds to, with bound variable list bvl. 
-   * Examples:
-   *
-   * (store (storeall (Array Int Int) 2) 0 1) 
-   * becomes
-   * ((lambda x. (ite (= x 0) 1 2))
-   *
-   * (store (storeall (Array Int (Array Int Int)) (storeall (Array Int Int) 4)) 0 (store (storeall (Array Int Int) 3) 1 2))
-   * becomes
-   * (lambda xy. (ite (= x 0) (ite (= x 1) 2 3) 4))
-   *
-   * (store (store (storeall (Array Int Bool) false) 2 true) 1 true)
-   * becomes
-   * (lambda x. (ite (= x 1) true (ite (= x 2) true false)))
-   *
-   * Notice that the return body of the lambda is rewritten to ensure that the representation is canonical. Hence the last
-   * example will in fact be returned as:
-   * (lambda x. (ite (= x 1) true (= x 2)))
+ /** Get function type for array type
+  *
+  * This returns the function type of terms returned by the function
+  * getLambdaForArrayRepresentation( t, bvl ),
+  * where t.getType()=atn.
+  *
+  * bvl should be a bound variable list whose variables correspond in-order
+  * to the index types of the (curried) Array type. For example, a bound
+  * variable list bvl whose variables have types (Int, Real) can be given as
+  * input when paired with atn = (Array Int (Array Real Bool)), or (Array Int
+  * (Array Real (Array Bool Bool))). This function returns (-> Int Real Bool)
+  * and (-> Int Real (Array Bool Bool)) respectively in these cases.
+  * On the other hand, the above bvl is not a proper input for
+  * atn = (Array Int (Array Bool Bool)) or (Array Int Int).
+  * If the types of bvl and atn do not match, we throw an assertion failure.
+  */
+ static TypeNode getFunctionTypeForArrayType(TypeNode atn, Node bvl);
+ /** Get array type for function type
+  *
+  * This returns the array type of terms returned by
+  * getArrayRepresentationForLambda( t ), where t.getType()=ftn.
+  */
+ static TypeNode getArrayTypeForFunctionType(TypeNode ftn);
+ /**
+  * Given an array constant a, returns a lambda expression that it corresponds
+  * to, with bound variable list bvl.
+  * Examples:
+  *
+  * (store (storeall (Array Int Int) 2) 0 1)
+  * becomes
+  * ((lambda x. (ite (= x 0) 1 2))
+  *
+  * (store (storeall (Array Int (Array Int Int)) (storeall (Array Int Int) 4)) 0
+  * (store (storeall (Array Int Int) 3) 1 2))
+  * becomes
+  * (lambda xy. (ite (= x 0) (ite (= x 1) 2 3) 4))
+  *
+  * (store (store (storeall (Array Int Bool) false) 2 true) 1 true)
+  * becomes
+  * (lambda x. (ite (= x 1) true (ite (= x 2) true false)))
+  *
+  * Notice that the return body of the lambda is rewritten to ensure that the
+  * representation is canonical. Hence the last
+  * example will in fact be returned as:
+  * (lambda x. (ite (= x 1) true (= x 2)))
+  */
+ static Node getLambdaForArrayRepresentation(TNode a, TNode bvl);
+ /** given a lambda expression n, returns an array term. reqConst is true if we
+  * require the return value to be a constant.
+   * This does the opposite direction of the examples described above.
    */
-  static Node getLambdaForArrayRepresentation( TNode a, TNode bvl );
-  /** given a lambda expression n, returns an array term. reqConst is true if we require the return value to be a constant. 
-    * This does the opposite direction of the examples described above.
-    */
-  static Node getArrayRepresentationForLambda( TNode n, bool reqConst = false );
+ static Node getArrayRepresentationForLambda(TNode n, bool reqConst = false);
 };/* class TheoryBuiltinRewriter */
 
 }/* CVC4::theory::builtin namespace */
diff --git a/src/theory/builtin/type_enumerator.cpp b/src/theory/builtin/type_enumerator.cpp
new file mode 100644
index 000000000..6435d3b4b
--- /dev/null
+++ b/src/theory/builtin/type_enumerator.cpp
@@ -0,0 +1,50 @@
+/*********************                                                        */
+/*! \file type_enumerator.cpp
+ ** \verbatim
+ ** Top contributors (to current version):
+ **   Andrew Reynolds
+ ** This file is part of the CVC4 project.
+ ** Copyright (c) 2009-2017 by the authors listed in the file AUTHORS
+ ** in the top-level source directory) and their institutional affiliations.
+ ** All rights reserved.  See the file COPYING in the top-level source
+ ** directory for licensing information.\endverbatim
+ **
+ ** \brief Enumerator for uninterpreted sorts and functions.
+ **
+ ** Enumerator for uninterpreted sorts and functions.
+ **/
+
+#include "theory/builtin/type_enumerator.h"
+
+namespace CVC4 {
+namespace theory {
+namespace builtin {
+
+FunctionEnumerator::FunctionEnumerator(TypeNode type,
+                                       TypeEnumeratorProperties* tep)
+    : TypeEnumeratorBase<FunctionEnumerator>(type),
+      d_arrayEnum(TheoryBuiltinRewriter::getArrayTypeForFunctionType(type), tep)
+{
+  Assert(type.getKind() == kind::FUNCTION_TYPE);
+  d_bvl = NodeManager::currentNM()->getBoundVarListForFunctionType(type);
+}
+
+Node FunctionEnumerator::operator*()
+{
+  if (isFinished())
+  {
+    throw NoMoreValuesException(getType());
+  }
+  Node a = *d_arrayEnum;
+  return TheoryBuiltinRewriter::getLambdaForArrayRepresentation(a, d_bvl);
+}
+
+FunctionEnumerator& FunctionEnumerator::operator++() throw()
+{
+  ++d_arrayEnum;
+  return *this;
+}
+
+} /* CVC4::theory::builtin namespace */
+} /* CVC4::theory namespace */
+} /* CVC4 namespace */
\ No newline at end of file
diff --git a/src/theory/builtin/type_enumerator.h b/src/theory/builtin/type_enumerator.h
index 1ab732710..e75443140 100644
--- a/src/theory/builtin/type_enumerator.h
+++ b/src/theory/builtin/type_enumerator.h
@@ -9,9 +9,9 @@
  ** All rights reserved.  See the file COPYING in the top-level source
  ** directory for licensing information.\endverbatim
  **
- ** \brief Enumerator for uninterpreted sorts
+ ** \brief Enumerator for uninterpreted sorts and functions.
  **
- ** Enumerator for uninterpreted sorts.
+ ** Enumerator for uninterpreted sorts and functions.
  **/
 
 #include "cvc4_private.h"
@@ -22,6 +22,7 @@
 #include "expr/kind.h"
 #include "expr/type_node.h"
 #include "expr/uninterpreted_constant.h"
+#include "theory/builtin/theory_builtin_rewriter.h"
 #include "theory/type_enumerator.h"
 #include "util/integer.h"
 
@@ -75,6 +76,30 @@ public:
 
 };/* class UninterpretedSortEnumerator */
 
+/** FunctionEnumerator
+* This enumerates function values, based on the enumerator for the
+* array type corresponding to the given function type.
+*/
+class FunctionEnumerator : public TypeEnumeratorBase<FunctionEnumerator>
+{
+ public:
+  FunctionEnumerator(TypeNode type, TypeEnumeratorProperties* tep = nullptr);
+  /** Get the current term of the enumerator. */
+  Node operator*() override;
+  /** Increment the enumerator. */
+  FunctionEnumerator& operator++() throw() override;
+  /** is the enumerator finished? */
+  bool isFinished() throw() override { return d_arrayEnum.isFinished(); }
+ private:
+  /** Enumerates arrays, which we convert to functions. */
+  TypeEnumerator d_arrayEnum;
+  /** The bound variable list for the function type we are enumerating.
+  * All terms output by this enumerator are of the form (LAMBDA d_bvl t) for
+  * some term t.
+  */
+  Node d_bvl;
+}; /* class FunctionEnumerator */
+
 }/* CVC4::theory::builtin namespace */
 }/* CVC4::theory namespace */
 }/* CVC4 namespace */