Print candidate rewrites in terms of original grammar (#1635)

1 files changed, 31 insertions, 6 deletions

diff --git a/src/theory/quantifiers/sygus_sampler.cpp b/src/theory/quantifiers/sygus_sampler.cpp
index fa0478ac7..0b9254818 100644
--- a/src/theory/quantifiers/sygus_sampler.cpp
+++ b/src/theory/quantifiers/sygus_sampler.cpp
@@ -64,13 +64,17 @@ Node LazyTrie::add(Node n,
   return Node::null();
 }
 
-SygusSampler::SygusSampler() : d_tds(nullptr), d_is_valid(false) {}
+SygusSampler::SygusSampler()
+    : d_tds(nullptr), d_use_sygus_type(false), d_is_valid(false)
+{
+}
 
 void SygusSampler::initialize(TypeNode tn,
                               std::vector<Node>& vars,
                               unsigned nsamples)
 {
   d_tds = nullptr;
+  d_use_sygus_type = false;
   d_is_valid = true;
   d_tn = tn;
   d_ftn = TypeNode::null();
@@ -105,9 +109,13 @@ void SygusSampler::initialize(TypeNode tn,
   initializeSamples(nsamples);
 }
 
-void SygusSampler::initializeSygus(TermDbSygus* tds, Node f, unsigned nsamples)
+void SygusSampler::initializeSygus(TermDbSygus* tds,
+                                   Node f,
+                                   unsigned nsamples,
+                                   bool useSygusType)
 {
   d_tds = tds;
+  d_use_sygus_type = useSygusType;
   d_is_valid = true;
   d_ftn = f.getType();
   Assert(d_ftn.isDatatype());
@@ -282,8 +290,23 @@ Node SygusSampler::registerTerm(Node n, bool forceKeep)
 {
   if (d_is_valid)
   {
-    Assert(n.getType() == d_tn);
-    return d_trie.add(n, this, 0, d_samples.size(), forceKeep);
+    Node bn = n;
+    // if this is a sygus type, get its builtin analog
+    if (d_use_sygus_type)
+    {
+      Assert(!d_ftn.isNull());
+      bn = d_tds->sygusToBuiltin(n);
+      bn = Rewriter::rewrite(bn);
+      d_builtin_to_sygus[bn] = n;
+    }
+    Assert(bn.getType() == d_tn);
+    Node res = d_trie.add(bn, this, 0, d_samples.size(), forceKeep);
+    if (d_use_sygus_type)
+    {
+      Assert(d_builtin_to_sygus.find(res) == d_builtin_to_sygus.end());
+      res = res != bn ? d_builtin_to_sygus[res] : n;
+    }
+    return res;
   }
   return n;
 }
@@ -645,9 +668,11 @@ void SygusSampler::registerSygusType(TypeNode tn)
 
 void SygusSamplerExt::initializeSygusExt(QuantifiersEngine* qe,
                                          Node f,
-                                         unsigned nsamples)
+                                         unsigned nsamples,
+                                         bool useSygusType)
 {
-  SygusSampler::initializeSygus(qe->getTermDatabaseSygus(), f, nsamples);
+  SygusSampler::initializeSygus(
+      qe->getTermDatabaseSygus(), f, nsamples, useSygusType);
 
   // initialize the dynamic rewriter
   std::stringstream ss;