represent storage in the bal encoding type as a custom struct (essentially revert the changes from master, as a stop-gap until we can create a cleaner solution)

core/types/bal/bal_encoding.go

@@ -19,6 +19,7 @@ package bal
 import (
 	"bytes"
 	"cmp"
+	"encoding/hex"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -31,6 +32,7 @@ import (
 	"io"
 	"maps"
 	"slices"
+	"strings"
 )
 
 //go:generate go run github.com/ethereum/go-ethereum/rlp/rlpgen -out bal_encoding_rlp_generated.go -type AccountAccess -decoder
@@ -142,10 +144,10 @@ func (e BlockAccessList) ValidateGasLimit(blockGasLimit uint64) error {
 		// Count unique storage keys across both reads and writes
 		uniqueSlots := make(map[common.Hash]struct{})
 		for _, sc := range account.StorageChanges {
-			uniqueSlots[sc.Slot.Bytes32()] = struct{}{}
+			uniqueSlots[sc.Slot.ToHash()] = struct{}{}
 		}
 		for _, sr := range account.StorageReads {
-			uniqueSlots[sr.Bytes32()] = struct{}{}
+			uniqueSlots[sr.ToHash()] = struct{}{}
 		}
 		balItems += uint64(len(uniqueSlots))
 	}
@@ -187,47 +189,6 @@ type encodingAccountNonce struct {
 	Nonce uint64 `json:"nonce"`
 }
 
-// encodingStorageWrite is the encoding format of StorageWrites.
-type encodingStorageWrite struct {
-	TxIdx      uint32       `json:"txIndex"`
-	ValueAfter *uint256.Int `json:"valueAfter"`
-}
-
-// encodingStorageWrite is the encoding format of SlotWrites.
-type encodingSlotWrites struct {
-	Slot     *uint256.Int           `json:"slot"`
-	Accesses []encodingStorageWrite `json:"accesses"`
-}
-
-// validate returns an instance of the encoding-representation slot writes in
-// working representation.
-func (e *encodingSlotWrites) validate(blockTxCount int) error {
-	if e.Slot == nil {
-		return errors.New("nil slot key")
-	}
-	if !slices.IsSortedFunc(e.Accesses, func(a, b encodingStorageWrite) int {
-		return cmp.Compare[uint32](a.TxIdx, b.TxIdx)
-	}) {
-		return errors.New("storage write tx indices not in order")
-	}
-	for i, access := range e.Accesses {
-		if access.ValueAfter == nil {
-			return errors.New("nil storage write post")
-		}
-		if i > 0 && e.Accesses[i-1].TxIdx == access.TxIdx {
-			return errors.New("duplicate storage write index")
-		}
-	}
-	// TODO: add test that covers there are actually storage modifications here
-	// if there aren't, it should be a bad block
-	if len(e.Accesses) == 0 {
-		return fmt.Errorf("empty storage writes")
-	} else if int(e.Accesses[len(e.Accesses)-1].TxIdx) >= blockTxCount+2 {
-		return fmt.Errorf("storage access reported index higher than allowed")
-	}
-	return nil
-}
-
 // encodingCodeChange contains the runtime bytecode deployed at an address
 // and the transaction index where the deployment took place.
 type encodingCodeChange struct {
@@ -239,7 +200,7 @@ type encodingCodeChange struct {
 type AccountAccess struct {
 	Address        common.Address          `json:"address,omitempty"`        // 20-byte Ethereum address
 	StorageChanges []encodingSlotWrites    `json:"storageChanges,omitempty"` // EncodedStorage changes (slot -> [tx_index -> new_value])
-	StorageReads   []*uint256.Int          `json:"storageReads,omitempty"`   // Read-only storage keys
+	StorageReads   []*EncodedStorage       `json:"storageReads,omitempty"`   // Read-only storage keys
 	BalanceChanges []encodingBalanceChange `json:"balanceChanges,omitempty"` // Balance changes ([tx_index -> post_balance])
 	NonceChanges   []encodingAccountNonce  `json:"nonceChanges,omitempty"`   // Nonce changes ([tx_index -> new_nonce])
 	CodeChanges    []encodingCodeChange    `json:"code,omitempty"`           // CodeChanges changes ([tx_index -> new_code])
@@ -251,7 +212,8 @@ type AccountAccess struct {
 func (e *AccountAccess) validate(blockTxCount int) error {
 	// Check the storage write slots are sorted in order
 	if !slices.IsSortedFunc(e.StorageChanges, func(a, b encodingSlotWrites) int {
-		return a.Slot.Cmp(b.Slot)
+		aHash, bHash := a.Slot.ToHash(), b.Slot.ToHash()
+		return bytes.Compare(aHash[:], bHash[:])
 	}) {
 		return errors.New("storage writes slots not in lexicographic order")
 	}
@@ -263,17 +225,17 @@ func (e *AccountAccess) validate(blockTxCount int) error {
 	readKeys := make(map[common.Hash]struct{})
 	writeKeys := make(map[common.Hash]struct{})
 	for _, readKey := range e.StorageReads {
-		if _, ok := readKeys[readKey.Bytes32()]; ok {
+		if _, ok := readKeys[readKey.ToHash()]; ok {
 			return errors.New("duplicate read key")
 		}
-		readKeys[readKey.Bytes32()] = struct{}{}
+		readKeys[readKey.ToHash()] = struct{}{}
 	}
 	for _, write := range e.StorageChanges {
 		writeKey := write.Slot
-		if _, ok := writeKeys[writeKey.Bytes32()]; ok {
+		if _, ok := writeKeys[writeKey.ToHash()]; ok {
 			return errors.New("duplicate write key")
 		}
-		writeKeys[writeKey.Bytes32()] = struct{}{}
+		writeKeys[writeKey.ToHash()] = struct{}{}
 	}
 
 	for readKey := range readKeys {
@@ -283,8 +245,9 @@ func (e *AccountAccess) validate(blockTxCount int) error {
 	}
 
 	// Check the storage read slots are sorted in order
-	if !slices.IsSortedFunc(e.StorageReads, func(a, b *uint256.Int) int {
-		return a.Cmp(b)
+	if !slices.IsSortedFunc(e.StorageReads, func(a, b *EncodedStorage) int {
+		aHash, bHash := a.ToHash(), b.ToHash()
+		return bytes.Compare(aHash[:], bHash[:])
 	}) {
 		return errors.New("storage read slots not in lexicographic order")
 	}
@@ -395,7 +358,7 @@ func (a *ConstructionAccountAccesses) toEncodingObj(addr common.Address) Account
 	slices.SortFunc(writeSlots, common.Hash.Cmp)
 	for _, slot := range writeSlots {
 		var obj encodingSlotWrites
-		obj.Slot = new(uint256.Int).SetBytes(slot[:])
+		obj.Slot = NewEncodedStorageFromHash(slot)
 
 		slotWrites := a.StorageWrites[slot]
 		obj.Accesses = make([]encodingStorageWrite, 0, len(slotWrites))
@@ -405,7 +368,7 @@ func (a *ConstructionAccountAccesses) toEncodingObj(addr common.Address) Account
 		for _, index := range indices {
 			obj.Accesses = append(obj.Accesses, encodingStorageWrite{
 				TxIdx:      index,
-				ValueAfter: new(uint256.Int).SetBytes(slotWrites[index].Bytes()),
+				ValueAfter: NewEncodedStorageFromHash(slotWrites[index]),
 			})
 		}
 		res.StorageChanges = append(res.StorageChanges, obj)
@@ -415,7 +378,7 @@ func (a *ConstructionAccountAccesses) toEncodingObj(addr common.Address) Account
 	readSlots := slices.Collect(maps.Keys(a.StorageReads))
 	slices.SortFunc(readSlots, common.Hash.Cmp)
 	for _, slot := range readSlots {
-		res.StorageReads = append(res.StorageReads, new(uint256.Int).SetBytes(slot.Bytes()))
+		res.StorageReads = append(res.StorageReads, NewEncodedStorageFromHash(slot))
 	}
 
 	// Convert balance changes
@@ -470,3 +433,113 @@ func (c *ConstructionBlockAccessList) ToEncodingObj() *BlockAccessList {
 }
 
 type ContractCode []byte
+
+// encodingStorageWrite is the encoding format of StorageWrites.
+type encodingStorageWrite struct {
+	TxIdx      uint32          `json:"txIndex"`
+	ValueAfter *EncodedStorage `json:"valueAfter"`
+}
+
+// EncodedStorage can represent either a storage key or value
+type EncodedStorage struct {
+	inner *uint256.Int
+}
+
+var _ rlp.Encoder = &EncodedStorage{}
+var _ rlp.Decoder = &EncodedStorage{}
+
+func (s *EncodedStorage) ToHash() common.Hash {
+	if s == nil {
+		return common.Hash{}
+	}
+	return s.inner.Bytes32()
+}
+
+func NewEncodedStorageFromHash(hash common.Hash) *EncodedStorage {
+	return &EncodedStorage{
+		new(uint256.Int).SetBytes(hash[:]),
+	}
+}
+
+func (s *EncodedStorage) UnmarshalJSON(b []byte) error {
+	var str string
+	if err := json.Unmarshal(b, &str); err != nil {
+		return err
+	}
+
+	str = strings.TrimLeft(str, "0x")
+	if len(str) == 0 {
+		return nil
+	}
+
+	if len(str)%2 == 1 {
+		str = "0" + str
+	}
+
+	val, err := hex.DecodeString(str)
+	if err != nil {
+		return err
+	}
+
+	if len(val) > 32 {
+		return fmt.Errorf("storage key/value cannot be greater than 32 bytes")
+	}
+
+	// TODO: check is s == nil ?? should be programmer error
+
+	*s = EncodedStorage{
+		inner: new(uint256.Int).SetBytes(val),
+	}
+	return nil
+}
+
+func (s EncodedStorage) MarshalJSON() ([]byte, error) {
+	return json.Marshal(s.inner.Hex())
+}
+
+func (s *EncodedStorage) EncodeRLP(_w io.Writer) error {
+	return s.inner.EncodeRLP(_w)
+}
+
+func (s *EncodedStorage) DecodeRLP(dec *rlp.Stream) error {
+	if s == nil {
+		*s = EncodedStorage{}
+	}
+	s.inner = uint256.NewInt(0)
+	return dec.ReadUint256(s.inner)
+}
+
+// encodingStorageWrite is the encoding format of SlotWrites.
+type encodingSlotWrites struct {
+	Slot     *EncodedStorage        `json:"slot"`
+	Accesses []encodingStorageWrite `json:"accesses"`
+}
+
+// validate returns an instance of the encoding-representation slot writes in
+// working representation.
+func (e *encodingSlotWrites) validate(blockTxCount int) error {
+	if e.Slot == nil {
+		return errors.New("nil slot key")
+	}
+	if !slices.IsSortedFunc(e.Accesses, func(a, b encodingStorageWrite) int {
+		return cmp.Compare[uint32](a.TxIdx, b.TxIdx)
+	}) {
+		return errors.New("storage write tx indices not in order")
+	}
+	for i, access := range e.Accesses {
+		if access.ValueAfter == nil {
+			return errors.New("nil storage write post")
+		}
+		if i > 0 && e.Accesses[i-1].TxIdx == access.TxIdx {
+			return errors.New("duplicate storage write index")
+		}
+	}
+	// TODO: add test that covers there are actually storage modifications here
+	// if there aren't, it should be a bad block
+	if len(e.Accesses) == 0 {
+		return fmt.Errorf("empty storage writes")
+	} else if int(e.Accesses[len(e.Accesses)-1].TxIdx) >= blockTxCount+2 {
+		return fmt.Errorf("storage access reported index higher than allowed")
+	}
+	return nil
+}

core/types/bal/bal_encoding_rlp_generated.go

@@ -14,19 +14,15 @@ func (obj *AccountAccess) EncodeRLP(_w io.Writer) error {
 	_tmp1 := w.List()
 	for _, _tmp2 := range obj.StorageChanges {
 		_tmp3 := w.List()
-		if _tmp2.Slot == nil {
-			w.Write(rlp.EmptyString)
-		} else {
-			w.WriteUint256(_tmp2.Slot)
+		if err := _tmp2.Slot.EncodeRLP(w); err != nil {
+			return err
 		}
 		_tmp4 := w.List()
 		for _, _tmp5 := range _tmp2.Accesses {
 			_tmp6 := w.List()
 			w.WriteUint64(uint64(_tmp5.TxIdx))
-			if _tmp5.ValueAfter == nil {
-				w.Write(rlp.EmptyString)
-			} else {
-				w.WriteUint256(_tmp5.ValueAfter)
+			if err := _tmp5.ValueAfter.EncodeRLP(w); err != nil {
+				return err
 			}
 			w.ListEnd(_tmp6)
 		}
@@ -36,10 +32,8 @@ func (obj *AccountAccess) EncodeRLP(_w io.Writer) error {
 	w.ListEnd(_tmp1)
 	_tmp7 := w.List()
 	for _, _tmp8 := range obj.StorageReads {
-		if _tmp8 == nil {
-			w.Write(rlp.EmptyString)
-		} else {
-			w.WriteUint256(_tmp8)
+		if err := _tmp8.EncodeRLP(w); err != nil {
+			return err
 		}
 	}
 	w.ListEnd(_tmp7)
@@ -99,11 +93,11 @@ func (obj *AccountAccess) DecodeRLP(dec *rlp.Stream) error {
 					return err
 				}
 				// Slot:
-				var _tmp4 uint256.Int
-				if err := dec.ReadUint256(&_tmp4); err != nil {
+				_tmp4 := new(EncodedStorage)
+				if err := _tmp4.DecodeRLP(dec); err != nil {
 					return err
 				}
-				_tmp3.Slot = &_tmp4
+				_tmp3.Slot = _tmp4
 				// Accesses:
 				var _tmp5 []encodingStorageWrite
 				if _, err := dec.List(); err != nil {
@@ -122,11 +116,11 @@ func (obj *AccountAccess) DecodeRLP(dec *rlp.Stream) error {
 						}
 						_tmp6.TxIdx = _tmp7
 						// ValueAfter:
-						var _tmp8 uint256.Int
-						if err := dec.ReadUint256(&_tmp8); err != nil {
+						_tmp8 := new(EncodedStorage)
+						if err := _tmp8.DecodeRLP(dec); err != nil {
 							return err
 						}
-						_tmp6.ValueAfter = &_tmp8
+						_tmp6.ValueAfter = _tmp8
 						if err := dec.ListEnd(); err != nil {
 							return err
 						}
@@ -148,16 +142,16 @@ func (obj *AccountAccess) DecodeRLP(dec *rlp.Stream) error {
 		}
 		_tmp0.StorageChanges = _tmp2
 		// StorageReads:
-		var _tmp9 []*uint256.Int
+		var _tmp9 []*EncodedStorage
 		if _, err := dec.List(); err != nil {
 			return err
 		}
 		for dec.MoreDataInList() {
-			var _tmp10 uint256.Int
-			if err := dec.ReadUint256(&_tmp10); err != nil {
+			_tmp10 := new(EncodedStorage)
+			if err := _tmp10.DecodeRLP(dec); err != nil {
 				return err
 			}
-			_tmp9 = append(_tmp9, &_tmp10)
+			_tmp9 = append(_tmp9, _tmp10)
 		}
 		if err := dec.ListEnd(); err != nil {
 			return err

core/types/bal/bal_test.go

@@ -116,15 +116,13 @@ func makeTestAccountAccess(sort bool) AccountAccess {
 		nonces        []encodingAccountNonce
 	)
 	for i := 0; i < 5; i++ {
-		slotBytes := testrand.Hash()
 		slot := encodingSlotWrites{
-			Slot: new(uint256.Int).SetBytes(slotBytes[:]),
+			Slot: NewEncodedStorageFromHash(testrand.Hash()),
 		}
 		for j := 0; j < 3; j++ {
-			valueBytes := testrand.Hash()
 			slot.Accesses = append(slot.Accesses, encodingStorageWrite{
 				TxIdx:      uint32(i*3 + j),
-				ValueAfter: new(uint256.Int).SetBytes(valueBytes[:]),
+				ValueAfter: NewEncodedStorageFromHash(testrand.Hash()),
 			})
 		}
 		if sort {
@@ -136,7 +134,7 @@ func makeTestAccountAccess(sort bool) AccountAccess {
 	}
 	if sort {
 		slices.SortFunc(storageWrites, func(a, b encodingSlotWrites) int {
-			return bytes.Compare(a.Slot.Bytes(), b.Slot.Bytes())
+			return bytes.Compare(a.Slot.inner.Bytes(), b.Slot.inner.Bytes())
 		})
 	}
 
@@ -173,9 +171,9 @@ func makeTestAccountAccess(sort bool) AccountAccess {
 		})
 	}
 
-	var encodedStorageReads []*uint256.Int
+	var encodedStorageReads []*EncodedStorage
 	for _, slot := range storageReads {
-		encodedStorageReads = append(encodedStorageReads, new(uint256.Int).SetBytes(slot[:]))
+		encodedStorageReads = append(encodedStorageReads, NewEncodedStorageFromHash(slot))
 	}
 	return AccountAccess{
 		Address:        [20]byte(testrand.Bytes(20)),
@@ -220,8 +218,7 @@ func TestBlockAccessListCopy(t *testing.T) {
 	// Make sure the mutations on copy won't affect the origin
 	for i := range *cpyCpy {
 		for j := 0; j < len((*cpyCpy)[i].StorageReads); j++ {
-			slotBytes := testrand.Hash()
-			(*cpyCpy)[i].StorageReads[j] = new(uint256.Int).SetBytes(slotBytes[:])
+			(*cpyCpy)[i].StorageReads[j] = NewEncodedStorageFromHash(testrand.Hash())
 		}
 	}
 	if !reflect.DeepEqual(list, *cpy) {