Use uint64 for Page valid bitlist + Reuse hasher and buffers

cannon/mipsevm/memory/binary_tree.go

@@ -40,6 +40,10 @@ func (m *BinaryTreeIndex) Invalidate(addr Word) {
 	gindex := (uint64(1) << (WordSize - PageAddrSize)) | uint64(addr>>PageAddrSize)
 
 	for gindex > 0 {
+		n := m.nodes[gindex]
+		if n != nil {
+			ReleaseByte32(n)
+		}
 		m.nodes[gindex] = nil
 		gindex >>= 1
 	}
@@ -70,9 +74,10 @@ func (m *BinaryTreeIndex) MerkleizeSubtree(gindex uint64) [32]byte {
 	}
 	left := m.MerkleizeSubtree(gindex << 1)
 	right := m.MerkleizeSubtree((gindex << 1) | 1)
-	r := HashPair(left, right)
-	m.nodes[gindex] = &r
-	return r
+	r := GetByte32()
+	HashPairNodes(r, &left, &right)
+	m.nodes[gindex] = r
+	return *r
 }
 
 func (m *BinaryTreeIndex) MerkleProof(addr Word) (out [MemProofSize]byte) {
@@ -112,6 +117,10 @@ func (m *BinaryTreeIndex) AddPage(pageIndex Word) {
 	// make nodes to root
 	k := (1 << PageKeySize) | uint64(pageIndex)
 	for k > 0 {
+		n := m.nodes[k]
+		if n != nil {
+			ReleaseByte32(n)
+		}
 		m.nodes[k] = nil
 		k >>= 1
 	}

cannon/mipsevm/memory/hasher.go

@@ -0,0 +1,74 @@
+package memory
+
+import (
+	"sync"
+
+	"github.com/ethereum/go-ethereum/crypto"
+)
+
+// byte32Pool is a sync.Pool for [32]byte slices
+var byte32Pool = sync.Pool{
+	New: func() interface{} {
+		var b [32]byte
+		return &b // Return a pointer to avoid extra allocations
+	},
+}
+
+// GetByte32 retrieves a *[32]byte from the pool
+func GetByte32() *[32]byte {
+	return byte32Pool.Get().(*[32]byte)
+}
+
+// ReleaseByte32 returns a *[32]byte to the pool
+func ReleaseByte32(b *[32]byte) {
+	// Optional: Zero the array before putting it back
+	*b = [32]byte{}
+	byte32Pool.Put(b)
+}
+
+var hashPool = sync.Pool{
+	New: func() interface{} {
+		return crypto.NewKeccakState()
+	},
+}
+
+func GetHasher() crypto.KeccakState {
+	return hashPool.Get().(crypto.KeccakState)
+}
+
+func PutHasher(h crypto.KeccakState) {
+	h.Reset()
+	hashPool.Put(h)
+}
+
+func HashPairNodes(out *[32]byte, left, right *[32]byte) {
+	h := GetHasher()
+	h.Write(left[:])
+	h.Write(right[:])
+	_, _ = h.Read(out[:])
+	PutHasher(h)
+}
+
+func HashData(out *[32]byte, data ...[]byte) {
+	h := GetHasher()
+	for _, b := range data {
+		h.Write(b)
+	}
+	_, _ = h.Read(out[:])
+	PutHasher(h)
+}
+
+func HashPair(left, right [32]byte) (out [32]byte) {
+	HashPairNodes(&out, &left, &right)
+	//fmt.Printf("0x%x 0x%x -> 0x%x\n", left, right, out)
+	return out
+}
+
+var zeroHashes = func() [256][32]byte {
+	// empty parts of the tree are all zero. Precompute the hash of each full-zero range sub-tree level.
+	var out [256][32]byte
+	for i := 1; i < 256; i++ {
+		out[i] = HashPair(out[i-1], out[i-1])
+	}
+	return out
+}()

cannon/mipsevm/memory/memory.go

@@ -9,7 +9,6 @@ import (
 	"sort"
 
 	"github.com/ethereum-optimism/optimism/cannon/mipsevm/arch"
-	"github.com/ethereum/go-ethereum/crypto"
 	"golang.org/x/exp/maps"
 )
 
@@ -28,21 +27,6 @@ const (
 
 type Word = arch.Word
 
-func HashPair(left, right [32]byte) [32]byte {
-	out := crypto.Keccak256Hash(left[:], right[:])
-	//fmt.Printf("0x%x 0x%x -> 0x%x\n", left, right, out)
-	return out
-}
-
-var zeroHashes = func() [256][32]byte {
-	// empty parts of the tree are all zero. Precompute the hash of each full-zero range sub-tree level.
-	var out [256][32]byte
-	for i := 1; i < 256; i++ {
-		out[i] = HashPair(out[i-1], out[i-1])
-	}
-	return out
-}()
-
 type Memory struct {
 	merkleIndex PageIndex
 	// Note: since we don't de-alloc Pages, we don't do ref-counting.
@@ -155,7 +139,7 @@ func (m *Memory) SetWord(addr Word, v Word) {
 		// Go may mmap relatively large ranges, but we only allocate the pages just in time.
 		p = m.AllocPage(pageIndex)
 	} else {
-		prevValid := p.Ok[1]
+		prevValid := p.getBit(1)
 		p.invalidate(pageAddr)
 		if prevValid {
 			m.merkleIndex.Invalidate(addr) // invalidate this branch of memory, now that the value changed
@@ -177,7 +161,7 @@ func (m *Memory) GetWord(addr Word) Word {
 		return 0
 	}
 	pageAddr := addr & PageAddrMask
-	return arch.ByteOrderWord.Word(p.Data[pageAddr : pageAddr+arch.WordSizeBytes])
+	return arch.ByteOrderWord.Word(p.Data[pageAddr : pageAddr+arch.WordSizeBytes : pageAddr+arch.WordSizeBytes])
 }
 
 func (m *Memory) AllocPage(pageIndex Word) *CachedPage {

cannon/mipsevm/memory/page.go

@@ -9,8 +9,6 @@ import (
 	"fmt"
 	"io"
 	"sync"
-
-	"github.com/ethereum/go-ethereum/crypto"
 )
 
 var zlibWriterPool = sync.Pool{
@@ -62,51 +60,80 @@ func (p *Page) UnmarshalText(dat []byte) error {
 	return err
 }
 
+// This is a compile time assertion to ensure that the PageSize is 4096
+// This is necessary because of the way we track intermediate nodes in CachedPage using OkLow and OkHigh
+var _ [0]struct{} = [PageSize - 4096]struct{}{}
+
 type CachedPage struct {
 	Data *Page
 	// intermediate nodes only
 	Cache [PageSize / 32][32]byte
-	// true if the intermediate node is valid
-	Ok [PageSize / 32]bool
+	// bit set to 1 if the intermediate node is valid
+	OkLow, OkHigh uint64 // size is PageSize / 32 == 64 + 64 == 128
+}
+
+func (p *CachedPage) getLowHighMask(k uint64) (
+	lowMask, highMask uint64,
+) {
+	mask := uint64(1) << (k & 63) // Bitmask
+	isHigh := k >> 6              // 1 if k >= 64, 0 if k < 64
+
+	return mask * (1 - isHigh),
+		mask * isHigh
 }
 
 func (p *CachedPage) invalidate(pageAddr Word) {
 	if pageAddr >= PageSize {
 		panic("invalid page addr")
 	}
-	k := (1 << PageAddrSize) | pageAddr
-	// first cache layer caches nodes that has two 32 byte leaf nodes.
+
+	k := uint64((1 << PageAddrSize) | pageAddr)
 	k >>= 5 + 1
+
 	for k > 0 {
-		p.Ok[k] = false
+		lowMask, highMask := p.getLowHighMask(k)
+		p.OkLow &^= lowMask
+		p.OkHigh &^= highMask
 		k >>= 1
 	}
 }
 
+func (p *CachedPage) getBit(k uint64) bool {
+	lowMask, highMask := p.getLowHighMask(k)
+	return (p.OkLow&lowMask | p.OkHigh&highMask) != 0
+}
+
+func (p *CachedPage) setBit(k uint64) {
+	lowMask, highMask := p.getLowHighMask(k)
+	p.OkLow |= lowMask
+	p.OkHigh |= highMask
+}
+
 func (p *CachedPage) InvalidateFull() {
-	p.Ok = [PageSize / 32]bool{} // reset everything to false
+	p.OkHigh = 0
+	p.OkLow = 0
 }
 
 func (p *CachedPage) MerkleRoot() [32]byte {
 	// hash the bottom layer
 	for i := uint64(0); i < PageSize; i += 64 {
 		j := PageSize/32/2 + i/64
-		if p.Ok[j] {
+		if p.getBit(j) {
 			continue
 		}
-		p.Cache[j] = crypto.Keccak256Hash(p.Data[i : i+64])
+		HashData(&p.Cache[j], p.Data[i:i+64])
 		//fmt.Printf("0x%x 0x%x -> 0x%x\n", p.Data[i:i+32], p.Data[i+32:i+64], p.Cache[j])
-		p.Ok[j] = true
+		p.setBit(j)
 	}
 
 	// hash the cache layers
 	for i := PageSize/32 - 2; i > 0; i -= 2 {
 		j := i >> 1
-		if p.Ok[j] {
+		if p.getBit(uint64(j)) {
 			continue
 		}
-		p.Cache[j] = HashPair(p.Cache[i], p.Cache[i+1])
-		p.Ok[j] = true
+		HashPairNodes(&p.Cache[j], &p.Cache[i], &p.Cache[i+1])
+		p.setBit(uint64(j))
 	}
 
 	return p.Cache[1]
@@ -120,7 +147,7 @@ func (p *CachedPage) MerkleizeSubtree(gindex uint64) [32]byte {
 		}
 		// it's pointing to a bottom node
 		nodeIndex := gindex & (PageAddrMask >> 5)
-		return *(*[32]byte)(p.Data[nodeIndex*32 : nodeIndex*32+32])
+		return *(*[32]byte)(p.Data[nodeIndex*32 : nodeIndex*32+32 : nodeIndex*32+32])
 	}
 	return p.Cache[gindex]
 }