Repository: tetratelabs/wazero
Pull request #139 (2022-01-17 07:25:11)

Adds binary encoding of TypeSection

codefromthecrypt: I redid this as a map, because it is the more natural type and by the time we start checking ranges, etc, the other approach started falling apart.
codefromthecrypt: Thanks for your patience and advice @mathetake!

Patch:
diff --git a/wasm/section.go b/wasm/section.go
index b2d165a9..9c45ec84 100644
--- a/wasm/section.go
+++ b/wasm/section.go
@@ -35,7 +35,7 @@ func (m *Module) encodeSections(buffer []byte) (bytes []byte) {
 		bytes = append(bytes, encodeCustomSection(name, data)...)
 	}
 	if len(m.TypeSection) > 0 {
-		panic("TODO: TypeSection")
+		bytes = append(bytes, encodeTypeSection(m.TypeSection)...)
 	}
 	if len(m.ImportSection) > 0 {
 		panic("TODO: ImportSection")
@@ -89,6 +89,25 @@ func encodeSection(sectionID SectionID, contents []byte) []byte {
 	return append([]byte{sectionID}, encodeSizePrefixed(contents)...)
 }
 
+// encodeTypeSection encodes a SectionIDType with any types in the Module.TypeSection
+// See https://www.w3.org/TR/wasm-core-1/#type-section%E2%91%A0
+func encodeTypeSection(ts []*FunctionType) []byte {
+	typeCount := len(ts)
+	data := leb128.EncodeUint32(uint32(typeCount))
+	for _, ft := range ts {
+		// Function types are encoded by the byte 0x60 followed by the respective vectors of parameter and result types.
+		// See https://www.w3.org/TR/wasm-core-1/#function-types%E2%91%A4
+		data = append(data, 0x60)
+		data = append(data, encodeValTypes(ft.Params)...)
+		data = append(data, encodeValTypes(ft.Results)...)
+	}
+
+	// Finally, make the header
+	dataSize := leb128.EncodeUint32(uint32(len(data)))
+	header := append([]byte{SectionIDType}, dataSize...)
+	return append(header, data...)
+}
+
 func (m *Module) readSections(r *reader) error {
 	for {
 		sectionID := make([]byte, 1)
diff --git a/wasm/type.go b/wasm/type.go
index c8b97e00..d67b4418 100644
--- a/wasm/type.go
+++ b/wasm/type.go
@@ -11,6 +11,24 @@ type FunctionType struct {
 	Params, Results []ValueType
 }
 
+var nullary = []byte{0x60, 0, 0}
+
+// encodedOneParam is a cache of FunctionType.encode values for param length 1 and result length 0
+var encodedOneParam = map[ValueType][]byte{
+	ValueTypeI32: {0x60, 1, ValueTypeI32, 0},
+	ValueTypeI64: {0x60, 1, ValueTypeI64, 0},
+	ValueTypeF32: {0x60, 1, ValueTypeF32, 0},
+	ValueTypeF64: {0x60, 1, ValueTypeF64, 0},
+}
+
+// encodedOneResult is a cache of FunctionType.encode values for param length 0 and result length 1
+var encodedOneResult = map[ValueType][]byte{
+	ValueTypeI32: {0x60, 0, 1, ValueTypeI32},
+	ValueTypeI64: {0x60, 0, 1, ValueTypeI64},
+	ValueTypeF32: {0x60, 0, 1, ValueTypeF32},
+	ValueTypeF64: {0x60, 0, 1, ValueTypeF64},
+}
+
 func (t *FunctionType) String() (ret string) {
 	for _, b := range t.Params {
 		ret += formatValueType(b)
@@ -28,6 +46,35 @@ func (t *FunctionType) String() (ret string) {
 	return
 }
 
+// encode returns a byte slice in WebAssembly 1.0 (MVP) Binary Format.
+//
+// Note: Function types are encoded by the byte 0x60 followed by the respective vectors of parameter and result types.
+// See https://www.w3.org/TR/wasm-core-1/#function-types%E2%91%A4
+func (t *FunctionType) encode() []byte {
+	paramCount, resultCount := len(t.Params), len(t.Results)
+	if paramCount == 0 && resultCount == 0 {
+		return nullary
+	}
+	if resultCount == 0 {
+		if paramCount == 1 {
+			if encoded, ok := encodedOneParam[t.Params[0]]; ok {
+				return encoded
+			}
+		}
+		return append(append([]byte{0x60}, encodeValTypes(t.Params)...), 0)
+	} else if resultCount == 1 {
+		if paramCount == 0 {
+			if encoded, ok := encodedOneResult[t.Results[0]]; ok {
+				return encoded
+			}
+		}
+		return append(append([]byte{0x60}, encodeValTypes(t.Params)...), 1, t.Results[0])
+	}
+	// This branch should never be reaches as WebAssembly 1.0 (MVP) supports at most 1 result
+	data := append([]byte{0x60}, encodeValTypes(t.Params)...)
+	return append(data, encodeValTypes(t.Results)...)
+}
+
 func readFunctionType(r io.Reader) (*FunctionType, error) {
 	b := make([]byte, 1)
 	if _, err := io.ReadFull(r, b); err != nil {
diff --git a/wasm/value.go b/wasm/value.go
index 53251914..0c6b8719 100644
--- a/wasm/value.go
+++ b/wasm/value.go
@@ -35,6 +35,38 @@ func formatValueType(t ValueType) (ret string) {
 	return
 }
 
+var noValType = []byte{0}
+
+// encodedValTypes is a cache of size prefixed binary encoding of known val types.
+var encodedValTypes = map[ValueType][]byte{
+	ValueTypeI32: {1, ValueTypeI32},
+	ValueTypeI64: {1, ValueTypeI64},
+	ValueTypeF32: {1, ValueTypeF32},
+	ValueTypeF64: {1, ValueTypeF64},
+}
+
+// encodeValTypes fast paths binary encoding of common value type lengths
+func encodeValTypes(vt []ValueType) []byte {
+	// Special case nullary and parameter lengths of wasi_snapshot_preview1 to avoid excess allocations
+	switch uint32(len(vt)) {
+	case 0: // nullary
+		return noValType
+	case 1: // ex $wasi_snapshot_preview1.fd_close or any result
+		if encoded, ok := encodedValTypes[vt[0]]; ok {
+			return encoded
+		}
+	case 2: // ex $wasi_snapshot_preview1.environ_sizes_get
+		return []byte{2, vt[0], vt[1]}
+	case 4: // ex $wasi_snapshot_preview1.fd_write
+		return []byte{4, vt[0], vt[1], vt[2], vt[3]}
+	case 9: // ex $wasi_snapshot_preview1.fd_write
+		return []byte{9, vt[0], vt[1], vt[2], vt[3], vt[4], vt[5], vt[6], vt[7], vt[8]}
+	}
+	// Slow path others until someone complains with a valid signature
+	count := leb128.EncodeUint32(uint32(len(vt)))
+	return append(count, vt...)
+}
+
 func readValueTypes(r io.Reader, num uint32) ([]ValueType, error) {
 	ret := make([]ValueType, num)
 	buf := make([]ValueType, num)


Test patch:
diff --git a/wasm/module_test.go b/wasm/module_test.go
index 5f1099e1..72e78702 100644
--- a/wasm/module_test.go
+++ b/wasm/module_test.go
@@ -7,6 +7,8 @@ import (
 )
 
 func TestModule_Encode(t *testing.T) {
+	i32 := ValueTypeI32
+
 	tests := []struct {
 		name     string
 		input    *Module
@@ -55,6 +57,23 @@ func TestModule_Encode(t *testing.T) {
 				0x06, // the Module name simple is 6 bytes long
 				's', 'i', 'm', 'p', 'l', 'e'),
 		},
+		{
+			name: "type section",
+			input: &Module{
+				TypeSection: []*FunctionType{
+					{},
+					{Params: []ValueType{i32, i32}, Results: []ValueType{i32}},
+					{Params: []ValueType{i32, i32, i32, i32}, Results: []ValueType{i32}},
+				},
+			},
+			expected: append(append(magic, version...),
+				SectionIDType, 0x12, // 18 bytes in this section
+				0x03,             // 3 types
+				0x60, 0x00, 0x00, // func=0x60 no param no result
+				0x60, 0x02, i32, i32, 0x01, i32, // func=0x60 2 params and 1 result
+				0x60, 0x04, i32, i32, i32, i32, 0x01, i32, // func=0x60 4 params and 1 result
+			),
+		},
 	}
 
 	for _, tt := range tests {
diff --git a/wasm/type_test.go b/wasm/type_test.go
index 22c23f78..dbc1cbb9 100644
--- a/wasm/type_test.go
+++ b/wasm/type_test.go
@@ -6,6 +6,85 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestFunctionType_Encode(t *testing.T) {
+	i32, i64 := ValueTypeI32, ValueTypeI64
+	tests := []struct {
+		name     string
+		input    *FunctionType
+		expected []byte
+	}{
+		{
+			name:     "empty",
+			input:    &FunctionType{},
+			expected: []byte{0x60, 0, 0},
+		},
+		{
+			name:     "one param no result",
+			input:    &FunctionType{Params: []ValueType{i32}},
+			expected: []byte{0x60, 1, i32, 0},
+		},
+		{
+			name:     "undefined param no result", // ensure future spec changes don't panic
+			input:    &FunctionType{Params: []ValueType{0x6f}},
+			expected: []byte{0x60, 1, 0x6f, 0},
+		},
+		{
+			name:     "no param one result",
+			input:    &FunctionType{Results: []ValueType{i32}},
+			expected: []byte{0x60, 0, 1, i32},
+		},
+		{
+			name:     "no param undefined result", // ensure future spec changes don't panic
+			input:    &FunctionType{Results: []ValueType{0x6f}},
+			expected: []byte{0x60, 0, 1, 0x6f},
+		},
+		{
+			name:     "one param one result",
+			input:    &FunctionType{Params: []ValueType{i64}, Results: []ValueType{i32}},
+			expected: []byte{0x60, 1, i64, 1, i32},
+		},
+		{
+			name:     "undefined param undefined result", // ensure future spec changes don't panic
+			input:    &FunctionType{Params: []ValueType{0x6f}, Results: []ValueType{0x6f}},
+			expected: []byte{0x60, 1, 0x6f, 1, 0x6f},
+		},
+		{
+			name:     "two params no result",
+			input:    &FunctionType{Params: []ValueType{i32, i64}},
+			expected: []byte{0x60, 2, i32, i64, 0},
+		},
+		{
+			name:     "no param two results", // this is just for coverage as WebAssembly 1.0 (MVP) does not allow it!
+			input:    &FunctionType{Results: []ValueType{i32, i64}},
+			expected: []byte{0x60, 0, 2, i32, i64},
+		},
+		{
+			name:     "one param two results", // this is just for coverage as WebAssembly 1.0 (MVP) does not allow it!
+			input:    &FunctionType{Params: []ValueType{i64}, Results: []ValueType{i32, i64}},
+			expected: []byte{0x60, 1, i64, 2, i32, i64},
+		},
+		{
+			name:     "two param one result",
+			input:    &FunctionType{Params: []ValueType{i32, i64}, Results: []ValueType{i32}},
+			expected: []byte{0x60, 2, i32, i64, 1, i32},
+		},
+		{
+			name:     "two param two results", // this is just for coverage as WebAssembly 1.0 (MVP) does not allow it!
+			input:    &FunctionType{Params: []ValueType{i32, i64}, Results: []ValueType{i32, i64}},
+			expected: []byte{0x60, 2, i32, i64, 2, i32, i64},
+		},
+	}
+
+	for _, tt := range tests {
+		tc := tt
+
+		t.Run(tc.name, func(t *testing.T) {
+			bytes := tc.input.encode()
+			require.Equal(t, tc.expected, bytes)
+		})
+	}
+}
+
 func TestFunctionType_String(t *testing.T) {
 	tp := FunctionType{}
 	require.Equal(t, "null_null", tp.String())
diff --git a/wasm/value_test.go b/wasm/value_test.go
new file mode 100644
index 00000000..cceed3b3
--- /dev/null
+++ b/wasm/value_test.go
@@ -0,0 +1,101 @@
+package wasm
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestEncodeValTypes(t *testing.T) {
+	i32, i64, f32, f64 := ValueTypeI32, ValueTypeI64, ValueTypeF32, ValueTypeF64
+	tests := []struct {
+		name     string
+		input    []ValueType
+		expected []byte
+	}{
+		{
+			name:     "empty",
+			input:    []ValueType{},
+			expected: []byte{0},
+		},
+		{
+			name:     "undefined", // ensure future spec changes don't panic
+			input:    []ValueType{0x6f},
+			expected: []byte{1, 0x6f},
+		},
+		{
+			name:     "i32",
+			input:    []ValueType{i32},
+			expected: []byte{1, i32},
+		},
+		{
+			name:     "i64",
+			input:    []ValueType{i64},
+			expected: []byte{1, i64},
+		},
+		{
+			name:     "f32",
+			input:    []ValueType{f32},
+			expected: []byte{1, f32},
+		},
+		{
+			name:     "f64",
+			input:    []ValueType{f64},
+			expected: []byte{1, f64},
+		},
+		{
+			name:     "i32i64",
+			input:    []ValueType{i32, i64},
+			expected: []byte{2, i32, i64},
+		},
+		{
+			name:     "i32i64f32",
+			input:    []ValueType{i32, i64, f32},
+			expected: []byte{3, i32, i64, f32},
+		},
+		{
+			name:     "i32i64f32f64",
+			input:    []ValueType{i32, i64, f32, f64},
+			expected: []byte{4, i32, i64, f32, f64},
+		},
+		{
+			name:     "i32i64f32f64i32",
+			input:    []ValueType{i32, i64, f32, f64, i32},
+			expected: []byte{5, i32, i64, f32, f64, i32},
+		},
+		{
+			name:     "i32i64f32f64i32i64",
+			input:    []ValueType{i32, i64, f32, f64, i32, i64},
+			expected: []byte{6, i32, i64, f32, f64, i32, i64},
+		},
+		{
+			name:     "i32i64f32f64i32i64f32",
+			input:    []ValueType{i32, i64, f32, f64, i32, i64, f32},
+			expected: []byte{7, i32, i64, f32, f64, i32, i64, f32},
+		},
+		{
+			name:     "i32i64f32f64i32i64f32f64",
+			input:    []ValueType{i32, i64, f32, f64, i32, i64, f32, f64},
+			expected: []byte{8, i32, i64, f32, f64, i32, i64, f32, f64},
+		},
+		{
+			name:     "i32i64f32f64i32i64f32f64i32",
+			input:    []ValueType{i32, i64, f32, f64, i32, i64, f32, f64, i32},
+			expected: []byte{9, i32, i64, f32, f64, i32, i64, f32, f64, i32},
+		},
+		{
+			name:     "i32i64f32f64i32i64f32f64i32i64",
+			input:    []ValueType{i32, i64, f32, f64, i32, i64, f32, f64, i32, i64},
+			expected: []byte{10, i32, i64, f32, f64, i32, i64, f32, f64, i32, i64},
+		},
+	}
+
+	for _, tt := range tests {
+		tc := tt
+
+		t.Run(tc.name, func(t *testing.T) {
+			bytes := encodeValTypes(tc.input)
+			require.Equal(t, tc.expected, bytes)
+		})
+	}
+}
diff --git a/wasm/wat/convert_test.go b/wasm/wat/convert_test.go
index 3d0349d6..1b99c1ef 100644
--- a/wasm/wat/convert_test.go
+++ b/wasm/wat/convert_test.go
@@ -207,6 +207,7 @@ func TestTextToBinary(t *testing.T) {
 						4: "hello",
 					},
 					LocalNames: map[uint32]map[uint32]string{
+						1: {0: "fd", 1: "iovs_ptr", 2: "iovs_len", 3: "nwritten_ptr"},
 						2: {0: "x", 1: "y"},
 						3: {0: "l", 1: "r"},
 					},
diff --git a/wasm/wat/parser_test.go b/wasm/wat/parser_test.go
index 7a47f21d..c04d4d95 100644
--- a/wasm/wat/parser_test.go
+++ b/wasm/wat/parser_test.go
@@ -20,7 +20,8 @@ import (
 var example = []byte(`(module $example
 	(type $i32i32_i32 (func (param i32 i32) (result i32)))
 	(import "wasi_snapshot_preview1" "arg_sizes_get" (func $runtime.arg_sizes_get (type $i32i32_i32)))
-	(import "wasi_snapshot_preview1" "fd_write" (func $runtime.fd_write (param i32 i32 i32 i32) (result i32)))
+	(import "wasi_snapshot_preview1" "fd_write" (func $runtime.fd_write
+		(param $fd i32) (param $iovs_ptr i32) (param $iovs_len i32) (param $nwritten_ptr i32) (result i32)))
 	(import "Math" "Mul" (func $mul (param $x f32) (param $y f32) (result f32)))
 	(import "Math" "Add" (func $add (type $i32i32_i32) (param $l i32) (param $r i32) (result i32)))
 	(type (func))
@@ -509,7 +510,13 @@ func TestParseModule(t *testing.T) {
 					{importIndex: 0, module: "wasi_snapshot_preview1", name: "arg_sizes_get", funcName: "runtime.arg_sizes_get",
 						typeIndex: &index{numeric: 0, line: 3, col: 86}},
 					{importIndex: 1, module: "wasi_snapshot_preview1", name: "fd_write", funcName: "runtime.fd_write",
-						typeIndex: &index{numeric: 2, line: 3, col: 81}},
+						typeIndex: &index{numeric: 2, line: 3, col: 81},
+						paramNames: paramNames{
+							&paramNameIndex{uint32(0), []byte("fd")},
+							&paramNameIndex{uint32(1), []byte("iovs_ptr")},
+							&paramNameIndex{uint32(2), []byte("iovs_len")},
+							&paramNameIndex{uint32(3), []byte("nwritten_ptr")},
+						}},
 					{importIndex: 2, module: "Math", name: "Mul", funcName: "mul",
 						typeIndex: &index{numeric: 3, line: 3, col: 81},
 						paramNames: paramNames{
@@ -518,10 +525,10 @@ func TestParseModule(t *testing.T) {
 						},
 					},
 					{importIndex: 3, module: "Math", name: "Add", funcName: "add",
-						typeIndex: &index{numeric: 0, line: 6, col: 40},
+						typeIndex: &index{numeric: 0, line: 7, col: 40},
 						typeInlined: &inlinedTypeFunc{
 							typeFunc: &typeFunc{name: "i32i32_i32", params: []wasm.ValueType{i32, i32}, result: i32},
-							line:     6, col: 35,
+							line:     7, col: 35,
 						},
 						paramNames: paramNames{
 							&paramNameIndex{uint32(0), []byte{'l'}},
@@ -529,9 +536,9 @@ func TestParseModule(t *testing.T) {
 						},
 					},
 					{importIndex: 4, module: "", name: "hello", funcName: "hello",
-						typeIndex: &index{numeric: 1, line: 8, col: 40}},
+						typeIndex: &index{numeric: 1, line: 9, col: 40}},
 				},
-				startFunction: &index{numeric: 4, line: 9, col: 9},
+				startFunction: &index{numeric: 4, line: 10, col: 9},
 			},
 		},
 	}