Eliminate partial operators within lambdas during grammar normalization (#2570)

2 files changed, 111 insertions, 26 deletions

diff --git a/src/theory/quantifiers/sygus/sygus_grammar_norm.cpp b/src/theory/quantifiers/sygus/sygus_grammar_norm.cpp
index 3d066e8dd..b014a30c6 100644
--- a/src/theory/quantifiers/sygus/sygus_grammar_norm.cpp
+++ b/src/theory/quantifiers/sygus/sygus_grammar_norm.cpp
@@ -67,6 +67,88 @@ bool OpPosTrie::getOrMakeType(TypeNode tn,
   return d_children[op_pos[ind]].getOrMakeType(tn, unres_tn, op_pos, ind + 1);
 }
 
+Kind SygusGrammarNorm::TypeObject::getEliminateKind(Kind ok)
+{
+  Kind nk = ok;
+  // We also must ensure that builtin operators which are eliminated
+  // during expand definitions are replaced by the proper operator.
+  if (ok == BITVECTOR_UDIV)
+  {
+    nk = BITVECTOR_UDIV_TOTAL;
+  }
+  else if (ok == BITVECTOR_UREM)
+  {
+    nk = BITVECTOR_UREM_TOTAL;
+  }
+  else if (ok == DIVISION)
+  {
+    nk = DIVISION_TOTAL;
+  }
+  else if (ok == INTS_DIVISION)
+  {
+    nk = INTS_DIVISION_TOTAL;
+  }
+  else if (ok == INTS_MODULUS)
+  {
+    nk = INTS_MODULUS_TOTAL;
+  }
+  return nk;
+}
+
+Node SygusGrammarNorm::TypeObject::eliminatePartialOperators(Node n)
+{
+  NodeManager* nm = NodeManager::currentNM();
+  std::unordered_map<TNode, Node, TNodeHashFunction> visited;
+  std::unordered_map<TNode, Node, TNodeHashFunction>::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[cur] = Node::null();
+      visit.push_back(cur);
+      for (const Node& cn : cur)
+      {
+        visit.push_back(cn);
+      }
+    }
+    else if (it->second.isNull())
+    {
+      Node ret = cur;
+      bool childChanged = false;
+      std::vector<Node> children;
+      if (cur.getMetaKind() == metakind::PARAMETERIZED)
+      {
+        children.push_back(cur.getOperator());
+      }
+      for (const Node& cn : cur)
+      {
+        it = visited.find(cn);
+        Assert(it != visited.end());
+        Assert(!it->second.isNull());
+        childChanged = childChanged || cn != it->second;
+        children.push_back(it->second);
+      }
+      Kind ok = cur.getKind();
+      Kind nk = getEliminateKind(ok);
+      if (nk != ok || childChanged)
+      {
+        ret = nm->mkNode(nk, children);
+      }
+      visited[cur] = ret;
+    }
+  } while (!visit.empty());
+  Assert(visited.find(n) != visited.end());
+  Assert(!visited.find(n)->second.isNull());
+  return visited[n];
+}
+
 void SygusGrammarNorm::TypeObject::addConsInfo(SygusGrammarNorm* sygus_norm,
                                                const DatatypeConstructor& cons)
 {
@@ -74,36 +156,15 @@ void SygusGrammarNorm::TypeObject::addConsInfo(SygusGrammarNorm* sygus_norm,
   /* Recover the sygus operator to not lose reference to the original
    * operator (NOT, ITE, etc) */
   Node sygus_op = Node::fromExpr(cons.getSygusOp());
+  Trace("sygus-grammar-normalize-debug")
+      << ".....operator is " << sygus_op << std::endl;
   Node exp_sop_n = Node::fromExpr(
       smt::currentSmtEngine()->expandDefinitions(sygus_op.toExpr()));
+  // if it is a builtin operator, convert to total version if necessary
   if (exp_sop_n.getKind() == kind::BUILTIN)
   {
     Kind ok = NodeManager::operatorToKind(sygus_op);
-    Kind nk = ok;
-    Trace("sygus-grammar-normalize-debug")
-        << "...builtin operator is " << ok << std::endl;
-    // We also must ensure that builtin operators which are eliminated
-    // during expand definitions are replaced by the proper operator.
-    if (ok == kind::BITVECTOR_UDIV)
-    {
-      nk = kind::BITVECTOR_UDIV_TOTAL;
-    }
-    else if (ok == kind::BITVECTOR_UREM)
-    {
-      nk = kind::BITVECTOR_UREM_TOTAL;
-    }
-    else if (ok == kind::DIVISION)
-    {
-      nk = kind::DIVISION_TOTAL;
-    }
-    else if (ok == kind::INTS_DIVISION)
-    {
-      nk = kind::INTS_DIVISION_TOTAL;
-    }
-    else if (ok == kind::INTS_MODULUS)
-    {
-      nk = kind::INTS_MODULUS_TOTAL;
-    }
+    Kind nk = getEliminateKind(ok);
     if (nk != ok)
     {
       Trace("sygus-grammar-normalize-debug")
@@ -111,7 +172,21 @@ void SygusGrammarNorm::TypeObject::addConsInfo(SygusGrammarNorm* sygus_norm,
       exp_sop_n = NodeManager::currentNM()->operatorOf(nk);
     }
   }
-  d_ops.push_back(Rewriter::rewrite(exp_sop_n));
+  else
+  {
+    exp_sop_n = Rewriter::rewrite(exp_sop_n);
+    Trace("sygus-grammar-normalize-debug")
+        << ".....operator (post-rewrite) is " << exp_sop_n << std::endl;
+    // eliminate all partial operators from it
+    exp_sop_n = eliminatePartialOperators(exp_sop_n);
+    Trace("sygus-grammar-normalize-debug")
+        << ".....operator (eliminate partial operators) is " << exp_sop_n
+        << std::endl;
+    // rewrite again
+    exp_sop_n = Rewriter::rewrite(exp_sop_n);
+  }
+
+  d_ops.push_back(exp_sop_n);
   Trace("sygus-grammar-normalize-defs")
       << "\tOriginal op: " << cons.getSygusOp()
       << "\n\tExpanded one: " << exp_sop_n
diff --git a/src/theory/quantifiers/sygus/sygus_grammar_norm.h b/src/theory/quantifiers/sygus/sygus_grammar_norm.h
index a0f81dcf3..993d41668 100644
--- a/src/theory/quantifiers/sygus/sygus_grammar_norm.h
+++ b/src/theory/quantifiers/sygus/sygus_grammar_norm.h
@@ -207,6 +207,16 @@ class SygusGrammarNorm
      */
     void addConsInfo(SygusGrammarNorm* sygus_norm,
                      const DatatypeConstructor& cons);
+    /**
+     * Returns the total version of Kind k if it is a partial operator, or
+     * otherwise k itself.
+     */
+    static Kind getEliminateKind(Kind k);
+    /**
+     * Returns a version of n where all partial functions such as bvudiv
+     * have been replaced by their total versions like bvudiv_total.
+     */
+    static Node eliminatePartialOperators(Node n);
 
     /** builds a datatype with the information in the type object
      *