Repository: tetratelabs/wazero
Pull request #303 (2022-03-01 03:04:24)

text: defers evaluation of current section index
The current section index is typically only read when there is an
unresolved reference. This pushes evaluation until when it is needed,
which simpilifies tracking and reduces arity of callbacks.



Patch:
diff --git a/internal/wasm/text/decoder.go b/internal/wasm/text/decoder.go
index 1a15f0bf..1aca36ab 100644
--- a/internal/wasm/text/decoder.go
+++ b/internal/wasm/text/decoder.go
@@ -138,9 +138,9 @@ func DecodeModule(source []byte) (result *wasm.Module, err error) {
 
 func newModuleParser(module *wasm.Module) *moduleParser {
 	p := moduleParser{module: module,
-		typeNamespace:   newIndexNamespace(),
-		funcNamespace:   newIndexNamespace(),
-		memoryNamespace: newIndexNamespace(),
+		typeNamespace:   newIndexNamespace(module.SectionElementCount),
+		funcNamespace:   newIndexNamespace(module.SectionElementCount),
+		memoryNamespace: newIndexNamespace(module.SectionElementCount),
 	}
 	p.typeParser = newTypeParser(p.typeNamespace, p.onTypeEnd)
 	p.typeUseParser = newTypeUseParser(module, p.typeNamespace)
@@ -183,7 +183,6 @@ func (p *moduleParser) beginModuleField(tok tokenType, tokenBytes []byte, _, _ u
 			return p.parseImportModule, nil
 		case wasm.ExternTypeFuncName:
 			p.pos = positionFunc
-			p.funcParser.currentIdx = p.module.SectionElementCount(wasm.SectionIDFunction)
 			return p.funcParser.begin, nil
 		case wasm.ExternTypeTableName:
 			return nil, fmt.Errorf("TODO: %s", tokenBytes)
@@ -369,12 +368,10 @@ func (p *moduleParser) addFunctionName(name string) {
 // Ex. If there is no signature `(import "" "main" (func))`
 //                     calls onImportFunc here ---^
 func (p *moduleParser) parseImportFunc(tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
-	idx := p.module.SectionElementCount(wasm.SectionIDImport)
 	if tok == tokenID { // Ex. (func $main $main)
 		return nil, fmt.Errorf("redundant ID %s", tokenBytes)
 	}
-
-	return p.typeUseParser.begin(wasm.SectionIDImport, idx, p.onImportFunc, tok, tokenBytes, line, col)
+	return p.typeUseParser.begin(wasm.SectionIDImport, p.onImportFunc, tok, tokenBytes, line, col)
 }
 
 // onImportFunc records the type index and local names of the current imported function, and increments
@@ -530,8 +527,7 @@ func (p *moduleParser) parseExportDesc(tok tokenType, tokenBytes []byte, line, c
 	default:
 		panic(fmt.Errorf("BUG: unhandled parsing state on parseExportDesc: %v", p.pos))
 	}
-	eIdx := p.module.SectionElementCount(wasm.SectionIDExport)
-	typeIdx, resolved, err := namespace.parseIndex(wasm.SectionIDExport, eIdx, 0, tok, tokenBytes, line, col)
+	typeIdx, resolved, err := namespace.parseIndex(wasm.SectionIDExport, 0, tok, tokenBytes, line, col)
 	if err != nil {
 		return nil, err
 	}
@@ -539,6 +535,7 @@ func (p *moduleParser) parseExportDesc(tok tokenType, tokenBytes []byte, line, c
 
 	// All sections in wasm.Module are numeric indices except exports. Hence, we have to special-case here
 	if !resolved {
+		eIdx := p.module.SectionElementCount(wasm.SectionIDExport)
 		if p.unresolvedExports == nil {
 			p.unresolvedExports = map[wasm.Index]*wasm.Export{eIdx: e}
 		} else {
@@ -580,7 +577,7 @@ func (p *moduleParser) parseExportEnd(tok tokenType, tokenBytes []byte, _, _ uin
 
 // parseStart returns parseStartEnd after recording the start function index, or errs if it couldn't be read.
 func (p *moduleParser) parseStart(tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
-	idx, _, err := p.funcNamespace.parseIndex(wasm.SectionIDStart, 0, 0, tok, tokenBytes, line, col)
+	idx, _, err := p.funcNamespace.parseIndex(wasm.SectionIDStart, 0, tok, tokenBytes, line, col)
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/wasm/text/func_parser.go b/internal/wasm/text/func_parser.go
index b24b11d8..af3e00ed 100644
--- a/internal/wasm/text/func_parser.go
+++ b/internal/wasm/text/func_parser.go
@@ -31,7 +31,6 @@ type funcParser struct {
 	// funcNamespace is described by moduleParser.funcNamespace
 	funcNamespace *indexNamespace
 
-	currentIdx        wasm.Index
 	currentName       string
 	currentTypeIdx    wasm.Index
 	currentParamNames wasm.NameMap
@@ -88,7 +87,7 @@ func (p *funcParser) parseFunc(tok tokenType, tokenBytes []byte, line, col uint3
 		return nil, fmt.Errorf("redundant ID %s", tokenBytes)
 	}
 
-	return p.typeUseParser.begin(wasm.SectionIDFunction, p.currentIdx, p.afterTypeUse, tok, tokenBytes, line, col)
+	return p.typeUseParser.begin(wasm.SectionIDFunction, p.afterTypeUse, tok, tokenBytes, line, col)
 }
 
 // afterTypeUse is a tokenParser that starts after a type use.
@@ -127,7 +126,7 @@ func sExpressionsUnsupported(tok tokenType, tokenBytes []byte, _, _ uint32) (tok
 	return nil, fmt.Errorf("TODO: s-expressions are not yet supported: %s", tokenBytes)
 }
 
-func (p *funcParser) beginFieldOrInstruction(tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
+func (p *funcParser) beginFieldOrInstruction(tok tokenType, tokenBytes []byte, _, _ uint32) (tokenParser, error) {
 	switch tok {
 	case tokenLParen:
 		return sExpressionsUnsupported, nil
@@ -174,7 +173,7 @@ func (p *funcParser) end() (tokenParser, error) {
 // placeholder byte(0) is added instead and will be resolved later.
 func (p *funcParser) parseFuncIndex(tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
 	bodyOffset := uint32(len(p.currentBody))
-	idx, resolved, err := p.funcNamespace.parseIndex(wasm.SectionIDCode, p.currentIdx, bodyOffset, tok, tokenBytes, line, col)
+	idx, resolved, err := p.funcNamespace.parseIndex(wasm.SectionIDCode, bodyOffset, tok, tokenBytes, line, col)
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/wasm/text/index_namespace.go b/internal/wasm/text/index_namespace.go
index 5684ddff..a4b70fd7 100644
--- a/internal/wasm/text/index_namespace.go
+++ b/internal/wasm/text/index_namespace.go
@@ -7,8 +7,9 @@ import (
 	wasm "github.com/tetratelabs/wazero/internal/wasm"
 )
 
-func newIndexNamespace() *indexNamespace {
-	return &indexNamespace{idToIdx: map[string]wasm.Index{}}
+// newIndexNamespace sectionElementCount parameter should be internalwasm.Module SectionElementCount unless testing.
+func newIndexNamespace(sectionElementCount func(wasm.SectionID) uint32) *indexNamespace {
+	return &indexNamespace{sectionElementCount: sectionElementCount, idToIdx: map[string]wasm.Index{}}
 }
 
 // indexNamespace contains the count in an index namespace and any association of symbolic IDs to numeric indices.
@@ -21,6 +22,8 @@ func newIndexNamespace() *indexNamespace {
 //
 // See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#text-context
 type indexNamespace struct {
+	sectionElementCount func(wasm.SectionID) uint32
+
 	unresolvedIndices []*unresolvedIndex
 
 	// count is the count of items in this namespace
@@ -57,7 +60,7 @@ func (i *indexNamespace) requireNoID(idToken []byte) (string, error) {
 }
 
 // parseIndex is a tokenParser called in a field that can only contain a symbolic identifier or raw numeric index.
-func (i *indexNamespace) parseIndex(section wasm.SectionID, idx wasm.Index, bodyOffset uint32, tok tokenType, tokenBytes []byte, line, col uint32) (targetIdx wasm.Index, resolved bool, err error) {
+func (i *indexNamespace) parseIndex(section wasm.SectionID, bodyOffset uint32, tok tokenType, tokenBytes []byte, line, col uint32) (targetIdx wasm.Index, resolved bool, err error) {
 	switch tok {
 	case tokenUN: // Ex. 2
 		if i, overflow := decodeUint32(tokenBytes); overflow {
@@ -69,12 +72,12 @@ func (i *indexNamespace) parseIndex(section wasm.SectionID, idx wasm.Index, body
 		if targetIdx < i.count {
 			resolved = true
 		} else {
-			i.recordOutOfRange(section, idx, bodyOffset, targetIdx, line, col)
+			i.recordOutOfRange(section, bodyOffset, targetIdx, line, col)
 		}
 	case tokenID: // Ex. $main
 		targetID := string(stripDollar(tokenBytes))
 		if targetIdx, resolved = i.idToIdx[targetID]; !resolved {
-			i.recordUnresolved(section, idx, bodyOffset, targetID, line, col)
+			i.recordUnresolved(section, bodyOffset, targetID, line, col)
 		}
 		return
 	case tokenRParen:
@@ -88,14 +91,16 @@ func (i *indexNamespace) parseIndex(section wasm.SectionID, idx wasm.Index, body
 // recordUnresolved records an ID, such as "main", is not yet resolvable.
 //
 // See unresolvedIndex for parameter descriptions
-func (i *indexNamespace) recordUnresolved(section wasm.SectionID, idx wasm.Index, bodyOffset uint32, targetID string, line, col uint32) {
+func (i *indexNamespace) recordUnresolved(section wasm.SectionID, bodyOffset uint32, targetID string, line, col uint32) {
+	idx := i.sectionElementCount(section)
 	i.unresolvedIndices = append(i.unresolvedIndices, &unresolvedIndex{section: section, idx: idx, bodyOffset: bodyOffset, targetID: targetID, line: line, col: col})
 }
 
 // recordUnresolved records numeric index is currently out of bounds.
 //
 // See unresolvedIndex for parameter descriptions
-func (i *indexNamespace) recordOutOfRange(section wasm.SectionID, idx wasm.Index, bodyOffset uint32, targetIdx wasm.Index, line, col uint32) {
+func (i *indexNamespace) recordOutOfRange(section wasm.SectionID, bodyOffset uint32, targetIdx wasm.Index, line, col uint32) {
+	idx := i.sectionElementCount(section)
 	i.unresolvedIndices = append(i.unresolvedIndices, &unresolvedIndex{section: section, idx: idx, bodyOffset: bodyOffset, targetIdx: targetIdx, line: line, col: col})
 }
 
diff --git a/internal/wasm/text/typeuse_parser.go b/internal/wasm/text/typeuse_parser.go
index bbf33dff..94daefa4 100644
--- a/internal/wasm/text/typeuse_parser.go
+++ b/internal/wasm/text/typeuse_parser.go
@@ -32,13 +32,12 @@ type onTypeUse func(typeIdx wasm.Index, paramNames wasm.NameMap, pos callbackPos
 // "type", "param" and "result" inner fields in the correct order.
 // Note: typeUseParser is reusable. The caller resets via begin.
 type typeUseParser struct {
-	// module during parsing is a read-only pointer to the TypeSection.
+	// module during parsing is a read-only pointer to the TypeSection and SectionElementCount
 	module *wasm.Module
 
 	typeNamespace *indexNamespace
 
 	section wasm.SectionID
-	idx     wasm.Index
 
 	// inlinedTypes are anonymous types defined by signature, which at the time of definition didn't match a
 	// module-defined type. Ex. `(param i32)` in `(import (func (param i32)))`
@@ -101,26 +100,25 @@ type typeUseParser struct {
 //      beginTypeParamOrResult starts here --^             ^
 //                                onTypeUse resumes here --+
 //
-func (p *typeUseParser) begin(section wasm.SectionID, idx wasm.Index, onTypeUse onTypeUse, tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
+func (p *typeUseParser) begin(section wasm.SectionID, onTypeUse onTypeUse, tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
 	pos := callbackPositionUnhandledToken
 	p.pos = positionInitial // to ensure errorContext reports properly
 	switch tok {
 	case tokenLParen:
 		p.section = section
-		p.idx = idx
 		p.onTypeUse = onTypeUse
 		return p.beginTypeParamOrResult, nil
 	case tokenRParen:
 		pos = callbackPositionEndField
 	}
-	return onTypeUse(p.emptyTypeIndex(section, idx), nil, pos, tok, tokenBytes, line, col)
+	return onTypeUse(p.emptyTypeIndex(section), nil, pos, tok, tokenBytes, line, col)
 }
 
 // v_v is a nullary function type (void -> void)
 var v_v = &wasm.FunctionType{}
 
 // inlinedTypeIndex searches for any existing empty type to re-use
-func (p *typeUseParser) emptyTypeIndex(section wasm.SectionID, idx wasm.Index) wasm.Index {
+func (p *typeUseParser) emptyTypeIndex(section wasm.SectionID) wasm.Index {
 	for i, t := range p.module.TypeSection {
 		if t == v_v {
 			return wasm.Index(i)
@@ -142,6 +140,7 @@ func (p *typeUseParser) emptyTypeIndex(section wasm.SectionID, idx wasm.Index) w
 	}
 
 	// typePos is not needed on empty as there's nothing to verify
+	idx := p.module.SectionElementCount(section)
 	i := &inlinedTypeIndex{section: section, idx: idx, inlinedIdx: inlinedIdx, typePos: nil}
 	p.inlinedTypeIndices = append(p.inlinedTypeIndices, i)
 	return 0 // substitute index that will be replaced later
@@ -168,7 +167,7 @@ func (p *typeUseParser) beginTypeParamOrResult(tok tokenType, tokenBytes []byte,
 // parseType parses a type index inside the "type" field. If not yet in the TypeSection, the position is recorded for
 // resolution later. Finally, this returns parseTypeEnd to finish the field.
 func (p *typeUseParser) parseType(tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
-	idx, resolved, err := p.typeNamespace.parseIndex(p.section, p.idx, 0, tok, tokenBytes, line, col)
+	idx, resolved, err := p.typeNamespace.parseIndex(p.section, 0, tok, tokenBytes, line, col)
 	if err != nil {
 		return nil, err
 	}
@@ -378,7 +377,7 @@ func (p *typeUseParser) end(pos callbackPosition, tok tokenType, tokenBytes []by
 			return nil, err
 		}
 	} else if p.currentInlinedType == nil { // no type was found
-		typeIdx = p.emptyTypeIndex(p.section, p.idx)
+		typeIdx = p.emptyTypeIndex(p.section)
 	} else { // There was no explicitly defined type, so search for any existing ones to re-use
 		typeIdx = p.inlinedTypeIndex()
 	}
@@ -444,7 +443,8 @@ type inlinedTypeIndex struct {
 }
 
 func (p *typeUseParser) recordInlinedType(inlinedIdx wasm.Index) {
-	i := &inlinedTypeIndex{section: p.section, idx: p.idx, inlinedIdx: inlinedIdx, typePos: p.currentTypeIndexUnresolved}
+	idx := p.module.SectionElementCount(p.section)
+	i := &inlinedTypeIndex{section: p.section, idx: idx, inlinedIdx: inlinedIdx, typePos: p.currentTypeIndexUnresolved}
 	p.inlinedTypeIndices = append(p.inlinedTypeIndices, i)
 }
 


Test patch:
diff --git a/internal/wasm/text/func_parser_test.go b/internal/wasm/text/func_parser_test.go
index d7b189a1..ff587d76 100644
--- a/internal/wasm/text/func_parser_test.go
+++ b/internal/wasm/text/func_parser_test.go
@@ -51,7 +51,9 @@ func TestFuncParser(t *testing.T) {
 				return parseErr, nil
 			}
 
-			require.NoError(t, parseFunc(newFuncParser(&typeUseParser{module: &wasm.Module{}}, newIndexNamespace(), setFunc), tc.source))
+			module := &wasm.Module{}
+			fp := newFuncParser(&typeUseParser{module: module}, newIndexNamespace(module.SectionElementCount), setFunc)
+			require.NoError(t, parseFunc(fp, tc.source))
 			require.Equal(t, tc.expected, parsedCode)
 		})
 	}
@@ -107,7 +109,8 @@ func TestFuncParser_Call_Unresolved(t *testing.T) {
 				return parseErr, nil
 			}
 
-			fp := newFuncParser(&typeUseParser{module: &wasm.Module{}}, newIndexNamespace(), setFunc)
+			module := &wasm.Module{}
+			fp := newFuncParser(&typeUseParser{module: module}, newIndexNamespace(module.SectionElementCount), setFunc)
 			require.NoError(t, parseFunc(fp, tc.source))
 			require.Equal(t, tc.expectedCode, parsedCode)
 			require.Equal(t, []*unresolvedIndex{tc.expectedUnresolvedIndex}, fp.funcNamespace.unresolvedIndices)
@@ -141,7 +144,10 @@ func TestFuncParser_Call_Resolved(t *testing.T) {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			funcNamespace := newIndexNamespace()
+			funcNamespace := newIndexNamespace(func(sectionID wasm.SectionID) uint32 {
+				require.Equal(t, wasm.SectionIDFunction, sectionID)
+				return 0
+			})
 			_, err := funcNamespace.setID([]byte("$not_main"))
 			require.NoError(t, err)
 			funcNamespace.count++
@@ -221,7 +227,8 @@ func TestFuncParser_Errors(t *testing.T) {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			fp := newFuncParser(&typeUseParser{module: &wasm.Module{}}, newIndexNamespace(), failOnFunc)
+			module := &wasm.Module{}
+			fp := newFuncParser(&typeUseParser{module: module}, newIndexNamespace(module.SectionElementCount), failOnFunc)
 			require.EqualError(t, parseFunc(fp, tc.source), tc.expectedErr)
 		})
 	}
diff --git a/internal/wasm/text/index_namespace_test.go b/internal/wasm/text/index_namespace_test.go
index 997b814f..080d04a2 100644
--- a/internal/wasm/text/index_namespace_test.go
+++ b/internal/wasm/text/index_namespace_test.go
@@ -10,7 +10,10 @@ import (
 )
 
 func TestIndexNamespace_SetId(t *testing.T) {
-	in := newIndexNamespace()
+	in := newIndexNamespace(func(_ wasm.SectionID) uint32 {
+		t.Fail()
+		return 0
+	})
 	t.Run("set when empty", func(t *testing.T) {
 		id, err := in.setID([]byte("$x"))
 		require.NoError(t, err)
diff --git a/internal/wasm/text/memory_parser_test.go b/internal/wasm/text/memory_parser_test.go
index f536ee28..3fbf421a 100644
--- a/internal/wasm/text/memory_parser_test.go
+++ b/internal/wasm/text/memory_parser_test.go
@@ -60,7 +60,10 @@ func TestMemoryParser(t *testing.T) {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			memoryNamespace := newIndexNamespace()
+			memoryNamespace := newIndexNamespace(func(sectionID wasm.SectionID) uint32 {
+				require.Equal(t, wasm.SectionIDMemory, sectionID)
+				return 0
+			})
 			parsed, tp, err := parseMemoryType(memoryNamespace, tc.input)
 			require.NoError(t, err)
 			require.Equal(t, tc.expected, parsed)
@@ -133,14 +136,21 @@ func TestMemoryParser_Errors(t *testing.T) {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			parsed, _, err := parseMemoryType(newIndexNamespace(), tc.input)
+			memoryNamespace := newIndexNamespace(func(sectionID wasm.SectionID) uint32 {
+				require.Equal(t, wasm.SectionIDMemory, sectionID)
+				return 0
+			})
+			parsed, _, err := parseMemoryType(memoryNamespace, tc.input)
 			require.EqualError(t, err, tc.expectedErr)
 			require.Nil(t, parsed)
 		})
 	}
 
 	t.Run("duplicate ID", func(t *testing.T) {
-		memoryNamespace := newIndexNamespace()
+		memoryNamespace := newIndexNamespace(func(sectionID wasm.SectionID) uint32 {
+			require.Equal(t, wasm.SectionIDMemory, sectionID)
+			return 0
+		})
 		_, err := memoryNamespace.setID([]byte("$mem"))
 		require.NoError(t, err)
 		memoryNamespace.count++
diff --git a/internal/wasm/text/type_parser_test.go b/internal/wasm/text/type_parser_test.go
index c872d305..2cebb338 100644
--- a/internal/wasm/text/type_parser_test.go
+++ b/internal/wasm/text/type_parser_test.go
@@ -100,7 +100,10 @@ func TestTypeParser(t *testing.T) {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			typeNamespace := newIndexNamespace()
+			typeNamespace := newIndexNamespace(func(sectionID wasm.SectionID) uint32 {
+				require.Equal(t, wasm.SectionIDType, sectionID)
+				return 0
+			})
 			parsed, tp, err := parseFunctionType(typeNamespace, tc.input)
 			require.NoError(t, err)
 			require.Equal(t, tc.expected, parsed)
@@ -228,14 +231,21 @@ func TestTypeParser_Errors(t *testing.T) {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			parsed, _, err := parseFunctionType(newIndexNamespace(), tc.input)
+			typeNamespace := newIndexNamespace(func(sectionID wasm.SectionID) uint32 {
+				require.Equal(t, wasm.SectionIDType, sectionID)
+				return 0
+			})
+			parsed, _, err := parseFunctionType(typeNamespace, tc.input)
 			require.EqualError(t, err, tc.expectedErr)
 			require.Nil(t, parsed)
 		})
 	}
 
 	t.Run("duplicate ID", func(t *testing.T) {
-		typeNamespace := newIndexNamespace()
+		typeNamespace := newIndexNamespace(func(sectionID wasm.SectionID) uint32 {
+			require.Equal(t, wasm.SectionIDType, sectionID)
+			return 0
+		})
 		_, err := typeNamespace.setID([]byte("$v_v"))
 		require.NoError(t, err)
 		typeNamespace.count++
diff --git a/internal/wasm/text/typeuse_parser_test.go b/internal/wasm/text/typeuse_parser_test.go
index efe112dd..ad7b3c77 100644
--- a/internal/wasm/text/typeuse_parser_test.go
+++ b/internal/wasm/text/typeuse_parser_test.go
@@ -81,7 +81,8 @@ func TestTypeUseParser_InlinesTypesWhenNotYetAdded(t *testing.T) {
 	}
 
 	runTypeUseParserTests(t, tests, func(tc *typeUseParserTest) (*typeUseParser, func(t *testing.T)) {
-		tp := newTypeUseParser(&wasm.Module{}, newIndexNamespace())
+		module := &wasm.Module{}
+		tp := newTypeUseParser(module, newIndexNamespace(module.SectionElementCount))
 		return tp, func(t *testing.T) {
 			// We should have inlined the type, and it is the first type use, which means the inlined index is zero
 			require.Zero(t, tp.inlinedTypeIndices[0].inlinedIdx)
@@ -116,7 +117,8 @@ func TestTypeUseParser_UnresolvedType(t *testing.T) {
 		},
 	}
 	runTypeUseParserTests(t, tests, func(tc *typeUseParserTest) (*typeUseParser, func(t *testing.T)) {
-		tp := newTypeUseParser(&wasm.Module{}, newIndexNamespace())
+		module := &wasm.Module{}
+		tp := newTypeUseParser(module, newIndexNamespace(module.SectionElementCount))
 		return tp, func(t *testing.T) {
 			require.NotNil(t, tp.typeNamespace.unresolvedIndices)
 			if tc.expectedInlinedType == nil {
@@ -207,10 +209,9 @@ func TestTypeUseParser_ReuseExistingType(t *testing.T) {
 		},
 	}
 	runTypeUseParserTests(t, tests, func(tc *typeUseParserTest) (*typeUseParser, func(t *testing.T)) {
-		typeNamespace := newIndexNamespace()
-
 		// Add types to cover the main ways types uses are declared
 		module := &wasm.Module{TypeSection: []*wasm.FunctionType{v_i32, v_v, i32_v, i32i64_i32}}
+		typeNamespace := newIndexNamespace(module.SectionElementCount)
 		_, err := typeNamespace.setID([]byte("$v_i32"))
 		require.NoError(t, err)
 		typeNamespace.count++
@@ -260,7 +261,8 @@ func TestTypeUseParser_ReuseExistingInlinedType(t *testing.T) {
 		},
 	}
 	runTypeUseParserTests(t, tests, func(tc *typeUseParserTest) (*typeUseParser, func(t *testing.T)) {
-		tp := newTypeUseParser(&wasm.Module{}, newIndexNamespace())
+		module := &wasm.Module{}
+		tp := newTypeUseParser(module, newIndexNamespace(module.SectionElementCount))
 		// inline a type that doesn't match the test
 		require.NoError(t, parseTypeUse(tp, "((param i32 i64))", ignoreTypeUse))
 		// inline the test type
@@ -305,7 +307,8 @@ func TestTypeUseParser_BeginResets(t *testing.T) {
 		},
 	}
 	runTypeUseParserTests(t, tests, func(tc *typeUseParserTest) (*typeUseParser, func(t *testing.T)) {
-		tp := newTypeUseParser(&wasm.Module{}, newIndexNamespace())
+		module := &wasm.Module{}
+		tp := newTypeUseParser(module, newIndexNamespace(module.SectionElementCount))
 		// inline a type that uses all fields
 		require.NoError(t, parseTypeUse(tp, "((type $i32i64_i32) (param $x i32) (param $y i64) (result i32))", ignoreTypeUse))
 		require.NoError(t, parseTypeUse(tp, tc.source, ignoreTypeUse))
@@ -482,17 +485,18 @@ func TestTypeUseParser_Errors(t *testing.T) {
 		tc := tt
 
 		t.Run(tc.name, func(t *testing.T) {
-			err := parseTypeUse(newTypeUseParser(&wasm.Module{}, newIndexNamespace()), tc.source, failOnTypeUse)
+			module := &wasm.Module{}
+			tp := newTypeUseParser(module, newIndexNamespace(module.SectionElementCount))
+			err := parseTypeUse(tp, tc.source, failOnTypeUse)
 			require.EqualError(t, err, tc.expectedErr)
 		})
 	}
 }
 
 func TestTypeUseParser_FailsMatch(t *testing.T) {
-	typeNamespace := newIndexNamespace()
-
 	// Add types to cover the main ways types uses are declared
 	module := &wasm.Module{TypeSection: []*wasm.FunctionType{v_v, i32i64_i32}}
+	typeNamespace := newIndexNamespace(module.SectionElementCount)
 	_, err := typeNamespace.setID([]byte("$v_v"))
 	require.NoError(t, err)
 	typeNamespace.count++
@@ -544,7 +548,7 @@ var failOnTypeUse onTypeUse = func(typeIdx wasm.Index, paramNames wasm.NameMap,
 
 func parseTypeUse(tp *typeUseParser, source string, onTypeUse onTypeUse) error {
 	var parser tokenParser = func(tok tokenType, tokenBytes []byte, line, col uint32) (tokenParser, error) {
-		return tp.begin(wasm.SectionIDFunction, 0, onTypeUse, tok, tokenBytes, line, col)
+		return tp.begin(wasm.SectionIDFunction, onTypeUse, tok, tokenBytes, line, col)
 	}
 
 	line, col, err := lex(skipTokens(1, parser), []byte(source))