feat: wazero fork upgrade to use uint64 for memory.Size()

.gitignore

@@ -26,3 +26,5 @@ tmp
 # node_modules used by polkadot.js/api tests
 tests/polkadotjs_test/node_modules
 !tests/polkadotjs_test/test/*.wasm
+
+*.json

go.mod

@@ -200,6 +200,6 @@ require (
 
 go 1.21
 
-replace github.com/tetratelabs/wazero => github.com/ChainSafe/wazero v0.0.0-20231114190045-1d874d099362
+replace github.com/tetratelabs/wazero => github.com/ChainSafe/wazero v0.0.0-20240306154750-130c2d1dc80b
 
 replace github.com/centrifuge/go-substrate-rpc-client/v4 => github.com/timwu20/go-substrate-rpc-client/v4 v4.0.0-20231110032757-3d8e441b7303

go.sum

@@ -10,8 +10,8 @@ git.apache.org/thrift.git v0.0.0-20180902110319-2566ecd5d999/go.mod h1:fPE2ZNJGy
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/ChainSafe/go-schnorrkel v1.1.0 h1:rZ6EU+CZFCjB4sHUE1jIu8VDoB/wRKZxoe1tkcO71Wk=
 github.com/ChainSafe/go-schnorrkel v1.1.0/go.mod h1:ABkENxiP+cvjFiByMIZ9LYbRoNNLeBLiakC1XeTFxfE=
-github.com/ChainSafe/wazero v0.0.0-20231114190045-1d874d099362 h1:hbvvSSB436JJalwq/2fRZwJpptvq9HMOLYVZX9oVHKM=
-github.com/ChainSafe/wazero v0.0.0-20231114190045-1d874d099362/go.mod h1:0U0G41+ochRKoPKCJlh0jMg1CHkyfK8kDqiirMmKY8A=
+github.com/ChainSafe/wazero v0.0.0-20240306154750-130c2d1dc80b h1:EqM7+3eUCn0aKpB4Zr9527/D8xt6LAxHPZoaJMIr25I=
+github.com/ChainSafe/wazero v0.0.0-20240306154750-130c2d1dc80b/go.mod h1:ytl6Zuh20R/eROuyDaGPkp82O9C/DJfXAwJfQ3X6/7Y=
 github.com/DataDog/zstd v1.4.1/go.mod h1:1jcaCB/ufaK+sKp1NBhlGmpz41jOoPQ35bpF36t7BBo=
 github.com/DataDog/zstd v1.4.5 h1:EndNeuB0l9syBZhut0wns3gV1hL8zX8LIu6ZiVHWLIQ=
 github.com/DataDog/zstd v1.4.5/go.mod h1:1jcaCB/ufaK+sKp1NBhlGmpz41jOoPQ35bpF36t7BBo=

lib/runtime/allocator/freeing_bump.go

@@ -37,7 +37,7 @@ const (
 	MaxPossibleAllocations uint32 = 33554432
 
 	PageSize     = 65536
-	MaxWasmPages = (4 * 1024 * 1024 * 1024 / PageSize) - 1
+	MaxWasmPages = 4 * 1024 * 1024 * 1024 / PageSize
 )
 
 var (
@@ -332,7 +332,7 @@ type FreeingBumpHeapAllocator struct {
 	bumper                 uint32
 	freeLists              *FreeLists
 	poisoned               bool
-	lastObservedMemorySize uint32
+	lastObservedMemorySize uint64
 	stats                  AllocationStats
 }
 
@@ -393,7 +393,7 @@ func (f *FreeingBumpHeapAllocator) Allocate(mem runtime.Memory, size uint32) (pt
 	link := f.freeLists.heads[order]
 	switch value := link.(type) {
 	case Ptr:
-		if uint64(value.headerPtr)+uint64(order.size())+uint64(HeaderSize) > uint64(mem.Size()) {
+		if uint64(value.headerPtr)+uint64(order.size())+uint64(HeaderSize) > mem.Size() {
 			return 0, fmt.Errorf("%w: pointer: %d, order size: %d",
 				ErrInvalidHeaderPointerDetected, value.headerPtr, order.size())
 		}
@@ -505,14 +505,17 @@ func (f *FreeingBumpHeapAllocator) Deallocate(mem runtime.Memory, ptr uint32) (e
 func bump(bumper *uint32, size uint32, mem runtime.Memory) (uint32, error) {
 	requiredSize := uint64(*bumper) + uint64(size)
 
-	if requiredSize > uint64(mem.Size()) {
+	if requiredSize > mem.Size() {
 		requiredPages, ok := pagesFromSize(requiredSize)
 		if !ok {
-			return 0, fmt.Errorf("%w: required size %d dont fit uint32",
-				ErrAllocatorOutOfSpace, requiredSize)
+			panic(fmt.Sprintf("cannot calculate number of pages from size %d", requiredSize))
+		}
+
+		currentPages, ok := pagesFromSize(mem.Size())
+		if !ok {
+			panic(fmt.Sprintf("cannot calculate current number of pages, current size: %d", mem.Size()))
 		}
 
-		currentPages := mem.Size() / PageSize
 		if currentPages >= requiredPages {
 			panic(fmt.Sprintf("current pages %d >= required pages %d", currentPages, requiredPages))
 		}
@@ -539,7 +542,7 @@ func bump(bumper *uint32, size uint32, mem runtime.Memory) (uint32, error) {
 				ErrCannotGrowLinearMemory, currentPages, nextPages)
 		}
 
-		pagesIncrease := (mem.Size() / PageSize) == nextPages
+		pagesIncrease := (mem.Size() / PageSize) == uint64(nextPages)
 		if !pagesIncrease {
 			panic(fmt.Sprintf("number of pages should have increased! previous: %d, desired: %d", currentPages, nextPages))
 		}

lib/runtime/allocator/memory_test.go

@@ -26,8 +26,8 @@ func (m *MemoryInstance) pages() uint32 {
 	return pages
 }
 
-func (m *MemoryInstance) Size() uint32 {
-	return m.pages() * PageSize
+func (m *MemoryInstance) Size() uint64 {
+	return uint64(m.pages() * PageSize)
 }
 
 func (m *MemoryInstance) Grow(pages uint32) (uint32, bool) {
@@ -54,7 +54,7 @@ func (m *MemoryInstance) WriteUint64Le(offset uint32, v uint64) bool {
 	copy(m.data[offset:offset+8], encoded)
 	return true
 }
-func (*MemoryInstance) Read(_, _ uint32) ([]byte, bool) {
+func (*MemoryInstance) Read(_ uint32, _ uint64) ([]byte, bool) {
 	return nil, false
 }
 

lib/runtime/memory.go

@@ -9,7 +9,7 @@ type Memory interface {
 	// has 1 page: 65536
 	//
 	// See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#-hrefsyntax-instr-memorymathsfmemorysize%E2%91%A0
-	Size() uint32
+	Size() uint64
 
 	// Grow increases memory by the delta in pages (65536 bytes per page).
 	// The return val is the previous memory size in pages, or false if the
@@ -66,7 +66,7 @@ type Memory interface {
 	// shared. Those who need a stable view must set Wasm memory min=max, or
 	// use wazero.RuntimeConfig WithMemoryCapacityPages to ensure max is always
 	// allocated.
-	Read(offset, byteCount uint32) ([]byte, bool)
+	Read(offset uint32, byteCount uint64) ([]byte, bool)
 
 	// WriteByte writes a single byte to the underlying buffer at the offset in or returns false if out of range.
 	WriteByte(offset uint32, v byte) bool //nolint:govet

lib/runtime/wazero/imports.go

@@ -51,8 +51,8 @@
 
 // splitPointerSize converts a 64bit pointer size to an
 // uint32 pointer and a uint32 size.
-func splitPointerSize(pointerSize uint64) (ptr, size uint32) {
-	return uint32(pointerSize), uint32(pointerSize >> 32)
+func splitPointerSize(pointerSize uint64) (ptr uint32, size uint64) {
+	return uint32(pointerSize), uint64(pointerSize >> 32)
 }
 
 // read will read from 64 bit pointer size and return a byte slice
@@ -2265,7 +2265,7 @@
 
 	var written uint
 	valueOutPtr, valueOutSize := splitPointerSize(valueOut)
-	if uint32(len(data)) <= valueOutSize {
+	if uint64(len(data)) <= valueOutSize {
 		written = uint(len(data))
 	} else {
 		written = uint(valueOutSize)

lib/runtime/wazero/instance.go

@@ -42,6 +42,7 @@ type Instance struct {
 	Module   api.Module
 	Context  *runtime.Context
 	codeHash common.Hash
+	heapBase uint32
 	sync.Mutex
 }
 
@@ -408,24 +409,23 @@ func NewInstance(code []byte, cfg Config) (instance *Instance, err error) {
 		return nil, err
 	}
 
-	global := mod.ExportedGlobal("__heap_base")
-	if global == nil {
+	encodedHeapBase := mod.ExportedGlobal("__heap_base")
+	if encodedHeapBase == nil {
 		return nil, fmt.Errorf("wazero error: nil global for __heap_base")
 	}
 
-	hb := api.DecodeU32(global.Get())
+	heapBase := api.DecodeU32(encodedHeapBase.Get())
 	// hb = runtime.DefaultHeapBase
 
 	mem := mod.Memory()
 	if mem == nil {
 		return nil, fmt.Errorf("wazero error: nil memory for module")
 	}
 
-	allocator := allocator.NewFreeingBumpHeapAllocator(hb)
 	instance = &Instance{
-		Runtime: rt,
+		heapBase: heapBase,
+		Runtime:  rt,
 		Context: &runtime.Context{
-			Allocator:       allocator,
 			Keystore:        cfg.Keystore,
 			Validator:       cfg.Role == common.AuthorityRole,
 			NodeStorage:     cfg.NodeStorage,
@@ -460,8 +460,11 @@ func (i *Instance) Exec(function string, data []byte) (result []byte, err error)
 	i.Lock()
 	defer i.Unlock()
 
+	// instantiate a new allocator on every execution func
+	allocator := allocator.NewFreeingBumpHeapAllocator(i.heapBase)
+
 	dataLength := uint32(len(data))
-	inputPtr, err := i.Context.Allocator.Allocate(i.Module.Memory(), dataLength)
+	inputPtr, err := allocator.Allocate(i.Module.Memory(), dataLength)
 	if err != nil {
 		return nil, fmt.Errorf("allocating input memory: %w", err)
 	}
@@ -471,18 +474,23 @@ func (i *Instance) Exec(function string, data []byte) (result []byte, err error)
 	if mem == nil {
 		panic("nil memory")
 	}
+
 	ok := mem.Write(inputPtr, data)
 	if !ok {
 		panic("write overflow")
 	}
 
+	i.Context.Allocator = allocator
+	defer func() {
+		i.Context.Allocator = nil
+	}()
+
 	runtimeFunc := i.Module.ExportedFunction(function)
 	if runtimeFunc == nil {
 		return nil, fmt.Errorf("%w: %s", ErrExportFunctionNotFound, function)
 	}
 
 	ctx := context.WithValue(context.Background(), runtimeContextKey, i.Context)
-
 	values, err := runtimeFunc.Call(ctx, api.EncodeU32(inputPtr), api.EncodeU32(dataLength))
 	if err != nil {
 		return nil, fmt.Errorf("running runtime function: %w", err)
@@ -497,6 +505,7 @@ func (i *Instance) Exec(function string, data []byte) (result []byte, err error)
 	if !ok {
 		panic("write overflow")
 	}
+
 	return result, nil
 }