From 11262aa81c5a356fd89d9666319b6558b88b087e Mon Sep 17 00:00:00 2001
From: mmsqe <mavis@crypto.com>
Date: Fri, 4 Oct 2024 14:00:33 +0800
Subject: [PATCH] Apply suggestions from code review

---
 x/evm/keeper/keeper.go           | 13 +++----------
 x/evm/keeper/state_transition.go |  6 ++++++
 x/evm/types/key.go               |  6 ++++++
 3 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/x/evm/keeper/keeper.go b/x/evm/keeper/keeper.go
index 0a8c4b0235..0b37664c9d 100644
--- a/x/evm/keeper/keeper.go
+++ b/x/evm/keeper/keeper.go
@@ -16,7 +16,6 @@
 package keeper
 
 import (
-	"encoding/binary"
 	"math/big"
 
 	errorsmod "cosmossdk.io/errors"
@@ -311,27 +310,21 @@ func (k Keeper) AddTransientGasUsed(ctx sdk.Context, gasUsed uint64) (uint64, er
 // SetHeaderHash stores the hash of the current block header in the store.
 func (k Keeper) SetHeaderHash(ctx sdk.Context) {
 	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixHeaderHash)
-	key := make([]byte, 8)
 	height, err := ethermint.SafeUint64(ctx.BlockHeight())
 	if err != nil {
 		panic(err)
 	}
-	binary.BigEndian.PutUint64(key, height)
-	store.Set(key, ctx.HeaderHash())
+	store.Set(types.GetHeaderHashKey(height), ctx.HeaderHash())
 }
 
 // GetHeaderHash retrieves the hash of a block header from the store by height.
 func (k Keeper) GetHeaderHash(ctx sdk.Context, height uint64) []byte {
 	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixHeaderHash)
-	key := make([]byte, 8)
-	binary.BigEndian.PutUint64(key, height)
-	return store.Get(key)
+	return store.Get(types.GetHeaderHashKey(height))
 }
 
 // DeleteHeaderHash removes the hash of a block header from the store by height
 func (k Keeper) DeleteHeaderHash(ctx sdk.Context, height uint64) {
 	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefixHeaderHash)
-	key := make([]byte, 8)
-	binary.BigEndian.PutUint64(key, height)
-	store.Delete(key)
+	store.Delete(types.GetHeaderHashKey(height))
 }
diff --git a/x/evm/keeper/state_transition.go b/x/evm/keeper/state_transition.go
index d24c138d40..c6d493f56a 100644
--- a/x/evm/keeper/state_transition.go
+++ b/x/evm/keeper/state_transition.go
@@ -103,6 +103,12 @@ func (k Keeper) GetHashFn(ctx sdk.Context) vm.GetHashFunc {
 		if ctx.BlockHeight() < h {
 			return common.Hash{}
 		}
+		if ctx.BlockHeight() == h {
+			headerHash := ctx.HeaderHash()
+			if len(headerHash) != 0 {
+				return common.BytesToHash(headerHash)
+			}
+		}
 		return common.BytesToHash(k.GetHeaderHash(ctx, height))
 	}
 }
diff --git a/x/evm/types/key.go b/x/evm/types/key.go
index 96e100be63..37a1c05206 100644
--- a/x/evm/types/key.go
+++ b/x/evm/types/key.go
@@ -112,3 +112,9 @@ func ObjectBloomKey(txIndex, msgIndex int) []byte {
 	binary.BigEndian.PutUint64(key[9:], value)
 	return key[:]
 }
+
+func GetHeaderHashKey(height uint64) []byte {
+	key := make([]byte, len(KeyPrefixHeaderHash)+8)
+	binary.BigEndian.PutUint64(key, height)
+	return key
+}