Finish iterator class

[satune.git] / src / AST / iterator.cc

diff --git a/src/AST/iterator.cc b/src/AST/iterator.cc
index ac81e96ad8fe8b91b1f86a1b4628e932725aa64c..2a18889dc2ad39564b40df6fc7e9222874b5e3ce 100644 (file)
--- 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;
+}