diff --git a/src/api/MainSolver.cc b/src/api/MainSolver.cc
index dd65cda5e..833cd1987 100644
--- a/src/api/MainSolver.cc
+++ b/src/api/MainSolver.cc
@@ -62,8 +62,8 @@ MainSolver::pop()
             for (int i = 0; i < partitionsToInvalidate.size(); ++i) {
                 PTRef part = partitionsToInvalidate[i];
                 auto index = pmanager.getPartitionIndex(part);
-                assert(index != -1);
-                opensmt::setbit(mask, static_cast<unsigned int>(index));
+                assert(index != PartIdx_Undef);
+                opensmt::setbit(mask, static_cast<unsigned int>(index.id));
             }
             pmanager.invalidatePartitions(mask);
         }
@@ -121,7 +121,7 @@ MainSolver::insertFormula(PTRef root, char** msg)
         //     thus we need the old value of count. TODO: Find a good interface for this so it cannot be broken this easily
         unsigned int partition_index = inserted_formulas_count++;
         pmanager.assignTopLevelPartitionIndex(partition_index, root);
-        assert(pmanager.getPartitionIndex(root) != -1);
+        assert(pmanager.getPartitionIndex(root) != PartIdx_Undef);
     }
     else {
         ++inserted_formulas_count;
@@ -156,22 +156,21 @@ sstat MainSolver::simplifyFormulas(char** err_msg)
             for (int j = 0; j < flas.size() && status != s_False; ++j) {
                 PTRef fla = flas[j];
                 if (fla == logic.getTerm_true()) { continue; }
-                assert(pmanager.getPartitionIndex(fla) != -1);
+                assert(pmanager.getPartitionIndex(fla) != PartIdx_Undef);
                 // Optimize the dag for cnfization
                 if (logic.isBooleanOperator(fla)) {
                     PTRef old = fla;
                     fla = rewriteMaxArity(fla);
                     pmanager.transferPartitionMembership(old, fla);
                 }
-                assert(pmanager.getPartitionIndex(fla) != -1);
+                assert(pmanager.getPartitionIndex(fla) != PartIdx_Undef);
                 pmanager.propagatePartitionMask(fla);
-                getSMTSolver().setPartition(pmanager.getPartitionIndex(fla));
-                status = giveToSolver(fla, frame.getId());
+                status = giveToSolver(fla, frame.getId(), pmanager.getPartitionIndex(fla));
             }
         } else {
             PTRef root = frame.root;
             if (logic.isFalse(root)) {
-                giveToSolver(getLogic().getTerm_false(), frame.getId());
+                giveToSolver(getLogic().getTerm_false(), frame.getId(), PartIdx_Undef);
                 status = s_False;
                 break;
             }
@@ -180,7 +179,7 @@ sstat MainSolver::simplifyFormulas(char** err_msg)
                 root = rewriteMaxArity(root);
             }
             root_instance.setRoot(root);
