Simplification of the preregister and register throught a NodeVisitor class. The theoryOf is not all in one place, theory::theoryOf. The uninterpreted sorts belong to the builtin theory and are dispatched to the apropriate theory (QF_UF, QF_AX) through theoryOf based on the setting in the Theory class.

8 files changed, 146 insertions, 250 deletions

diff --git a/src/theory/arith/theory_arith.cpp b/src/theory/arith/theory_arith.cpp
index 93c02942d..3831536e9 100644
--- a/src/theory/arith/theory_arith.cpp
+++ b/src/theory/arith/theory_arith.cpp
@@ -281,6 +281,13 @@ void TheoryArith::preRegisterTerm(TNode n) {
       if(!left.getAttribute(Slack())){
         setupSlack(left);
       }
+    } else {
+      if (theoryOf(left) != THEORY_ARITH && !d_arithvarNodeMap.hasArithVar(left)) {
+        // The only way not to get it through pre-register is if it's a foreign term
+        ++(d_statistics.d_statUserVariables);
+        ArithVar av = requestArithVar(left, false);
+        setupInitialValue(av);
+      } 
     }
   }
   Debug("arith_preregister") << "end arith::preRegisterTerm("<< n <<")"<< endl;
@@ -306,7 +313,7 @@ ArithVar TheoryArith::requestArithVar(TNode x, bool basic){
   return varX;
 }
 
