Fix: find_element(): No value found for key (#378)

* fix declare then deploy blocks

* add blocks to CI

---------

Co-authored-by: Herman Obst Demaestri <70286869+HermanObst@users.noreply.github.com>

crates/bin/prove_block/src/state_utils.rs

@@ -52,7 +52,7 @@ pub(crate) async fn get_formatted_state_update(
         state_diff.declared_classes.iter().map(|declared_item| declared_item.class_hash).collect();
 
     // TODO: Handle deprecated classes
-    let mut class_hash_to_compiled_class_hash: HashMap<Felt252, Felt252> = format_declared_classes(&state_diff);
+    let mut class_hash_to_compiled_class_hash: HashMap<Felt252, Felt252> = HashMap::new();
     let (compiled_contract_classes, deprecated_compiled_contract_classes, declared_class_hash_component_hashes) =
         build_compiled_class_and_maybe_update_class_hash_to_compiled_class_hash(
             rpc_client,
@@ -65,6 +65,9 @@ pub(crate) async fn get_formatted_state_update(
         )
         .await?;
 
+    // OS will expect a Zero in compiled_class_hash for new classes. Overwrite the needed entries.
+    format_declared_classes(&state_diff, &mut class_hash_to_compiled_class_hash);
+
     Ok((
         FormattedStateUpdate {
             class_hash_to_compiled_class_hash,
@@ -232,15 +235,18 @@ async fn build_compiled_class_and_maybe_update_class_hash_to_compiled_class_hash
     Ok((compiled_contract_classes, deprecated_compiled_contract_classes, declared_class_hash_to_component_hashes))
 }
 
-/// Format StateDiff's DeclaredClassItem to a HashMap<class_hash, compiled_class_hash>
-fn format_declared_classes(state_diff: &StateDiff) -> HashMap<Felt252, Felt252> {
+fn format_declared_classes(state_diff: &StateDiff, class_hash_to_compiled_class_hash: &mut HashMap<Felt252, Felt252>) {
     // The comment below explicits that the value should be 0 for new classes:
     // From execute_transactions.cairo
     // Note that prev_value=0 enforces that a class may be declared only once.
     // dict_update{dict_ptr=contract_class_changes}(
     //     key=[class_hash_ptr], prev_value=0, new_value=compiled_class_hash
     // );
-    let class_hash_to_compiled_class_hash =
-        state_diff.declared_classes.iter().map(|class| (class.class_hash, Felt::ZERO)).collect();
-    class_hash_to_compiled_class_hash
+
+    // class_hash_to_compiled_class_hash is already populated. However, for classes
+    // that are defined in state_diff.declared_classes, we need to set the
+    // compiled_class_hashes to zero as it was explain above
+    for class in state_diff.declared_classes.iter() {
+        class_hash_to_compiled_class_hash.insert(class.class_hash, Felt::ZERO);
+    }
 }

crates/bin/prove_block/tests/prove_block.rs

@@ -15,6 +15,7 @@ use rstest::rstest;
 // # * 97581, 101556, 102076 deploy account txns
 // # * 155016 / 125622 fix writes to zero (storage value not included in the tree)
 // # * 160035: EvalCircuit hint used
+// # * 164333 / 169203: Declare and Deploy on the same block
 #[rstest]
 #[case::small_block_with_only_invoke_txs(76793)]
 #[case::additional_basic_blocks_1(76766)]
@@ -37,6 +38,8 @@ use rstest::rstest;
 #[case::deploy_account_many_txs(102076)]
 #[case::edge_bottom_not_found(155016)]
 #[case::eval_circuit(160035)]
+#[case::declare_and_deploy_in_same_block(164333)]
+#[case::declare_and_deploy_in_same_block(169206)]
 #[ignore = "Requires a running Pathfinder node"]
 #[tokio::test(flavor = "multi_thread")]
 async fn test_prove_selected_blocks(#[case] block_number: u64) {