-            status = giveToSolver(root, frame.getId());
+            status = giveToSolver(root, frame.getId(), PartIdx_Undef);
         }
     }
     if (status == s_False) {
diff --git a/src/api/MainSolver.h b/src/api/MainSolver.h
index 7bf463d6a..2e785981c 100644
--- a/src/api/MainSolver.h
+++ b/src/api/MainSolver.h
@@ -133,8 +133,8 @@ class MainSolver
     sstat simplifyFormulas(char** msg);
     sstat solve           ();
 
-    sstat giveToSolver(PTRef root, FrameId push_id) {
-        if (ts.cnfizeAndGiveToSolver(root, push_id) == l_False) return s_False;
+    sstat giveToSolver(PTRef root, FrameId push_id, PartIdx partitionIndex) {
+        if (ts.cnfizeAndGiveToSolver(root, push_id, partitionIndex) == l_False) return s_False;
         return s_Undef; }
 
     PTRef rewriteMaxArity(PTRef);
diff --git a/src/api/PartitionManager.h b/src/api/PartitionManager.h
index ec8177783..676844f56 100644
--- a/src/api/PartitionManager.h
+++ b/src/api/PartitionManager.h
@@ -58,7 +58,7 @@ class PartitionManager {
         partitionInfo.transferPartitionMembership(old, new_ptref);
     }
 
-    int getPartitionIndex(PTRef ref) const {
+    PartIdx getPartitionIndex(PTRef ref) const {
         return partitionInfo.getPartitionIndex(ref);
     }
 };
diff --git a/src/cnfizers/Cnfizer.cc b/src/cnfizers/Cnfizer.cc
index 77316d604..50af144d2 100644
--- a/src/cnfizers/Cnfizer.cc
+++ b/src/cnfizers/Cnfizer.cc
@@ -53,11 +53,11 @@ void Cnfizer::initialize()
     vec<Lit> c;
     Lit l = this->getOrCreateLiteralFor (logic.getTerm_true());
     c.push (l);
-    addClause(c);
+    addClause(c, PartIdx_Default); // Check what's the "default" partition!
     c.pop();
     l = this->getOrCreateLiteralFor (logic.getTerm_false());
     c.push (~l);
-    addClause(c);
+    addClause(c, PartIdx_Default); // Check what's the "default" partition!
 }
 
 lbool
@@ -132,7 +132,7 @@ void Cnfizer::setFrameTerm(FrameId frame_id)
 //
 // Main Routine. Examine formula and give it to the solver
 //
-lbool Cnfizer::cnfizeAndGiveToSolver(PTRef formula, FrameId frame_id)
+lbool Cnfizer::cnfizeAndGiveToSolver(PTRef formula, FrameId frame_id, PartIdx partitionIndex)
 {
     // Get the variable for the incrementality.
     setFrameTerm(frame_id);
@@ -174,7 +174,7 @@ lbool Cnfizer::cnfizeAndGiveToSolver(PTRef formula, FrameId frame_id)
 #ifdef PEDANTIC_DEBUG
             cerr << " => Already in CNF" << endl;
 #endif
-            res = giveToSolver (f);
+            res = giveToSolver (f, partitionIndex);
         }
 
         // Check whether it can be rewritten using deMorgan laws
@@ -184,14 +184,14 @@ lbool Cnfizer::cnfizeAndGiveToSolver(PTRef formula, FrameId frame_id)
 #ifdef PEDANTIC_DEBUG
             cout << " => Will be de Morganized" << endl;
 #endif
-            res = deMorganize (f);
+            res = deMorganize (f, partitionIndex);
         }
         else
         {
 #ifdef PEDANTIC_DEBUG
             cout << " => proper cnfization" << endl;
 #endif // PEDANTIC_DEBUG
-            res = cnfizeAndAssert (f); // Perform actual cnfization (implemented in subclasses)
+            res = cnfizeAndAssert (f, partitionIndex); // Perform actual cnfization (implemented in subclasses)
         }
         alreadyAsserted.insert(f, frame_term);
     }
@@ -199,7 +199,7 @@ lbool Cnfizer::cnfizeAndGiveToSolver(PTRef formula, FrameId frame_id)
     if (res) {
         vec<PTRef> nestedBoolRoots = getNestedBoolRoots(formula);
         for (int i = 0; i < nestedBoolRoots.size(); ++i) {
-            res &= cnfize(nestedBoolRoots[i]); // cnfize the formula without asserting the top level
+            res &= cnfize(nestedBoolRoots[i], partitionIndex); // cnfize the formula without asserting the top level
         }
         assert(res);
         declareVars(logic.propFormulasAppearingInUF);
@@ -208,7 +208,7 @@ lbool Cnfizer::cnfizeAndGiveToSolver(PTRef formula, FrameId frame_id)
     return res == false ? l_False : l_Undef;
 }
 
-bool Cnfizer::cnfizeAndAssert(PTRef formula) {
+bool Cnfizer::cnfizeAndAssert(PTRef formula, PartIdx partitionIndex) {
     assert(formula != PTRef_Undef);
     // Top level formula must not be and anymore
     assert(!logic.isAnd(formula));
@@ -216,9 +216,9 @@ bool Cnfizer::cnfizeAndAssert(PTRef formula) {
     // Add the top level literal as a unit to solver.
     vec<Lit> clause;
     clause.push(this->getOrCreateLiteralFor(formula));
-    res &= addClause(clause);
+    res &= addClause(clause, partitionIndex);
 
-    res &= cnfize(formula);
+    res &= cnfize(formula, partitionIndex);
     return res;
 }
 
@@ -234,7 +234,7 @@ void Cnfizer::declareVars(vec<PTRef>& vars)
 // Apply simple de Morgan laws to the formula
 //
 
-bool Cnfizer::deMorganize ( PTRef formula )
+bool Cnfizer::deMorganize ( PTRef formula, PartIdx partitionIndex )
 {
     assert (!logic.isAnd (formula));
     Pterm &pt = logic.getPterm (formula);
@@ -259,7 +259,7 @@ bool Cnfizer::deMorganize ( PTRef formula )
 #endif
         }
 
-        rval = addClause(clause);
+        rval = addClause(clause, partitionIndex);
     }
 
     return rval;
@@ -408,7 +408,7 @@ bool Cnfizer::checkPureConj (PTRef e, Map<PTRef, bool, PTRefHash, Equal<PTRef> >
     return true;
 }
 
-bool Cnfizer::addClause(const vec<Lit> & c_in)
+bool Cnfizer::addClause(const vec<Lit> & c_in, PartIdx partitionIndex)
 {
     vec<Lit> c;
     c_in.copyTo(c);
@@ -430,14 +430,14 @@ bool Cnfizer::addClause(const vec<Lit> & c_in)
     }
 
 #endif
-    bool res = solver.addOriginalSMTClause(c);
+    bool res = solver.addOriginalSMTClause(c, partitionIndex);
     return res;
 }
 //
 // Give the formula to the solver
 //
 
