Utility function for getting component types (#3703)

2 files changed, 33 insertions, 0 deletions

diff --git a/src/expr/node_algorithm.cpp b/src/expr/node_algorithm.cpp
index eda4dadd2..9721845f7 100644
--- a/src/expr/node_algorithm.cpp
+++ b/src/expr/node_algorithm.cpp
@@ -18,6 +18,7 @@
 #include "expr/node_algorithm.h"
 
 #include "expr/attribute.h"
+#include "expr/dtype.h"
 
 namespace CVC4 {
 namespace expr {
@@ -522,5 +523,27 @@ Node substituteCaptureAvoiding(TNode n,
   return visited[n];
 }
 
+void getComponentTypes(
+    TypeNode t, std::unordered_set<TypeNode, TypeNodeHashFunction>& types)
+{
+  std::vector<TypeNode> toProcess;
+  toProcess.push_back(t);
+  do
+  {
+    TypeNode curr = toProcess.back();
+    toProcess.pop_back();
+    // if not already visited
+    if (types.find(t) == types.end())
+    {
+      types.insert(t);
+      // get component types from the children
+      for (unsigned i = 0, nchild = t.getNumChildren(); i < nchild; i++)
+      {
+        toProcess.push_back(t[i]);
+      }
+    }
+  } while (!toProcess.empty());
+}
+
 }  // namespace expr
 }  // namespace CVC4
diff --git a/src/expr/node_algorithm.h b/src/expr/node_algorithm.h
index e5a21d565..03d77a2f9 100644
--- a/src/expr/node_algorithm.h
+++ b/src/expr/node_algorithm.h
@@ -151,6 +151,16 @@ Node substituteCaptureAvoiding(TNode n,
                                std::vector<Node>& src,
                                std::vector<Node>& dest);
 
+/**
+ * Get component types in type t. This adds all types that are subterms of t
+ * when viewed as a term. For example, if t is (Array T1 T2), then
+ * (Array T1 T2), T1 and T2 are component types of t.
+ * @param t The type node under investigation
+ * @param types The set which component types are added to.
+ */
+void getComponentTypes(
+    TypeNode t, std::unordered_set<TypeNode, TypeNodeHashFunction>& types);
+
 }  // namespace expr
 }  // namespace CVC4