From db935e5fb2eea6bbbd16acc5a1d7f1e2ad9249fc Mon Sep 17 00:00:00 2001
From: ajreynol <andrew.j.reynolds@gmail.com>
Date: Tue, 14 Jan 2025 17:45:05 -0600
Subject: [PATCH] Trying assertion rewrite

---
 src/rewriter/basic_rewrite_rcons.cpp |  2 +-
 src/smt/assertions.cpp               | 25 +++++++++++++++++++++----
 src/smt/assertions.h                 |  5 +++++
 test/regress/cli/CMakeLists.txt      |  1 +
 4 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/src/rewriter/basic_rewrite_rcons.cpp b/src/rewriter/basic_rewrite_rcons.cpp
index 1db452bbd5c..0bc27554e70 100644
--- a/src/rewriter/basic_rewrite_rcons.cpp
+++ b/src/rewriter/basic_rewrite_rcons.cpp
@@ -1600,8 +1600,8 @@ bool BasicRewriteRCons::ensureProofMacroLambdaAppElimShadow(CDProof* cdp,
   size_t i = 0;
   while (i < matchConds.size())
   {
-    Assert(mc.getKind() == Kind::EQUAL);
     Node mc = matchConds[i];
+    Assert(mc.getKind() == Kind::EQUAL);
     i++;
     if (mc[0].getKind() == mc[1].getKind() && mc[0].isClosure())
     {
diff --git a/src/smt/assertions.cpp b/src/smt/assertions.cpp
index 7ac0ea07e57..39d4339a464 100644
--- a/src/smt/assertions.cpp
+++ b/src/smt/assertions.cpp
@@ -26,6 +26,7 @@
 #include "smt/env.h"
 #include "theory/trust_substitutions.h"
 #include "util/result.h"
+#include "proof/proof.h"
 
 using namespace cvc5::internal::theory;
 using namespace cvc5::internal::kind;
@@ -120,12 +121,28 @@ void Assertions::addFormula(TNode n,
   if (isFunDef)
   {
     // if a non-recursive define-fun, just add as a top-level substitution
-    if (n.getKind() == Kind::EQUAL && n[0].isVar())
+    if (n.getKind() == Kind::EQUAL && n[0].isVar() && n[1].getKind()==Kind::LAMBDA)
     {
-      // A define-fun is an assumption in the overall proof, thus
-      // we justify the substitution with ASSUME here.
+      Node defRew = nm->mkNode(Kind::LAMBDA, n[1][0], rewrite(n[1][1]);
+      if (d_env.isProofProducing())
+      {
+        if (d_defFunRewPf==nullptr)
+        {
+          d_defFunRewPf = std::make_shared<CDProof>(d_env);
+        }
+        if (defRew!=n[1])
+        {
+          // A define-fun is an assumption in the overall proof, thus
+          // we justify the substitution with ASSUME here.
+          d_defFunRewPf->addStep(n, ProofRule::ASSUME, {}, {n});
+          Node eqRew = n[1].eqNode(defRew);
+          d_defFunRewPf->addStep(eqRew, ProofRule::MACRO_SR_EQ_INTRO, {}, {n[1]});
+          Node eqFinal = n[0].eqNode(defRew);
+          d_defFunRewPf->addStep(eqFinal, ProofRule::TRANS, {n, eqRew}, {});
+        }
+      }
       d_env.getTopLevelSubstitutions().addSubstitution(
-          n[0], n[1], ProofRule::ASSUME, {}, {n});
+          n[0], defRew, d_defFunRewPf.get());
       return;
     }
   }
diff --git a/src/smt/assertions.h b/src/smt/assertions.h
index b01084b8cf4..1cfd4d82aaf 100644
--- a/src/smt/assertions.h
+++ b/src/smt/assertions.h
@@ -26,6 +26,9 @@
 #include "smt/env_obj.h"
 
 namespace cvc5::internal {
+
+class CDProof;
+
 namespace smt {
 
 class AbstractValues;
@@ -141,6 +144,8 @@ class Assertions : protected EnvObj
    * The list of assumptions from the previous call to checkSatisfiability.
    */
   std::vector<Node> d_assumptions;
+  /** Proof generator storing proofs of rewriting for defined functions */
+  std::shared_ptr<CDProof> d_defFunRewPf;
 };
 
 }  // namespace smt
diff --git a/test/regress/cli/CMakeLists.txt b/test/regress/cli/CMakeLists.txt
index 3fdcb6e3a12..8e2f57f63a4 100644
--- a/test/regress/cli/CMakeLists.txt
+++ b/test/regress/cli/CMakeLists.txt
@@ -1342,6 +1342,7 @@ set(regress_0_tests
   regress0/proofs/dd_ic_macro-elim-shadow.smt2
   regress0/proofs/dd_ic_pf_764.smt2
   regress0/proofs/dd_int_check_bvnot_lam_value_norm.smt2
+  regress0/proofs/dd_int_check_rec_elim_shadow.smt2
   regress0/proofs/dd_issue3481_beta_red.smt2
   regress0/proofs/dd_M301_array_norm_op.smt2
   regress0/proofs/dd_N627_fun_const.smt2