From 014a84b541cf8626e18e9830d8490da23a5397c2 Mon Sep 17 00:00:00 2001
From: Brice Dobry <brice@hiro.so>
Date: Wed, 11 Dec 2024 10:30:47 -0500
Subject: [PATCH] fix: set epoch to 3.1 in Clarity DB

---
 stackslib/src/chainstate/stacks/db/blocks.rs |  3 +--
 stackslib/src/clarity_vm/clarity.rs          | 26 ++++++++++++++++++++
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/stackslib/src/chainstate/stacks/db/blocks.rs b/stackslib/src/chainstate/stacks/db/blocks.rs
index 791f2064dc..10d691b40f 100644
--- a/stackslib/src/chainstate/stacks/db/blocks.rs
+++ b/stackslib/src/chainstate/stacks/db/blocks.rs
@@ -4104,8 +4104,7 @@ impl StacksChainState {
                         current_epoch = StacksEpochId::Epoch30;
                     }
                     StacksEpochId::Epoch30 => {
-                        // no special initialization is needed, since only the coinbase emission
-                        // schedule is changing.
+                        receipts.append(&mut clarity_tx.block.initialize_epoch_3_1()?);
                         current_epoch = StacksEpochId::Epoch31;
                     }
                     StacksEpochId::Epoch31 => {
diff --git a/stackslib/src/clarity_vm/clarity.rs b/stackslib/src/clarity_vm/clarity.rs
index c89679f414..a412a4d304 100644
--- a/stackslib/src/clarity_vm/clarity.rs
+++ b/stackslib/src/clarity_vm/clarity.rs
@@ -1524,6 +1524,32 @@ impl<'a, 'b> ClarityBlockConnection<'a, 'b> {
         })
     }
 
+    pub fn initialize_epoch_3_1(&mut self) -> Result<Vec<StacksTransactionReceipt>, Error> {
+        // use the `using!` statement to ensure that the old cost_tracker is placed
+        //  back in all branches after initialization
+        using!(self.cost_track, "cost tracker", |old_cost_tracker| {
+            // epoch initialization is *free*.
+            // NOTE: this also means that cost functions won't be evaluated.
+            self.cost_track.replace(LimitedCostTracker::new_free());
+            self.epoch = StacksEpochId::Epoch31;
+            self.as_transaction(|tx_conn| {
+                // bump the epoch in the Clarity DB
+                tx_conn
+                    .with_clarity_db(|db| {
+                        db.set_clarity_epoch_version(StacksEpochId::Epoch31)?;
+                        Ok(())
+                    })
+                    .unwrap();
+
+                // require 3.1 rules henceforth in this connection as well
+                tx_conn.epoch = StacksEpochId::Epoch31;
+            });
+
+            debug!("Epoch 3.1 initialized");
+            (old_cost_tracker, Ok(vec![]))
+        })
+    }
+
     pub fn start_transaction_processing<'c>(&'c mut self) -> ClarityTransactionConnection<'c, 'a> {
         let store = &mut self.datastore;
         let cost_track = &mut self.cost_track;