Get operators in node (#3094)

This commit adds a function to node_algorithm.{h,cpp} that returns the operators that occur in a given node.

2 files changed, 73 insertions, 1 deletions

diff --git a/src/expr/node_algorithm.cpp b/src/expr/node_algorithm.cpp
index 841f9ea28..c20dddbcc 100644
--- a/src/expr/node_algorithm.cpp
+++ b/src/expr/node_algorithm.cpp
@@ -304,6 +304,47 @@ void getSymbols(TNode n,
   } while (!visit.empty());
 }
 
+void getOperatorsMap(
+    TNode n,
+    std::map<TypeNode, std::unordered_set<Node, NodeHashFunction>>& ops)
+{
+  std::unordered_set<TNode, TNodeHashFunction> visited;
+  getOperatorsMap(n, ops, visited);
+}
+
+void getOperatorsMap(
+    TNode n,
+    std::map<TypeNode, std::unordered_set<Node, NodeHashFunction>>& ops,
+    std::unordered_set<TNode, TNodeHashFunction>& visited)
+{
+  // nodes that we still need to visit
+  std::vector<TNode> visit;
+  // current node
+  TNode cur;
+  visit.push_back(n);
+  do
+  {
+    cur = visit.back();
+    visit.pop_back();
+    // if cur is in the cache, do nothing
+    if (visited.find(cur) == visited.end())
+    {
+      // fetch the correct type
+      TypeNode tn = cur.getType();
+      // add the current operator to the result
+      if (cur.hasOperator())
+      {
+        ops[tn].insert(NodeManager::currentNM()->operatorOf(cur.getKind()));
+      }
+      // add children to visit in the future
+      for (TNode cn : cur)
+      {
+        visit.push_back(cn);
+      }
+    }
+  } while (!visit.empty());
+}
+
 Node substituteCaptureAvoiding(TNode n, Node src, Node dest)
 {
   if (n == src)
diff --git a/src/expr/node_algorithm.h b/src/expr/node_algorithm.h
index 727f5ba75..3686b6686 100644
--- a/src/expr/node_algorithm.h
+++ b/src/expr/node_algorithm.h
@@ -25,6 +25,7 @@
 #include <vector>
 
 #include "expr/node.h"
+#include "expr/type_node.h"
 
 namespace CVC4 {
 namespace expr {
@@ -87,10 +88,40 @@ bool getVariables(TNode n, std::unordered_set<TNode, TNodeHashFunction>& vs);
  * @param syms The set which the symbols of n are added to
  */
 void getSymbols(TNode n, std::unordered_set<Node, NodeHashFunction>& syms);
-/** Same as above, with a visited cache */
+
+/**
+ * For term n, this function collects the symbols that occur as a subterms
+ * of n. A symbol is a variable that does not have kind BOUND_VARIABLE.
+ * @param n The node under investigation
+ * @param syms The set which the symbols of n are added to
+ * @param visited A cache to be used for visited nodes.
+ */
 void getSymbols(TNode n,
                 std::unordered_set<Node, NodeHashFunction>& syms,
                 std::unordered_set<TNode, TNodeHashFunction>& visited);
+
+/**
+ * For term n, this function collects the operators that occur in n.
+ * @param n The node under investigation
+ * @param ops The map (from each type to operators of that type) which the
+ * operators of n are added to
+ */
+void getOperatorsMap(
+    TNode n,
+    std::map<TypeNode, std::unordered_set<Node, NodeHashFunction>>& ops);
+
+/**
+ * For term n, this function collects the operators that occur in n.
+ * @param n The node under investigation
+ * @param ops The map (from each type to operators of that type) which the
+ * operators of n are added to
+ * @param visited A cache to be used for visited nodes.
+ */
+void getOperatorsMap(
+    TNode n,
+    std::map<TypeNode, std::unordered_set<Node, NodeHashFunction>>& ops,
+    std::unordered_set<TNode, TNodeHashFunction>& visited);
+
 /**
  * Substitution of Nodes in a capture avoiding way.
  */