1 files changed, 82 insertions, 27 deletions

diff --git a/src/theory/datatypes/type_enumerator.h b/src/theory/datatypes/type_enumerator.h
index fe6a54ba7..76d49ecee 100644
--- a/src/theory/datatypes/type_enumerator.h
+++ b/src/theory/datatypes/type_enumerator.h
@@ -31,60 +31,115 @@ namespace theory {
 namespace datatypes {
 
 class DatatypesEnumerator : public TypeEnumeratorBase<DatatypesEnumerator> {
+  /** The datatype we're enumerating */
   const Datatype& d_datatype;
-  size_t d_count;
+  /** The datatype constructor we're currently enumerating */
   size_t d_ctor;
-  size_t d_numArgs;
-  size_t *d_argDistance;
+  /** The "first" constructor to consider; it's non-recursive */
+  size_t d_zeroCtor;
+  /** Delegate enumerators for the arguments of the current constructor */
   TypeEnumerator** d_argEnumerators;
 
+  /** Allocate and initialize the delegate enumerators */
+  void newEnumerators() {
+    d_argEnumerators = new TypeEnumerator*[d_datatype[d_ctor].getNumArgs()];
+    for(size_t a = 0; a < d_datatype[d_ctor].getNumArgs(); ++a) {
+      d_argEnumerators[a] = NULL;
+    }
+  }
+
+  /** Delete the delegate enumerators */
+  void deleteEnumerators() {
+    if(d_argEnumerators != NULL) {
+      for(size_t a = 0; a < d_datatype[d_ctor].getNumArgs(); ++a) {
+        delete d_argEnumerators[a];
+      }
+      delete [] d_argEnumerators;
+      d_argEnumerators = NULL;
+    }
+  }
+
 public:
 
   DatatypesEnumerator(TypeNode type) throw() :
     TypeEnumeratorBase(type),
     d_datatype(DatatypeType(type.toType()).getDatatype()),
-    d_count(0),
     d_ctor(0),
-    d_numArgs(0),
-    d_argDistance(NULL),
+    d_zeroCtor(0),
     d_argEnumerators(NULL) {
-    for(size_t c = 0; c < d_datatype.getNumConstructors(); ++c) {
-      if(d_datatype[c].getNumArgs() > d_numArgs) {
-        d_numArgs = d_datatype[c].getNumArgs();
+
+    /* find the "zero" constructor (the first non-recursive one) */
+    /* FIXME: this isn't sufficient for mutually-recursive datatypes! */
+    while(d_zeroCtor < d_datatype.getNumConstructors()) {
+      bool recursive = false;
+      for(size_t a = 0; a < d_datatype[d_zeroCtor].getNumArgs() && !recursive; ++a) {
+        recursive = (Node::fromExpr(d_datatype[d_zeroCtor][a].getSelector()).getType()[1] == type);
       }
+      if(!recursive) {
+        break;
+      }
+      ++d_zeroCtor;
     }
-    d_argDistance = new size_t[d_numArgs];
-    d_argEnumerators = new TypeEnumerator*[d_numArgs];
-    size_t a;
-    for(a = 0; a < d_datatype[0].getNumArgs(); ++a) {
-      d_argDistance[a] = 0;
-      d_argEnumerators[a] = new TypeEnumerator(TypeNode::fromType(d_datatype[0][a].getType()));
-    }
-    while(a < d_numArgs) {
-      d_argDistance[a] = 0;
-      d_argEnumerators[a++] = NULL;
-    }
+
+    /* start with the non-recursive constructor */
+    d_ctor = d_zeroCtor;
+
+    /* allocate space for the enumerators */
+    newEnumerators();
   }
 
   ~DatatypesEnumerator() throw() {
-    delete [] d_argDistance;
-    for(size_t a = 0; a < d_numArgs; ++a) {
-      delete d_argEnumerators[a];
-    }
-    delete [] d_argEnumerators;
+    deleteEnumerators();
   }
 
   Node operator*() throw(NoMoreValuesException) {
     if(d_ctor < d_datatype.getNumConstructors()) {
       DatatypeConstructor ctor = d_datatype[d_ctor];
-      return NodeManager::currentNM()->mkNode(kind::APPLY_CONSTRUCTOR, ctor.getConstructor());
+      NodeBuilder<> b(kind::APPLY_CONSTRUCTOR);
+      b << ctor.getConstructor();
+      try {
+        for(size_t a = 0; a < d_datatype[d_ctor].getNumArgs(); ++a) {
+          if(d_argEnumerators[a] == NULL) {
+            d_argEnumerators[a] = new TypeEnumerator(Node::fromExpr(d_datatype[d_ctor][a].getSelector()).getType()[1]);
+          }
+          b << **d_argEnumerators[a];
+        }
+      } catch(NoMoreValuesException&) {
+        InternalError();
+      }
+      return Node(b);
     } else {
       throw NoMoreValuesException(getType());
     }
   }
 
   DatatypesEnumerator& operator++() throw() {
-    ++d_ctor;
+    if(d_ctor < d_datatype.getNumConstructors()) {
+      for(size_t a = d_datatype[d_ctor].getNumArgs(); a > 0; --a) {
+        try {
+          *++*d_argEnumerators[a - 1];
+          return *this;
+        } catch(NoMoreValuesException&) {
+          *d_argEnumerators[a - 1] = TypeEnumerator(Node::fromExpr(d_datatype[d_ctor][a - 1].getSelector()).getType()[1]);
+        }
+      }
+
+      // Here, we need to step from the current constructor to the next one
+
+      // first, delete the current delegate enumerators
+      deleteEnumerators();
+
+      // Find the next constructor (only complicated by the notion of the "zero" constructor
+      d_ctor = (d_ctor == d_zeroCtor) ? 0 : d_ctor + 1;
+      if(d_ctor == d_zeroCtor) {
+        ++d_ctor;
+      }
+
+      // If we aren't out of constructors, allocate space for the new delegate enumerators
+      if(d_ctor < d_datatype.getNumConstructors()) {
+        newEnumerators();
+      }
+    }
     return *this;
   }