Adds Memory.IndexByte and memory grow example (#489)

This adds Memory.IndexByte which allows efficent scanning for a delimiter, ex NUL(0) in null-terminated strings. This also adds an ad-hoc test to ensure we can export memory functions such as grow. While this is implicitly in the spectests, it is easier to find in the ad-hoc tests. Signed-off-by: Adrian Cole <adrian@tetrate.io>
2022-04-20 20:33:15 +08:00
parent 043f67ab84
commit d238a004a8
7 changed files with 165 additions and 23 deletions
--- a/api/wasm.go
+++ b/api/wasm.go
@@ -179,7 +179,11 @@ type Memory interface {
 	// See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#-hrefsyntax-instr-memorymathsfmemorysize%E2%91%A0
 	Size() uint32

-	// ReadByte reads a single byte from the underlying buffer at the offset in or returns false if out of range.
+	// IndexByte returns the index of the first instance of c in the underlying buffer at the offset or returns false if
+	// not found or out of range.
+	IndexByte(offset uint32, c byte) (uint32, bool)
+
+	// ReadByte reads a single byte from the underlying buffer at the offset or returns false if out of range.
 	ReadByte(offset uint32) (byte, bool)

 	// ReadUint32Le reads a uint32 in little-endian encoding from the underlying buffer at the offset in or returns
--- a/internal/integration_test/engine/adhoc_test.go
+++ b/internal/integration_test/engine/adhoc_test.go
@@ -27,6 +27,7 @@ var tests = map[string]func(t *testing.T, r wazero.Runtime){
 	"host function with numeric parameter":    testHostFunctionNumericParameter,
 	"close module with in-flight calls":       testCloseInFlight,
 	"multiple instantiation from same source": testMultipleInstantiation,
+	"exported function that grows memory":     testMemOps,
 }

 func TestEngineJIT(t *testing.T) {
@@ -374,6 +375,42 @@ func testCloseInFlight(t *testing.T, r wazero.Runtime) {
 	}
 }

+func testMemOps(t *testing.T, r wazero.Runtime) {
+	// Instantiate a module that manages its memory
+	memory, err := r.InstantiateModuleFromCode(testCtx, []byte(`(module $memory
+  (func $grow (param $delta i32) (result (;previous_size;) i32) local.get 0 memory.grow)
+  (func $size (result (;size;) i32) memory.size)
+
+  (memory 0)
+
+  (export "size" (func $size))
+  (export "grow" (func $grow))
+  (export "memory" (memory 0))
+)`))
+	require.NoError(t, err)
+	defer memory.Close()
+
+	// Check the export worked
+	require.Equal(t, memory.Memory(), memory.ExportedMemory("memory"))
+
+	// Check the size command worked
+	results, err := memory.ExportedFunction("size").Call(testCtx)
+	require.NoError(t, err)
+	require.Zero(t, results[0])
+	require.Zero(t, memory.ExportedMemory("memory").Size())
+
+	// Try to grow the memory by one page
+	results, err = memory.ExportedFunction("grow").Call(testCtx, 1)
+	require.NoError(t, err)
+	require.Zero(t, results[0]) // should succeed and return the old size in pages.
+
+	// Check the size command works!
+	results, err = memory.ExportedFunction("size").Call(testCtx)
+	require.NoError(t, err)
+	require.Equal(t, uint64(1), results[0])                 // 1 page
+	require.Equal(t, uint32(65536), memory.Memory().Size()) // 64KB
+}
+
 func testMultipleInstantiation(t *testing.T, r wazero.Runtime) {
 	compiled, err := r.CompileModule(testCtx, []byte(`(module $test
 		(memory 1)
--- a/internal/wasm/memory.go
+++ b/internal/wasm/memory.go
@@ -1,9 +1,12 @@
 package wasm

 import (
+	"bytes"
 	"encoding/binary"
 	"fmt"
 	"math"
+
+	"github.com/tetratelabs/wazero/api"
 )

 const (
@@ -18,6 +21,9 @@ const (
 	MemoryPageSizeInBits = 16
 )

+// compile-time check to ensure MemoryInstance implements api.Memory
+var _ api.Memory = &MemoryInstance{}
+
 // MemoryInstance represents a memory instance in a store, and implements api.Memory.
 //
 // Note: In WebAssembly 1.0 (20191205), there may be up to one Memory per store, which means the precise memory is always
@@ -28,7 +34,7 @@ type MemoryInstance struct {
 	Min, Max uint32
 }

-// Size implements api.Memory Size
+// Size implements the same method as documented on api.Memory.
 func (m *MemoryInstance) Size() uint32 {
 	return uint32(len(m.Buffer))
 }
@@ -38,7 +44,20 @@ func (m *MemoryInstance) hasSize(offset uint32, sizeInBytes uint32) bool {
 	return uint64(offset)+uint64(sizeInBytes) <= uint64(m.Size()) // uint64 prevents overflow on add
 }

-// ReadByte implements api.Memory ReadByte
+// IndexByte implements the same method as documented on api.Memory.
+func (m *MemoryInstance) IndexByte(offset uint32, c byte) (uint32, bool) {
+	if offset >= m.Size() {
+		return 0, false
+	}
+	b := m.Buffer[offset:]
+	if result := bytes.IndexByte(b, c); result == -1 {
+		return 0, false
+	} else {
+		return uint32(result) + offset, true
+	}
+}
+
+// ReadByte implements the same method as documented on api.Memory.
 func (m *MemoryInstance) ReadByte(offset uint32) (byte, bool) {
 	if offset >= m.Size() {
 		return 0, false
@@ -46,7 +65,7 @@ func (m *MemoryInstance) ReadByte(offset uint32) (byte, bool) {
 	return m.Buffer[offset], true
 }

-// ReadUint32Le implements api.Memory ReadUint32Le
+// ReadUint32Le implements the same method as documented on api.Memory.
 func (m *MemoryInstance) ReadUint32Le(offset uint32) (uint32, bool) {
 	if !m.hasSize(offset, 4) {
 		return 0, false
@@ -54,7 +73,7 @@ func (m *MemoryInstance) ReadUint32Le(offset uint32) (uint32, bool) {
 	return binary.LittleEndian.Uint32(m.Buffer[offset : offset+4]), true
 }

-// ReadFloat32Le implements api.Memory ReadFloat32Le
+// ReadFloat32Le implements the same method as documented on api.Memory.
 func (m *MemoryInstance) ReadFloat32Le(offset uint32) (float32, bool) {
 	v, ok := m.ReadUint32Le(offset)
 	if !ok {
@@ -63,7 +82,7 @@ func (m *MemoryInstance) ReadFloat32Le(offset uint32) (float32, bool) {
 	return math.Float32frombits(v), true
 }

-// ReadUint64Le implements api.Memory ReadUint64Le
+// ReadUint64Le implements the same method as documented on api.Memory.
 func (m *MemoryInstance) ReadUint64Le(offset uint32) (uint64, bool) {
 	if !m.hasSize(offset, 8) {
 		return 0, false
@@ -71,7 +90,7 @@ func (m *MemoryInstance) ReadUint64Le(offset uint32) (uint64, bool) {
 	return binary.LittleEndian.Uint64(m.Buffer[offset : offset+8]), true
 }

-// ReadFloat64Le implements api.Memory ReadFloat64Le
+// ReadFloat64Le implements the same method as documented on api.Memory.
 func (m *MemoryInstance) ReadFloat64Le(offset uint32) (float64, bool) {
 	v, ok := m.ReadUint64Le(offset)
 	if !ok {
@@ -80,7 +99,7 @@ func (m *MemoryInstance) ReadFloat64Le(offset uint32) (float64, bool) {
 	return math.Float64frombits(v), true
 }

-// Read implements api.Memory Read
+// Read implements the same method as documented on api.Memory.
 func (m *MemoryInstance) Read(offset, byteCount uint32) ([]byte, bool) {
 	if !m.hasSize(offset, byteCount) {
 		return nil, false
@@ -88,7 +107,7 @@ func (m *MemoryInstance) Read(offset, byteCount uint32) ([]byte, bool) {
 	return m.Buffer[offset : offset+byteCount], true
 }

-// WriteByte implements api.Memory WriteByte
+// WriteByte implements the same method as documented on api.Memory.
 func (m *MemoryInstance) WriteByte(offset uint32, v byte) bool {
 	if offset >= m.Size() {
 		return false
@@ -97,7 +116,7 @@ func (m *MemoryInstance) WriteByte(offset uint32, v byte) bool {
 	return true
 }

-// WriteUint32Le implements api.Memory WriteUint32Le
+// WriteUint32Le implements the same method as documented on api.Memory.
 func (m *MemoryInstance) WriteUint32Le(offset, v uint32) bool {
 	if !m.hasSize(offset, 4) {
 		return false
@@ -106,12 +125,12 @@ func (m *MemoryInstance) WriteUint32Le(offset, v uint32) bool {
 	return true
 }

-// WriteFloat32Le implements api.Memory WriteFloat32Le
+// WriteFloat32Le implements the same method as documented on api.Memory.
 func (m *MemoryInstance) WriteFloat32Le(offset uint32, v float32) bool {
 	return m.WriteUint32Le(offset, math.Float32bits(v))
 }

-// WriteUint64Le implements api.Memory WriteUint64Le
+// WriteUint64Le implements the same method as documented on api.Memory.
 func (m *MemoryInstance) WriteUint64Le(offset uint32, v uint64) bool {
 	if !m.hasSize(offset, 8) {
 		return false
@@ -120,12 +139,12 @@ func (m *MemoryInstance) WriteUint64Le(offset uint32, v uint64) bool {
 	return true
 }

-// WriteFloat64Le implements api.Memory WriteFloat64Le
+// WriteFloat64Le implements the same method as documented on api.Memory.
 func (m *MemoryInstance) WriteFloat64Le(offset uint32, v float64) bool {
 	return m.WriteUint64Le(offset, math.Float64bits(v))
 }

-// Write implements api.Memory Write
+// Write implements the same method as documented on api.Memory.
 func (m *MemoryInstance) Write(offset uint32, val []byte) bool {
 	if !m.hasSize(offset, uint32(len(val))) {
 		return false
--- a/internal/wasm/memory_test.go
+++ b/internal/wasm/memory_test.go
@@ -45,6 +45,19 @@ func TestMemoryInstance_Grow_Size(t *testing.T) {
 	require.Equal(t, max, m.PageSize())
 }

+func TestIndexByte(t *testing.T) {
+	var mem = &MemoryInstance{Buffer: []byte{0, 0, 0, 0, 16, 0, 0, 0}, Min: 1}
+	v, ok := mem.IndexByte(4, 16)
+	require.True(t, ok)
+	require.Equal(t, uint32(4), v)
+
+	_, ok = mem.IndexByte(5, 16)
+	require.False(t, ok)
+
+	_, ok = mem.IndexByte(9, 16)
+	require.False(t, ok)
+}
+
 func TestReadByte(t *testing.T) {
 	var mem = &MemoryInstance{Buffer: []byte{0, 0, 0, 0, 0, 0, 0, 16}, Min: 1}
 	v, ok := mem.ReadByte(7)
--- a/internal/wasm/text/func_parser.go
+++ b/internal/wasm/text/func_parser.go
@@ -158,16 +158,26 @@ func (p *funcParser) beginInstruction(tokenBytes []byte) (next tokenParser, err
 	case wasm.OpcodeI32ConstName: // See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#syntax-instr-numeric
 		opCode = wasm.OpcodeI32Const
 		next = p.parseI32
+	case wasm.OpcodeI32LoadName: // See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#memory-instructions%E2%91%A8
+		return p.encodeMemArgOp(wasm.OpcodeI32Load, alignment32)
+	case wasm.OpcodeI32StoreName: // See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#memory-instructions%E2%91%A8
+		return p.encodeMemArgOp(wasm.OpcodeI32Store, alignment32)
 	case wasm.OpcodeI64ConstName: // See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#syntax-instr-numeric
 		opCode = wasm.OpcodeI64Const
 		next = p.parseI64
 	case wasm.OpcodeI64LoadName: // See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#memory-instructions%E2%91%A8
-		return p.encodeI64Instruction(wasm.OpcodeI64Load)
+		return p.encodeMemArgOp(wasm.OpcodeI64Load, alignment64)
 	case wasm.OpcodeI64StoreName: // See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#memory-instructions%E2%91%A8
-		return p.encodeI64Instruction(wasm.OpcodeI64Store)
+		return p.encodeMemArgOp(wasm.OpcodeI64Store, alignment64)
 	case wasm.OpcodeLocalGetName: // See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#variable-instructions%E2%91%A0
 		opCode = wasm.OpcodeLocalGet
 		next = p.parseLocalIndex
+	case wasm.OpcodeMemoryGrowName: // See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#memory-instructions%E2%91%A6
+		p.currentBody = append(p.currentBody, wasm.OpcodeMemoryGrow, 0x00) // reserved arg0
+		return p.beginFieldOrInstruction, nil
+	case wasm.OpcodeMemorySizeName: // See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#memory-instructions%E2%91%A6
+		p.currentBody = append(p.currentBody, wasm.OpcodeMemorySize, 0x00) // reserved arg0
+		return p.beginFieldOrInstruction, nil

 		// Next are sign-extension-ops
 		// See https://github.com/WebAssembly/spec/blob/main/proposals/sign-extension-ops/Overview.md
@@ -202,13 +212,16 @@ func (p *funcParser) beginInstruction(tokenBytes []byte) (next tokenParser, err
 	return next, nil
 }

-func (p *funcParser) encodeI64Instruction(oc wasm.Opcode) (tokenParser, error) {
-	p.currentBody = append(
-		p.currentBody,
-		oc,
-		3, // alignment=3 (natural alignment) because 2^3 = size of I64 (8 bytes)
-		0, // offset=0 because that's the default
-	)
+const (
+	// alignment32 is because it is 32bit is 2^2 bytes
+	alignment32 = 2
+	// alignment64 is because it is 64bit is 2^3 bytes
+	alignment64 = 3
+)
+
+func (p *funcParser) encodeMemArgOp(oc wasm.Opcode, alignment byte) (tokenParser, error) {
+	offset := byte(0) // offset=0 because that's the default
+	p.currentBody = append(p.currentBody, oc, alignment, offset)
 	return p.beginFieldOrInstruction, nil
 }

--- a/internal/wasm/text/func_parser_test.go
+++ b/internal/wasm/text/func_parser_test.go
@@ -62,6 +62,25 @@ func TestFuncParser(t *testing.T) {
 				wasm.OpcodeI32Const, 0x02, wasm.OpcodeI32Const, 0x01, wasm.OpcodeI32Sub, wasm.OpcodeEnd,
 			}},
 		},
+		{
+			name:   "i32.load",
+			source: "(func i32.const 8 i32.load)",
+			expected: &wasm.Code{Body: []byte{
+				wasm.OpcodeI32Const, 8, // dynamic memory offset to load
+				wasm.OpcodeI32Load, 0x2, 0x0, // load alignment=2 (natural alignment) staticOffset=0
+				wasm.OpcodeEnd,
+			}},
+		},
+		{
+			name:   "i32.store",
+			source: "(func i32.const 8 i32.const 37 i32.store)",
+			expected: &wasm.Code{Body: []byte{
+				wasm.OpcodeI32Const, 8, // dynamic memory offset to store
+				wasm.OpcodeI32Const, 37, // value to store
+				wasm.OpcodeI32Store, 0x2, 0x0, // load alignment=2 (natural alignment) staticOffset=0
+				wasm.OpcodeEnd,
+			}},
+		},
 		{
 			name:     "i64.const",
 			source:   "(func i64.const 356)",
@@ -86,6 +105,25 @@ func TestFuncParser(t *testing.T) {
 				wasm.OpcodeEnd,
 			}},
 		},
+		{
+			name:   "memory.grow",
+			source: "(func i32.const 2 memory.grow drop)",
+			expected: &wasm.Code{Body: []byte{
+				wasm.OpcodeI32Const, 2, // how many pages to grow
+				wasm.OpcodeMemoryGrow, 0, // memory index zero
+				wasm.OpcodeDrop, // drop the previous page count (or -1 if grow failed)
+				wasm.OpcodeEnd,
+			}},
+		},
+		{
+			name:   "memory.size",
+			source: "(func memory.size drop)",
+			expected: &wasm.Code{Body: []byte{
+				wasm.OpcodeMemorySize, 0, // memory index zero
+				wasm.OpcodeDrop, // drop the page count
+				wasm.OpcodeEnd,
+			}},
+		},

 		// Below are changes to test/core/i32 and i64.wast from the commit that added "sign-extension-ops" support.
 		// See https://github.com/WebAssembly/spec/commit/e308ca2ae04d5083414782e842a81f931138cf2e
--- a/internal/wazeroir/compiler_test.go
+++ b/internal/wazeroir/compiler_test.go
@@ -58,6 +58,24 @@ func TestCompile(t *testing.T) {
 				Signature:    &wasm.FunctionType{Params: []wasm.ValueType{wasm.ValueTypeI32}, Results: []wasm.ValueType{wasm.ValueTypeI32}},
 			},
 		},
+		{
+			name: "memory.grow", // Ex to expose ops to grow memory
+			module: requireModuleText(t, `(module
+  (func (param $delta i32) (result (;previous_size;) i32) local.get 0 memory.grow)
+)`),
+			expected: &CompilationResult{
+				Operations: []Operation{ // begin with params: [$delta]
+					&OperationPick{Depth: 0},                                 // [$delta, $delta]
+					&OperationMemoryGrow{},                                   // [$delta, $old_size]
+					&OperationDrop{Depth: &InclusiveRange{Start: 1, End: 1}}, // [$old_size]
+					&OperationBr{Target: &BranchTarget{}},                    // return!
+				},
+				LabelCallers: map[string]uint32{},
+				Types:        []*wasm.FunctionType{{Params: []wasm.ValueType{i32}, Results: []wasm.ValueType{i32}}},
+				Functions:    []uint32{0},
+				Signature:    &wasm.FunctionType{Params: []wasm.ValueType{wasm.ValueTypeI32}, Results: []wasm.ValueType{wasm.ValueTypeI32}},
+			},
+		},
 	}

 	for _, tt := range tests {