1 files changed, 132 insertions, 62 deletions
diff --git a/src/expr/node_algorithm.cpp b/src/expr/node_algorithm.cpp
index 13265db71..badf0b75a 100644
--- a/src/expr/node_algorithm.cpp
+++ b/src/expr/node_algorithm.cpp
@@ -281,64 +281,32 @@ bool hasBoundVar(TNode n)
   return n.getAttribute(HasBoundVarAttr());
 }
 
-bool hasFreeVar(TNode n)
-{
-  // optimization for variables and constants
-  if (n.getNumChildren() == 0)
-  {
-    return n.getKind() == kind::BOUND_VARIABLE;
-  }
-  std::unordered_set<Node> fvs;
-  return getFreeVariables(n, fvs, false);
-}
-
-struct HasClosureTag
-{
-};
-struct HasClosureComputedTag
-{
-};
-/** Attribute true for expressions with closures in them */
-typedef expr::Attribute<HasClosureTag, bool> HasClosureAttr;
-typedef expr::Attribute<HasClosureComputedTag, bool> HasClosureComputedAttr;
-
-bool hasClosure(Node n)
-{
-  if (!n.getAttribute(HasClosureComputedAttr()))
-  {
-    bool hasC = false;
-    if (n.isClosure())
-    {
-      hasC = true;
-    }
-    else
-    {
-      for (auto i = n.begin(); i != n.end() && !hasC; ++i)
-      {
-        hasC = hasClosure(*i);
-      }
-    }
-    if (!hasC && n.hasOperator())
-    {
-      hasC = hasClosure(n.getOperator());
-    }
-    n.setAttribute(HasClosureAttr(), hasC);
-    n.setAttribute(HasClosureComputedAttr(), true);
-    return hasC;
-  }
-  return n.getAttribute(HasClosureAttr());
-}
-
-bool getFreeVariables(TNode n, std::unordered_set<Node>& fvs, bool computeFv)
-{
-  std::unordered_set<TNode> scope;
-  return getFreeVariablesScope(n, fvs, scope, computeFv);
-}
-
-bool getFreeVariablesScope(TNode n,
-                           std::unordered_set<Node>& fvs,
-                           std::unordered_set<TNode>& scope,
-                           bool computeFv)
+/**
+ * Check variables internal, which is used as a helper to implement many of the
+ * methods in this file.
+ *
+ * This computes the free variables in n, that is, the subterms of n of kind
+ * BOUND_VARIABLE that are not bound in n or occur in scope, adds these to fvs
+ * if computeFv is true.
+ *
+ * @param n The node under investigation
+ * @param fvs The set which free variables are added to
+ * @param scope The scope we are considering.
+ * @param wasShadow Flag set to true if variable shadowing was encountered.
+ * Only computed if checkShadow is true.
+ * @param computeFv If this flag is false, then we only return true/false and
+ * do not add to fvs.
+ * @param checkShadow If this flag is true, we immediately return true if a
+ * variable is shadowing. If this flag is false, we give an assertion failure
+ * when this occurs.
+ * @return true iff this node contains a free variable.
+ */
+bool checkVariablesInternal(TNode n,
+                            std::unordered_set<Node>& fvs,
+                            std::unordered_set<TNode>& scope,
+                            bool& wasShadow,
+                            bool computeFv = true,
+                            bool checkShadow = false)
 {
   std::unordered_set<TNode> visited;
   std::vector<TNode> visit;
@@ -376,13 +344,26 @@ bool getFreeVariablesScope(TNode n,
         // add to scope
         for (const TNode& cn : cur[0])
         {
-          // should not shadow
-          Assert(scope.find(cn) == scope.end())
-              << "Shadowed variable " << cn << " in " << cur << "\n";
+          if (checkShadow)
+          {
+            if (scope.find(cn) != scope.end())
+            {
+              wasShadow = true;
+              return true;
+            }
+          }
+          else
+          {
+            // should not shadow
+            Assert(scope.find(cn) == scope.end())
+                << "Shadowed variable " << cn << " in " << cur << "\n";
+          }
           scope.insert(cn);
         }
         // must make recursive call to use separate cache
-        if (getFreeVariablesScope(cur[1], fvs, scope, computeFv) && !computeFv)
+        if (checkVariablesInternal(
+                cur[1], fvs, scope, wasShadow, computeFv, checkShadow)
+            && !computeFv)
         {
           return true;
         }
@@ -406,6 +387,95 @@ bool getFreeVariablesScope(TNode n,
   return !fvs.empty();
 }
 
+/** Same as above, without checking for shadowing */
+bool getVariablesInternal(TNode n,
+                          std::unordered_set<Node>& fvs,
+                          std::unordered_set<TNode>& scope,
+                          bool computeFv = true)
+{
+  bool wasShadow = false;
+  return checkVariablesInternal(n, fvs, scope, wasShadow, computeFv, false);
+}
+
+bool hasFreeVar(TNode n)
+{
+  // optimization for variables and constants
+  if (n.getNumChildren() == 0)
+  {
+    return n.getKind() == kind::BOUND_VARIABLE;
+  }
+  std::unordered_set<Node> fvs;
+  std::unordered_set<TNode> scope;
+  return getVariablesInternal(n, fvs, scope, false);
+}
+
+bool hasFreeOrShadowedVar(TNode n, bool& wasShadow)
+{
+  // optimization for variables and constants
+  if (n.getNumChildren() == 0)
+  {
+    return n.getKind() == kind::BOUND_VARIABLE;
+  }
+  std::unordered_set<Node> fvs;
+  std::unordered_set<TNode> scope;
+  return checkVariablesInternal(n, fvs, scope, wasShadow, false, true);
+}
+
+struct HasClosureTag
+{
+};
+struct HasClosureComputedTag
+{
+};
+/** Attribute true for expressions with closures in them */
+typedef expr::Attribute<HasClosureTag, bool> HasClosureAttr;
+typedef expr::Attribute<HasClosureComputedTag, bool> HasClosureComputedAttr;
+
+bool hasClosure(Node n)
+{
+  if (!n.getAttribute(HasClosureComputedAttr()))
+  {
+    bool hasC = false;
+    if (n.isClosure())
+    {
+      hasC = true;
+    }
+    else
+    {
+      for (auto i = n.begin(); i != n.end() && !hasC; ++i)
+      {
+        hasC = hasClosure(*i);
+      }
+    }
+    if (!hasC && n.hasOperator())
+    {
+      hasC = hasClosure(n.getOperator());
+    }
+    n.setAttribute(HasClosureAttr(), hasC);
+    n.setAttribute(HasClosureComputedAttr(), true);
+    return hasC;
+  }
+  return n.getAttribute(HasClosureAttr());
+}
+
+bool getFreeVariables(TNode n, std::unordered_set<Node>& fvs)
+{
+  std::unordered_set<TNode> scope;
+  return getVariablesInternal(n, fvs, scope);
+}
+
+bool getFreeVariablesScope(TNode n,
+                           std::unordered_set<Node>& fvs,
+                           std::unordered_set<TNode>& scope)
+{
+  return getVariablesInternal(n, fvs, scope);
+}
+bool hasFreeVariablesScope(TNode n, std::unordered_set<TNode>& scope)
+{
+  std::unordered_set<Node> fvs;
+  return getVariablesInternal(n, fvs, scope, false);
+}
+
 bool getVariables(TNode n, std::unordered_set<TNode>& vs)
 {
   std::unordered_set<TNode> visited;