From 18aa5690b5afa41ce2498a28adba27b37df7816d Mon Sep 17 00:00:00 2001
From: bdemsky <bdemsky@uci.edu>
Date: Thu, 7 Sep 2017 18:41:59 -0700
Subject: [PATCH] Finish iterator class

---
 src/AST/iterator.cc            | 81 +++++++++++++++++++++++++++++++++-
 src/AST/iterator.h             | 15 ++++++-
 src/ASTTransform/preprocess.cc |  2 +-
 src/ASTTransform/preprocess.h  | 20 ++++-----
 src/Collections/structs.h      |  1 +
 5 files changed, 106 insertions(+), 13 deletions(-)

diff --git a/src/AST/iterator.cc b/src/AST/iterator.cc
index ac81e96..2a18889 100644
--- a/src/AST/iterator.cc
+++ b/src/AST/iterator.cc
@@ -1,5 +1,6 @@
 #include "iterator.h"
 #include "boolean.h"
+#include "element.h"
 #include "csolver.h"
 
 BooleanIterator::BooleanIterator(CSolver * _solver) :
@@ -50,7 +51,7 @@ void BooleanIterator::updateNext() {
 			index.push(topindex+1);
 			Boolean *newchild=logicop->inputs.get(topindex).getBoolean();
 			if (discovered.add(newchild)) {
-				boolean.push(logicop->inputs.get(topindex).getBoolean());
+				boolean.push(newchild);
 				index.push(0);
 			}
 			break;
@@ -66,3 +67,81 @@ Boolean * BooleanIterator::next() {
 	updateNext();
 	return b;
 }
+
+ElementIterator::ElementIterator(CSolver *_solver) :
+	bit(_solver),
+	base(NULL),
+	baseindex(0) {
+	
+}
+
+ElementIterator::~ElementIterator() {
+}
+
+bool ElementIterator::hasNext() {
+	return element.getSize() != 0;
+}
+
+void ElementIterator::updateNext() {
+	if (element.getSize() != 0) {
+		element.pop();
+		index.pop();
+	}
+	
+	while(true) {
+		if (element.getSize() == 0) {
+			if (base != NULL) {
+				if (baseindex == base->inputs.getSize()) {
+					base = NULL;
+					continue;
+				} else {
+					Element * e = base->inputs.get(baseindex);
+					baseindex++;
+					if (discovered.add(e)) {
+						element.push(e);
+						index.push(0);
+					}
+				}
+			} else {
+				if (bit.hasNext()) {
+					Boolean *b=bit.next();
+					if (b->type == PREDICATEOP) {
+						base = (BooleanPredicate *)b;
+						baseindex=0;
+					}
+					continue;
+				} else
+					return;
+			}
+		}
+		Element *topelement=element.last();
+		uint topindex=index.last();
+		switch(topelement->type) {
+		case ELEMSET:
+		case ELEMCONST:
+			return;
+		case ELEMFUNCRETURN: {
+			ElementFunction * func=(ElementFunction*) topelement;
+			uint size=func->inputs.getSize();
+			if (topindex == size)
+				return;
+			index.pop();
+			index.push(topindex+1);
+			Element *newchild=func->inputs.get(topindex);
+			if (discovered.add(newchild)) {
+				element.push(newchild);
+				index.push(0);
+			}
+			break;
+		}
+		default:
+			ASSERT(0);
+		}
+	}
+}
+
+Element * ElementIterator::next() {
+	Element * e = element.last();
+	updateNext();
+	return e;
+}
diff --git a/src/AST/iterator.h b/src/AST/iterator.h
index b82ee37..0e7303f 100644
--- a/src/AST/iterator.h
+++ b/src/AST/iterator.h
@@ -10,6 +10,7 @@ class BooleanIterator {
 	bool hasNext();
 	Boolean * next();
 	CMEMALLOC;
+
  private:
 	SetIteratorBooleanEdge * solverit;
 	HashsetBoolean discovered;
@@ -20,10 +21,22 @@ class BooleanIterator {
 
 class ElementIterator {
  public:
+	ElementIterator(CSolver *_solver);
+	~ElementIterator();
+	bool hasNext();
+	Element * next();
 	CMEMALLOC;
+
  private:
 	BooleanIterator bit;
-};
+	BooleanPredicate *base;
+	uint baseindex;
 
+	HashsetElement discovered;
+
+	Vector<Element *> element;
+	Vector<uint> index;
+	void updateNext();
+};
 
 #endif
diff --git a/src/ASTTransform/preprocess.cc b/src/ASTTransform/preprocess.cc
index a23d918..1e04fd0 100644
--- a/src/ASTTransform/preprocess.cc
+++ b/src/ASTTransform/preprocess.cc
@@ -5,7 +5,7 @@
 #include "iterator.h"
 
 Preprocess::Preprocess(CSolver *_solver)
-        : Transform(_solver)
+	: Transform(_solver)
 {
 }
 
diff --git a/src/ASTTransform/preprocess.h b/src/ASTTransform/preprocess.h
index 1272266..ea9800d 100644
--- a/src/ASTTransform/preprocess.h
+++ b/src/ASTTransform/preprocess.h
@@ -4,16 +4,16 @@
 #include "transform.h"
 
 class Preprocess : public Transform {
-	public:
-        Preprocess(CSolver *_solver);
-        ~Preprocess();
-        void doTransform();
-
-        CMEMALLOC;
-private:
-				HashsetBoolean toremove;
-				void processBooleanVar(BooleanVar * b);
-				void resolveBooleanVars();
+ public:
+	Preprocess(CSolver *_solver);
+	~Preprocess();
+	void doTransform();
+	
+	CMEMALLOC;
+ private:
+	HashsetBoolean toremove;
+	void processBooleanVar(BooleanVar * b);
+	void resolveBooleanVars();
 };
 
 #endif
diff --git a/src/Collections/structs.h b/src/Collections/structs.h
index 801cb87..e6f451e 100644
--- a/src/Collections/structs.h
+++ b/src/Collections/structs.h
@@ -24,6 +24,7 @@ typedef Hashset<OrderNode *, uintptr_t, PTRSHIFT, order_node_hash_function, orde
 typedef Hashset<OrderEdge *, uintptr_t, PTRSHIFT, order_edge_hash_function, order_edge_equals> HashsetOrderEdge;
 typedef Hashset<OrderElement *, uintptr_t, PTRSHIFT, order_element_hash_function, order_element_equals> HashsetOrderElement;
 typedef Hashset<Boolean *, uintptr_t, PTRSHIFT> HashsetBoolean;
+typedef Hashset<Element *, uintptr_t, PTRSHIFT> HashsetElement;
 typedef SetIterator<Boolean *, uintptr_t, PTRSHIFT> SetIteratorBoolean;
 
 typedef Hashtable<OrderNode *, HashsetOrderNode *, uintptr_t, PTRSHIFT> HashtableNodeToNodeSet;
-- 
2.34.1