From 0065cd8102dab5f1dff44927be44801e0212bef0 Mon Sep 17 00:00:00 2001
From: Yoav Gross <yoav@starkware.co>
Date: Mon, 20 May 2024 14:07:16 +0300
Subject: [PATCH] fix: support txs without signature

---
 src/block_hash/transaction_commitment.rs      | 20 +++++++----
 src/block_hash/transaction_commitment_test.rs | 34 +++++++++++++------
 2 files changed, 37 insertions(+), 17 deletions(-)

diff --git a/src/block_hash/transaction_commitment.rs b/src/block_hash/transaction_commitment.rs
index b56b7fb..0347e48 100644
--- a/src/block_hash/transaction_commitment.rs
+++ b/src/block_hash/transaction_commitment.rs
@@ -10,24 +10,32 @@ mod transaction_commitment_test;
 
 /// The elements used to calculate a leaf in the transactions Patricia tree.
 #[derive(Clone)]
-pub struct TransactionLeafElements {
-    transaction_hash: TransactionHash,
-    transaction_signature: TransactionSignature,
+pub struct TransactionLeafElement {
+    pub transaction_hash: TransactionHash,
+    pub transaction_signature: Option<TransactionSignature>,
 }
 
 /// Returns the root of a Patricia tree where each leaf is
 /// Poseidon(transaction_hash, transaction_signature).
+/// The leaf of a transaction types without a signature field is: Poseidon(transaction_hash, 0).
 pub fn calculate_transactions_commitment<H: HashFunction>(
-    transaction_leaf_elements: &[TransactionLeafElements],
+    transaction_leaf_elements: &[TransactionLeafElement],
 ) -> TransactionCommitment {
     let transaction_leaves =
         transaction_leaf_elements.iter().map(calculate_transaction_leaf).collect();
     TransactionCommitment(calculate_root::<H>(transaction_leaves))
 }
 
-fn calculate_transaction_leaf(transaction_leaf_elements: &TransactionLeafElements) -> StarkFelt {
+fn calculate_transaction_leaf(transaction_leaf_elements: &TransactionLeafElement) -> StarkFelt {
     HashChain::new()
         .chain(&transaction_leaf_elements.transaction_hash.0)
-        .chain_iter(transaction_leaf_elements.transaction_signature.0.iter())
+        .chain_iter(
+            transaction_leaf_elements
+                .transaction_signature
+                .as_ref()
+                .unwrap_or(&TransactionSignature(vec![StarkFelt::ZERO]))
+                .0
+                .iter(),
+        )
         .get_poseidon_hash()
 }
diff --git a/src/block_hash/transaction_commitment_test.rs b/src/block_hash/transaction_commitment_test.rs
index e4f0623..1f7131e 100644
--- a/src/block_hash/transaction_commitment_test.rs
+++ b/src/block_hash/transaction_commitment_test.rs
@@ -1,5 +1,6 @@
+use super::TransactionLeafElement;
 use crate::block_hash::transaction_commitment::{
-    calculate_transaction_leaf, calculate_transactions_commitment, TransactionLeafElements,
+    calculate_transaction_leaf, calculate_transactions_commitment,
 };
 use crate::core::TransactionCommitment;
 use crate::hash::{PoseidonHashCalculator, StarkFelt};
@@ -7,11 +8,7 @@ use crate::transaction::{TransactionHash, TransactionSignature};
 
 #[test]
 fn test_transaction_leaf_regression() {
-    let transaction_hash = TransactionHash(StarkFelt::ONE);
-    let transaction_signature = TransactionSignature(vec![StarkFelt::TWO, StarkFelt::THREE]);
-    let transaction_leaf_elements =
-        TransactionLeafElements { transaction_hash, transaction_signature };
-
+    let transaction_leaf_elements = get_transaction_leaf_element();
     let expected_leaf =
         StarkFelt::try_from("0x2f0d8840bcf3bc629598d8a6cc80cb7c0d9e52d93dab244bbf9cd0dca0ad082")
             .unwrap();
@@ -20,12 +17,21 @@ fn test_transaction_leaf_regression() {
 }
 
 #[test]
-fn test_transactions_commitment_regression() {
-    let transaction_hash = TransactionHash(StarkFelt::ONE);
-    let transaction_signature = TransactionSignature(vec![StarkFelt::TWO, StarkFelt::THREE]);
-    let transaction_leaf_elements =
-        TransactionLeafElements { transaction_hash, transaction_signature };
+fn test_transaction_leaf_without_signature_regression() {
+    let transaction_leaf_elements = TransactionLeafElement {
+        transaction_hash: TransactionHash(StarkFelt::ONE),
+        transaction_signature: None,
+    };
+    let expected_leaf =
+        StarkFelt::try_from("0x00a93bf5e58b9378d093aa86ddc2f61a3295a1d1e665bd0ef3384dd07b30e033")
+            .unwrap();
 
+    assert_eq!(expected_leaf, calculate_transaction_leaf(&transaction_leaf_elements));
+}
+
+#[test]
+fn test_transactions_commitment_regression() {
+    let transaction_leaf_elements = get_transaction_leaf_element();
     let expected_root =
         StarkFelt::try_from("0x0282b635972328bd1cfa86496fe920d20bd9440cd78ee8dc90ae2b383d664dcf")
             .unwrap();
@@ -38,3 +44,9 @@ fn test_transactions_commitment_regression() {
         ],)
     );
 }
+
+fn get_transaction_leaf_element() -> TransactionLeafElement {
+    let transaction_hash = TransactionHash(StarkFelt::ONE);
+    let transaction_signature = TransactionSignature(vec![StarkFelt::TWO, StarkFelt::THREE]);
+    TransactionLeafElement { transaction_hash, transaction_signature: Some(transaction_signature) }
+}