-bool Cnfizer::giveToSolver ( PTRef f )
+bool Cnfizer::giveToSolver ( PTRef f, PartIdx partitionIndex )
 {
     vec<Lit> clause;
 
@@ -447,7 +447,7 @@ bool Cnfizer::giveToSolver ( PTRef f )
     if (logic.isLit (f))
     {
         clause.push (this->getOrCreateLiteralFor (f));
-        return addClause(clause);
+        return addClause(clause, partitionIndex);
     }
 
     //
@@ -462,7 +462,7 @@ bool Cnfizer::giveToSolver ( PTRef f )
         for (int i = 0; i < lits.size(); i++)
             clause.push (this->getOrCreateLiteralFor (lits[i]));
 
-        return addClause(clause);
+        return addClause(clause, partitionIndex);
     }
 
     //
@@ -474,7 +474,7 @@ bool Cnfizer::giveToSolver ( PTRef f )
     bool result = true;
 
     for (unsigned i = 0; i < conj.size_( ) && result; i++)
-        result = giveToSolver (conj[i]);
+        result = giveToSolver (conj[i], partitionIndex);
     return result;
 }
 
diff --git a/src/cnfizers/Cnfizer.h b/src/cnfizers/Cnfizer.h
index 33f7cfab9..6fbdaa7fc 100644
--- a/src/cnfizers/Cnfizer.h
+++ b/src/cnfizers/Cnfizer.h
@@ -74,7 +74,7 @@ class Cnfizer
 
     virtual ~Cnfizer( ) { }
 
-    lbool cnfizeAndGiveToSolver (PTRef, FrameId frame_id); // Main routine
+    lbool cnfizeAndGiveToSolver (PTRef, FrameId frame_id, PartIdx partitionIndex); // Main routine
 
     lbool  getTermValue(PTRef) const;
 
@@ -86,9 +86,9 @@ class Cnfizer
 
 protected:
 
-    virtual bool cnfizeAndAssert        ( PTRef );       // Cnfize and assert the top-level.
-    virtual bool cnfize                 ( PTRef ) = 0;   // Actual cnfization. To be implemented in derived classes
-    bool     deMorganize                ( PTRef );       // Apply deMorgan rules whenever feasible
+    virtual bool cnfizeAndAssert        ( PTRef, PartIdx partitionIndex );       // Cnfize and assert the top-level.
+    virtual bool cnfize                 ( PTRef, PartIdx partitionIndex ) = 0;   // Actual cnfization. To be implemented in derived classes
+    bool     deMorganize                ( PTRef, PartIdx partitionIndex );       // Apply deMorgan rules whenever feasible
     void     declareVars                (vec<PTRef>&);   // Declare a set of Boolean atoms to the solver (without asserting them)
 
 public:
@@ -98,10 +98,10 @@ class Cnfizer
     void     retrieveTopLevelFormulae   ( PTRef, vec<PTRef> & );         // Retrieves the list of top-level formulae
 protected:
 
-    bool     giveToSolver               ( PTRef );                              // Gives formula to the SAT solver
+    bool     giveToSolver               ( PTRef, PartIdx partitionIndex );        // Gives formula to the SAT solver
 
 
-    bool addClause(const vec<Lit> &);
+    bool addClause(const vec<Lit> &, PartIdx partitionIndex);
 
     void  retrieveClause             ( PTRef, vec<PTRef> & );         // Retrieve a clause from a formula
     void  retrieveConjuncts          ( PTRef, vec<PTRef> & );         // Retrieve the list of conjuncts
