diff --git a/src/smt/theory_array_full.cpp b/src/smt/theory_array_full.cpp
index 3b566cd3e..cc13ca707 100644
--- a/src/smt/theory_array_full.cpp
+++ b/src/smt/theory_array_full.cpp
@@ -29,8 +29,8 @@ namespace smt {
 
     theory_array_full::theory_array_full(context& ctx) : 
         theory_array(ctx),
-        m_sort2epsilon0(ctx.get_manager()),
-        m_sort2epsilon1(ctx.get_manager()) {}
+        m_sort2epsilon(ctx.get_manager()),
+        m_sort2diag(ctx.get_manager()) {}
 
     theory_array_full::~theory_array_full() {
         std::for_each(m_var_data_full.begin(), m_var_data_full.end(), delete_proc<var_data_full>());
@@ -591,7 +591,7 @@ namespace smt {
         return true;        
     }
 
-    bool theory_array_full::has_large_domain(app* array_term) {
+   bool theory_array_full::has_large_domain(app* array_term) {
         SASSERT(is_array_sort(array_term));
         sort* s = m.get_sort(array_term);
         unsigned dim = get_dimension(s);
@@ -691,56 +691,66 @@ namespace smt {
 
         m_stats.m_num_default_store_axiom++;
 
-        app_ref def1(m), def2(m);
+        expr_ref def1(m), def2(m);
 
         TRACE("array", tout << mk_bounded_pp(store_app, m) << "\n";);
 
+        unsigned num_args = store_app->get_num_args();
+
+        def1 = mk_default(store_app);
         def2 = mk_default(store_app->get_arg(0));
+
         bool is_new = false;
-        if (!has_large_domain(store_app) && !has_unitary_domain(store_app)) {
+
+        if (has_unitary_domain(store_app)) {
+            def2 = store_app->get_arg(num_args - 1);
+        }        
+        else if (!has_large_domain(store_app)) {
             //
             // let A = store(B, i, v)
             // 
             // Add:
             //   default(A) = ite(epsilon1 = i, v, default(B))
-            //   A[epsilon2] = B[epsilon2]
+            //   A[diag(i)] = B[diag(i)]
             // 
-            expr_ref_vector eqs(m), args1(m), args2(m);
+            expr_ref_vector eqs(m);
+            expr_ref_vector args1(m), args2(m);
             args1.push_back(store_app->get_arg(0));
             args2.push_back(store_app);
-            unsigned num_args = store_app->get_num_args();
+
             for (unsigned i = 1; i + 1 < num_args; ++i) {
-                sort* srt = m.get_sort(store_app->get_arg(i));
+                expr* arg = store_app->get_arg(i);
+                sort* srt = m.get_sort(arg);
                 auto ep = mk_epsilon(srt);
-                eqs.push_back(m.mk_eq(ep.first, store_app->get_arg(i)));
-                args1.push_back(ep.second);
-                args2.push_back(ep.second);
-            }
-            
+                eqs.push_back(m.mk_eq(ep.first, arg));
+                args1.push_back(m.mk_app(ep.second, arg));
+                args2.push_back(m.mk_app(ep.second, arg));
+            }            
             expr_ref eq(mk_and(eqs), m);
-            def2 = m.mk_ite(eq, store_app->get_arg(num_args-1), def2); 
-            
+            def2 = m.mk_ite(eq, store_app->get_arg(num_args-1), def2);             
             app_ref sel1(m), sel2(m);
             sel1 = mk_select(args1);
             sel2 = mk_select(args2);
             is_new = try_assign_eq(sel1, sel2);
         }
 
-        def1 = mk_default(store_app);
         ctx.internalize(def1, false);
         ctx.internalize(def2, false);
         return try_assign_eq(def1, def2) || is_new;
     }
 
-    std::pair<app*,app*> theory_array_full::mk_epsilon(sort* s) {
-        app* eps0 = nullptr, *eps1 = nullptr;
-        if (!m_sort2epsilon0.find(s, eps0) || !m_sort2epsilon1.find(s, eps1)) {
-            eps0 = m.mk_fresh_const("epsilon", s);
-            eps1 = m.mk_fresh_const("epsilon", s);
-            m_trail_stack.push(ast2ast_trail<theory_array, sort, app>(m_sort2epsilon0, s, eps0));   
-            m_trail_stack.push(ast2ast_trail<theory_array, sort, app>(m_sort2epsilon1, s, eps1));  
-        } 
-        return std::make_pair(eps0, eps1);
+    std::pair<app*,func_decl*> theory_array_full::mk_epsilon(sort* s) {
+        app* eps = nullptr;
+        func_decl* diag = nullptr;
+        if (!m_sort2epsilon.find(s, eps)) {
+            eps = m.mk_fresh_const("epsilon", s);
+            m_trail_stack.push(ast2ast_trail<theory_array, sort, app>(m_sort2epsilon, s, eps));   
+        }
+        if (!m_sort2diag.find(s, diag)) {
+            diag = m.mk_fresh_func_decl("diag", 1, &s, s);
+            m_trail_stack.push(ast2ast_trail<theory_array, sort, func_decl>(m_sort2diag, s, diag));   
+        }
+        return std::make_pair(eps, diag);
     }
 
     final_check_status theory_array_full::assert_delayed_axioms() {        
diff --git a/src/smt/theory_array_full.h b/src/smt/theory_array_full.h
index 56e46e352..e5bb2668d 100644
--- a/src/smt/theory_array_full.h
+++ b/src/smt/theory_array_full.h
@@ -34,7 +34,8 @@ namespace smt {
 
         ptr_vector<var_data_full> m_var_data_full;
 
-        ast2ast_trailmap<sort, app> m_sort2epsilon0, m_sort2epsilon1;
+        ast2ast_trailmap<sort, app> m_sort2epsilon;
+        ast2ast_trailmap<sort, func_decl> m_sort2diag;
         obj_pair_map<expr, expr, bool> m_eqs;
         
         static unsigned const m_default_map_fingerprint = UINT_MAX - 112;
@@ -80,7 +81,7 @@ namespace smt {
 
         bool has_large_domain(app* array_term);
         bool has_unitary_domain(app* array_term);
-        std::pair<app*,app*> mk_epsilon(sort* s);
+        std::pair<app*,func_decl*> mk_epsilon(sort* s);
 
         bool instantiate_select_const_axiom(enode* select, enode* cnst);
         bool instantiate_select_as_array_axiom(enode* select, enode* arr);