From 0b7e50aa1f0f2e6a77fb7e2a1f48b6af8ce5b91d Mon Sep 17 00:00:00 2001
From: Andrew Reynolds <andrew.j.reynolds@gmail.com>
Date: Fri, 21 May 2021 13:04:49 -0500
Subject: [PATCH] Add utility to get all types occurring in a term (#6588)

Required for external proof conversions.
---
 src/expr/node_algorithm.cpp | 28 ++++++++++++++++++++++++++++
 src/expr/node_algorithm.h   | 18 ++++++++++++++++++
 2 files changed, 46 insertions(+)

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
-- 
2.30.2