Add utility to get all types occurring in a term (#6588)

Required for external proof conversions.

2 files changed, 46 insertions, 0 deletions

diff --git a/src/expr/node_algorithm.cpp b/src/expr/node_algorithm.cpp
index 19955443a..2220caff9 100644
--- a/src/expr/node_algorithm.cpp
+++ b/src/expr/node_algorithm.cpp
@@ -647,6 +647,34 @@ Node substituteCaptureAvoiding(TNode n,
   return visited[n];
 }
 
+void getTypes(TNode n, std::unordered_set<TypeNode>& types)
+{
+  std::unordered_set<TNode> visited;
+  getTypes(n, types, visited);
+}
+
+void getTypes(TNode n,
+              std::unordered_set<TypeNode>& types,
+              std::unordered_set<TNode>& visited)
+{
+  std::unordered_set<TNode>::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())
+    {
+      visited.insert(cur);
+      types.insert(cur.getType());
+      visit.insert(visit.end(), cur.begin(), cur.end());
+    }
+  } while (!visit.empty());
+}
+
 void getComponentTypes(TypeNode t, std::unordered_set<TypeNode>& types)
 {
   std::vector<TypeNode> toProcess;
diff --git a/src/expr/node_algorithm.h b/src/expr/node_algorithm.h
index b8576f787..72ea03176 100644
--- a/src/expr/node_algorithm.h
+++ b/src/expr/node_algorithm.h
@@ -198,6 +198,24 @@ Node substituteCaptureAvoiding(TNode n, Node src, Node dest);
 Node substituteCaptureAvoiding(TNode n,
                                std::vector<Node>& src,
                                std::vector<Node>& dest);
+/**
+ * Collect all types in n, which adds to types all types for which a subterm
+ * of n has that type. Operators are not considered in the traversal.
+ * @param n The node under investigation
+ * @param types The set of types
+ */
+void getTypes(TNode n, std::unordered_set<TypeNode>& types);
+
+/**
+ * Collect all types in n, which adds to types all types for which a subterm
+ * of n has that type. Operators are not considered in the traversal.
+ * @param n The node under investigation
+ * @param types The set of types
+ * @param visited A cache of nodes we have already visited
+ */
+void getTypes(TNode n,
+              std::unordered_set<TypeNode>& types,
+              std::unordered_set<TNode>& visited);
 
 /**
  * Get component types in type t. This adds all types that are subterms of t