From f77037e9a56adadf1b6c988f10c6e9aceb00a7fe Mon Sep 17 00:00:00 2001
From: Nikolaj Bjorner <nbjorner@microsoft.com>
Date: Wed, 25 May 2022 20:23:43 -0400
Subject: [PATCH] expand select/store when I/J are values #6053

Signed-off-by: Nikolaj Bjorner <nbjorner@microsoft.com>
---
 src/ast/array_decl_plugin.cpp       | 19 +++++++++++++++++++
 src/ast/array_decl_plugin.h         |  3 +++
 src/ast/rewriter/array_rewriter.cpp | 18 ++++++++++++++++--
 3 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/src/ast/array_decl_plugin.cpp b/src/ast/array_decl_plugin.cpp
index 31bbaac1e..b336dd2e5 100644
--- a/src/ast/array_decl_plugin.cpp
+++ b/src/ast/array_decl_plugin.cpp
@@ -630,6 +630,24 @@ bool array_decl_plugin::is_fully_interp(sort * s) const {
     return m_manager->is_fully_interp(get_array_range(s));
 }
 
+bool array_decl_plugin::is_value(app * _e) const {
+    expr* e = _e;
+    array_util u(*m_manager);
+    while (true) {
+        if (u.is_const(e, e))
+            return m_manager->is_value(e);
+        if (u.is_store(e)) {            
+            for (unsigned i = 1; i < to_app(e)->get_num_args(); ++i)
+                if (!m_manager->is_value(to_app(e)->get_arg(i)))
+                    return false;
+            e = to_app(e)->get_arg(0);
+            continue;
+        }
+        return false;
+    }
+}
+
+
 func_decl * array_recognizers::get_as_array_func_decl(expr * n) const { 
     SASSERT(is_as_array(n)); 
     return to_func_decl(to_app(n)->get_decl()->get_parameter(0).get_ast()); 
@@ -704,3 +722,4 @@ func_decl* array_util::mk_array_ext(sort *domain, unsigned i) {
     parameter p(i);
     return m_manager.mk_func_decl(m_fid, OP_ARRAY_EXT, 1, &p, 2, domains);
 }
+
diff --git a/src/ast/array_decl_plugin.h b/src/ast/array_decl_plugin.h
index dd8443025..51f380243 100644
--- a/src/ast/array_decl_plugin.h
+++ b/src/ast/array_decl_plugin.h
@@ -137,6 +137,9 @@ class array_decl_plugin : public decl_plugin {
     expr * get_some_value(sort * s) override;
 
     bool is_fully_interp(sort * s) const override;
+
+    bool is_value(app * e) const override;
+
 };
 
 class array_recognizers {
diff --git a/src/ast/rewriter/array_rewriter.cpp b/src/ast/rewriter/array_rewriter.cpp
index 88a0ee28e..b8b30a206 100644
--- a/src/ast/rewriter/array_rewriter.cpp
+++ b/src/ast/rewriter/array_rewriter.cpp
@@ -179,8 +179,21 @@ br_status array_rewriter::mk_select_core(unsigned num_args, expr * const * args,
             result = m().mk_app(get_fid(), OP_SELECT, num_args, new_args.data());
             return BR_REWRITE1;
         }
-        default:
-            if (m_blast_select_store || (m_expand_select_store && to_app(args[0])->get_arg(0)->get_ref_count() == 1)) {
+        default: {
+            auto are_values = [&]() {
+                for (unsigned i = 1; i < num_args; ++i) {
+                    if (!m().is_value(args[i]))
+                        return false;
+                    if (!m().is_value(to_app(args[0])->get_arg(i)))
+                        return false;
+                }
+                return true;
+            };
+            bool should_expand =
+                m_blast_select_store ||
+                are_values() ||
+                (m_expand_select_store && to_app(args[0])->get_arg(0)->get_ref_count() == 1);
+            if (should_expand) {
                 // select(store(a, I, v), J) --> ite(I=J, v, select(a, J))
                 ptr_buffer<expr> new_args;
                 new_args.push_back(to_app(args[0])->get_arg(0));
@@ -203,6 +216,7 @@ br_status array_rewriter::mk_select_core(unsigned num_args, expr * const * args,
             }
             return BR_FAILED;
         }
+        }
     }
 
     if (m_util.is_const(args[0])) {