From 88a00e7cf4ed152f363a3c3e9d384c7de940c160 Mon Sep 17 00:00:00 2001
From: greged93 <82421016+greged93@users.noreply.github.com>
Date: Wed, 20 Mar 2024 11:41:11 +0100
Subject: [PATCH] dev: small fix (#876)

small fix
---
 src/test_utils/hive/mod.rs | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/test_utils/hive/mod.rs b/src/test_utils/hive/mod.rs
index ec9f0bc76..dd14f83de 100644
--- a/src/test_utils/hive/mod.rs
+++ b/src/test_utils/hive/mod.rs
@@ -2,7 +2,6 @@ use std::collections::HashMap;
 
 use ef_testing::evm_sequencer::account::{AccountType, KakarotAccount};
 use ethers::types::U256 as EthersU256;
-use eyre::eyre;
 use katana_primitives::{
     contract::ContractAddress,
     genesis::json::{GenesisContractJson, GenesisJson},
@@ -61,8 +60,8 @@ impl HiveGenesisConfig {
         let contract_account_class_hash = ClassHash(builder.contract_account_class_hash()?.into());
 
         // Fetch the contracts from the alloc field.
-        let mut additional_kakarot_storage = HashMap::new();
-        let mut fee_token_storage = HashMap::new();
+        let mut additional_kakarot_storage = HashMap::with_capacity(self.alloc.len()); // 1 mapping per contract
+        let mut fee_token_storage = HashMap::with_capacity(2 * self.alloc.len()); // 2 allowances per contract
         let contracts = self
             .alloc
             .into_iter()
@@ -119,9 +118,10 @@ impl HiveGenesisConfig {
         let kakarot_address = ContractAddress::new(kakarot_address);
         let mut genesis = builder.build()?;
 
-        let kakarot_contract =
-            genesis.contracts.get_mut(&kakarot_address).ok_or(eyre!("Kakarot contract not found"))?;
-        kakarot_contract.storage.get_or_insert_with(HashMap::new).extend(additional_kakarot_storage);
+        let kakarot_contract = genesis.contracts.entry(kakarot_address);
+        kakarot_contract.and_modify(|contract| {
+            contract.storage.get_or_insert_with(HashMap::new).extend(additional_kakarot_storage)
+        });
 
         genesis.fee_token.storage.get_or_insert_with(HashMap::new).extend(fee_token_storage);