Repository: tetratelabs/wazero
Pull request #177 (2022-02-01 08:01:07)

Consolidate memory consts and grow operation
This commit consolidates the logic and constants around memory pages,
and memory grow, size operations. Prior to this, memory.grow and size were
implemented in each engine and embedded deeply inside of runtime, which
made it impossible for us to write unit tests on it.



Patch:
diff --git a/RATIONALE.md b/RATIONALE.md
index 02afadfd..69eeafd5 100644
--- a/RATIONALE.md
+++ b/RATIONALE.md
@@ -13,21 +13,21 @@ Note: `microwasm` was never specified formally, and only exists in a historical
 https://github.com/bytecodealliance/wasmtime/blob/v0.29.0/crates/lightbeam/src/microwasm.rs
 
 
-### Size limitations
-#### Number of functions
+## Size limitations
+### Number of functions
 
 The possible number of function instances in [a store](https://www.w3.org/TR/wasm-core-1/#store%E2%91%A0) is not specified in the WebAssembly specifications since [`funcaddr`](https://www.w3.org/TR/wasm-core-1/#syntax-funcaddr) corresponding to a function instance can be arbitrary number. 
 In wazero, we choose to use `uint32` to represent `funcaddr`. Therefore the maximum number of function instances a store can instantiate is limited to 2^32. 
 
 That is because not only we _believe_ that all use cases are fine with the limitation, but also we have no way to test wazero runtimes under these unusual circumstances.
 
-#### Number of function types
+### Number of function types
 
 There's no limitation on the number of function types in [a store](https://www.w3.org/TR/wasm-core-1/#store%E2%91%A0) according to the spec. In wazero implementation, we assign each function type to a unique ID, and choose to use `uint32` to represent the IDs.
 Therefore the maximum number of function types a store can have is limited to 2^32. 
 
 This is due to the same reason for the limitation on the number of functions above.
 
-### JIT engine implementation
+## JIT engine implementation
 
 See [wasm/jit/RATIONALE.md](wasm/jit/RATIONALE.md).
diff --git a/wasm/binary/types.go b/wasm/binary/types.go
index 576b6132..8abe0ab9 100644
--- a/wasm/binary/types.go
+++ b/wasm/binary/types.go
@@ -64,13 +64,13 @@ func decodeMemoryType(r io.Reader) (*wasm.MemoryType, error) {
 	if err != nil {
 		return nil, err
 	}
-	if ret.Min > uint32(wasm.PageSize) {
+	if ret.Min > uint32(wasm.MemoryPageSize) {
 		return nil, fmt.Errorf("memory min must be at most 65536 pages (4GiB)")
 	}
 	if ret.Max != nil {
 		if *ret.Max < ret.Min {
 			return nil, fmt.Errorf("memory size minimum must not be greater than maximum")
-		} else if *ret.Max > uint32(wasm.PageSize) {
+		} else if *ret.Max > uint32(wasm.MemoryPageSize) {
 			return nil, fmt.Errorf("memory max must be at most 65536 pages (4GiB)")
 		}
 	}
diff --git a/wasm/engine.go b/wasm/engine.go
index eb274dca..5f319bdb 100644
--- a/wasm/engine.go
+++ b/wasm/engine.go
@@ -1,7 +1,5 @@
 package wasm
 
-const PageSize uint64 = 65536
-
 // Engine is the interface implemented by interpreters.
 type Engine interface {
 	// Call invokes a function instance f with given parameters.
diff --git a/wasm/interpreter/interpreter.go b/wasm/interpreter/interpreter.go
index e62630c5..951b9dc5 100644
--- a/wasm/interpreter/interpreter.go
+++ b/wasm/interpreter/interpreter.go
@@ -745,24 +745,14 @@ func (it *interpreter) callNativeFunc(f *interpreterFunction) {
 			}
 		case wazeroir.OperationKindMemorySize:
 			{
-				v := uint64(len(memoryInst.Buffer)) / wasm.PageSize
-				it.push(v)
+				it.push(uint64(memoryInst.PageSize()))
 				frame.pc++
 			}
 		case wazeroir.OperationKindMemoryGrow:
 			{
 				n := it.pop()
-				max := uint64(math.MaxUint32)
-				if memoryInst.Max != nil {
-					max = uint64(*memoryInst.Max) * wasm.PageSize
-				}
-				if uint64(n*wasm.PageSize+uint64(len(memoryInst.Buffer))) > max {
-					v := int32(-1)
-					it.push(uint64(v))
-				} else {
-					it.push(uint64(len(memoryInst.Buffer)) / wasm.PageSize)
-					memoryInst.Buffer = append(memoryInst.Buffer, make([]byte, n*wasm.PageSize)...)
-				}
+				res := memoryInst.Grow(uint32(n))
+				it.push(uint64(res))
 				frame.pc++
 			}
 		case wazeroir.OperationKindConstI32, wazeroir.OperationKindConstI64,
diff --git a/wasm/jit/engine.go b/wasm/jit/engine.go
index b10289a1..09e31f88 100644
--- a/wasm/jit/engine.go
+++ b/wasm/jit/engine.go
@@ -617,25 +617,14 @@ func (e *engine) builtinFunctionGrowCallFrameStack() {
 
 func (e *engine) builtinFunctionMemoryGrow(mem *wasm.MemoryInstance) {
 	newPages := e.popValue()
-	max := uint64(math.MaxUint32)
-	if mem.Max != nil {
-		max = uint64(*mem.Max) * wasm.PageSize
-	}
 
-	currentLen := uint64(len(mem.Buffer))
-	// If exceeds the max of memory size, we push -1 according to the spec.
-	if uint64(newPages*wasm.PageSize+currentLen) > max {
-		v := int32(-1)
-		e.pushValue(uint64(v))
-	} else {
-		e.pushValue(currentLen / wasm.PageSize) // Grow returns the prior memory size on change.
-		mem.Buffer = append(mem.Buffer, make([]byte, newPages*wasm.PageSize)...)
+	res := mem.Grow(uint32(newPages))
+	e.pushValue(uint64(res))
 
-		// Update the moduleContext's fields as they become stale after the update ^^.
-		bufSliceHeader := (*reflect.SliceHeader)(unsafe.Pointer(&mem.Buffer))
-		e.moduleContext.memorySliceLen = uint64(bufSliceHeader.Len)
-		e.moduleContext.memoryElement0Address = bufSliceHeader.Data
-	}
+	// Update the moduleContext's fields as they become stale after the update ^^.
+	bufSliceHeader := (*reflect.SliceHeader)(unsafe.Pointer(&mem.Buffer))
+	e.moduleContext.memorySliceLen = uint64(bufSliceHeader.Len)
+	e.moduleContext.memoryElement0Address = bufSliceHeader.Data
 }
 
 func (e *engine) Compile(f *wasm.FunctionInstance) (err error) {
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index cd017d73..45587a02 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -1263,8 +1263,6 @@ func (c *amd64Compiler) compilePick(o *wazeroir.OperationPick) error {
 		prog.To.Type = obj.TYPE_REG
 		prog.To.Reg = reg
 		c.addInstruction(prog)
-	} else if pickTarget.onConditionalRegister() {
-		panic("TODO: compilePick for targets on a conditonal register")
 	}
 	// Now we already placed the picked value on the register,
 	// so push the location onto the stack.
@@ -5202,7 +5200,7 @@ func (c *amd64Compiler) readInstructionAddress(destinationRegister int16, before
 		binary.LittleEndian.PutUint32(code[readInstructionAddress.Pc+3:], offset)
 
 		// See the comment at readInstructionAddress.From.Reg above. Here we drop the most significant bit of the third byte of the LEA instruction.
-		code[readInstructionAddress.Pc+2] = code[readInstructionAddress.Pc+2] & 0b01111111
+		code[readInstructionAddress.Pc+2] &= 0b01111111
 		return nil
 	})
 }
diff --git a/wasm/mempage.go b/wasm/mempage.go
new file mode 100644
index 00000000..c4388a13
--- /dev/null
+++ b/wasm/mempage.go
@@ -0,0 +1,51 @@
+package wasm
+
+const (
+	// MemoryPageSize is the unit of memory length in WebAssembly,
+	// and is defined as 2^16 = 65536.
+	// See https://www.w3.org/TR/wasm-core-1/#memory-instances%E2%91%A0
+	MemoryPageSize = 65536
+	// memoryMaxPages is maximum number of pages defined (2^16).
+	// See https://www.w3.org/TR/wasm-core-1/#grow-mem
+	memoryMaxPages = MemoryPageSize
+	// memoryPageSizeInBit satisfies the relation: "1 << memoryPageSizeInBit == memoryPageSize".
+	memoryPageSizeInBit = 16
+)
+
+// MemoryPagesToBytesNum converts the given pages into the number of bytes contained in these pages.
+func memoryPagesToBytesNum(pages uint32) (bytesNum uint64) {
+	return uint64(pages) << memoryPageSizeInBit
+}
+
+// MemoryPagesToBytesNum converts the given nuber of bytes into the number of pages.
+func memoryBytesNumToPages(bytesNum uint64) (pages uint32) {
+	return uint32(bytesNum >> memoryPageSizeInBit)
+}
+
+// Grow extends the memory buffer by "newPages" * memoryPageSize.
+// The logic here is described in https://www.w3.org/TR/wasm-core-1/#grow-mem.
+//
+// Returns -1 if the operation resulted in excedding the maximum memory pages.
+// Otherwise, returns the prior memory size after growing the memory buffer.
+func (m *MemoryInstance) Grow(newPages uint32) (result uint32) {
+	currentPages := memoryBytesNumToPages(uint64(len(m.Buffer)))
+
+	maxPages := uint32(memoryMaxPages)
+	if m.Max != nil {
+		maxPages = *m.Max
+	}
+
+	// If exceeds the max of memory size, we push -1 according to the spec.
+	if currentPages+newPages > maxPages {
+		return 0xffffffff // = -1 in signed 32 bit integer.
+	} else {
+		// Otherwise, grow the memory.
+		m.Buffer = append(m.Buffer, make([]byte, memoryPagesToBytesNum(newPages))...)
+		return currentPages
+	}
+}
+
+// PageSize returns the current memory buffer size in pages.
+func (m *MemoryInstance) PageSize() (result uint32) {
+	return memoryBytesNumToPages(uint64(len(m.Buffer)))
+}
diff --git a/wasm/store.go b/wasm/store.go
index ed7613d7..4d75bb01 100644
--- a/wasm/store.go
+++ b/wasm/store.go
@@ -640,7 +640,7 @@ func (s *Store) buildMemoryInstances(module *Module, target *ModuleInstance) (ro
 			return rollbackFuncs, fmt.Errorf("multiple memories not supported")
 		}
 		target.Memory = &MemoryInstance{
-			Buffer: make([]byte, uint64(memSec.Min)*PageSize),
+			Buffer: make([]byte, memoryPagesToBytesNum(memSec.Min)),
 			Min:    memSec.Min,
 			Max:    memSec.Max,
 		}
@@ -670,11 +670,11 @@ func (s *Store) buildMemoryInstances(module *Module, target *ModuleInstance) (ro
 		}
 
 		size := uint64(offset) + uint64(len(d.Init))
-		max := uint64(math.MaxUint32)
+		maxPage := uint32(memoryMaxPages)
 		if int(d.MemoryIndex) < len(module.MemorySection) && module.MemorySection[d.MemoryIndex].Max != nil {
-			max = uint64(*module.MemorySection[d.MemoryIndex].Max)
+			maxPage = *module.MemorySection[d.MemoryIndex].Max
 		}
-		if size > max*PageSize {
+		if size > memoryPagesToBytesNum(maxPage) {
 			return rollbackFuncs, fmt.Errorf("memory size out of limit %d * 64Ki", int(*(module.MemorySection[d.MemoryIndex].Max)))
 		}
 
@@ -1834,7 +1834,7 @@ func (s *Store) AddTableInstance(moduleName, name string, min uint32, max *uint3
 
 func (s *Store) AddMemoryInstance(moduleName, name string, min uint32, max *uint32) error {
 	memory := &MemoryInstance{
-		Buffer: make([]byte, uint64(min)*PageSize),
+		Buffer: make([]byte, memoryPagesToBytesNum(min)),
 		Min:    min,
 		Max:    max,
 	}


Test patch:
diff --git a/wasm/jit/jit_amd64_test.go b/wasm/jit/jit_amd64_test.go
index 44eb4bf9..b3a576e2 100644
--- a/wasm/jit/jit_amd64_test.go
+++ b/wasm/jit/jit_amd64_test.go
@@ -147,7 +147,7 @@ func newJITEnvironment() *jitEnv {
 	return &jitEnv{
 		eng: newEngine(),
 		moduleInstance: &wasm.ModuleInstance{
-			Memory:  &wasm.MemoryInstance{Buffer: make([]byte, wasm.PageSize*defaultMemoryPageNumInTest)},
+			Memory:  &wasm.MemoryInstance{Buffer: make([]byte, wasm.MemoryPageSize*defaultMemoryPageNumInTest)},
 			Tables:  []*wasm.TableInstance{{}},
 			Globals: []*wasm.GlobalInstance{},
 		},
diff --git a/wasm/mempage_test.go b/wasm/mempage_test.go
new file mode 100644
index 00000000..31052726
--- /dev/null
+++ b/wasm/mempage_test.go
@@ -0,0 +1,58 @@
+package wasm
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestMemoryPageConsts(t *testing.T) {
+	a := int32(-1)
+	fmt.Printf("%x", uint32(a))
+	require.Equal(t, MemoryPageSize, 1<<memoryPageSizeInBit)
+	require.Equal(t, MemoryPageSize, memoryMaxPages)
+	require.Equal(t, MemoryPageSize, 1<<16)
+}
+
+func Test_MemoryPagesToBytesNum(t *testing.T) {
+	for _, numPage := range []uint32{0, 1, 5, 10} {
+		require.Equal(t, uint64(numPage)*MemoryPageSize, memoryPagesToBytesNum(numPage))
+	}
+}
+
+func Test_MemoryBytesNumToPages(t *testing.T) {
+	for _, numbytes := range []uint64{0, MemoryPageSize * 1, MemoryPageSize * 10} {
+		require.Equal(t, uint32(numbytes/MemoryPageSize), memoryBytesNumToPages(numbytes))
+	}
+}
+
+func TestMemryInstance_Grow_Size(t *testing.T) {
+	t.Run("with max", func(t *testing.T) {
+		max := uint32(10)
+		m := &MemoryInstance{Max: &max, Buffer: make([]byte, 0)}
+		require.Equal(t, uint32(0), m.Grow(5))
+		require.Equal(t, uint32(5), m.PageSize())
+		// Zero page grow is well-defined, should return the current page correctly.
+		require.Equal(t, uint32(5), m.Grow(0))
+		require.Equal(t, uint32(5), m.PageSize())
+		require.Equal(t, uint32(5), m.Grow(4))
+		require.Equal(t, uint32(9), m.PageSize())
+		// At this point, the page size equal 9,
+		// so trying to grow two pages should result in failure.
+		require.Equal(t, int32(-1), int32(m.Grow(2)))
+		require.Equal(t, uint32(9), m.PageSize())
+		// But growing one page is still perimitted.
+		require.Equal(t, uint32(9), m.Grow(1))
+		// Ensure that the current page size equals the max.
+		require.Equal(t, max, m.PageSize())
+	})
+	t.Run("without max", func(t *testing.T) {
+		m := &MemoryInstance{Buffer: make([]byte, 0)}
+		require.Equal(t, uint32(0), m.Grow(1))
+		require.Equal(t, uint32(1), m.PageSize())
+		// Trying to grow above memoryMaxPages, the operation should fail.
+		require.Equal(t, int32(-1), int32(m.Grow(memoryMaxPages)))
+		require.Equal(t, uint32(1), m.PageSize())
+	})
+}