Depurification

2 files changed, 96 insertions, 36 deletions

diff --git a/src/theory/strings/skolem_cache.cpp b/src/theory/strings/skolem_cache.cpp
index 7efedc310..05cb82a7c 100644
--- a/src/theory/strings/skolem_cache.cpp
+++ b/src/theory/strings/skolem_cache.cpp
@@ -60,11 +60,41 @@ Node SkolemCache::mkTypedSkolemCached(
     std::tie(id, a, b, c) = normalizeStringSkolem(id, a, b, c);
   }
 
-  std::map<SkolemId, Node>::iterator it = d_skolemCache[a][b][c].find(id);
-  if (it == d_skolemCache[a][b][c].end())
+  std::tuple<Node, Node, Node, SkolemId> args = std::make_tuple(a, b, c, id);
+  const auto& it = d_skolemCache.find(args);
+  if (it == d_skolemCache.end())
   {
-    Node sk = mkTypedSkolem(tn, name);
-    d_skolemCache[a][b][c][id] = sk;
+    NodeManager* nm = NodeManager::currentNM();
+    Node sk;
+    if (id == SK_PREFIX)
+    {
+      Node da = depurify(a);
+      Node db = depurify(b);
+      sk = mkSkolemCached(
+          Rewriter::rewrite(nm->mkNode(STRING_SUBSTR, da, d_zero, db)),
+          SK_PURIFY,
+          "pur");
+    }
+    else if (id == SK_SUFFIX_REM)
+    {
+      Node da = depurify(a);
+      Node db = depurify(b);
+      sk = mkSkolemCached(
+          Rewriter::rewrite(
+              nm->mkNode(STRING_SUBSTR,
+                         da,
+                         db,
+                         nm->mkNode(MINUS, nm->mkNode(STRING_LENGTH, da), db))),
+          SK_PURIFY,
+          "pur");
+    }
+    else
+    {
+      sk = mkTypedSkolem(tn, name);
+    }
+
+    d_skolemCache[args] = sk;
+    d_skolemToArgs[sk] = args;
     return sk;
   }
   return it->second;
@@ -94,6 +124,57 @@ bool SkolemCache::isSkolem(Node n) const
   return d_allSkolems.find(n) != d_allSkolems.end();
 }
 
+Node SkolemCache::depurify(Node n)
+{
+  std::vector<Node> toVisit;
+  std::unordered_map<Node, Node, NodeHashFunction> result;
+
+  toVisit.push_back(n);
+  while (!toVisit.empty())
+  {
+    Node curr = toVisit.back();
+    toVisit.pop_back();
+
+    if (result.find(curr) == result.end())
+    {
+      if (curr.getNumChildren() == 0)
+      {
+        const auto& it = d_skolemToArgs.find(curr);
+        if (it != d_skolemToArgs.end()
+            && std::get<3>(it->second) == SkolemCache::SK_PURIFY)
+        {
+          result[curr] = std::get<0>(it->second);
+        } else {
+          result[curr] = curr;
+        }
+      }
+      else
+      {
+        toVisit.push_back(curr);
+        if (curr.getMetaKind() == metakind::PARAMETERIZED) {
+          toVisit.push_back(curr.getOperator());
+        }
+        toVisit.insert(toVisit.end(), curr.begin(), curr.end());
+        result[curr] = Node::null();
+      }
+    }
+    else if (result[curr].isNull())
+    {
+      NodeBuilder<> nb(curr.getKind());
+      if (curr.getMetaKind() == metakind::PARAMETERIZED) {
+        nb << result[curr.getOperator()];
+      }
+      for (const Node& ncurr : curr)
+      {
+        nb << result[ncurr];
+      }
+      result[curr] = nb;
+    }
+  }
+
+  return result[n];
+}
+
 std::tuple<SkolemCache::SkolemId, Node, Node, Node>
 SkolemCache::normalizeStringSkolem(SkolemId id, Node a, Node b, Node c)
 {
@@ -279,12 +360,12 @@ SkolemCache::normalizeStringSkolem(SkolemId id, Node a, Node b, Node c)
     if (!len.isNull())
     {
       id = SK_PURIFY;
-      a = nm->mkNode(
-          STRING_SUBSTR,
-          a,
-          len,
-          nm->mkNode(
-              MINUS, nm->mkNode(STRING_LENGTH, a), nm->mkNode(PLUS, len, len)));
+      a = Rewriter::rewrite(nm->mkNode(STRING_SUBSTR,
+                                       a,
+                                       len,
+                                       nm->mkNode(MINUS,
+                                                  nm->mkNode(STRING_LENGTH, a),
+                                                  nm->mkNode(PLUS, len, len))));
       b = Node::null();
     }
   }
@@ -299,30 +380,6 @@ SkolemCache::normalizeStringSkolem(SkolemId id, Node a, Node b, Node c)
     }
   }
 
-  if (id == SK_PURIFY && a.getKind() == kind::STRING_SUBSTR)
-  {
-    Node s = a[0];
-    Node n = a[1];
-    Node m = a[2];
-
-    if (n == d_zero)
-    {
-      // SK_PURIFY((str.substr x 0 m)) ---> SK_PREFIX(x, m)
-      id = SK_PREFIX;
-      a = s;
-      b = m;
-    }
-    else if (TheoryStringsRewriter::checkEntailArith(
-                 nm->mkNode(PLUS, n, m), nm->mkNode(STRING_LENGTH, s)))
-    {
-      // SK_PURIFY((str.substr x n m)) ---> SK_SUFFIX_REM(x, n)
-      // if n + m >= (str.len x)
-      id = SK_SUFFIX_REM;
-      a = s;
-      b = n;
-    }
-  }
-
   if (id == SK_PREFIX && b.getKind() == kind::STRING_STRIDOF && a == b[0]
       && b[2] == d_zero)
   {
diff --git a/src/theory/strings/skolem_cache.h b/src/theory/strings/skolem_cache.h
index 08df744de..1085ab43c 100644
--- a/src/theory/strings/skolem_cache.h
+++ b/src/theory/strings/skolem_cache.h
@@ -158,6 +158,8 @@ class SkolemCache
   bool isSkolem(Node n) const;
 
  private:
+  Node depurify(Node n);
+
   /**
    * Simplifies the arguments for a string skolem used for indexing into the
    * cache. In certain cases, we can share skolems with similar arguments e.g.
@@ -182,8 +184,9 @@ class SkolemCache
   /** Constant node zero */
   Node d_zero;
   /** map from node pairs and identifiers to skolems */
-  std::map<Node, std::map<Node, std::map<Node, std::map<SkolemId, Node>>>>
-      d_skolemCache;
+  std::map<std::tuple<Node, Node, Node, SkolemId>, Node> d_skolemCache;
+
+  std::map<Node, std::tuple<Node, Node, Node, SkolemId>> d_skolemToArgs;
   /** the set of all skolems we have generated */
   std::unordered_set<Node, NodeHashFunction> d_allSkolems;
 };