cannon: allocator preheap instrumentation [rebased from ec2#2]

.circleci/config.yml

@@ -2229,7 +2229,7 @@ workflows:
           notify: true
           mentions: "@proofs-team"
           no_output_timeout: 90m
-          test_timeout: 240m
+          test_timeout: 480m
           resource_class: ethereum-optimism/latitude-fps-1
           context:
             - slack

cannon/cmd/run.go

@@ -395,7 +395,7 @@ func Run(ctx *cli.Context) error {
 		}
 	}
 
-	state, err := versions.LoadStateFromFile(ctx.Path(RunInputFlag.Name))
+	state, err := versions.LoadStateFromFileWithLargeICache(ctx.Path(RunInputFlag.Name))
 	if err != nil {
 		return fmt.Errorf("failed to load state: %w", err)
 	}

cannon/mipsevm/arch/arch64.go

@@ -15,11 +15,12 @@ const (
 	ExtMask     = 0x7
 
 	// Ensure virtual address is limited to 48-bits as many user programs assume such to implement packed pointers
-	// limit          0x00_00_FF_FF_FF_FF_FF_FF
-	HeapStart       = 0x00_00_10_00_00_00_00_00
-	HeapEnd         = 0x00_00_60_00_00_00_00_00
-	ProgramBreak    = 0x00_00_40_00_00_00_00_00
-	HighMemoryStart = 0x00_00_7F_FF_FF_FF_F0_00
+	Limit            = 0x00_00_FF_FF_FF_FF_FF_FF
+	ProgramHeapStart = 0x00_00_00_c0_00_00_00_00
+	HeapStart        = 0x00_00_10_00_00_00_00_00
+	HeapEnd          = 0x00_00_60_00_00_00_00_00
+	ProgramBreak     = 0x00_00_40_00_00_00_00_00
+	HighMemoryStart  = 0x00_00_7F_FF_FF_FF_F0_00
 )
 
 // MIPS64 syscall table - https://github.com/torvalds/linux/blob/3efc57369a0ce8f76bf0804f7e673982384e4ac9/arch/mips/kernel/syscalls/syscall_n64.tbl. Generate the syscall numbers using the Makefile in that directory.

cannon/mipsevm/memory/binary_tree.go

@@ -1,7 +1,10 @@
 package memory
 
 import (
+	"fmt"
 	"math/bits"
+
+	"github.com/ethereum-optimism/optimism/cannon/mipsevm/arch"
 )
 
 // BinaryTreeIndex is a representation of the state of the memory in a binary merkle tree.
@@ -12,13 +15,40 @@ type BinaryTreeIndex struct {
 	pageTable map[Word]*CachedPage
 }
 