diff --git a/src/cnfizers/Tseitin.cc b/src/cnfizers/Tseitin.cc
index cae09fbb5..292bf1534 100644
--- a/src/cnfizers/Tseitin.cc
+++ b/src/cnfizers/Tseitin.cc
@@ -33,7 +33,7 @@ WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 //
 
 
-bool Tseitin::cnfize(PTRef term) {
+bool Tseitin::cnfize(PTRef term, PartIdx partitionIndex) {
     bool res = true;
     vec<PTRef> unprocessed_terms {term};
 
@@ -55,15 +55,15 @@ bool Tseitin::cnfize(PTRef term) {
         // by calling findLit
         int sz = logic.getPterm(ptr).size();
         if (logic.isAnd(ptr))
-            res &= cnfizeAnd(ptr);
+            res &= cnfizeAnd(ptr, partitionIndex);
         else if (logic.isOr(ptr))
-            res &= cnfizeOr(ptr);
+            res &= cnfizeOr(ptr, partitionIndex);
         else if (logic.isXor(ptr))
-            res &= cnfizeXor(ptr);
+            res &= cnfizeXor(ptr, partitionIndex);
         else if (logic.isIff(ptr))
-            res &= cnfizeIff(ptr);
+            res &= cnfizeIff(ptr, partitionIndex);
         else if (logic.isImplies(ptr))
-            res &= cnfizeImplies(ptr);
+            res &= cnfizeImplies(ptr, partitionIndex);
         // Ites are handled through the ite manager system and treated here as variables
 //        else if (logic.isIte(ptr))
 //            res &= cnfizeIfthenelse(ptr);
@@ -84,7 +84,7 @@ bool Tseitin::cnfize(PTRef term) {
 }
 
 
-bool Tseitin::cnfizeAnd(PTRef and_term)
+bool Tseitin::cnfizeAnd(PTRef and_term, PartIdx partitionIndex)
 {
 //  assert( list );
 //  assert( list->isList( ) );
@@ -107,16 +107,16 @@ bool Tseitin::cnfizeAnd(PTRef and_term)
         PTRef arg = logic.getPterm(and_term)[i];
         little_clause.push( this->getOrCreateLiteralFor(arg) );
         big_clause   .push(~this->getOrCreateLiteralFor(arg));
-        res &= addClause(little_clause);        // Adds a little clause to the solver
+        res &= addClause(little_clause, partitionIndex);        // Adds a little clause to the solver
         little_clause.pop();
     }
-    res &= addClause(big_clause);                    // Adds a big clause to the solver
+    res &= addClause(big_clause, partitionIndex);                    // Adds a big clause to the solver
     return res;
 }
 
 
 
-bool Tseitin::cnfizeOr(PTRef or_term)
+bool Tseitin::cnfizeOr(PTRef or_term, PartIdx partitionIndex)
 {
     //
     // ( a_0 | ... | a_{n-1} )
@@ -136,16 +136,16 @@ bool Tseitin::cnfizeOr(PTRef or_term)
         little_clause.push(~arg);
         big_clause   .push( arg);
 
-        res &= addClause(little_clause);        // Adds a little clause to the solver
+        res &= addClause(little_clause, partitionIndex);        // Adds a little clause to the solver
 
         little_clause.pop();
     }
-    res &= addClause(big_clause);                    // Adds a big clause to the solver
+    res &= addClause(big_clause, partitionIndex);                    // Adds a big clause to the solver
     return res;
 }
 
 
-bool Tseitin::cnfizeXor(PTRef xor_term)
+bool Tseitin::cnfizeXor(PTRef xor_term, PartIdx partitionIndex)
 {
     //
     // ( a_0 xor a_1 )
@@ -169,7 +169,7 @@ bool Tseitin::cnfizeXor(PTRef xor_term)
     clause.push(arg0);
     clause.push(arg1);
 
-    res &= addClause(clause);
+    res &= addClause(clause, partitionIndex);
     clause.pop();
     clause.pop();
 
@@ -177,7 +177,7 @@ bool Tseitin::cnfizeXor(PTRef xor_term)
     clause.push(~arg0);
     clause.push(~arg1);
 
-    res &= addClause(clause);
+    res &= addClause(clause, partitionIndex);
     clause.pop();
     clause.pop();
 
@@ -188,7 +188,7 @@ bool Tseitin::cnfizeXor(PTRef xor_term)
     clause.push(~arg0);
     clause.push( arg1);
 
-    res &= addClause(clause);
+    res &= addClause(clause, partitionIndex);
 
     clause.pop();
     clause.pop();
@@ -197,11 +197,11 @@ bool Tseitin::cnfizeXor(PTRef xor_term)
     clause.push( arg0);
     clause.push(~arg1);
 
-    res &= addClause(clause);
+    res &= addClause(clause, partitionIndex);
     return res;
 }
 
-bool Tseitin::cnfizeIff(PTRef eq_term)
+bool Tseitin::cnfizeIff(PTRef eq_term, PartIdx partitionIndex)
 {
 
     //
@@ -225,7 +225,7 @@ bool Tseitin::cnfizeIff(PTRef eq_term)
     clause.push( arg0);
     clause.push(~arg1);
 
-    res &= addClause(clause);
+    res &= addClause(clause, partitionIndex);
 
     clause.pop();
     clause.pop();
@@ -234,7 +234,7 @@ bool Tseitin::cnfizeIff(PTRef eq_term)
     clause.push(~arg0);
     clause.push( arg1);
 
-    res &= addClause(clause);
+    res &= addClause(clause, partitionIndex);
 
     clause.pop();
     clause.pop();
@@ -246,7 +246,7 @@ bool Tseitin::cnfizeIff(PTRef eq_term)
     clause.push(arg0);
     clause.push(arg1);
 
-    res &= addClause(clause);
+    res &= addClause(clause, partitionIndex);
 
     clause.pop();
     clause.pop();
@@ -255,12 +255,12 @@ bool Tseitin::cnfizeIff(PTRef eq_term)
     clause.push(~arg0);
     clause.push(~arg1);
 
-    res &= addClause(clause);
+    res &= addClause(clause, partitionIndex);
     return res;
 }
 
 
-bool Tseitin::cnfizeIfthenelse(PTRef ite_term)
+bool Tseitin::cnfizeIfthenelse(PTRef ite_term, PartIdx partitionIndex)
 {
     //  (!a | !i | t) & (!a | i | e) & (a | !i | !t) & (a | i | !e)
     //
@@ -286,25 +286,25 @@ bool Tseitin::cnfizeIfthenelse(PTRef ite_term)
 
     clause.push(~v); clause.push(~a0); clause.push(a1);
 
-    res &= addClause(clause);
+    res &= addClause(clause, partitionIndex);
 
     clause.clear();
 
     clause.push(~v); clause.push(a0); clause.push(a2);
-    res &= addClause(clause);
+    res &= addClause(clause, partitionIndex);
     clause.clear();
 
     clause.push(v); clause.push(~a0); clause.push(~a1);
-    res &= addClause(clause);
+    res &= addClause(clause, partitionIndex);
     clause.clear();
 
     clause.push(v); clause.push(a0); clause.push(~a2);
-    res &= addClause(clause);
+    res &= addClause(clause, partitionIndex);
     return res;
 }
 
 
-bool Tseitin::cnfizeImplies(PTRef impl_term)
+bool Tseitin::cnfizeImplies(PTRef impl_term, PartIdx partitionIndex)
 {
     // ( a_0 => a_1 )
     //
@@ -329,19 +329,19 @@ bool Tseitin::cnfizeImplies(PTRef impl_term)
 
     clause.push(a0);
 
-    res &= addClause(clause);
+    res &= addClause(clause, partitionIndex);
 
     clause.pop();
 
     clause.push(~a1);
 
-    res &= addClause(clause);
+    res &= addClause(clause, partitionIndex);
 
     clause.clear();
 
     clause.push(~v); clause.push(~a0); clause.push(a1);
 
-    res &= addClause(clause);
+    res &= addClause(clause, partitionIndex);
     return res;
 }
 
diff --git a/src/cnfizers/Tseitin.h b/src/cnfizers/Tseitin.h
index 8add4dddb..00a2d86ec 100644
--- a/src/cnfizers/Tseitin.h
+++ b/src/cnfizers/Tseitin.h
@@ -55,13 +55,13 @@ class Tseitin : public Cnfizer
     // Cache of already cnfized terms. Note that this is different from Cnfizer cache of already processed top-level flas
     Cache alreadyCnfized;
 
-    bool cnfize             (PTRef);                          // Cnfize the given term
-    bool cnfizeAnd          (PTRef);                          // Cnfize conjunctions
-    bool cnfizeOr           (PTRef);                          // Cnfize disjunctions
-    bool cnfizeIff          (PTRef);                          // Cnfize iffs
-    bool cnfizeXor          (PTRef);                          // Cnfize xors
-    bool cnfizeIfthenelse   (PTRef);                          // Cnfize if then elses
-    bool cnfizeImplies      (PTRef);                          // Cnfize if then elses
+    bool cnfize             (PTRef, PartIdx);                          // Cnfize the given term
+    bool cnfizeAnd          (PTRef, PartIdx);                          // Cnfize conjunctions
+    bool cnfizeOr           (PTRef, PartIdx);                          // Cnfize disjunctions
+    bool cnfizeIff          (PTRef, PartIdx);                          // Cnfize iffs
+    bool cnfizeXor          (PTRef, PartIdx);                          // Cnfize xors
+    bool cnfizeIfthenelse   (PTRef, PartIdx);                          // Cnfize if then elses
+    bool cnfizeImplies      (PTRef, PartIdx);                          // Cnfize if then elses
 //    void copyArgsWithCache(PTRef, vec<PTRef>&, Map<PTRef, PTRef, PTRefHash>&);
 };
 
