Repository: tetratelabs/wazero
Pull request #299 (2022-02-28 05:26:13)

Adds import count functions to unlock shared text decoding
Right now, it is a lot of copy/paste to decode inlined import
abbreviation with the same parser as module-defined ones. The primary
reason for this is that the import section is only guaranteed before
module-defined *after* expanding abbreviations. For example,

`(func (import "foo" "bar") ...` follows imports and may be followed by
any number of abbreviations. This means you can't know the end of
imports until you see the first `(func` which doesn't abbreviate one.

While certain external types, namely memory and table, can only be
imported once, this also applies to globals.

This change makes counting the current imports a `func() uint32` so that
it can be decoupled and used for any external type. Other PRs can
simplify memory and table which can only be imported once, but the
general signatures here will still work.



Patch:
diff --git a/internal/wasm/counts.go b/internal/wasm/counts.go
new file mode 100644
index 00000000..96865ea2
--- /dev/null
+++ b/internal/wasm/counts.go
@@ -0,0 +1,81 @@
+package internalwasm
+
+import "fmt"
+
+// ImportFuncCount returns the possibly empty count of imported functions. This plus SectionElementCount of
+// SectionIDFunction is the size of the function index namespace.
+func (m *Module) ImportFuncCount() uint32 {
+	return m.importCount(ExternTypeFunc)
+}
+
+// ImportTableCount returns the possibly empty count of imported tables. This plus SectionElementCount of SectionIDTable
+// is the size of the table index namespace.
+func (m *Module) ImportTableCount() uint32 {
+	return m.importCount(ExternTypeTable) // TODO: once validation happens on decode, this is zero or one.
+}
+
+// ImportMemoryCount returns the possibly empty count of imported memories. This plus SectionElementCount of
+// SectionIDMemory is the size of the memory index namespace.
+func (m *Module) ImportMemoryCount() uint32 {
+	return m.importCount(ExternTypeMemory) // TODO: once validation happens on decode, this is zero or one.
+}
+
+// ImportGlobalCount returns the possibly empty count of imported globals. This plus SectionElementCount of
+// SectionIDGlobal is the size of the global index namespace.
+func (m *Module) ImportGlobalCount() uint32 {
+	return m.importCount(ExternTypeGlobal)
+}
+
+// importCount returns the count of a specific type of import. This is important because it is easy to mistake the
+// length of the import section with the count of a specific kind of import.
+func (m *Module) importCount(et ExternType) (res uint32) {
+	for _, im := range m.ImportSection {
+		if im.Type == et {
+			res++
+		}
+	}
+	return
+}
+
+// SectionElementCount returns the count of elements in a given section ID
+//
+// For example...
+// * SectionIDType returns the count of FunctionType
+// * SectionIDCustom returns one if the NameSection is present
+// * SectionIDExport returns the count of unique export names
+func (m *Module) SectionElementCount(sectionID SectionID) uint32 { // element as in vector elements!
+	switch sectionID {
+	case SectionIDCustom:
+		if m.NameSection != nil {
+			return 1
+		}
+		return 0
+	case SectionIDType:
+		return uint32(len(m.TypeSection))
+	case SectionIDImport:
+		return uint32(len(m.ImportSection))
+	case SectionIDFunction:
+		return uint32(len(m.FunctionSection))
+	case SectionIDTable:
+		return uint32(len(m.TableSection))
+	case SectionIDMemory:
+		return uint32(len(m.MemorySection))
+	case SectionIDGlobal:
+		return uint32(len(m.GlobalSection))
+	case SectionIDExport:
+		return uint32(len(m.ExportSection))
+	case SectionIDStart:
+		if m.StartSection != nil {
+			return 1
+		}
+		return 0
+	case SectionIDElement:
+		return uint32(len(m.ElementSection))
+	case SectionIDCode:
+		return uint32(len(m.CodeSection))
+	case SectionIDData:
+		return uint32(len(m.DataSection))
+	default:
+		panic(fmt.Errorf("BUG: unknown section: %d", sectionID))
+	}
+}
diff --git a/internal/wasm/module.go b/internal/wasm/module.go
index 1485fe2e..79bb142c 100644
--- a/internal/wasm/module.go
+++ b/internal/wasm/module.go
@@ -139,8 +139,9 @@ type Module struct {
 	NameSection *NameSection
 }
 
