Disconnect two-way peg data at the connecting block height

lib/state/two_way_peg_data.rs

@@ -1049,7 +1049,7 @@ pub fn disconnect(
             latest_withdrawal_bundle_event_block_hash,
             last_withdrawal_bundle_event_block_hash
         );
-        assert_eq!(block_height - 1, last_withdrawal_bundle_event_block_height);
+        assert_eq!(block_height, last_withdrawal_bundle_event_block_height);
         if !state
             .withdrawal_bundle_event_blocks
             .delete(rwtxn, &last_withdrawal_bundle_event_block_seq_idx)?
@@ -1062,14 +1062,14 @@ pub fn disconnect(
         .map(|(height, _bundle)| height)
         .unwrap_or_default();
     if block_height - last_withdrawal_bundle_failure_height
-        > WITHDRAWAL_BUNDLE_FAILURE_GAP
+        >= WITHDRAWAL_BUNDLE_FAILURE_GAP
         && let Some(bundle_m6id) =
             state.pending_withdrawal_bundle.try_get(rwtxn, &())?
         && let (bundle, bundle_status) = state
             .withdrawal_bundles
             .try_get(rwtxn, &bundle_m6id)?
             .ok_or(error::PendingWithdrawalBundleUnknown(bundle_m6id))?
-        && bundle_status.latest().height == block_height - 1
+        && bundle_status.latest().height == block_height
     {
         state.pending_withdrawal_bundle.delete(rwtxn, &())?;
         if let (Some(bundle_status), _latest_bundle_status) =
@@ -1094,7 +1094,7 @@ pub fn disconnect(
             .last(rwtxn)?
             .ok_or(Error::NoDepositBlock)?;
         assert_eq!(latest_deposit_block_hash, last_deposit_block_hash);
-        assert_eq!(block_height - 1, last_deposit_block_height);
+        assert_eq!(block_height, last_deposit_block_height);
         if !state
             .deposit_blocks
             .delete(rwtxn, &last_deposit_block_seq_idx)?
@@ -1255,12 +1255,12 @@ mod tests {
             .unwrap();
         state
             .withdrawal_bundle_event_blocks
-            .put(&mut rwtxn, &0, &(event_block_hash, block_height - 1))
+            .put(&mut rwtxn, &0, &(event_block_hash, block_height))
             .unwrap();
         // a deposit record at the same sequence index that must survive
         state
             .deposit_blocks
-            .put(&mut rwtxn, &0, &(deposit_block_hash, block_height - 1))
+            .put(&mut rwtxn, &0, &(deposit_block_hash, block_height))
             .unwrap();
         rwtxn.commit().unwrap();
 
@@ -1403,4 +1403,98 @@ mod tests {
         );
         Ok(())
     }
+
+    // connecting a deposit then disconnecting it on a reorg must round-trip
+    #[test]
+    fn deposit_reorg_round_trips() {
+        use crate::types::proto::mainchain::Deposit;
+        use crate::types::{Body, FilledTransaction, Header};
+
+        let env = temp_env("deposit_reorg_round_trips");
+        let state = State::new(&env).unwrap();
+
+        let empty_body = Body {
+            coinbase: Vec::new(),
+            transactions: Vec::new(),
+            authorizations: Vec::new(),
+        };
+        let no_txs: &[FilledTransaction] = &[];
+        let merkle_root = Body::compute_merkle_root(&[], no_txs).unwrap();
+        let main0 = bitcoin::BlockHash::from_byte_array([10; 32]);
+        let main1 = bitcoin::BlockHash::from_byte_array([11; 32]);
+
+        let genesis = Header {
+            merkle_root,
+            prev_side_hash: None,
+            prev_main_hash: main0,
+            roots: Vec::new(),
+        };
+        {
+            let mut rwtxn = env.write_txn().unwrap();
+            state
+                .apply_block(&mut rwtxn, &genesis, &empty_body)
+                .unwrap();
+            state
+                .connect_two_way_peg_data(&mut rwtxn, &TwoWayPegData::default())
+                .unwrap();
+            rwtxn.commit().unwrap();
+        }
+
+        let block1 = Header {
+            merkle_root,
+            prev_side_hash: Some(genesis.hash()),
+            prev_main_hash: main1,
+            roots: Vec::new(),
+        };
+        let deposit_outpoint = bitcoin::OutPoint {
+            txid: bitcoin::Txid::from_byte_array([2; 32]),
+            vout: 0,
+        };
+        let deposit_key =
+            OutPointKey::from(&OutPoint::Deposit(deposit_outpoint));
+        let deposit_twpd = {
+            let mut block_info = LinkedHashMap::new();
+            block_info.insert(
+                main1,
+                BlockInfo {
+                    bmm_commitment: None,
+                    events: vec![BlockEvent::Deposit(Deposit {
+                        tx_index: 0,
+                        outpoint: deposit_outpoint,
+                        output: Output {
+                            address: Address::ALL_ZEROS,
+                            content: OutputContent::Value(
+                                bitcoin::Amount::from_sat(1000),
+                            ),
+                        },
+                    })],
+                },
+            );
+            TwoWayPegData { block_info }
+        };
+        {
+            let mut rwtxn = env.write_txn().unwrap();
+            state.apply_block(&mut rwtxn, &block1, &empty_body).unwrap();
+            state
+                .connect_two_way_peg_data(&mut rwtxn, &deposit_twpd)
+                .unwrap();
+            assert!(
+                state.utxos.try_get(&rwtxn, &deposit_key).unwrap().is_some()
+            );
+            assert!(state.deposit_blocks.last(&rwtxn).unwrap().is_some());
+            rwtxn.commit().unwrap();
+        }
+
+        {
+            let mut rwtxn = env.write_txn().unwrap();
+            state
+                .disconnect_two_way_peg_data(&mut rwtxn, &deposit_twpd)
+                .unwrap();
+            assert!(
+                state.utxos.try_get(&rwtxn, &deposit_key).unwrap().is_none()
+            );
+            assert!(state.deposit_blocks.last(&rwtxn).unwrap().is_none());
+            rwtxn.commit().unwrap();
+        }
+    }
 }