Optimize make_aggsig_final_message() (#682)

* initial commit * extend_from_slice in test
Chia-Network · Aug 29, 2024 · c134fed · c134fed
1 parent a5bf674
commit c134fed
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 33 deletions.
diff --git a/crates/chia-consensus/src/gen/make_aggsig_final_message.rs b/crates/chia-consensus/src/gen/make_aggsig_final_message.rs
@@ -9,50 +9,46 @@ use chia_protocol::Coin;
 
 pub fn make_aggsig_final_message(
     opcode: ConditionOpcode,
-    msg: &[u8],
+    msg: &mut Vec<u8>,
     spend: &OwnedSpendConditions,
     constants: &ConsensusConstants,
-) -> Vec<u8> {
-    let mut result = Vec::<u8>::with_capacity(msg.len() + 96);
-    result.extend(msg);
+) {
     match opcode {
         AGG_SIG_PARENT => {
-            result.extend(spend.parent_id.as_slice());
-            result.extend(constants.agg_sig_parent_additional_data.as_slice());
+            msg.extend(spend.parent_id.as_slice());
+            msg.extend(constants.agg_sig_parent_additional_data.as_slice());
         }
         AGG_SIG_PUZZLE => {
-            result.extend(spend.puzzle_hash.as_slice());
-            result.extend(constants.agg_sig_puzzle_additional_data.as_slice());
+            msg.extend(spend.puzzle_hash.as_slice());
+            msg.extend(constants.agg_sig_puzzle_additional_data.as_slice());
         }
         AGG_SIG_AMOUNT => {
-            result.extend(u64_to_bytes(spend.coin_amount).as_slice());
-            result.extend(constants.agg_sig_amount_additional_data.as_slice());
+            msg.extend(u64_to_bytes(spend.coin_amount).as_slice());
+            msg.extend(constants.agg_sig_amount_additional_data.as_slice());
         }
         AGG_SIG_PUZZLE_AMOUNT => {
-            result.extend(spend.puzzle_hash.as_slice());
-            result.extend(u64_to_bytes(spend.coin_amount).as_slice());
-            result.extend(constants.agg_sig_puzzle_amount_additional_data.as_slice());
+            msg.extend(spend.puzzle_hash.as_slice());
+            msg.extend(u64_to_bytes(spend.coin_amount).as_slice());
+            msg.extend(constants.agg_sig_puzzle_amount_additional_data.as_slice());
         }
         AGG_SIG_PARENT_AMOUNT => {
-            result.extend(spend.parent_id.as_slice());
-            result.extend(u64_to_bytes(spend.coin_amount).as_slice());
-            result.extend(constants.agg_sig_parent_amount_additional_data.as_slice());
+            msg.extend(spend.parent_id.as_slice());
+            msg.extend(u64_to_bytes(spend.coin_amount).as_slice());
+            msg.extend(constants.agg_sig_parent_amount_additional_data.as_slice());
         }
         AGG_SIG_PARENT_PUZZLE => {
-            result.extend(spend.parent_id.as_slice());
-            result.extend(spend.puzzle_hash.as_slice());
-            result.extend(constants.agg_sig_parent_puzzle_additional_data.as_slice());
+            msg.extend(spend.parent_id.as_slice());
+            msg.extend(spend.puzzle_hash.as_slice());
+            msg.extend(constants.agg_sig_parent_puzzle_additional_data.as_slice());
         }
         AGG_SIG_ME => {
             let coin: Coin = Coin::new(spend.parent_id, spend.puzzle_hash, spend.coin_amount);
 
-            result.extend(coin.coin_id().as_slice());
-            result.extend(constants.agg_sig_me_additional_data.as_slice());
+            msg.extend(coin.coin_id().as_slice());
+            msg.extend(constants.agg_sig_me_additional_data.as_slice());
         }
-        _ => return result,
-    };
-
-    result
+        _ => {}
+    }
 }
 
 fn u64_to_bytes(val: u64) -> Bytes {
@@ -120,10 +116,10 @@ mod tests {
             hex!("4444444444444444444444444444444444444444444444444444444444444444").into();
         let puzzle_hash: Vec<u8> =
             hex!("3333333333333333333333333333333333333333333333333333333333333333").into();
-        let msg = b"message";
+        let mut msg = b"message".to_vec();
 
         let mut expected_result = Vec::<u8>::new();
-        expected_result.extend(msg);
+        expected_result.extend_from_slice(msg.as_slice());
 
         let coin = Coin::new(
             Bytes32::try_from(parent_id.clone()).expect("test should pass"),
@@ -188,7 +184,7 @@ mod tests {
 
         let spend = OwnedSpendConditions::from(&a, spend);
 
-        let result = make_aggsig_final_message(opcode, msg, &spend, &TEST_CONSTANTS);
-        assert_eq!(result, expected_result);
+        make_aggsig_final_message(opcode, &mut msg, &spend, &TEST_CONSTANTS);
+        assert_eq!(msg, expected_result);
     }
 }
diff --git a/crates/chia-consensus/src/spendbundle_validation.rs b/crates/chia-consensus/src/spendbundle_validation.rs
@@ -37,6 +37,9 @@ pub fn validate_clvm_and_signature(
     // Collect all pairs in a single vector to avoid multiple iterations
     let mut pairs = Vec::new();
 
+    let mut aug_msg = Vec::<u8>::new();
+    let mut final_msg = Vec::<u8>::new();
+
     for spend in &npcresult.spends {
         let condition_items_pairs = [
             (AGG_SIG_PARENT, &spend.agg_sig_parent),
@@ -50,12 +53,15 @@ pub fn validate_clvm_and_signature(
 
         for (condition, items) in condition_items_pairs {
             for (pk, msg) in items {
-                let mut aug_msg = pk.to_bytes().to_vec();
-                let msg = make_aggsig_final_message(condition, msg.as_slice(), spend, constants);
-                aug_msg.extend_from_slice(msg.as_ref());
+                aug_msg.clear();
+                final_msg.clear();
+                final_msg.extend_from_slice(msg.as_slice());
+                aug_msg.extend_from_slice(&pk.to_bytes());
+                make_aggsig_final_message(condition, &mut final_msg, spend, constants);
+                aug_msg.extend(&final_msg);
                 let aug_hash = hash_to_g2(&aug_msg);
                 let pairing = aug_hash.pair(pk);
-                pairs.push((hash_pk_and_msg(&pk.to_bytes(), &msg), pairing));
+                pairs.push((hash_pk_and_msg(&pk.to_bytes(), &final_msg), pairing));
             }
         }
     }