-// TypeOfFunction returns the wasm.SectionIDType index for the given function namespace index or nil.
+// TypeOfFunction returns the internalwasm.SectionIDType index for the given function namespace index or nil.
 // Note: The function index namespace is preceded by imported functions.
+// TODO: Returning nil should be impossible when decode results are validated. Validate decode before backfilling tests.
 func (m *Module) TypeOfFunction(funcIdx Index) *FunctionType {
 	typeSectionLength := uint32(len(m.TypeSection))
 	if typeSectionLength == 0 {
@@ -373,49 +374,6 @@ func (m *Module) allDeclarations() (functions []Index, globals []*GlobalType, me
 	return
 }
 
-// SectionElementCount returns the count of elements in a given section ID
-//
-// For example...
-// * SectionIDType returns the count of FunctionType
-// * SectionIDCustom returns one if the NameSection is present
-// * SectionIDExport returns the count of unique export names
-func (m *Module) SectionElementCount(sectionID SectionID) uint32 { // element as in vector elements!
-	switch sectionID {
-	case SectionIDCustom:
-		if m.NameSection != nil {
-			return 1
-		}
-		return 0
-	case SectionIDType:
-		return uint32(len(m.TypeSection))
-	case SectionIDImport:
-		return uint32(len(m.ImportSection))
-	case SectionIDFunction:
-		return uint32(len(m.FunctionSection))
-	case SectionIDTable:
-		return uint32(len(m.TableSection))
-	case SectionIDMemory:
-		return uint32(len(m.MemorySection))
-	case SectionIDGlobal:
-		return uint32(len(m.GlobalSection))
-	case SectionIDExport:
-		return uint32(len(m.ExportSection))
-	case SectionIDStart:
-		if m.StartSection != nil {
-			return 1
-		}
-		return 0
-	case SectionIDElement:
-		return uint32(len(m.ElementSection))
-	case SectionIDCode:
-		return uint32(len(m.CodeSection))
-	case SectionIDData:
-		return uint32(len(m.DataSection))
-	default:
-		panic(fmt.Errorf("BUG: unknown section: %d", sectionID))
-	}
-}
-
 // SectionID identifies the sections of a Module in the WebAssembly 1.0 (20191205) Binary Format.
 //
 // Note: these are defined in the wasm package, instead of the binary package, as a key per section is needed regardless


Test patch:
diff --git a/internal/wasm/counts_test.go b/internal/wasm/counts_test.go
new file mode 100644
index 00000000..c0312384
--- /dev/null
+++ b/internal/wasm/counts_test.go
@@ -0,0 +1,298 @@
+package internalwasm
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestModule_ImportFuncCount(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    *Module
+		expected uint32
+	}{
+		{
+			name:  "none",
+			input: &Module{},
+		},
+		{
+			name:  "none with function section",
+			input: &Module{FunctionSection: []Index{0}},
+		},
+		{
+			name:     "one",
+			input:    &Module{ImportSection: []*Import{{Type: ExternTypeFunc}}},
+			expected: 1,
+		},
+		{
+			name:     "one with function section",
+			input:    &Module{ImportSection: []*Import{{Type: ExternTypeFunc}}, FunctionSection: []Index{0}},
+			expected: 1,
+		},
+		{
+			name:     "one with other imports",
+			input:    &Module{ImportSection: []*Import{{Type: ExternTypeFunc}, {Type: ExternTypeMemory}}},
+			expected: 1,
+		},
+		{
+			name:     "two",
+			input:    &Module{ImportSection: []*Import{{Type: ExternTypeFunc}, {Type: ExternTypeFunc}}},
+			expected: 2,
+		},
+	}
+
+	for _, tt := range tests {
+		tc := tt
+
+		t.Run(tc.name, func(t *testing.T) {
+			require.Equal(t, tc.expected, tc.input.ImportFuncCount())
+		})
+	}
+}
+
+// TODO: once we fix up-front validation, this only needs to check zero or one
+func TestModule_ImportTableCount(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    *Module
+		expected uint32
+	}{
+		{
+			name:  "none",
+			input: &Module{},
+		},
+		{
+			name:  "none with table section",
+			input: &Module{TableSection: []*TableType{{0x70, &LimitsType{1, nil}}}},
+		},
+		{
+			name:     "one",
+			input:    &Module{ImportSection: []*Import{{Type: ExternTypeTable}}},
+			expected: 1,
+		},
+		{
+			name: "one with table section",
+			input: &Module{
+				ImportSection: []*Import{{Type: ExternTypeTable}},
+				TableSection:  []*TableType{{0x70, &LimitsType{1, nil}}},
+			},
+			expected: 1,
+		},
+		{
+			name:     "one with other imports",
+			input:    &Module{ImportSection: []*Import{{Type: ExternTypeTable}, {Type: ExternTypeMemory}}},
+			expected: 1,
+		},
+		{
+			name:     "two",
+			input:    &Module{ImportSection: []*Import{{Type: ExternTypeTable}, {Type: ExternTypeTable}}},
+			expected: 2,
+		},
+	}
+
+	for _, tt := range tests {
+		tc := tt
+
+		t.Run(tc.name, func(t *testing.T) {
+			require.Equal(t, tc.expected, tc.input.ImportTableCount())
+		})
+	}
+}
+
+// TODO: once we fix up-front validation, this only needs to check zero or one
+func TestModule_ImportMemoryCount(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    *Module
+		expected uint32
+	}{
+		{
+			name:  "none",
+			input: &Module{},
+		},
+		{
+			name:  "none with memory section",
+			input: &Module{MemorySection: []*MemoryType{{Min: 1}}},
+		},
+		{
+			name:     "one",
+			input:    &Module{ImportSection: []*Import{{Type: ExternTypeMemory}}},
+			expected: 1,
+		},
+		{
+			name: "one with memory section",
+			input: &Module{
+				ImportSection: []*Import{{Type: ExternTypeMemory}},
+				MemorySection: []*MemoryType{{Min: 1}},
+			},
+			expected: 1,
+		},
+		{
+			name:     "one with other imports",
+			input:    &Module{ImportSection: []*Import{{Type: ExternTypeMemory}, {Type: ExternTypeTable}}},
+			expected: 1,
+		},
+		{
+			name:     "two",
+			input:    &Module{ImportSection: []*Import{{Type: ExternTypeMemory}, {Type: ExternTypeMemory}}},
+			expected: 2,
+		},
+	}
+
+	for _, tt := range tests {
+		tc := tt
+
+		t.Run(tc.name, func(t *testing.T) {
+			require.Equal(t, tc.expected, tc.input.ImportMemoryCount())
+		})
+	}
+}
+
+func TestModule_ImportGlobalCount(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    *Module
+		expected uint32
+	}{
+		{
+			name:  "none",
+			input: &Module{},
+		},
+		{
+			name:  "none with global section",
+			input: &Module{GlobalSection: []*Global{{Type: &GlobalType{ValType: ValueTypeI64}}}},
+		},
+		{
+			name:     "one",
+			input:    &Module{ImportSection: []*Import{{Type: ExternTypeGlobal}}},
+			expected: 1,
+		},
+		{
+			name: "one with global section",
+			input: &Module{
+				ImportSection: []*Import{{Type: ExternTypeGlobal}},
+				GlobalSection: []*Global{{Type: &GlobalType{ValType: ValueTypeI64}}},
+			},
+			expected: 1,
+		},
+		{
+			name:     "one with other imports",
+			input:    &Module{ImportSection: []*Import{{Type: ExternTypeGlobal}, {Type: ExternTypeMemory}}},
+			expected: 1,
+		},
+		{
+			name:     "two",
+			input:    &Module{ImportSection: []*Import{{Type: ExternTypeGlobal}, {Type: ExternTypeGlobal}}},
+			expected: 2,
+		},
+	}
+
+	for _, tt := range tests {
+		tc := tt
+
+		t.Run(tc.name, func(t *testing.T) {
+			require.Equal(t, tc.expected, tc.input.ImportGlobalCount())
+		})
+	}
+}
+
+func TestModule_SectionElementCount(t *testing.T) {
+	i32, f32 := ValueTypeI32, ValueTypeF32
+	zero := uint32(0)
+	empty := &ConstantExpression{Opcode: OpcodeI32Const, Data: []byte{0x00}}
+
+	tests := []struct {
+		name     string
+		input    *Module
+		expected map[string]uint32
+	}{
+		{
+			name:     "empty",
+			input:    &Module{},
+			expected: map[string]uint32{},
+		},
+		{
+			name:     "only name section",
+			input:    &Module{NameSection: &NameSection{ModuleName: "simple"}},
+			expected: map[string]uint32{"custom": 1},
+		},
+		{
+			name: "type section",
+			input: &Module{
+				TypeSection: []*FunctionType{
+					{},
+					{Params: []ValueType{i32, i32}, Results: []ValueType{i32}},
+					{Params: []ValueType{i32, i32, i32, i32}, Results: []ValueType{i32}},
+				},
+			},
+			expected: map[string]uint32{"type": 3},
+		},
+		{
+			name: "type and import section",
+			input: &Module{
+				TypeSection: []*FunctionType{
+					{Params: []ValueType{i32, i32}, Results: []ValueType{i32}},
+					{Params: []ValueType{f32, f32}, Results: []ValueType{f32}},
+				},
+				ImportSection: []*Import{
+					{
+						Module: "Math", Name: "Mul",
+						Type:     ExternTypeFunc,
+						DescFunc: 1,
+					}, {
+						Module: "Math", Name: "Add",
+						Type:     ExternTypeFunc,
+						DescFunc: 0,
+					},
+				},
+			},
+			expected: map[string]uint32{"import": 2, "type": 2},
+		},
+		{
+			name: "type function and start section",
+			input: &Module{
+				TypeSection:     []*FunctionType{{}},
+				FunctionSection: []Index{0},
+				CodeSection: []*Code{
+					{Body: []byte{OpcodeLocalGet, 0, OpcodeLocalGet, 1, OpcodeI32Add, OpcodeEnd}},
+				},
+				ExportSection: map[string]*Export{
+					"AddInt": {Name: "AddInt", Type: ExternTypeFunc, Index: Index(0)},
+				},
+				StartSection: &zero,
+			},
+			expected: map[string]uint32{"code": 1, "export": 1, "function": 1, "start": 1, "type": 1},
+		},
+		{
+			name: "memory and data",
+			input: &Module{
+				MemorySection: []*MemoryType{{Min: 1}},
+				DataSection:   []*DataSegment{{MemoryIndex: 0, OffsetExpression: empty}},
+			},
+			expected: map[string]uint32{"data": 1, "memory": 1},
+		},
+		{
+			name: "table and element",
+			input: &Module{
+				TableSection:   []*TableType{{ElemType: 0x70, Limit: &LimitsType{Min: 1}}},
+				ElementSection: []*ElementSegment{{TableIndex: 0, OffsetExpr: empty}},
+			},
+			expected: map[string]uint32{"element": 1, "table": 1},
+		},
+	}
+
+	for _, tt := range tests {
+		tc := tt
+
+		t.Run(tc.name, func(t *testing.T) {
+			actual := map[string]uint32{}
+			for i := SectionID(0); i <= SectionIDData; i++ {
+				if size := tc.input.SectionElementCount(i); size > 0 {
+					actual[SectionIDName(i)] = size
+				}
+			}
+			require.Equal(t, tc.expected, actual)
+		})
+	}
+}