Add getFreeVariables method to node algorithm (#2852)

2 files changed, 31 insertions, 2 deletions

diff --git a/src/expr/node_algorithm.cpp b/src/expr/node_algorithm.cpp
index 3905ad5c9..6923efec2 100644
--- a/src/expr/node_algorithm.cpp
+++ b/src/expr/node_algorithm.cpp
@@ -159,6 +159,14 @@ bool hasBoundVar(TNode n)
 
 bool hasFreeVar(TNode n)
 {
+  std::unordered_set<Node, NodeHashFunction> fvs;
+  return getFreeVariables(n, fvs, false);
+}
+
+bool getFreeVariables(TNode n,
+                      std::unordered_set<Node, NodeHashFunction>& fvs,
+                      bool computeFv)
+{
   std::unordered_set<TNode, TNodeHashFunction> bound_var;
   std::unordered_map<TNode, bool, TNodeHashFunction> visited;
   std::vector<TNode> visit;
@@ -184,7 +192,14 @@ bool hasFreeVar(TNode n)
       {
         if (bound_var.find(cur) == bound_var.end())
         {
-          return true;
+          if (computeFv)
+          {
+            fvs.insert(cur);
+          }
+          else
+          {
+            return true;
+          }
         }
       }
       else if (isQuant)
@@ -218,7 +233,8 @@ bool hasFreeVar(TNode n)
       visited[cur] = true;
     }
   } while (!visit.empty());
-  return false;
+
+  return !fvs.empty();
 }
 
 void getSymbols(TNode n, std::unordered_set<Node, NodeHashFunction>& syms)
diff --git a/src/expr/node_algorithm.h b/src/expr/node_algorithm.h
index d825d7f57..bf2cb5877 100644
--- a/src/expr/node_algorithm.h
+++ b/src/expr/node_algorithm.h
@@ -60,6 +60,19 @@ bool hasBoundVar(TNode n);
 bool hasFreeVar(TNode n);
 
 /**
+ * Get the free variables in n, that is, the subterms of n of kind
+ * BOUND_VARIABLE that are not bound in n, adds these to fvs.
+ * @param n The node under investigation
+ * @param fvs The set which free variables are added to
+ * @param computeFv If this flag is false, then we only return true/false and
+ * do not add to fvs.
+ * @return true iff this node contains a free variable.
+ */
+bool getFreeVariables(TNode n,
+                      std::unordered_set<Node, NodeHashFunction>& fvs,
+                      bool computeFv = true);
+
+/**
  * 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