Improvements to rewrite rules from inputs (#2625)

2 files changed, 60 insertions, 0 deletions

diff --git a/src/expr/node_algorithm.cpp b/src/expr/node_algorithm.cpp
index 4bbfb5df8..3905ad5c9 100644
--- a/src/expr/node_algorithm.cpp
+++ b/src/expr/node_algorithm.cpp
@@ -63,6 +63,61 @@ bool hasSubterm(TNode n, TNode t, bool strict)
   return false;
 }
 
+bool hasSubtermMulti(TNode n, TNode t)
+{
+  std::unordered_map<TNode, bool, TNodeHashFunction> visited;
+  std::unordered_map<TNode, bool, TNodeHashFunction> contains;
+  std::unordered_map<TNode, bool, TNodeHashFunction>::iterator it;
+  std::vector<TNode> visit;
+  TNode cur;
+  visit.push_back(n);
+  do
+  {
+    cur = visit.back();
+    visit.pop_back();
+    it = visited.find(cur);
+
+    if (it == visited.end())
+    {
+      if (cur == t)
+      {
+        visited[cur] = true;
+        contains[cur] = true;
+      }
+      else
+      {
+        visited[cur] = false;
+        visit.push_back(cur);
+        for (const Node& cc : cur)
+        {
+          visit.push_back(cc);
+        }
+      }
+    }
+    else if (!it->second)
+    {
+      bool doesContain = false;
+      for (const Node& cn : cur)
+      {
+        it = contains.find(cn);
+        Assert(it != visited.end());
+        if (it->second)
+        {
+          if (doesContain)
+          {
+            // two children have t, return true
+            return true;
+          }
+          doesContain = true;
+        }
+      }
+      contains[cur] = doesContain;
+      visited[cur] = true;
+    }
+  } while (!visit.empty());
+  return false;
+}
+
 struct HasBoundVarTag
 {
 };
diff --git a/src/expr/node_algorithm.h b/src/expr/node_algorithm.h
index 7453bc292..d825d7f57 100644
--- a/src/expr/node_algorithm.h
+++ b/src/expr/node_algorithm.h
@@ -39,6 +39,11 @@ namespace expr {
 bool hasSubterm(TNode n, TNode t, bool strict = false);
 
 /**
+ * Check if the node n has >1 occurrences of a subterm t.
+ */
+bool hasSubtermMulti(TNode n, TNode t);
+
+/**
  * Returns true iff the node n contains a bound variable, that is a node of
  * kind BOUND_VARIABLE. This bound variable may or may not be free.
  * @param n The node under investigation