-void TheoryArith::asVectors(Polynomial& p, std::vector<Rational>& coeffs, std::vector<ArithVar>& variables) const{
+void TheoryArith::asVectors(Polynomial& p, std::vector<Rational>& coeffs, std::vector<ArithVar>& variables) {
   for(Polynomial::iterator i = p.begin(), end = p.end(); i != end; ++i){
     const Monomial& mono = *i;
     const Constant& constant = mono.getConstant();
@@ -317,10 +324,19 @@ void TheoryArith::asVectors(Polynomial& p, std::vector<Rational>& coeffs, std::v
     Debug("rewriter") << "should be var: " << n << endl;
 
     Assert(isLeaf(n));
-    Assert(d_arithvarNodeMap.hasArithVar(n));
-
-    ArithVar av = d_arithvarNodeMap.asArithVar(n);
+    Assert(theoryOf(n) != THEORY_ARITH || d_arithvarNodeMap.hasArithVar(n));
 
+    ArithVar av;
+    if (theoryOf(n) != THEORY_ARITH && !d_arithvarNodeMap.hasArithVar(n)) {
+      // The only way not to get it through pre-register is if it's a foreign term
+      ++(d_statistics.d_statUserVariables);
+      av = requestArithVar(n,false);
+      setupInitialValue(av);
+    } else {
+      // Otherwise, we already have it's variable
+      av = d_arithvarNodeMap.asArithVar(n);
+    }
+    
     coeffs.push_back(constant.getValue());
     variables.push_back(av);
   }
@@ -334,8 +350,6 @@ void TheoryArith::setupSlack(TNode left){
 
   d_rowHasBeenAdded = true;
 
-  ArithVar varSlack = requestArithVar(left, true);
-
   Polynomial polyLeft = Polynomial::parsePolynomial(left);
 
   vector<ArithVar> variables;
@@ -343,8 +357,8 @@ void TheoryArith::setupSlack(TNode left){
 
   asVectors(polyLeft, coefficients, variables);
 
+  ArithVar varSlack = requestArithVar(left, true);
   d_tableau.addRow(varSlack, coefficients, variables);
-
   setupInitialValue(varSlack);
 }
 
diff --git a/src/theory/arith/theory_arith.h b/src/theory/arith/theory_arith.h
index 255e2a304..8213dd47a 100644
--- a/src/theory/arith/theory_arith.h
+++ b/src/theory/arith/theory_arith.h
@@ -240,7 +240,7 @@ private:
 
   void asVectors(Polynomial& p,
                  std::vector<Rational>& coeffs,
-                 std::vector<ArithVar>& variables) const;
+                 std::vector<ArithVar>& variables);
 
   /** Routine for debugging. Print the assertions the theory is aware of. */
   void debugPrintAssertions();
diff --git a/src/theory/builtin/kinds b/src/theory/builtin/kinds
index f5195e256..d170469e0 100644
--- a/src/theory/builtin/kinds
+++ b/src/theory/builtin/kinds
@@ -228,6 +228,22 @@ sort BUILTIN_OPERATOR_TYPE \
     not-well-founded \
     "Built in type for built in operators"
 
+# Justified because we can have an unbounded-but-finite number of
+# sorts.  Assuming we have |Z| is probably ok for now..
+sort KIND_TYPE \
+    Cardinality::INTEGERS \
+    not-well-founded \
+    "Uninterpreted Sort"
+
+variable SORT_TAG "sort tag"
+parameterized SORT_TYPE SORT_TAG 0: "sort type"
+# This is really "unknown" cardinality, but maybe this will be good
+# enough (for now) ?  Once we support quantifiers, maybe reconsider
+# this..
+cardinality SORT_TYPE "Cardinality(Cardinality::INTEGERS)"
+well-founded SORT_TYPE false
+
+
 # A kind representing "inlined" operators defined with OPERATOR
 # Conceptually, (EQUAL a b) is actually an (APPLY EQUAL a b), but it's
 # not stored that way.  If you ask for the operator of (EQUAL a b),
diff --git a/src/theory/theory.cpp b/src/theory/theory.cpp
index b1eb195c7..576e0188b 100644
--- a/src/theory/theory.cpp
+++ b/src/theory/theory.cpp
@@ -26,6 +26,9 @@ using namespace std;
 namespace CVC4 {
 namespace theory {
 
+/** Default value for the uninterpreted sorts is the UF theory */
+TheoryId Theory::d_uninterpretedSortOwner = THEORY_UF;
+
 std::ostream& operator<<(std::ostream& os, Theory::Effort level){
   switch(level){
   case Theory::MIN_EFFORT:
diff --git a/src/theory/theory.h b/src/theory/theory.h
index 93d78f57c..1d8797d1f 100644
--- a/src/theory/theory.h
+++ b/src/theory/theory.h
@@ -40,11 +40,6 @@ class TheoryEngine;
 
 namespace theory {
 
-/** Tag for the "newFact()-has-been-called-in-this-context" flag on Nodes */
-struct AssertedAttrTag {};
-/** The "newFact()-has-been-called-in-this-context" flag on Nodes */
-typedef CVC4::expr::CDAttribute<AssertedAttrTag, bool> Asserted;
-
 /**
  * Base class for T-solvers.  Abstract DPLL(T).
  *
@@ -178,31 +173,51 @@ protected:
     return d_facts.end();
   }
 
+  /**
+   * The theory that owns the uninterpreted sort.
+   */
+  static TheoryId d_uninterpretedSortOwner;
+
 public:
 
   /**
    * Return the ID of the theory responsible for the given type.
    */
   static inline TheoryId theoryOf(TypeNode typeNode) {
+    TheoryId id;
     if (typeNode.getKind() == kind::TYPE_CONSTANT) {
-      return typeConstantToTheoryId(typeNode.getConst<TypeConstant>());
+      id = typeConstantToTheoryId(typeNode.getConst<TypeConstant>());
     } else {
-      return kindToTheoryId(typeNode.getKind());
+      id = kindToTheoryId(typeNode.getKind());
+    }
+    if (id == THEORY_BUILTIN) {
+      return d_uninterpretedSortOwner;
     }
+    return id;
   }
 
+
   /**
    * Returns the ID of the theory responsible for the given node.
    */
   static inline TheoryId theoryOf(TNode node) {
-    if (node.getMetaKind() == kind::metakind::VARIABLE ||
-        node.getMetaKind() == kind::metakind::CONSTANT) {
-      // Constants, variables, 0-ary constructors
+    // Constants, variables, 0-ary constructors
+    if (node.getMetaKind() == kind::metakind::VARIABLE || node.getMetaKind() == kind::metakind::CONSTANT) {
       return theoryOf(node.getType());
-    } else {
-      // Regular nodes
-      return kindToTheoryId(node.getKind());
     }
+    // Equality is owned by the theory that owns the domain
+    if (node.getKind() == kind::EQUAL) {
+      return theoryOf(node[0].getType());
+    }
+    // Regular nodes are owned by the kind
+    return kindToTheoryId(node.getKind());
+  }
+
+  /**
+   * Set the owner of the uninterpreted sort.
+   */
+  static void setUninterpretedSortOwner(TheoryId theory) {
+    d_uninterpretedSortOwner = theory;
   }
 
   /**
diff --git a/src/theory/theory_engine.cpp b/src/theory/theory_engine.cpp
index b1b4d67dd..833c03e2f 100644
--- a/src/theory/theory_engine.cpp
+++ b/src/theory/theory_engine.cpp
@@ -29,6 +29,7 @@
 #include "theory/rewriter.h"
 #include "theory/theory_traits.h"
 
+#include "util/node_visitor.h"
 #include "util/ite_removal.h"
 
 using namespace std;
@@ -36,98 +37,28 @@ using namespace std;
 using namespace CVC4;
 using namespace CVC4::theory;
 
-namespace CVC4 {
-
-namespace theory {
-
 /** Tag for the "registerTerm()-has-been-called" flag on Nodes */
 struct RegisteredAttrTag {};
 /** The "registerTerm()-has-been-called" flag on Nodes */
-typedef CVC4::expr::CDAttribute<RegisteredAttrTag, bool> RegisteredAttr;
+typedef expr::CDAttribute<RegisteredAttrTag, bool> RegisteredAttr;
 
+/** Tag for the "preregisterTerm()-has-benn-called" flag on nodes */
 struct PreRegisteredAttrTag {};
-typedef expr::Attribute<PreRegisteredAttrTag, bool> PreRegistered;
-
-}/* CVC4::theory namespace */
+/** The "preregisterTerm()-has-been-called" flag on Nodes */
+typedef expr::Attribute<PreRegisteredAttrTag, bool> PreRegisteredAttr;
 
 void TheoryEngine::EngineOutputChannel::newFact(TNode fact) {
   TimerStat::CodeTimer codeTimer(d_newFactTimer);
 
-  //FIXME: Assert(fact.isLiteral(), "expected literal");
   if (fact.getKind() == kind::NOT) {
     // No need to register negations - only atoms
     fact = fact[0];
-// FIXME: In future, might want to track disequalities in shared term manager
-//     if (fact.getKind() == kind::EQUAL) {
-//       d_engine->getSharedTermManager()->addDiseq(fact);
-//     }
   }
-  else if (fact.getKind() == kind::EQUAL) {
-    // Automatically track all asserted equalities in the shared term manager
-    d_engine->getSharedTermManager()->addEq(fact);
-  }
-
-  if(Options::current()->theoryRegistration &&
-     !fact.getAttribute(RegisteredAttr()) &&
-     d_engine->d_needRegistration) {
-    list<TNode> toReg;
-    toReg.push_back(fact);
-
-    Trace("theory") << "Theory::get(): registering new atom" << endl;
-
-    /* Essentially this is doing a breadth-first numbering of
-     * non-registered subterms with children.  Any non-registered
-     * leaves are immediately registered. */
-    for(list<TNode>::iterator workp = toReg.begin();
-        workp != toReg.end();
-        ++workp) {
-
-      TNode n = *workp;
-      Theory* thParent = d_engine->theoryOf(n);
-
-      for(TNode::iterator i = n.begin(); i != n.end(); ++i) {
-        TNode c = *i;
-        Theory* thChild = d_engine->theoryOf(c);
 
-        if (thParent != thChild) {
-          d_engine->getSharedTermManager()->addTerm(c, thParent, thChild);
-        }
-        if(! c.getAttribute(RegisteredAttr())) {
-          if(c.getNumChildren() == 0) {
-            c.setAttribute(RegisteredAttr(), true);
-            thChild->registerTerm(c);
-          } else {
-            toReg.push_back(c);
-          }
-        }
-      }
-    }
-
-    /* Now register the list of terms in reverse order.  Between this
-     * and the above registration of leaves, this should ensure that
-     * all subterms in the entire tree were registered in
-     * reverse-topological order. */
-    for(list<TNode>::reverse_iterator i = toReg.rbegin();
-        i != toReg.rend();
-        ++i) {
-
-      TNode n = *i;
-
-      /* Note that a shared TNode in the DAG rooted at "fact" could
-       * appear twice on the list, so we have to avoid hitting it
-       * twice. */
-      // FIXME when ExprSets are online, use one of those to avoid
-      // duplicates in the above?
-      // Actually, that doesn't work because you have to make sure
-      // that the *last* occurrence is the one that gets processed first @CB
-      // This could be a big performance problem though because it requires
-      // traversing a DAG as a tree and that can really blow up @CB
-      if(! n.getAttribute(RegisteredAttr())) {
-        n.setAttribute(RegisteredAttr(), true);
-        d_engine->theoryOf(n)->registerTerm(n);
-      }
-    }
-  }/* Options::current()->theoryRegistration && !fact.getAttribute(RegisteredAttr()) */
+  if(Options::current()->theoryRegistration && !fact.getAttribute(RegisteredAttr()) && d_engine->d_needRegistration) {
+    RegisterVisitor visitor(*d_engine);
+    NodeVisitor<RegisterVisitor, RegisteredAttr>::run(visitor, fact);
+  }
 }
 
 TheoryEngine::TheoryEngine(context::Context* ctxt) :
@@ -162,6 +93,12 @@ TheoryEngine::~TheoryEngine() {
   delete d_sharedTermManager;
 }
 
+void TheoryEngine::setLogic(std::string logic) {
+  Assert(d_logic.empty());
+  // Set the logic
+  d_logic = logic;
+}
+
 struct preregister_stack_element {
   TNode node;
   bool children_added;
@@ -170,65 +107,8 @@ struct preregister_stack_element {
 };/* struct preprocess_stack_element */
 
 void TheoryEngine::preRegister(TNode preprocessed) {
-  // If we are pre-registered already we are done
-  if (preprocessed.getAttribute(PreRegistered())) {
-    return;
-  }
-
-  // Do a topological sort of the subexpressions and preregister them
-  vector<preregister_stack_element> toVisit;
-  toVisit.push_back((TNode) preprocessed);
-  while (!toVisit.empty()) {
-    preregister_stack_element& stackHead = toVisit.back();
-    // The current node we are processing
-    TNode current = stackHead.node;
-    // If we already added all the children its time to register or just
-    // pop from the stack
-    if (stackHead.children_added || current.getAttribute(PreRegistered())) {
-      if (!current.getAttribute(PreRegistered())) {
-        // Mark it as registered
-        current.setAttribute(PreRegistered(), true);
-        // Register this node
-        if (current.getKind() == kind::EQUAL) {
-          if(d_logic == "QF_AX") {
-            d_theoryTable[theory::THEORY_ARRAY]->preRegisterTerm(current);
-          } else {
-            Theory* theory = theoryOf(current);
-            TheoryId theoryLHS = theory->getId();
-            Trace("register") << "preregistering " << current
-                              << " with " << theoryLHS << std::endl;
-            markActive(theoryLHS);
-            theory->preRegisterTerm(current);
-          }
-        } else {
-          TheoryId theory = theoryIdOf(current);
-          Trace("register") << "preregistering " << current
-                            << " with " << theory << std::endl;
-          markActive(theory);
-          d_theoryTable[theory]->preRegisterTerm(current);
-          TheoryId typeTheory = theoryIdOf(current.getType());
-          if (theory != typeTheory) {
-            Trace("register") << "preregistering " << current
-                              << " with " << typeTheory << std::endl;
-            markActive(typeTheory);
-            d_theoryTable[typeTheory]->preRegisterTerm(current);
-          }
-        }
-      }
-      // Done with this node, remove from the stack
-      toVisit.pop_back();
-    } else {
-      // Mark that we have added the children
-      stackHead.children_added = true;
-      // We need to add the children
-      for(TNode::iterator child_it = current.begin(); child_it != current.end(); ++ child_it) {
-        TNode childNode = *child_it;
-        if (!childNode.getAttribute(PreRegistered())) {
-          toVisit.push_back(childNode);
-        }
-      }
-    }
-  }
+  PreRegisterVisitor visitor(*this);
+  NodeVisitor<PreRegisterVisitor, PreRegisteredAttr>::run(visitor, preprocessed);
 }
 
 /**
@@ -409,7 +289,7 @@ Node TheoryEngine::preprocess(TNode assertion) {
     }
 
     // If this is an atom, we preprocess it with the theory
-    if (theoryIdOf(current) != THEORY_BOOL) {
+    if (Theory::theoryOf(current) != THEORY_BOOL) {
       d_atomPreprocessingCache[current] = theoryOf(current)->preprocess(current);
       continue;
     }
@@ -455,5 +335,3 @@ Node TheoryEngine::preprocess(TNode assertion) {
   return d_atomPreprocessingCache[assertion];
 }
 
-
-}/* CVC4 namespace */
diff --git a/src/theory/theory_engine.h b/src/theory/theory_engine.h
index 1e5653564..c472a1c41 100644
--- a/src/theory/theory_engine.h
+++ b/src/theory/theory_engine.h
@@ -233,10 +233,11 @@ class TheoryEngine {
 
   bool hasRegisterTerm(theory::TheoryId th) const;
 
-public:
   /** The logic of the problem */
   std::string d_logic;
 
+public:
+
   /** Constructs a theory engine */
   TheoryEngine(context::Context* ctxt);
 
@@ -249,12 +250,16 @@ public:
    */
   template <class TheoryClass>
   void addTheory() {
-    TheoryClass* theory =
-      new TheoryClass(d_context, d_theoryOut, theory::Valuation(this));
+    TheoryClass* theory = new TheoryClass(d_context, d_theoryOut, theory::Valuation(this));
     d_theoryTable[theory->getId()] = theory;
     d_sharedTermManager->registerTheory(static_cast<TheoryClass*>(theory));
   }
 
+  /**
+   * Set's the logic (smt-lib format). All theory specific setup/hacks should go in here.
+   */
+  void setLogic(std::string logic);
+
   SharedTermManager* getSharedTermManager() {
     return d_sharedTermManager;
   }
@@ -312,45 +317,7 @@ public:
    * of built-in type.
    */
   inline theory::Theory* theoryOf(TNode node) {
-    if (node.getKind() == kind::EQUAL) {
-      return d_theoryTable[theoryIdOf(node[0])];
-    } else {
-      return d_theoryTable[theoryIdOf(node)];
-    }
-  }
-
-  /**
-   * Wrapper for theory::Theory::theoryOf() that implements the
-   * array/EUF hack.
-   */
-  inline theory::TheoryId theoryIdOf(TNode node) {
-    theory::TheoryId id = theory::Theory::theoryOf(node);
-    if(d_logic == "QF_AX" && id == theory::THEORY_UF) {
-      id = theory::THEORY_ARRAY;
-    }
-    return id;
-  }
-
-  /**
-   * Get the theory associated to a given Node.
-   *
-   * @returns the theory, or NULL if the TNode is
-   * of built-in type.
-   */
-  inline theory::Theory* theoryOf(const TypeNode& typeNode) {
-    return d_theoryTable[theoryIdOf(typeNode)];
-  }
-
-  /**
-   * Wrapper for theory::Theory::theoryOf() that implements the
-   * array/EUF hack.
-   */
-  inline theory::TheoryId theoryIdOf(const TypeNode& typeNode) {
-    theory::TheoryId id = theory::Theory::theoryOf(typeNode);
-    if(d_logic == "QF_AX" && id == theory::THEORY_UF) {
-      id = theory::THEORY_ARRAY;
-    }
-    return id;
+    return d_theoryTable[theory::Theory::theoryOf(node)];
   }
 
   /**
@@ -380,30 +347,12 @@ public:
   inline void assertFact(TNode node) {
     Trace("theory") << "TheoryEngine::assertFact(" << node << ")" << std::endl;
 
-    // Mark it as asserted in this context
-    //
-    // [MGD] removed for now, this appears to have a nontrivial
-    // performance penalty
-    // node.setAttribute(theory::Asserted(), true);
-
     // Get the atom
     TNode atom = node.getKind() == kind::NOT ? node[0] : node;
 
-    // Again, equality is a special case
-    if (atom.getKind() == kind::EQUAL) {
-      if(d_logic == "QF_AX") {
-        Trace("theory") << "TheoryEngine::assertFact QF_AX logic; everything goes to Arrays" << std::endl;
-        d_theoryTable[theory::THEORY_ARRAY]->assertFact(node);
-      } else {
-        theory::Theory* theory = theoryOf(atom);
-        Trace("theory") << "asserting " << node << " to " << theory->getId() << std::endl;
-        theory->assertFact(node);
-      }
-    } else {
-      theory::Theory* theory = theoryOf(atom);
-      Trace("theory") << "asserting " << node << " to " << theory->getId() << std::endl;
-      theory->assertFact(node);
-    }
+    theory::Theory* theory = theoryOf(atom);
+    Trace("theory") << "asserting " << node << " to " << theory->getId() << std::endl;
+    theory->assertFact(node);
   }
 
   /**
@@ -458,16 +407,7 @@ public:
   inline Node getExplanation(TNode node) {
     d_theoryOut.d_explanationNode = Node::null();
     TNode atom = node.getKind() == kind::NOT ? node[0] : node;
-    if (atom.getKind() == kind::EQUAL) {
-      if(d_logic == "QF_AX") {
-        Trace("theory") << "TheoryEngine::assertFact QF_AX logic; everything goes to Arrays" << std::endl;
-        d_theoryTable[theory::THEORY_ARRAY]->explain(node);
-      } else {
-        theoryOf(atom[0])->explain(node);
-      }
-    } else {
-      theoryOf(atom)->explain(node);
-    }
+    theoryOf(atom)->explain(node);
     Assert(!d_theoryOut.d_explanationNode.get().isNull());
     return d_theoryOut.d_explanationNode;
   }
@@ -501,6 +441,52 @@ private:
   };/* class TheoryEngine::Statistics */
   Statistics d_statistics;
 
+  ///////////////////////////
+  // Visitors
+  ///////////////////////////
+
+  /**
+   * Visitor that calls the apropriate theory to preregister the term.
+   */
+  class PreRegisterVisitor {
+
+    /** The engine */
+    TheoryEngine& d_engine;
+
+  public:
+
+    PreRegisterVisitor(TheoryEngine& engine): d_engine(engine) {}
+
+    void operator () (TNode current) {
+      // Get the theory of the term and pre-register it with the owning theory
+      theory::TheoryId currentTheoryId = theory::Theory::theoryOf(current);
+      Debug("register") << "preregistering " << current << " with " << currentTheoryId << std::endl;
+      d_engine.markActive(currentTheoryId);
+      theory::Theory* currentTheory = d_engine.theoryOf(current);
+      currentTheory->preRegisterTerm(current);
+    }
+  };
+
+  /**
+   * Visitor that calls the apropriate theory to preregister the term.
+   */
+  class RegisterVisitor {
+
+    /** The engine */
+    TheoryEngine& d_engine;
+
+  public:
+
+    RegisterVisitor(TheoryEngine& engine): d_engine(engine) {}
+
+    void operator () (TNode current) {
+      // Register this node to it's theory
+      theory::Theory* currentTheory = d_engine.theoryOf(current);
+      Debug("register") << "registering " << current << " with " << currentTheory->getId() << std::endl;
+      currentTheory->registerTerm(current);
+    }
+  };
+
 };/* class TheoryEngine */
 
 }/* CVC4 namespace */
diff --git a/src/theory/uf/kinds b/src/theory/uf/kinds
index 9498fa6fb..6cec23385 100644
--- a/src/theory/uf/kinds
+++ b/src/theory/uf/kinds
@@ -10,22 +10,6 @@ properties stable-infinite
 properties check propagate staticLearning presolve notifyRestart
 
 rewriter ::CVC4::theory::uf::TheoryUfRewriter "theory/uf/theory_uf_rewriter.h"
-
-# Justified because we can have an unbounded-but-finite number of
-# sorts.  Assuming we have |Z| is probably ok for now..
-sort KIND_TYPE \
-    Cardinality::INTEGERS \
-    not-well-founded \
-    "Uninterpreted Sort"
-
 parameterized APPLY_UF VARIABLE 1: "uninterpreted function application"
 
-variable SORT_TAG "sort tag"
-parameterized SORT_TYPE SORT_TAG 0: "sort type"
-# This is really "unknown" cardinality, but maybe this will be good
-# enough (for now) ?  Once we support quantifiers, maybe reconsider
-# this..
-cardinality SORT_TYPE "Cardinality(Cardinality::INTEGERS)"
-well-founded SORT_TYPE false
-
 endtheory