1 files changed, 19 insertions, 9 deletions
diff --git a/src/theory/difficulty_manager.cpp b/src/theory/difficulty_manager.cpp
index 0a9eba60d..7714abfa9 100644
--- a/src/theory/difficulty_manager.cpp
+++ b/src/theory/difficulty_manager.cpp
@@ -26,20 +26,30 @@ namespace cvc5 {
 namespace theory {
 
 DifficultyManager::DifficultyManager(context::Context* c, Valuation val)
-    : d_val(val), d_dfmap(c)
+    : d_input(c), d_val(val), d_dfmap(c)
 {
 }
 
+void DifficultyManager::notifyInputAssertions(
+    const std::vector<Node>& assertions)
+{
+  for (const Node& a : assertions)
+  {
+    d_input.insert(a);
+  }
+}
+
 void DifficultyManager::getDifficultyMap(std::map<Node, Node>& dmap)
 {
   NodeManager* nm = NodeManager::currentNM();
   for (const std::pair<const Node, uint64_t> p : d_dfmap)
   {
-    dmap[p.first] = nm->mkConst(CONST_RATIONAL, Rational(p.second));
+    dmap[p.first] = nm->mkConstInt(Rational(p.second));
   }
 }
 
-void DifficultyManager::notifyLemma(const std::map<TNode, TNode>& rse, Node n)
+void DifficultyManager::notifyLemma(const context::CDHashMap<Node, Node>& rse,
+                                    Node n)
 {
   if (options::difficultyMode() != options::DifficultyMode::LEMMA_LITERAL)
   {
@@ -63,7 +73,7 @@ void DifficultyManager::notifyLemma(const std::map<TNode, TNode>& rse, Node n)
   {
     litsToCheck.push_back(n);
   }
-  std::map<TNode, TNode>::const_iterator it;
+  context::CDHashMap<Node, Node>::const_iterator it;
   for (TNode nc : litsToCheck)
   {
     bool pol = nc.getKind() != kind::NOT;
@@ -72,23 +82,23 @@ void DifficultyManager::notifyLemma(const std::map<TNode, TNode>& rse, Node n)
     Trace("diff-man-debug")
         << "Check literal: " << atom << ", has reason = " << (it != rse.end())
         << std::endl;
-    if (it != rse.end())
+    // must be an input assertion
+    if (it != rse.end() && d_input.find(it->second) != d_input.end())
     {
       incrementDifficulty(it->second);
     }
   }
 }
 
-void DifficultyManager::notifyCandidateModel(const NodeList& input,
-                                             TheoryModel* m)
+void DifficultyManager::notifyCandidateModel(TheoryModel* m)
 {
   if (options::difficultyMode() != options::DifficultyMode::MODEL_CHECK)
   {
     return;
   }
   Trace("diff-man") << "DifficultyManager::notifyCandidateModel, #input="
-                    << input.size() << std::endl;
-  for (const Node& a : input)
+                    << d_input.size() << std::endl;
+  for (const Node& a : d_input)
   {
     // should have miniscoped the assertions upstream
     Assert(a.getKind() != kind::AND);