-func NewBinaryTreeMemory() *Memory {
-	pages := make(map[Word]*CachedPage)
+func NewBinaryTreeMemory(codeSize, heapSize arch.Word) *Memory {
+	pages := make(map[arch.Word]*CachedPage)
 	index := NewBinaryTreeIndex(pages)
+
+	// Default values (2 GiB) if not provided
+	if codeSize == 0 {
+		codeSize = 1 << 31 // 2 GiB
+	}
+	if heapSize == 0 {
+		heapSize = 1 << 31 // 2 GiB
+	}
+
+	// Defensive bounds: code region must not overlap heap start
+	if codeSize > arch.ProgramHeapStart {
+		panic(fmt.Sprintf("codeSize (0x%x) overlaps heap start (0x%x)", codeSize, arch.ProgramHeapStart))
+	}
+
+	indexedRegions := make([]MappedMemoryRegion, 2)
+	indexedRegions[0] = MappedMemoryRegion{
+		startAddr: 0,
+		endAddr:   codeSize,
+		Data:      make([]byte, codeSize),
+	}
+	indexedRegions[1] = MappedMemoryRegion{
+		startAddr: arch.ProgramHeapStart,
+		endAddr:   arch.ProgramHeapStart + heapSize,
+		Data:      make([]byte, heapSize),
+	}
+
 	return &Memory{
-		merkleIndex:  index,
-		pageTable:    pages,
-		lastPageKeys: [2]Word{^Word(0), ^Word(0)}, // default to invalid keys, to not match any pages
+		merkleIndex:   index,
+		pageTable:     pages,
+		lastPageKeys:  [2]arch.Word{^arch.Word(0), ^arch.Word(0)},
+		MappedRegions: indexedRegions,
 	}
 }
 

cannon/mipsevm/memory/memory.go

@@ -27,6 +27,20 @@ const (
 
 type Word = arch.Word
 
+type MappedMemoryRegion struct {
+	startAddr Word
+	endAddr   Word
+	Data      []byte
+}
+
+func (m *MappedMemoryRegion) AddrInRegion(addr Word) bool {
+	return addr >= m.startAddr && addr < m.endAddr
+}
+
+func (m *MappedMemoryRegion) PageIndexInRegion(pageIndex Word) bool {
+	return pageIndex >= m.startAddr>>PageAddrSize && pageIndex < m.endAddr>>PageAddrSize
+}
+
 type Memory struct {
 	merkleIndex PageIndex
 	// Note: since we don't de-alloc Pages, we don't do ref-counting.
@@ -38,6 +52,8 @@ type Memory struct {
 	// this prevents map lookups each instruction
 	lastPageKeys [2]Word
 	lastPage     [2]*CachedPage
+
+	MappedRegions []MappedMemoryRegion
 }
 
 type PageIndex interface {
@@ -50,8 +66,53 @@ type PageIndex interface {
 	New(pages map[Word]*CachedPage) PageIndex
 }
 
+func NewMemoryWithLargeRegions() *Memory {
+	const codeSize = 1 << 31
+	const heapSize = 1 << 31
+	return NewBinaryTreeMemory(codeSize, heapSize)
+}
+
 func NewMemory() *Memory {
-	return NewBinaryTreeMemory()
+	return NewBinaryTreeMemory(4096, 4096)
+}
+
+// start end size gap
+func (m *Memory) GetAllocatedRanges() [][4]Word {
+	var ranges [][4]Word
+	if len(m.pageTable) == 0 {
+		return ranges
+	}
+
+	// Extract and sort page addresses
+	keys := make([]Word, 0, len(m.pageTable))
+	for key := range m.pageTable {
+		keys = append(keys, key)
+	}
+	sort.Slice(keys, func(i, j int) bool { return keys[i] < keys[j] })
+
+	// Find contiguous ranges and gaps
+	start := keys[0]
+	prev := start
+	var lastEnd Word = start - 1
+
+	for i := 1; i < len(keys); i++ {
+		if keys[i] != prev+1 {
+			gap := start - lastEnd - 1 // Gap is calculated from end of prev range to start of new one
+			ranges = append(ranges, [4]Word{start, prev, prev - start + 1, gap})
+			lastEnd = prev
+			start = keys[i]
+		}
+		prev = keys[i]
+	}
+
+	// Append last range
+	gap := start - lastEnd - 1
+	ranges = append(ranges, [4]Word{start, prev, prev - start + 1, gap})
+	for i := 0; i < len(ranges); i++ {
+		ranges[i][0] <<= PageAddrSize
+		ranges[i][1] <<= PageAddrSize
+	}
+	return ranges
 }
 
 func (m *Memory) MerkleRoot() [32]byte {
@@ -66,15 +127,14 @@ func (m *Memory) PageCount() int {
 	return len(m.pageTable)
 }
 
-func (m *Memory) ForEachPage(fn func(pageIndex Word, page *Page) error) error {
+func (m *Memory) ForEachPage(fn func(pageIndex Word, page Page) error) error {
 	for pageIndex, cachedPage := range m.pageTable {
 		if err := fn(pageIndex, cachedPage.Data); err != nil {
 			return err
 		}
 	}
 	return nil
 }
-
 func (m *Memory) MerkleizeSubtree(gindex uint64) [32]byte {
 	return m.merkleIndex.MerkleizeSubtree(gindex)
 }
@@ -155,7 +215,15 @@ func (m *Memory) GetWord(addr Word) Word {
 	if addr&arch.ExtMask != 0 {
 		panic(fmt.Errorf("unaligned memory access: %x", addr))
 	}
+	for _, region := range m.MappedRegions {
+		if ok := region.AddrInRegion(addr); ok {
+			offset := addr - region.startAddr
+			return arch.ByteOrderWord.Word(region.Data[offset : offset+arch.WordSizeBytes : offset+arch.WordSizeBytes])
+		}
+	}
+
 	pageIndex := addr >> PageAddrSize
+
 	p, ok := m.PageLookup(pageIndex)
 	if !ok {
 		return 0
@@ -165,7 +233,17 @@ func (m *Memory) GetWord(addr Word) Word {
 }
 
 func (m *Memory) AllocPage(pageIndex Word) *CachedPage {
-	p := &CachedPage{Data: new(Page)}
+	p := new(CachedPage)
+	for _, region := range m.MappedRegions {
+		if region.PageIndexInRegion(pageIndex) {
+			indexAdjusted := pageIndex - region.startAddr>>PageAddrSize
+			p.Data = region.Data[indexAdjusted*PageSize : (indexAdjusted+1)*PageSize : (indexAdjusted+1)*PageSize]
+			break
+		}
+	}
+	if p.Data == nil {
+		p.Data = make(Page, PageSize)
+	}
 	m.pageTable[pageIndex] = p
 	m.merkleIndex.AddPage(pageIndex)
 	return p
@@ -237,8 +315,9 @@ func (m *Memory) Copy() *Memory {
 	}
 
 	for k, page := range m.pageTable {
-		data := new(Page)
-		*data = *page.Data
+		data := make(Page, PageSize)
+		// *data = *page.Data
+		copy(data, page.Data)
 		out.AllocPage(k).Data = data
 	}
 	return out
@@ -287,20 +366,23 @@ func (m *Memory) Deserialize(in io.Reader) error {
 			return err
 		}
 	}
+
 	return nil
 }
 
 type pageEntry struct {
-	Index Word  `json:"index"`
-	Data  *Page `json:"data"`
+	Index Word            `json:"index"`
+	Data  *[PageSize]byte `json:"data"`
 }
 
 func (m *Memory) MarshalJSON() ([]byte, error) { // nosemgrep
 	pages := make([]pageEntry, 0, len(m.pageTable))
 	for k, p := range m.pageTable {
+		data := new([PageSize]byte)
+		copy(data[:], p.Data)
 		pages = append(pages, pageEntry{
 			Index: k,
-			Data:  p.Data,
+			Data:  data,
 		})
 	}
 	sort.Slice(pages, func(i, j int) bool {
@@ -318,7 +400,8 @@ func (m *Memory) UnmarshalJSON(data []byte) error {
 		if _, ok := m.pageTable[p.Index]; ok {
 			return fmt.Errorf("cannot load duplicate page, entry %d, page index %d", i, p.Index)
 		}
-		m.AllocPage(p.Index).Data = p.Data
+		page := m.AllocPage(p.Index)
+		copy(page.Data, p.Data[:])
 	}
 	return nil
 }

cannon/mipsevm/memory/memory64_benchmark_test.go

@@ -9,9 +9,11 @@ import (
 )
 
 const (
-	smallDataset  = 12_500_000
-	mediumDataset = 100_000_000
-	largeDataset  = 400_000_000
+	smallDataset          = 12_500_000
+	mediumDataset         = 100_000_000
+	largeDataset          = 400_000_000
+	defaultCodeRegionSize = 4096
+	defaultHeapSize       = 4096
 )
 
 func BenchmarkMemoryOperations(b *testing.B) {
@@ -36,7 +38,7 @@ func BenchmarkMemoryOperations(b *testing.B) {
 	for _, bm := range benchmarks {
 		b.Run("BinaryTree", func(b *testing.B) {
 			b.Run(bm.name, func(b *testing.B) {
-				m := NewBinaryTreeMemory()
+				m := NewBinaryTreeMemory(defaultCodeRegionSize, defaultHeapSize)
 				b.ResetTimer()
 				bm.fn(b, m)
 			})