Add getSubtermKinds to node algorithm (#5398)

Required for a new algorithm for nested quantifier elimination.

diff --git a/src/expr/node_algorithm.cpp b/src/expr/node_algorithm.cpp
index 9b7b2b036b5efcba4bc2b80e99aee82217a4fd63..9d1c6ab41c4b7bec1181a068620cb5c87109efad 100644 (file)
--- a/src/expr/node_algorithm.cpp
+++ b/src/expr/node_algorithm.cpp
@@ -374,10 +374,7 @@ bool getFreeVariables(TNode n,
       {
         visit.push_back(cur.getOperator());
       }
-      for (const TNode& cn : cur)
-      {
-        visit.push_back(cn);
-      }
+      visit.insert(visit.end(), cur.begin(), cur.end());
     }
     else if (!itv->second)
     {
@@ -413,10 +410,7 @@ bool getVariables(TNode n, std::unordered_set<TNode, TNodeHashFunction>& vs)
       }
       else
       {
-        for (const TNode& cn : cur)
-        {
-          visit.push_back(cn);
-        }
+        visit.insert(visit.end(), cur.begin(), cur.end());
       }
       visited.insert(cur);
     }
@@ -453,10 +447,41 @@ void getSymbols(TNode n,
       {
         visit.push_back(cur.getOperator());
       }
-      for (TNode cn : cur)
+      visit.insert(visit.end(), cur.begin(), cur.end());
+    }
+  } while (!visit.empty());
+}
+
+void getKindSubterms(TNode n,
+                     Kind k,
+                     bool topLevel,
+                     std::unordered_set<Node, NodeHashFunction>& ts)
+{
+  std::unordered_set<TNode, TNodeHashFunction> visited;
+  std::vector<TNode> visit;
+  TNode cur;
+  visit.push_back(n);
+  do
+  {
+    cur = visit.back();
+    visit.pop_back();
+    if (visited.find(cur) == visited.end())
+    {
+      visited.insert(cur);
+      if (cur.getKind() == k)
+      {
+        ts.insert(cur);
+        if (topLevel)
+        {
+          // only considering top-level applications
+          continue;
+        }
+      }
+      if (cur.hasOperator())
       {
-        visit.push_back(cn);
+        visit.push_back(cur.getOperator());
       }
+      visit.insert(visit.end(), cur.begin(), cur.end());
     }
   } while (!visit.empty());
 }
@@ -500,10 +525,7 @@ void getOperatorsMap(
         ops[tn].insert(o);
       }
       // add children to visit in the future
-      for (TNode cn : cur)
-      {
-        visit.push_back(cn);
-      }
+      visit.insert(visit.end(), cur.begin(), cur.end());
     }
   } while (!visit.empty());
 }

diff --git a/src/expr/node_algorithm.h b/src/expr/node_algorithm.h
index 97a0cb6e258091ac24b65f7567e00f099367eb90..7b6dd4f8c3ac9a05f3bb4d3e41de8607236317a1 100644 (file)
--- a/src/expr/node_algorithm.h
+++ b/src/expr/node_algorithm.h
@@ -133,6 +133,18 @@ void getSymbols(TNode n,
                 std::unordered_set<Node, NodeHashFunction>& syms,
                 std::unordered_set<TNode, TNodeHashFunction>& visited);
 
+/**
+ * For term n, this function collects the subterms of n whose kind is k.
+ * @param n The node under investigation
+ * @param k The kind we are considering
+ * @param topLevel If true, we collect only the top-level subterms with kind k.
+ * @param ts The set which the symbols of n are added to
+ */
+void getKindSubterms(TNode n,
+                     Kind k,
+                     bool topLevel,
+                     std::unordered_set<Node, NodeHashFunction>& ts);
+
 /**
  * For term n, this function collects the operators that occur in n.
  * @param n The node under investigation