From 6317b0dd0e2c877a66e3c032cc00a3923572923a Mon Sep 17 00:00:00 2001
From: Kirill <paltsev.kir@gmail.com>
Date: Thu, 14 Mar 2024 21:21:38 +0400
Subject: [PATCH] Fix fee estimation in simulateTransactions, fix
 DataAvailability prices in traces (#1745)

---
 rpc/handlers.go    | 11 +++--------
 vm/rust/src/lib.rs |  2 +-
 vm/trace.go        | 14 ++++++++++++++
 3 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/rpc/handlers.go b/rpc/handlers.go
index 24c24defc7..2cb5a78971 100644
--- a/rpc/handlers.go
+++ b/rpc/handlers.go
@@ -1664,10 +1664,7 @@ func (h *Handler) simulateTransactions(id BlockID, transactions []BroadcastedTra
 		if !v0_6Response {
 			trace := traces[i]
 			executionResources := trace.TotalExecutionResources()
-			executionResources.DataAvailability = &vm.DataAvailability{
-				L1Gas:     gasConsumed.Uint64(),
-				L1DataGas: dataGasConsumed[i].Uint64(),
-			}
+			executionResources.DataAvailability = vm.NewDataAvailability(gasConsumed, dataGasConsumed[i], header.L1DAMode)
 			traces[i].ExecutionResources = executionResources
 		}
 
@@ -1819,10 +1816,8 @@ func (h *Handler) traceBlockTransactions(ctx context.Context, block *core.Block,
 			gasConsumed = gasConsumed.Div(gasConsumed, gasPrice) // division by zero felt is zero felt
 
 			executionResources := trace.TotalExecutionResources()
-			executionResources.DataAvailability = &vm.DataAvailability{
-				L1Gas:     gasConsumed.Uint64(),
-				L1DataGas: dataGasConsumed[index].Uint64(),
-			}
+			executionResources.DataAvailability = vm.NewDataAvailability(gasConsumed, dataGasConsumed[index],
+				header.L1DAMode)
 			traces[index].ExecutionResources = executionResources
 		}
 		result = append(result, TracedBlockTransaction{
diff --git a/vm/rust/src/lib.rs b/vm/rust/src/lib.rs
index 662b2351d5..4640b35a92 100644
--- a/vm/rust/src/lib.rs
+++ b/vm/rust/src/lib.rs
@@ -282,7 +282,7 @@ pub extern "C" fn cairoVMExecute(
                 }
 
                 // we are estimating fee, override actual fee calculation
-                if !charge_fee {
+                if  t.actual_fee.0 == 0 {
                     t.actual_fee = calculate_tx_fee(&t.actual_resources, &block_context, &fee_type).unwrap();
                 }
 
diff --git a/vm/trace.go b/vm/trace.go
index dd9269b18d..bd284767b3 100644
--- a/vm/trace.go
+++ b/vm/trace.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"slices"
 
+	"github.com/NethermindEth/juno/core"
 	"github.com/NethermindEth/juno/core/felt"
 )
 
@@ -242,3 +243,16 @@ type ExecutionResources struct {
 	ComputationResources
 	DataAvailability *DataAvailability `json:"data_availability,omitempty"`
 }
+
+func NewDataAvailability(gasConsumed, dataGasConsumed *felt.Felt, mode core.L1DAMode) *DataAvailability {
+	da := &DataAvailability{}
+
+	switch mode {
+	case core.Calldata:
+		da.L1Gas = gasConsumed.Uint64()
+	case core.Blob:
+		da.L1DataGas = dataGasConsumed.Uint64()
+	}
+
+	return da
+}