From 23df33f277f12e43df2c0b676578f90d53dba7bb Mon Sep 17 00:00:00 2001
From: Andres Noetzli <andres.noetzli@gmail.com>
Date: Fri, 18 Feb 2022 12:18:04 -0800
Subject: [PATCH] Fix `STRINGS_ARRAY_NTH_UPDATE` inference (#8121)

Our `STRINGS_ARRAY_NTH_UPDATE` inference was missing an explanation that
establishes the relation between the update and the nth term. On the
current set of benchmarks, this does not cause an issue, but there
should be cases where it does. The commit also changes the array core solver
to take better advantage of `d_indexMap`.
---
 src/theory/strings/array_core_solver.cpp | 35 ++++++++++++++----------
 1 file changed, 21 insertions(+), 14 deletions(-)

diff --git a/src/theory/strings/array_core_solver.cpp b/src/theory/strings/array_core_solver.cpp
index c79664a32..e31a5e5a4 100644
--- a/src/theory/strings/array_core_solver.cpp
+++ b/src/theory/strings/array_core_solver.cpp
@@ -159,23 +159,21 @@ void ArrayCoreSolver::checkUpdate(const std::vector<Node>& updateTerms)
           exp, lem, InferenceId::STRINGS_ARRAY_UPDATE_BOUND, false, true);
     }
 
-    for (const auto& nthIdxs : d_indexMap)
+    Node rn = d_state.getRepresentative(n);
+    Node rs = d_state.getRepresentative(n[0]);
+    for (const Node& r : {rn, rs})
     {
       // Enumerate n-th terms for sequences that are related to the current
       // update term
-      Node seq = nthIdxs.first;
-      if (!d_state.areEqual(seq, n) && !d_state.areEqual(seq, n[0]))
-      {
-        continue;
-      }
-
-      const std::set<Node>& indexes = nthIdxs.second;
-      Trace("seq-array-core-debug") << "  check nth for " << seq
+      const std::set<Node>& indexes = d_indexMap[r];
+      Trace("seq-array-core-debug") << "  check nth for " << r
                                     << " with indices " << indexes << std::endl;
       Node i = n[1];
       for (Node j : indexes)
       {
-        // nth(update(s, n, t), m)
+        // nth(x, m)
+        // y = update(s, n, t), m)
+        // x = y or x = s
         // ------------------------
         // nth(update(s, n, t)) =
         //   ite(0 <= m < len(s),
@@ -197,6 +195,15 @@ void ArrayCoreSolver::checkUpdate(const std::vector<Node>& updateTerms)
 
         std::vector<Node> exp;
         d_im.addToExplanation(termProxy, n, exp);
+        if (d_state.areEqual(r, n))
+        {
+          d_im.addToExplanation(r, n, exp);
+        }
+        else
+        {
+          Assert(d_state.areEqual(r, n[0]));
+          d_im.addToExplanation(r, n[0], exp);
+        }
         sendInference(exp, lem, InferenceId::STRINGS_ARRAY_NTH_UPDATE);
       }
     }
@@ -232,10 +239,10 @@ void ArrayCoreSolver::check(const std::vector<Node>& nthTerms,
   {
     // (seq.nth n[0] n[1])
     Node r = d_state.getRepresentative(n[0]);
-    Trace("seq-update") << "- " << r << ": " << n[1] << " -> " << n
-                        << std::endl;
-    d_writeModel[r][n[1]] = n;
-    d_indexMap[r].insert(n[1]);
+    Node ri = d_state.getRepresentative(n[1]);
+    Trace("seq-update") << "- " << r << ": " << ri << " -> " << n << std::endl;
+    d_writeModel[r][ri] = n;
+    d_indexMap[r].insert(ri);
 
     if (n[0].getKind() == STRING_REV)
     {
-- 
2.30.2