From 7095aa33b4e867db12955870426a9dd109afdc5c Mon Sep 17 00:00:00 2001
From: lmittmann <lmittmann@users.noreply.github.com>
Date: Sat, 29 Jul 2023 16:37:25 +0200
Subject: [PATCH] added WETH storage slot utils

---
 w3vm/util.go      | 38 +++++++++++++++++++++++++++++++
 w3vm/util_test.go | 58 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 96 insertions(+)
 create mode 100644 w3vm/util_test.go

diff --git a/w3vm/util.go b/w3vm/util.go
index 6b07a11c..8d9d9c1c 100644
--- a/w3vm/util.go
+++ b/w3vm/util.go
@@ -2,8 +2,10 @@ package w3vm
 
 import (
 	"crypto/rand"
+	"math/big"
 
 	"github.com/ethereum/go-ethereum/common"
+	"github.com/lmittmann/w3/internal/crypto"
 )
 
 // zero values
@@ -19,6 +21,42 @@ func RandA() (addr common.Address) {
 	return addr
 }
 
+var (
+	weth9BalancePos   = big.NewInt(3)
+	weth9AllowancePos = big.NewInt(4)
+)
+
+// WETHBalanceSlot returns the storage slot that stores the WETH balance of
+// the given addr.
+func WETHBalanceSlot(addr common.Address) common.Hash {
+	return slot(weth9BalancePos, addr)
+}
+
+// WETHAllowanceSlot returns the storage slot that stores the WETH allowance
+// of the given owner and spender.
+func WETHAllowanceSlot(owner, spender common.Address) common.Hash {
+	return slot2(weth9AllowancePos, owner, spender)
+}
+
+func slot(pos *big.Int, acc common.Address) common.Hash {
+	data := make([]byte, 64)
+	copy(data[12:32], acc[:])
+	pos.FillBytes(data[32:])
+
+	return crypto.Keccak256Hash(data)
+}
+
+func slot2(pos *big.Int, acc, acc2 common.Address) common.Hash {
+	data := make([]byte, 64)
+	copy(data[12:32], acc[:])
+	pos.FillBytes(data[32:])
+
+	copy(data[32:], crypto.Keccak256(data))
+	copy(data[12:32], acc2[:])
+
+	return crypto.Keccak256Hash(data)
+}
+
 // nilToZero converts sets a pointer to the zero value if it is nil.
 func nilToZero[T any](ptr *T) *T {
 	if ptr == nil {
diff --git a/w3vm/util_test.go b/w3vm/util_test.go
new file mode 100644
index 00000000..073a9144
--- /dev/null
+++ b/w3vm/util_test.go
@@ -0,0 +1,58 @@
+package w3vm_test
+
+import (
+	"strconv"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/lmittmann/w3"
+	"github.com/lmittmann/w3/w3vm"
+)
+
+func TestWETHBalanceSlot(t *testing.T) {
+	tests := []struct {
+		Addr     common.Address
+		WantSlot common.Hash
+	}{
+		{
+			Addr:     w3.A("0x000000000000000000000000000000000000dEaD"),
+			WantSlot: w3.H("0x262bb27bbdd95c1cdc8e16957e36e38579ea44f7f6413dd7a9c75939def06b2c"),
+		},
+		{
+			Addr:     w3.A("0x000000000000000000000000000000000000c0Fe"),
+			WantSlot: w3.H("0xf68b260b81af177c0bf1a03b5d62b15aea1b486f8df26c77f33aed7538cfeb2c"),
+		},
+	}
+
+	for i, test := range tests {
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			gotSlot := w3vm.WETHBalanceSlot(test.Addr)
+			if test.WantSlot != gotSlot {
+				t.Fatalf("want %s, got %s", test.WantSlot, gotSlot)
+			}
+		})
+	}
+}
+
+func TestWETHAllowanceSlot(t *testing.T) {
+	tests := []struct {
+		Owner    common.Address
+		Spender  common.Address
+		WantSlot common.Hash
+	}{
+		{
+			Owner:    w3.A("0x000000000000000000000000000000000000dEaD"),
+			Spender:  w3.A("0x000000000000000000000000000000000000c0Fe"),
+			WantSlot: w3.H("0xea3c5e9cf6f5b7aba5d41ca731cf1f8cb1373e841a1cb336cc4bfeddc27c7f8b"),
+		},
+	}
+
+	for i, test := range tests {
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			gotSlot := w3vm.WETHAllowanceSlot(test.Owner, test.Spender)
+			if test.WantSlot != gotSlot {
+				t.Fatalf("want %s, got %s", test.WantSlot, gotSlot)
+			}
+		})
+	}
+}