diff --git a/src/common/FlaPartitionMap.cc b/src/common/FlaPartitionMap.cc
index 9d864e931..f6c381759 100644
--- a/src/common/FlaPartitionMap.cc
+++ b/src/common/FlaPartitionMap.cc
@@ -25,14 +25,14 @@ void FlaPartitionMap::transferPartitionMembership(PTRef old, PTRef new_ptref) {
     }
 }
 
-int FlaPartitionMap::getPartitionIndex(PTRef ref) const {
+PartIdx FlaPartitionMap::getPartitionIndex(PTRef ref) const {
     auto it = top_level_flas.find(ref);
     if (it != top_level_flas.end()) {
-        return static_cast<int>(it->second);
+        return PartIdx{static_cast<int>(it->second)};
     }
     auto other_it = other_flas.find(ref);
     if (other_it != other_flas.end()) {
-        return static_cast<int>(other_it->second);
+        return PartIdx{static_cast<int>(other_it->second)};
     }
-    return -1;
+    return PartIdx_Undef;
 }
diff --git a/src/common/FlaPartitionMap.h b/src/common/FlaPartitionMap.h
index 444d7b598..d2800d520 100644
--- a/src/common/FlaPartitionMap.h
+++ b/src/common/FlaPartitionMap.h
@@ -10,6 +10,14 @@
 #include <map>
 #include <vector>
 
+struct PartIdx {
+    int id;
+    bool operator== (const PartIdx o) const { return o.id == id; }
+    bool operator!= (const PartIdx o) const { return o.id != id; }
+};
+const struct PartIdx PartIdx_Default = {0};
+const struct PartIdx PartIdx_Undef   = {INT32_MAX};
+
 class FlaPartitionMap {
     std::map<PTRef, unsigned int> top_level_flas;
     std::map<PTRef, unsigned int> other_flas;
@@ -20,7 +28,7 @@ class FlaPartitionMap {
     std::vector<PTRef> get_top_level_flas () const;
     unsigned getNoOfPartitions() const {return get_top_level_flas().size();}
     void transferPartitionMembership(PTRef old, PTRef new_ptref);
-    int getPartitionIndex(PTRef ref) const;
+    PartIdx getPartitionIndex(PTRef ref) const;
 };
 
 
diff --git a/src/common/PartitionInfo.h b/src/common/PartitionInfo.h
index b99ebc70e..f83765112 100644
--- a/src/common/PartitionInfo.h
+++ b/src/common/PartitionInfo.h
@@ -29,7 +29,7 @@ class PartitionInfo {
     inline std::vector<PTRef> getTopLevelFormulas() const { return flaPartitionMap.get_top_level_flas(); }
     inline unsigned int getNoOfPartitions() const {return flaPartitionMap.getNoOfPartitions(); }
     inline void transferPartitionMembership(PTRef o, PTRef n) { return flaPartitionMap.transferPartitionMembership(o, n); }
-    inline int getPartitionIndex(PTRef p) const { return flaPartitionMap.getPartitionIndex(p); }
+    inline PartIdx getPartitionIndex(PTRef p) const { return flaPartitionMap.getPartitionIndex(p); }
 
     void invalidatePartitions(ipartitions_t const & p);
 };
diff --git a/src/smtsolvers/CoreSMTSolver.cc b/src/smtsolvers/CoreSMTSolver.cc
index 3ce2b7a7c..29abf74aa 100644
--- a/src/smtsolvers/CoreSMTSolver.cc
+++ b/src/smtsolvers/CoreSMTSolver.cc
@@ -269,7 +269,7 @@ Var CoreSMTSolver::newVar(bool sign, bool dvar)
     return v;
 }
 
-bool CoreSMTSolver::addOriginalClause_(const vec<Lit> & _ps)
+bool CoreSMTSolver::addOriginalClause_(const vec<Lit> & _ps, PartIdx partitionIndex)
 {
     assert(decisionLevel() == 0);
     if (!isOK()) { return false; }
@@ -309,7 +309,7 @@ bool CoreSMTSolver::addOriginalClause_(const vec<Lit> & _ps)
         CRef inputClause = ca.alloc(original, false);
         clauseToInsert = resolvedUnits.empty() ? inputClause :
                 ps.size() == 0 ? CRef_Undef : ca.alloc(ps, false);
-        proof->newOriginalClause(inputClause);
+        proof->newOriginalClause(inputClause, partitionIndex);
         if (!resolvedUnits.empty()) {
             proof->beginChain( inputClause );
             for(Lit l : resolvedUnits) {
@@ -1035,12 +1035,6 @@ void CoreSMTSolver::finalizeProof(CRef finalConflict) {
     proof->endChain(CRef_Undef);
 }
 
-void CoreSMTSolver::setPartition(std::size_t partitionIndex) {
-    if (proof) {
-        proof->setPartition(partitionIndex);
-    }
-}
-
 /*_________________________________________________________________________________________________
   |
   |  analyzeFinal : (p : Lit)  ->  [void]
@@ -2165,7 +2159,7 @@ bool CoreSMTSolver::createSplit_scatter(bool last)
 
 bool CoreSMTSolver::excludeAssumptions(vec<Lit>& neg_constrs)
 {
-    addOriginalClause(neg_constrs);
+    addOriginalClause(neg_constrs, PartIdx_Undef);
     simplify();
     return ok;
 }
diff --git a/src/smtsolvers/CoreSMTSolver.h b/src/smtsolvers/CoreSMTSolver.h
index cc6346fa6..f156788bc 100644
--- a/src/smtsolvers/CoreSMTSolver.h
+++ b/src/smtsolvers/CoreSMTSolver.h
@@ -248,13 +248,13 @@ class CoreSMTSolver
     virtual Var newVar(bool polarity, bool dvar);//    (bool polarity = true, bool dvar = true); // Add a new variable with parameters specifying variable mode.
 public:
     void    addVar(Var v); // Anounce the existence of a variable to the solver
-    bool    addOriginalClause(const vec<Lit> & ps);
+    bool    addOriginalClause(const vec<Lit> & ps, PartIdx partitionIndex);
     bool    addEmptyClause();                                   // Add the empty clause, making the solver contradictory.
-    bool    addOriginalClause(Lit p);                                  // Add a unit clause to the solver.
-    bool    addOriginalClause(Lit p, Lit q);                           // Add a binary clause to the solver.
-    bool    addOriginalClause(Lit p, Lit q, Lit r);                    // Add a ternary clause to the solver.
+    bool    addOriginalClause(Lit p, PartIdx partitionIndex);       // Add a unit clause to the solver.
+    bool    addOriginalClause(Lit p, Lit q, PartIdx partitionIndex);// Add a binary clause to the solver.
+    bool    addOriginalClause(Lit p, Lit q, Lit r, PartIdx partitionIndex);// Add a ternary clause to the solver.
 protected:
-    bool addOriginalClause_(const vec<Lit> & _ps);                                          // Add a clause to the solver
+    bool addOriginalClause_(const vec<Lit> & _ps, PartIdx partitionIndex); // Add a clause to the solver
 public:
     // Solving:
     //
@@ -754,7 +754,6 @@ class CoreSMTSolver
 public:
 
     void    printTrace() const;
-    void    setPartition(std::size_t partitionIndex);
 
 protected:
     virtual inline void clausesPublish() {};
@@ -844,35 +843,35 @@ inline bool     CoreSMTSolver::enqueue         (Lit p, CRef from)
     return value(p) != l_Undef ? value(p) != l_False : (uncheckedEnqueue(p, from), true);
 }
 
-inline bool     CoreSMTSolver::addOriginalClause(const vec<Lit> & ps)
+inline bool     CoreSMTSolver::addOriginalClause(const vec<Lit> & ps, PartIdx partitionIndex)
 {
-    return addOriginalClause_(ps);
+    return addOriginalClause_(ps, partitionIndex);
 }
 inline bool     CoreSMTSolver::addEmptyClause  ()
 {
     add_tmp.clear();
-    return addOriginalClause_(add_tmp);
+    return addOriginalClause_(add_tmp, PartIdx_Default);
 }
-inline bool     CoreSMTSolver::addOriginalClause(Lit p)
+inline bool     CoreSMTSolver::addOriginalClause(Lit p, PartIdx partitionIndex)
 {
     add_tmp.clear();
     add_tmp.push(p);
-    return addOriginalClause_(add_tmp);
+    return addOriginalClause_(add_tmp, partitionIndex);
 }
-inline bool     CoreSMTSolver::addOriginalClause(Lit p, Lit q)
+inline bool     CoreSMTSolver::addOriginalClause(Lit p, Lit q, PartIdx partitionIndex)
 {
     add_tmp.clear();
     add_tmp.push(p);
     add_tmp.push(q);
-    return addOriginalClause_(add_tmp);
+    return addOriginalClause_(add_tmp, partitionIndex);
 }
-inline bool     CoreSMTSolver::addOriginalClause(Lit p, Lit q, Lit r)
+inline bool     CoreSMTSolver::addOriginalClause(Lit p, Lit q, Lit r, PartIdx partitionIndex)
 {
     add_tmp.clear();
     add_tmp.push(p);
     add_tmp.push(q);
     add_tmp.push(r);
-    return addOriginalClause_(add_tmp);
+    return addOriginalClause_(add_tmp, partitionIndex);
 }
 
 
diff --git a/src/smtsolvers/Proof.cc b/src/smtsolvers/Proof.cc
index d151c1e22..e32c504f3 100644
--- a/src/smtsolvers/Proof.cc
+++ b/src/smtsolvers/Proof.cc
@@ -35,11 +35,11 @@ WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
   , cl_al		( cl )
 { }
 
-void Proof::newOriginalClause(CRef c) {
-    assert(currentPartition != static_cast<std::size_t>(-1));
+void Proof::newOriginalClause(CRef c, PartIdx partitionIndex) {
+    assert(partitionIndex != PartIdx_Undef);
     newLeafClause(c, clause_type::CLA_ORIG);
     ipartitions_t partition = 0;
-    setbit(partition, currentPartition);
+    setbit(partition, partitionIndex.id);
     clause_class[c] |= partition;
 }
 
diff --git a/src/smtsolvers/Proof.h b/src/smtsolvers/Proof.h
index b994db68e..83f24dd19 100644
--- a/src/smtsolvers/Proof.h
+++ b/src/smtsolvers/Proof.h
@@ -31,6 +31,7 @@ WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 #include <unordered_map>
 #include <iosfwd>
 #include <functional>
+#include "FlaPartitionMap.h"
 
 //=================================================================================================
 
@@ -90,19 +91,16 @@ class Proof
     ClauseAllocator&            cl_al;
     std::unordered_map<Lit, CRef, LitHash> assumed_literals;
     std::unordered_map<CRef, ipartitions_t> clause_class;
-    std::size_t currentPartition = 0;
 
 public:
 
     Proof ( ClauseAllocator& cl );
     ~Proof( ) = default;
 
-    void setPartition(std::size_t partition) { currentPartition = partition; }
-
     ipartitions_t getClauseClassMask(CRef c) const { assert(clause_class.count(c) != 0); return clause_class.at(c); }
 
     // Notifies the proof about a new original clause.
-    void newOriginalClause(CRef c);
+    void newOriginalClause(CRef c, PartIdx partitionIndex);
 
     // Notifies the proof about a new T-clause.
     void newTheoryClause(CRef c) { newLeafClause(c, clause_type::CLA_THEORY); }
diff --git a/src/smtsolvers/SimpSMTSolver.cc b/src/smtsolvers/SimpSMTSolver.cc
index 0112906a2..b878572d7 100644
--- a/src/smtsolvers/SimpSMTSolver.cc
+++ b/src/smtsolvers/SimpSMTSolver.cc
@@ -184,7 +184,7 @@ lbool SimpSMTSolver::solve_(bool do_simp, bool turn_off_simp)
 //=================================================================================================
 // Added code
 
-bool SimpSMTSolver::addOriginalSMTClause(const vec<Lit> & smt_clause)
+bool SimpSMTSolver::addOriginalSMTClause(const vec<Lit> & smt_clause, PartIdx partitionIndex)
 {
     assert( config.sat_preprocess_theory == 0 );
 
@@ -212,7 +212,7 @@ bool SimpSMTSolver::addOriginalSMTClause(const vec<Lit> & smt_clause)
         cerr << "XXX skipped handling of unary theory literal?" << endl;
     }
     int nclauses = clauses.size();
-    if (!CoreSMTSolver::addOriginalClause_(smt_clause))
+    if (!CoreSMTSolver::addOriginalClause_(smt_clause, partitionIndex))
         return false;
 
     if (use_simplification && clauses.size() == nclauses + 1)
@@ -621,7 +621,7 @@ bool SimpSMTSolver::eliminateVar(Var v)
     vec<Lit>& resolvent = add_tmp;
     for (int i = 0; i < pos.size(); i++) {
         for (int j = 0; j < neg.size(); j++) {
-            if (merge(ca[pos[i]], ca[neg[j]], v, resolvent) && !addOriginalSMTClause(resolvent))
+            if (merge(ca[pos[i]], ca[neg[j]], v, resolvent) && !addOriginalSMTClause(resolvent, PartIdx_Undef))
                 return false;
         }
     }
@@ -663,7 +663,7 @@ bool SimpSMTSolver::substitute(Var v, Lit x)
 
         removeClause(cls[i]);
 
-        if (!addOriginalSMTClause(subst_clause))
+        if (!addOriginalSMTClause(subst_clause, PartIdx_Undef))
             return ok = false;
     }
 
diff --git a/src/smtsolvers/SimpSMTSolver.h b/src/smtsolvers/SimpSMTSolver.h
index 78ec5f392..c551e4db3 100644
--- a/src/smtsolvers/SimpSMTSolver.h
+++ b/src/smtsolvers/SimpSMTSolver.h
@@ -63,7 +63,7 @@ class SimpSMTSolver : public CoreSMTSolver
     //
     Var     newVar    (bool polarity = true, bool dvar = true) override;
 
-    bool addOriginalSMTClause(const vec<Lit> & smt_clause);
+    bool addOriginalSMTClause(const vec<Lit> & smt_clause, PartIdx partitionIndex);
 public:
 
     bool    substitute(Var v, Lit x);  // Replace all occurences of v with x (may cause a contradiction).
diff --git a/src/tsolvers/bvsolver/BitBlaster.cc b/src/tsolvers/bvsolver/BitBlaster.cc
index 1ba7a063b..22eff1487 100644
--- a/src/tsolvers/bvsolver/BitBlaster.cc
+++ b/src/tsolvers/bvsolver/BitBlaster.cc
@@ -1466,7 +1466,7 @@ BitBlaster::bbDistinct(PTRef tr)
 bool
 BitBlaster::addClause(vec<Lit> & c)
 {
-    return solverP.addOriginalClause(c);
+    return solverP.addOriginalClause(c, PartIdx_Undef);
 }
 
 //=============================================================================