From 904c6e21b11cafe451ac8d5bfeaeb99fb56705ad Mon Sep 17 00:00:00 2001
From: Nikolaj Bjorner <nbjorner@microsoft.com>
Date: Mon, 16 Aug 2021 03:28:21 -0700
Subject: [PATCH] modify #5454

---
 src/sat/smt/array_axioms.cpp | 5 ++++-
 src/sat/smt/array_solver.cpp | 3 +--
 src/sat/smt/array_solver.h   | 3 ++-
 3 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/sat/smt/array_axioms.cpp b/src/sat/smt/array_axioms.cpp
index 372e05c75..99c7ebd18 100644
--- a/src/sat/smt/array_axioms.cpp
+++ b/src/sat/smt/array_axioms.cpp
@@ -175,6 +175,9 @@ namespace array {
         ptr_buffer<expr> sel1_args, sel2_args;
         unsigned num_args = select->get_num_args();
 
+        expr* arg = select->get_arg(0);
+        if (arg != store && !can_beta_reduce(arg) && expr2enode(arg)->get_root() == expr2enode(store)->get_root()) 
+            return false;
         bool has_diff = false;
         for (unsigned i = 1; i < num_args; i++) 
             has_diff |= expr2enode(select->get_arg(i))->get_root() != expr2enode(store->get_arg(i))->get_root();
@@ -182,7 +185,7 @@ namespace array {
             return false;
 
         sel1_args.push_back(store);
-        sel2_args.push_back(store->get_arg(0));        
+        sel2_args.push_back(arg);       
         
         for (unsigned i = 1; i < num_args; i++) {
             sel1_args.push_back(select->get_arg(i));
diff --git a/src/sat/smt/array_solver.cpp b/src/sat/smt/array_solver.cpp
index 47ccfcc62..15978747c 100644
--- a/src/sat/smt/array_solver.cpp
+++ b/src/sat/smt/array_solver.cpp
@@ -272,8 +272,7 @@ namespace array {
         return !get_config().m_array_delay_exp_axiom && d.m_prop_upward;
     }
 
-    bool solver::can_beta_reduce(euf::enode* n) const {
-        expr* c = n->get_expr();
+    bool solver::can_beta_reduce(expr* c) const {
         return a.is_const(c) || a.is_as_array(c) || a.is_store(c) || is_lambda(c) || a.is_map(c);
     }
 }
diff --git a/src/sat/smt/array_solver.h b/src/sat/smt/array_solver.h
index 6cd5274c9..26f2d9022 100644
--- a/src/sat/smt/array_solver.h
+++ b/src/sat/smt/array_solver.h
@@ -209,7 +209,8 @@ namespace array {
         unsigned get_lambda_equiv_size(var_data const& d) const;
         bool should_set_prop_upward(var_data const& d) const;
         bool should_prop_upward(var_data const& d) const;
-        bool can_beta_reduce(euf::enode* n) const;
+        bool can_beta_reduce(euf::enode* n) const { return can_beta_reduce(n->get_expr()); }
+        bool can_beta_reduce(expr* e) const;
 
         var_data& get_var_data(euf::enode* n) { return get_var_data(n->get_th_var(get_id())); }
         var_data& get_var_data(theory_var v) { return *m_var_data[v]; }