mipsevm: improve memory code, implement tests

mipsevm/memory.go

@@ -23,7 +23,9 @@ const (
 )
 
 func HashPair(left, right [32]byte) [32]byte {
-	return crypto.Keccak256Hash(left[:], right[:])
+	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 {
@@ -53,20 +55,23 @@ func NewMemory() *Memory {
 	}
 }
 
-func (m *Memory) Invalidate(addr uint32, count uint32) {
-	// we invalidate nodes of 32 bytes at a time
-	minGindex := ((uint64(1) << 32) | uint64(addr)) >> 5
-	count >>= 5
+func (m *Memory) Invalidate(addr uint32) {
+	// addr must be aligned to 4 bytes
+	if addr&0x3 != 0 {
+		panic(fmt.Errorf("unaligned memory access: %x", addr))
+	}
+
+	// find page, and invalidate addr within it
+	if p, ok := m.Pages[addr>>pageAddrSize]; ok {
+		p.Invalidate(addr & pageAddrMask)
+	}
 
-	for minGindex > 0 {
-		for i := minGindex; i < minGindex+uint64(count); i++ {
-			m.Nodes[i] = nil
-		}
-		minGindex >>= 1
-		count >>= 1
-		if count == 0 {
-			count = 1
-		}
+	// find the gindex of the first page covering the address
+	gindex := ((uint64(1) << 32) | uint64(addr)) >> pageAddrSize
+
+	for gindex > 0 {
+		m.Nodes[gindex] = nil
+		gindex >>= 1
 	}
 }
 
@@ -76,7 +81,7 @@ func (m *Memory) MerkleizeSubtree(gindex uint64) [32]byte {
 		panic("gindex too deep")
 	}
 	if l > pageKeySize {
-		depthIntoPage := l - pageKeySize
+		depthIntoPage := l - 1 - pageKeySize
 		pageIndex := (gindex >> depthIntoPage) & pageKeyMask
 		if p, ok := m.Pages[uint32(pageIndex)]; ok {
 			pageGindex := (1 << depthIntoPage) | (gindex & ((1 << depthIntoPage) - 1))
@@ -113,17 +118,17 @@ func (m *Memory) MerkleProof(addr uint32) (out [28 * 32]byte) {
 }
 
 func (m *Memory) traverseBranch(parent uint64, addr uint32, depth uint8) (proof [][32]byte) {
-	if depth == 28 {
+	if depth == 32-5 {
 		proof = make([][32]byte, 0, 32-5+1)
 		proof = append(proof, m.MerkleizeSubtree(parent))
 		return
 	}
-	if depth > 28 {
+	if depth > 32-5 {
 		panic("traversed too deep")
 	}
 	self := parent << 1
 	sibling := self | 1
-	if addr&(1<<depth) == 1 {
+	if addr&(1<<(31-depth)) != 0 {
 		self, sibling = sibling, self
 	}
 	proof = m.traverseBranch(self, addr, depth+1)
@@ -150,7 +155,7 @@ func (m *Memory) SetMemory(addr uint32, v uint32) {
 		// Go may mmap relatively large ranges, but we only allocate the pages just in time.
 		p = m.AllocPage(pageIndex)
 	} else {
-		m.Invalidate(addr, 4) // invalidate this branch of memory, now that the value changed
+		m.Invalidate(addr) // invalidate this branch of memory, now that the value changed
 	}
 	binary.BigEndian.PutUint32(p.Data[pageAddr:pageAddr+4], v)
 }

mipsevm/memory_test.go

+package mipsevm
+
+import (
+	"bytes"
+	"crypto/rand"
+	"encoding/binary"
+	"encoding/json"
+	"io"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestMemoryMerkleProof(t *testing.T) {
+	t.Run("nearly empty tree", func(t *testing.T) {
+		m := NewMemory()
+		m.SetMemory(0x10000, 0xaabbccdd)
+		proof := m.MerkleProof(0x10000)
+		require.Equal(t, uint32(0xaabbccdd), binary.BigEndian.Uint32(proof[:4]))
+		for i := 0; i < 32-5; i++ {
+			require.Equal(t, zeroHashes[i][:], proof[32+i*32:32+i*32+32], "empty siblings")
+		}
+	})
+	t.Run("fuller tree", func(t *testing.T) {
+		m := NewMemory()
+		m.SetMemory(0x10000, 0xaabbccdd)
+		m.SetMemory(0x80004, 42)
+		m.SetMemory(0x13370000, 123)
+		root := m.MerkleRoot()
+		proof := m.MerkleProof(0x80004)
+		require.Equal(t, uint32(42), binary.BigEndian.Uint32(proof[4:8]))
+		node := *(*[32]byte)(proof[:32])
+		path := uint32(0x80004) >> 5
+		for i := 32; i < len(proof); i += 32 {
+			sib := *(*[32]byte)(proof[i : i+32])
+			if path&1 != 0 {
+				node = HashPair(sib, node)
+			} else {
+				node = HashPair(node, sib)
+			}
+			path >>= 1
+		}
+		require.Equal(t, root, node, "proof must verify")
+	})
+}
+
+func TestMemoryMerkleRoot(t *testing.T) {
+	t.Run("empty", func(t *testing.T) {
+		m := NewMemory()
+		root := m.MerkleRoot()
+		require.Equal(t, zeroHashes[32-5], root, "fully zeroed memory should have expected zero hash")
+	})
+	t.Run("empty page", func(t *testing.T) {
+		m := NewMemory()
+		m.SetMemory(0xF000, 0)
+		root := m.MerkleRoot()
+		require.Equal(t, zeroHashes[32-5], root, "fully zeroed memory should have expected zero hash")
+	})
+	t.Run("single page", func(t *testing.T) {
+		m := NewMemory()
+		m.SetMemory(0xF000, 1)
+		root := m.MerkleRoot()
+		require.NotEqual(t, zeroHashes[32-5], root, "non-zero memory")
+	})
+	t.Run("repeat zero", func(t *testing.T) {
+		m := NewMemory()
+		m.SetMemory(0xF000, 0)
+		m.SetMemory(0xF004, 0)
+		root := m.MerkleRoot()
+		require.Equal(t, zeroHashes[32-5], root, "zero still")
+	})
+	t.Run("two empty pages", func(t *testing.T) {
+		m := NewMemory()
+		m.SetMemory(pageSize*3, 0)
+		m.SetMemory(pageSize*10, 0)
+		root := m.MerkleRoot()
+		require.Equal(t, zeroHashes[32-5], root, "zero still")
+	})
+	t.Run("random few pages", func(t *testing.T) {
+		m := NewMemory()
+		m.SetMemory(pageSize*3, 1)
+		m.SetMemory(pageSize*5, 42)
+		m.SetMemory(pageSize*6, 123)
+		p3 := m.MerkleizeSubtree((1 << pageKeySize) | 3)
+		p5 := m.MerkleizeSubtree((1 << pageKeySize) | 5)
+		p6 := m.MerkleizeSubtree((1 << pageKeySize) | 6)
+		z := zeroHashes[pageAddrSize-5]
+		r1 := HashPair(
+			HashPair(
+				HashPair(z, z),  // 0,1
+				HashPair(z, p3), // 2,3
+			),
+			HashPair(
+				HashPair(z, p5), // 4,5
+				HashPair(p6, z), // 6,7
+			),
+		)
+		r2 := m.MerkleizeSubtree(1 << (pageKeySize - 3))
+		require.Equal(t, r1, r2, "expecting manual page combination to match subtree merkle func")
+	})
+	t.Run("invalidate page", func(t *testing.T) {
+		m := NewMemory()
+		m.SetMemory(0xF000, 0)
+		require.Equal(t, zeroHashes[32-5], m.MerkleRoot(), "zero at first")
+		m.SetMemory(0xF004, 1)
+		require.NotEqual(t, zeroHashes[32-5], m.MerkleRoot(), "non-zero")
+		m.SetMemory(0xF004, 0)
+		require.Equal(t, zeroHashes[32-5], m.MerkleRoot(), "zero again")
+	})
+}
+
+func TestMemoryReadWrite(t *testing.T) {
+
+	t.Run("large random", func(t *testing.T) {
+		m := NewMemory()
+		data := make([]byte, 20_000)
+		_, err := rand.Read(data[:])
+		require.NoError(t, err)
+		require.NoError(t, m.SetMemoryRange(0, bytes.NewReader(data)))
+		for _, i := range []uint32{0, 4, 1000, 20_000 - 4} {
+			v := m.GetMemory(i)
+			expected := binary.BigEndian.Uint32(data[i : i+4])
+			require.Equalf(t, expected, v, "read at %d", i)
+		}
+	})
+
+	t.Run("repeat range", func(t *testing.T) {
+		m := NewMemory()
+		data := []byte(strings.Repeat("under the big bright yellow sun ", 40))
+		require.NoError(t, m.SetMemoryRange(0x1337, bytes.NewReader(data)))
+		res, err := io.ReadAll(m.ReadMemoryRange(0x1337-10, uint32(len(data)+20)))
+		require.NoError(t, err)
+		require.Equal(t, make([]byte, 10), res[:10], "empty start")
+		require.Equal(t, data, res[10:len(res)-10], "result")
+		require.Equal(t, make([]byte, 10), res[len(res)-10:], "empty end")
+	})
+
+	t.Run("read-write", func(t *testing.T) {
+		m := NewMemory()
+		m.SetMemory(12, 0xAABBCCDD)
+		require.Equal(t, uint32(0xAABBCCDD), m.GetMemory(12))
+		m.SetMemory(12, 0xAABB1CDD)
+		require.Equal(t, uint32(0xAABB1CDD), m.GetMemory(12))
+	})
+
+	t.Run("unaligned read", func(t *testing.T) {
+		m := NewMemory()
+		m.SetMemory(12, 0xAABBCCDD)
+		m.SetMemory(16, 0x11223344)
+		require.Panics(t, func() {
+			m.GetMemory(13)
+		})
+		require.Panics(t, func() {
+			m.GetMemory(14)
+		})
+		require.Panics(t, func() {
+			m.GetMemory(15)
+		})
+		require.Equal(t, uint32(0x11223344), m.GetMemory(16))
+		require.Equal(t, uint32(0), m.GetMemory(20))
+		require.Equal(t, uint32(0xAABBCCDD), m.GetMemory(12))
+	})
+
+	t.Run("unaligned write", func(t *testing.T) {
+		m := NewMemory()
+		m.SetMemory(12, 0xAABBCCDD)
+		require.Panics(t, func() {
+			m.SetMemory(13, 0x11223344)
+		})
+		require.Panics(t, func() {
+			m.SetMemory(14, 0x11223344)
+		})
+		require.Panics(t, func() {
+			m.SetMemory(15, 0x11223344)
+		})
+		require.Equal(t, uint32(0xAABBCCDD), m.GetMemory(12))
+	})
+}
+
+func TestMemoryJSON(t *testing.T) {
+	m := NewMemory()
+	m.SetMemory(8, 123)
+	dat, err := json.Marshal(m)
+	require.NoError(t, err)
+	var res Memory
+	require.NoError(t, json.Unmarshal(dat, &res))
+	require.Equal(t, uint32(123), res.GetMemory(8))
+}

mipsevm/memory_page.go → mipsevm/page.go

@@ -60,7 +60,7 @@ func (p *CachedPage) MerkleRoot() [32]byte {
 	}
 
 	// hash the cache layers
-	for i := pageSize/32 - 2; i > 0; i++ {
+	for i := pageSize/32 - 2; i > 0; i -= 2 {
 		j := i >> 1
 		if p.Ok[j] {
 			continue
@@ -79,7 +79,7 @@ func (p *CachedPage) MerkleizeSubtree(gindex uint64) [32]byte {
 			panic("gindex too deep")
 		}
 		// it's pointing to a bottom node
-		nodeIndex := gindex & pageAddrMask
+		nodeIndex := gindex & (pageAddrMask >> 5)
 		return *(*[32]byte)(p.Data[nodeIndex*32 : nodeIndex*32+32])
 	}
 	return p.Cache[gindex]

mipsevm/page_test.go

+package mipsevm
+
+import (
+	"testing"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/stretchr/testify/require"
+)
+
+func TestCachedPage(t *testing.T) {
+	p := &CachedPage{Data: new(Page)}
+	p.Data[42] = 0xab
+
+	gindex := ((uint64(1) << pageAddrSize) | 42) >> 5
+	node := common.Hash(p.MerkleizeSubtree(gindex))
+	expectedLeaf := common.Hash{10: 0xab}
+	require.Equal(t, expectedLeaf, node, "leaf nodes should not be hashed")
+
+	node = p.MerkleizeSubtree(gindex >> 1)
+	expectedParent := common.Hash(HashPair(zeroHashes[0], expectedLeaf))
+	require.Equal(t, expectedParent, node, "can get the parent node")
+
+	node = p.MerkleizeSubtree(gindex >> 2)
+	expectedParentParent := common.Hash(HashPair(expectedParent, zeroHashes[1]))
+	require.Equal(t, expectedParentParent, node, "and the parent of the parent")
+
+	pre := p.MerkleRoot()
+	p.Data[42] = 0xcd
+	post := p.MerkleRoot()
+	require.Equal(t, pre, post, "no change expected until cache is invalidated")
+
+	p.Invalidate(42)
+	post2 := p.MerkleRoot()
+	require.NotEqual(t, post, post2, "change after cache invalidation")
+
+	p.Data[2000] = 0xef
+	p.Invalidate(42)
+	post3 := p.MerkleRoot()
+	require.Equal(t, post2, post3, "local invalidation is not global invalidation")
+
+	p.Invalidate(2000)
+	post4 := p.MerkleRoot()
+	require.NotEqual(t, post3, post4, "can see the change now")
+
+	p.Data[1000] = 0xff
+	p.InvalidateFull()
+	post5 := p.MerkleRoot()
+	require.NotEqual(t, post4, post5, "and global invalidation works regardless of changed data")
+}