Repository: tetratelabs/wazero
Pull request #707 (2022-07-19 03:21:12)

Supports mix of wasm and go funcs in the same module
This removes the constraint of a module being exclusively wasm or host
functions. Later pull requests can optimize special imports to be
implemented in wasm, particularly useful for disabled logging callbacks.


Patch:
diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go
index c06a89b0..e08a7301 100644
--- a/internal/engine/compiler/engine.go
+++ b/internal/engine/compiler/engine.go
@@ -417,43 +417,32 @@ func (e *engine) CompileModule(ctx context.Context, module *wasm.Module) error {
 
 	funcs := make([]*code, 0, len(module.FunctionSection))
 
-	if module.IsHostModule() {
-		for funcIndex := range module.HostFunctionSection {
-			compiled, err := compileHostFunction(module.TypeSection[module.FunctionSection[funcIndex]])
-			if err != nil {
+	irs, err := wazeroir.CompileFunctions(ctx, e.enabledFeatures, module)
+	if err != nil {
+		return err
+	}
+	for funcIndex, ir := range irs {
+		var compiled *code
+		if ir.GoFunc != nil {
+			sig := module.TypeSection[module.FunctionSection[funcIndex]]
+			if compiled, err = compileHostFunction(sig); err != nil {
 				def := module.FunctionDefinitionSection[uint32(funcIndex)+module.ImportFuncCount()]
 				return fmt.Errorf("error compiling host func[%s]: %w", def.DebugName(), err)
 			}
-
-			// As this uses mmap, we need a finalizer in case moduleEngine.Close was never called. Regardless, we need a
-			// finalizer due to how moduleEngine.Close is implemented.
-			e.setFinalizer(compiled, releaseCode)
-
-			compiled.indexInModule = wasm.Index(funcIndex)
-			compiled.sourceModule = module
-			funcs = append(funcs, compiled)
-		}
-	} else {
-		irs, err := wazeroir.CompileFunctions(ctx, e.enabledFeatures, module)
-		if err != nil {
-			return err
-		}
-
-		for funcIndex := range module.FunctionSection {
-			compiled, err := compileWasmFunction(e.enabledFeatures, irs[funcIndex])
-			if err != nil {
+		} else {
+			if compiled, err = compileWasmFunction(e.enabledFeatures, ir); err != nil {
 				def := module.FunctionDefinitionSection[uint32(funcIndex)+module.ImportFuncCount()]
 				return fmt.Errorf("error compiling wasm func[%s]: %w", def.DebugName(), err)
 			}
+		}
 
-			// As this uses mmap, we need to munmap on the compiled machine code when it's GCed.
-			e.setFinalizer(compiled, releaseCode)
+		// As this uses mmap, we need to munmap on the compiled machine code when it's GCed.
+		e.setFinalizer(compiled, releaseCode)
 
-			compiled.indexInModule = wasm.Index(funcIndex)
-			compiled.sourceModule = module
+		compiled.indexInModule = wasm.Index(funcIndex)
+		compiled.sourceModule = module
 
-			funcs = append(funcs, compiled)
-		}
+		funcs = append(funcs, compiled)
 	}
 	e.addCodes(module, funcs)
 	return nil
diff --git a/internal/engine/interpreter/interpreter.go b/internal/engine/interpreter/interpreter.go
index 283163eb..a06cb13e 100644
--- a/internal/engine/interpreter/interpreter.go
+++ b/internal/engine/interpreter/interpreter.go
@@ -220,24 +220,21 @@ func (e *engine) CompileModule(ctx context.Context, module *wasm.Module) error {
 	}
 
 	funcs := make([]*code, 0, len(module.FunctionSection))
-	if module.IsHostModule() {
-		// If this is the host module, there's nothing to do as the runtime representation of
+	irs, err := wazeroir.CompileFunctions(ctx, e.enabledFeatures, module)
+	if err != nil {
+		return err
+	}
+	for i, ir := range irs {
+		// If this is the host function, there's nothing to do as the runtime representation of
 		// host function in interpreter is its Go function itself as opposed to Wasm functions,
 		// which need to be compiled down to wazeroir.
-		for _, hf := range module.HostFunctionSection {
-			funcs = append(funcs, &code{hostFn: hf})
-		}
-	} else {
-		irs, err := wazeroir.CompileFunctions(ctx, e.enabledFeatures, module)
-		if err != nil {
-			return err
-		}
-		for i, ir := range irs {
-			compiled, err := e.lowerIR(ir)
-			if err != nil {
-				def := module.FunctionDefinitionSection[uint32(i)+module.ImportFuncCount()]
-				return fmt.Errorf("failed to lower func[%s] to wazeroir: %w", def.DebugName(), err)
-			}
+		if ir.GoFunc != nil {
+			funcs = append(funcs, &code{hostFn: ir.GoFunc})
+			continue
+		} else if compiled, err := e.lowerIR(ir); err != nil {
+			def := module.FunctionDefinitionSection[uint32(i)+module.ImportFuncCount()]
+			return fmt.Errorf("failed to lower func[%s] to wazeroir: %w", def.DebugName(), err)
+		} else {
 			funcs = append(funcs, compiled)
 		}
 	}
diff --git a/internal/wasm/binary/code.go b/internal/wasm/binary/code.go
index 8264c222..a2c3b2db 100644
--- a/internal/wasm/binary/code.go
+++ b/internal/wasm/binary/code.go
@@ -86,6 +86,10 @@ func decodeCode(r *bytes.Reader) (*wasm.Code, error) {
 //
 // See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#binary-code
 func encodeCode(c *wasm.Code) []byte {
+	if c.GoFunc != nil {
+		panic("BUG: GoFunc is not encodable")
+	}
+
 	// local blocks compress locals while preserving index order by grouping locals of the same type.
 	// https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#code-section%E2%91%A0
 	localBlockCount := uint32(0) // how many blocks of locals with the same type (types can repeat!)
diff --git a/internal/wasm/binary/encoder.go b/internal/wasm/binary/encoder.go
index 74890ead..65f6029a 100644
--- a/internal/wasm/binary/encoder.go
+++ b/internal/wasm/binary/encoder.go
@@ -10,10 +10,6 @@ var sizePrefixedName = []byte{4, 'n', 'a', 'm', 'e'}
 // Note: If saving to a file, the conventional extension is wasm
 // See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#binary-format%E2%91%A0
 func EncodeModule(m *wasm.Module) (bytes []byte) {
-	if m.SectionElementCount(wasm.SectionIDHostFunction) > 0 {
-		// TODO: See if there's a way to serialize reflect.Value references, potentially by name lookup in a store.
-		panic("BUG: HostFunctionSection is not encodable")
-	}
 	bytes = append(Magic, version...)
 	if m.SectionElementCount(wasm.SectionIDType) > 0 {
 		bytes = append(bytes, encodeTypeSection(m.TypeSection)...)
diff --git a/internal/wasm/counts.go b/internal/wasm/counts.go
index 82288712..5f10c8ea 100644
--- a/internal/wasm/counts.go
+++ b/internal/wasm/counts.go
@@ -79,8 +79,6 @@ func (m *Module) SectionElementCount(sectionID SectionID) uint32 { // element as
 		return uint32(len(m.CodeSection))
 	case SectionIDData:
 		return uint32(len(m.DataSection))
-	case SectionIDHostFunction:
-		return uint32(len(m.HostFunctionSection))
 	default:
 		panic(fmt.Errorf("BUG: unknown section: %d", sectionID))
 	}
diff --git a/internal/wasm/function_definition.go b/internal/wasm/function_definition.go
index 45d6c6dc..81a108b0 100644
--- a/internal/wasm/function_definition.go
+++ b/internal/wasm/function_definition.go
@@ -64,13 +64,11 @@ func (m *Module) BuildFunctionDefinitions() {
 		importFuncIdx++
 	}
 
-	// At the moment, a module can either be solely wasm or host functions.
-	isHostFunction := m.HostFunctionSection != nil
 	for codeIndex, typeIndex := range m.FunctionSection {
 		m.FunctionDefinitionSection = append(m.FunctionDefinitionSection, &FunctionDefinition{
 			index:          Index(codeIndex) + importCount,
 			funcType:       m.TypeSection[typeIndex],
-			isHostFunction: isHostFunction,
+			isHostFunction: m.CodeSection[codeIndex].GoFunc != nil,
 		})
 	}
 
diff --git a/internal/wasm/host.go b/internal/wasm/host.go
index a9906839..5d0094ca 100644
--- a/internal/wasm/host.go
+++ b/internal/wasm/host.go
@@ -74,10 +74,6 @@ func NewHostModule(
 	return
 }
 
-func (m *Module) IsHostModule() bool {
-	return len(m.HostFunctionSection) > 0
-}
-
 func addFuncs(
 	m *Module,
 	nameToGoFunc map[string]interface{},
@@ -92,7 +88,7 @@ func addFuncs(
 	moduleName := m.NameSection.ModuleName
 	m.NameSection.FunctionNames = make([]*NameAssoc, 0, funcCount)
 	m.FunctionSection = make([]Index, 0, funcCount)
-	m.HostFunctionSection = make([]*reflect.Value, 0, funcCount)
+	m.CodeSection = make([]*Code, 0, funcCount)
 	m.FunctionDefinitionSection = make([]*FunctionDefinition, 0, funcCount)
 
 	// Sort names for consistent iteration
@@ -116,7 +112,7 @@ func addFuncs(
 		}
 
 		m.FunctionSection = append(m.FunctionSection, m.maybeAddType(functionType))
-		m.HostFunctionSection = append(m.HostFunctionSection, &fn)
+		m.CodeSection = append(m.CodeSection, &Code{GoFunc: &fn})
 		m.ExportSection = append(m.ExportSection, &Export{Type: ExternTypeFunc, Name: exportName, Index: idx})
 		if namesLen > 0 {
 			m.NameSection.FunctionNames = append(m.NameSection.FunctionNames, &NameAssoc{Index: idx, Name: names[0]})
diff --git a/internal/wasm/module.go b/internal/wasm/module.go
index 26d83859..2aeb0e8a 100644
--- a/internal/wasm/module.go
+++ b/internal/wasm/module.go
@@ -38,7 +38,7 @@ type EncodeModule func(m *Module) (bytes []byte)
 // Differences from the specification:
 // * NameSection is the only key ("name") decoded from the SectionIDCustom.
 // * ExportSection is represented as a map for lookup convenience.
-// * HostFunctionSection is a custom section that contains any go `func`s. It may be present when CodeSection is not.
+// * Code.GoFunc is contains any go `func`. It may be present when Code.Body is not.
 type Module struct {
 	// TypeSection contains the unique FunctionType of functions imported or defined in this module.
 	//
@@ -133,8 +133,10 @@ type Module struct {
 	// Note: In the Binary Format, this is SectionIDElement.
 	ElementSection []*ElementSegment
 
-	// CodeSection is index-correlated with FunctionSection and contains each function's locals and body.
-	// When present, the HostFunctionSection must be nil.
+	// CodeSection is index-correlated with FunctionSection and contains each
+	// function's locals and body.
+	//
+	// When present, the HostFunctionSection of the same index must be nil.
 	//
 	// Note: In the Binary Format, this is SectionIDCode.
 	//
@@ -153,13 +155,6 @@ type Module struct {
 	// See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#custom-section%E2%91%A0
 	NameSection *NameSection
 
-	// HostFunctionSection is index-correlated with FunctionSection and contains a host function defined in Go.
-	// When present, the CodeSection must be nil.
-	//
-	// Note: This section currently has no serialization format, so is not encodable.
-	// See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#host-functions%E2%91%A2
-	HostFunctionSection []*reflect.Value
-
 	// validatedActiveElementSegments are built on Validate when
 	// SectionIDElement is non-empty and all inputs are valid.
 	//
@@ -236,10 +231,6 @@ func (m *Module) Validate(enabledFeatures Features) error {
 		return err
 	}
 
-	if m.SectionElementCount(SectionIDCode) > 0 && m.SectionElementCount(SectionIDHostFunction) > 0 {
-		return errors.New("cannot mix functions and host functions in the same module")
-	}
-
 	functions, globals, memory, tables, err := m.AllDeclarations()
 	if err != nil {
 		return err
@@ -588,24 +579,19 @@ func (m *Module) buildGlobals(importedGlobals []*GlobalInstance) (globals []*Glo
 //	  enginetest.go.
 func (m *ModuleInstance) BuildFunctions(mod *Module, listeners []experimental.FunctionListener) (fns []*FunctionInstance) {
 	fns = make([]*FunctionInstance, 0, len(mod.FunctionDefinitionSection))
-	if mod.IsHostModule() {
-		for i := range mod.HostFunctionSection {
-			fn := mod.HostFunctionSection[i]
-			fns = append(fns, &FunctionInstance{
-				Kind:   kind(fn.Type()),
-				GoFunc: fn,
-			})
-		}
-	} else {
-		for i := range mod.FunctionSection {
-			code := mod.CodeSection[i]
-			fns = append(fns, &FunctionInstance{
-				Kind:       FunctionKindWasm,
-				Body:       code.Body,
-				TypeID:     m.TypeIDs[mod.FunctionSection[i]],
-				LocalTypes: code.LocalTypes,
-			})
+	for i := range mod.FunctionSection {
+		code := mod.CodeSection[i]
+		fnKind := FunctionKindWasm
+		if fn := code.GoFunc; fn != nil {
+			fnKind = kind(fn.Type())
 		}
+		fns = append(fns, &FunctionInstance{
+			Kind:       fnKind,
+			Body:       code.Body,
+			GoFunc:     code.GoFunc,
+			TypeID:     m.TypeIDs[mod.FunctionSection[i]],
+			LocalTypes: code.LocalTypes,
+		})
 	}
 
 	importCount := mod.ImportFuncCount()
@@ -808,6 +794,15 @@ type Export struct {
 // Code is an entry in the Module.CodeSection containing the locals and body of the function.
 // See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#binary-code
 type Code struct {
+
+	// GoFunc is a host function defined in Go.
+	//
+	// When present, LocalTypes and Body must be nil.
+	//
+	// Note: This has no serialization format, so is not encodable.
+	// See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#host-functions%E2%91%A2
+	GoFunc *reflect.Value
+
 	// LocalTypes are any function-scoped variables in insertion order.
 	// See https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#binary-local
 	LocalTypes []ValueType
@@ -954,11 +949,6 @@ const (
 	SectionIDDataCount
 )
 
-// SectionIDHostFunction is a pseudo-section ID for host functions.
-//
-// Note: This is not defined in the WebAssembly 1.0 (20191205) Binary Format.
-const SectionIDHostFunction = SectionID(0xff)
-
 // SectionIDName returns the canonical name of a module section.
 // https://www.w3.org/TR/2019/REC-wasm-core-1-20191205/#sections%E2%91%A0
 func SectionIDName(sectionID SectionID) string {
@@ -987,8 +977,6 @@ func SectionIDName(sectionID SectionID) string {
 		return "code"
 	case SectionIDData:
 		return "data"
-	case SectionIDHostFunction:
-		return "host_function"
 	case SectionIDDataCount:
 		return "data_count"
 	}
diff --git a/internal/wasm/store.go b/internal/wasm/store.go
index d0611265..391375ab 100644
--- a/internal/wasm/store.go
+++ b/internal/wasm/store.go
@@ -367,14 +367,6 @@ func (s *Store) instantiate(
 	}
 	globals, memory := module.buildGlobals(importedGlobals), module.buildMemory()
 
-	// If there are no module-defined functions, assume this is a host module.
-	var funcSection SectionID
-	if module.HostFunctionSection == nil {
-		funcSection = SectionIDFunction
-	} else {
-		funcSection = SectionIDHostFunction
-	}
-
 	m := &ModuleInstance{Name: name, TypeIDs: typeIDs}
 	functions := m.BuildFunctions(module, listeners)
 
@@ -420,7 +412,7 @@ func (s *Store) instantiate(
 		if exitErr, ok := err.(*sys.ExitError); ok { // Don't wrap an exit error!
 			return nil, exitErr
 		} else if err != nil {
-			return nil, fmt.Errorf("start %s failed: %w", module.funcDesc(funcSection, funcIdx), err)
+			return nil, fmt.Errorf("start %s failed: %w", module.funcDesc(SectionIDFunction, funcIdx), err)
 		}
 	}
 
diff --git a/internal/wazeroir/compiler.go b/internal/wazeroir/compiler.go
index 45513cef..3129a016 100644
--- a/internal/wazeroir/compiler.go
+++ b/internal/wazeroir/compiler.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"math"
 	"os"
+	"reflect"
 	"strings"
 
 	"github.com/tetratelabs/wazero/internal/buildoptions"
@@ -176,6 +177,10 @@ func (c *compiler) resetUnreachable() {
 }
 
 type CompilationResult struct {
+	// GoFunc is present when the result is a host function.
+	// In this case, other fields can be ignored.
+	GoFunc *reflect.Value
+
 	// Operations holds wazeroir operations compiled from Wasm instructions in a Wasm function.
 	Operations []Operation
 	// LabelCallers maps Label.String() to the number of callers to that label.
@@ -233,6 +238,10 @@ func CompileFunctions(_ context.Context, enabledFeatures wasm.Features, module *
 		typeID := module.FunctionSection[funcIndex]
 		sig := module.TypeSection[typeID]
 		code := module.CodeSection[funcIndex]
+		if code.GoFunc != nil {
+			ret = append(ret, &CompilationResult{GoFunc: code.GoFunc})
+			continue
+		}
 		r, err := compile(enabledFeatures, sig, code.Body, code.LocalTypes, module.TypeSection, functions, globals)
 		if err != nil {
 			def := module.FunctionDefinitionSection[uint32(funcIndex)+module.ImportFuncCount()]


Test patch:
diff --git a/builder_test.go b/builder_test.go
index 3dbd8646..f57f70d3 100644
--- a/builder_test.go
+++ b/builder_test.go
@@ -53,8 +53,8 @@ func TestNewModuleBuilder_Compile(t *testing.T) {
 				TypeSection: []*wasm.FunctionType{
 					{Params: []api.ValueType{i32}, Results: []api.ValueType{i32}, ParamNumInUint64: 1, ResultNumInUint64: 1},
 				},
-				FunctionSection:     []wasm.Index{0},
-				HostFunctionSection: []*reflect.Value{&fnUint32_uint32},
+				FunctionSection: []wasm.Index{0},
+				CodeSection:     []*wasm.Code{{GoFunc: &fnUint32_uint32}},
 				ExportSection: []*wasm.Export{
 					{Name: "1", Type: wasm.ExternTypeFunc, Index: 0},
 				},
@@ -72,8 +72,8 @@ func TestNewModuleBuilder_Compile(t *testing.T) {
 				TypeSection: []*wasm.FunctionType{
 					{Params: []api.ValueType{i32}, Results: []api.ValueType{i32}, ParamNumInUint64: 1, ResultNumInUint64: 1},
 				},
-				FunctionSection:     []wasm.Index{0},
-				HostFunctionSection: []*reflect.Value{&fnUint32_uint32},
+				FunctionSection: []wasm.Index{0},
+				CodeSection:     []*wasm.Code{{GoFunc: &fnUint32_uint32}},
 				ExportSection: []*wasm.Export{
 					{Name: "1", Type: wasm.ExternTypeFunc, Index: 0},
 				},
@@ -92,8 +92,8 @@ func TestNewModuleBuilder_Compile(t *testing.T) {
 				TypeSection: []*wasm.FunctionType{
 					{Params: []api.ValueType{i64}, Results: []api.ValueType{i32}, ParamNumInUint64: 1, ResultNumInUint64: 1},
 				},
-				FunctionSection:     []wasm.Index{0},
-				HostFunctionSection: []*reflect.Value{&fnUint64_uint32},
+				FunctionSection: []wasm.Index{0},
+				CodeSection:     []*wasm.Code{{GoFunc: &fnUint64_uint32}},
 				ExportSection: []*wasm.Export{
 					{Name: "1", Type: wasm.ExternTypeFunc, Index: 0},
 				},
@@ -113,8 +113,8 @@ func TestNewModuleBuilder_Compile(t *testing.T) {
 					{Params: []api.ValueType{i32}, Results: []api.ValueType{i32}, ParamNumInUint64: 1, ResultNumInUint64: 1},
 					{Params: []api.ValueType{i64}, Results: []api.ValueType{i32}, ParamNumInUint64: 1, ResultNumInUint64: 1},
 				},
-				FunctionSection:     []wasm.Index{0, 1},
-				HostFunctionSection: []*reflect.Value{&fnUint32_uint32, &fnUint64_uint32},
+				FunctionSection: []wasm.Index{0, 1},
+				CodeSection:     []*wasm.Code{{GoFunc: &fnUint32_uint32}, {GoFunc: &fnUint64_uint32}},
 				ExportSection: []*wasm.Export{
 					{Name: "1", Type: wasm.ExternTypeFunc, Index: 0},
 					{Name: "2", Type: wasm.ExternTypeFunc, Index: 1},
@@ -137,8 +137,8 @@ func TestNewModuleBuilder_Compile(t *testing.T) {
 					{Params: []api.ValueType{i32}, Results: []api.ValueType{i32}, ParamNumInUint64: 1, ResultNumInUint64: 1},
 					{Params: []api.ValueType{i64}, Results: []api.ValueType{i32}, ParamNumInUint64: 1, ResultNumInUint64: 1},
 				},
-				FunctionSection:     []wasm.Index{0, 1},
-				HostFunctionSection: []*reflect.Value{&fnUint32_uint32, &fnUint64_uint32},
+				FunctionSection: []wasm.Index{0, 1},
+				CodeSection:     []*wasm.Code{{GoFunc: &fnUint32_uint32}, {GoFunc: &fnUint64_uint32}},
 				ExportSection: []*wasm.Export{
 					{Name: "1", Type: wasm.ExternTypeFunc, Index: 0},
 					{Name: "2", Type: wasm.ExternTypeFunc, Index: 1},
@@ -162,8 +162,8 @@ func TestNewModuleBuilder_Compile(t *testing.T) {
 					{Params: []api.ValueType{i32}, Results: []api.ValueType{i32}, ParamNumInUint64: 1, ResultNumInUint64: 1},
 					{Params: []api.ValueType{i64}, Results: []api.ValueType{i32}, ParamNumInUint64: 1, ResultNumInUint64: 1},
 				},
-				FunctionSection:     []wasm.Index{0, 1},
-				HostFunctionSection: []*reflect.Value{&fnUint32_uint32, &fnUint64_uint32},
+				FunctionSection: []wasm.Index{0, 1},
+				CodeSection:     []*wasm.Code{{GoFunc: &fnUint32_uint32}, {GoFunc: &fnUint64_uint32}},
 				ExportSection: []*wasm.Export{
 					{Name: "1", Type: wasm.ExternTypeFunc, Index: 0},
 					{Name: "2", Type: wasm.ExternTypeFunc, Index: 1},
@@ -457,13 +457,12 @@ func requireHostModuleEquals(t *testing.T, expected, actual *wasm.Module) {
 	require.Equal(t, expected.ExportSection, actual.ExportSection)
 	require.Equal(t, expected.StartSection, actual.StartSection)
 	require.Equal(t, expected.ElementSection, actual.ElementSection)
-	require.Zero(t, len(actual.CodeSection)) // Host functions are implemented in Go, not Wasm!
 	require.Equal(t, expected.DataSection, actual.DataSection)
 	require.Equal(t, expected.NameSection, actual.NameSection)
 
 	// Special case because reflect.Value can't be compared with Equals
-	require.Equal(t, len(expected.HostFunctionSection), len(actual.HostFunctionSection))
-	for i := range expected.HostFunctionSection {
-		require.Equal(t, (*expected.HostFunctionSection[i]).Type(), (*actual.HostFunctionSection[i]).Type())
+	require.Equal(t, len(expected.CodeSection), len(actual.CodeSection))
+	for _, c := range expected.CodeSection {
+		require.Equal(t, c.GoFunc.Type(), c.GoFunc.Type())
 	}
 }
diff --git a/experimental/log_listener_test.go b/experimental/log_listener_test.go
index d1ef1244..abf287ba 100644
--- a/experimental/log_listener_test.go
+++ b/experimental/log_listener_test.go
@@ -284,7 +284,9 @@ func Test_loggingListener(t *testing.T) {
 			}
 
 			if tc.isHostFunc {
-				m.HostFunctionSection = []*reflect.Value{&fnV}
+				m.CodeSection = []*wasm.Code{{GoFunc: &fnV}}
+			} else {
+				m.CodeSection = []*wasm.Code{{Body: []byte{wasm.OpcodeEnd}}}
 			}
 			m.BuildFunctionDefinitions()
 			l := lf.NewListener(m.FunctionDefinitionSection[0])
@@ -314,6 +316,7 @@ func Test_loggingListener_indentation(t *testing.T) {
 	m := &wasm.Module{
 		TypeSection:     []*wasm.FunctionType{{}},
 		FunctionSection: []wasm.Index{0, 0},
+		CodeSection:     []*wasm.Code{{Body: []byte{wasm.OpcodeEnd}}, {Body: []byte{wasm.OpcodeEnd}}},
 		NameSection: &wasm.NameSection{
 			ModuleName:    "test",
 			FunctionNames: wasm.NameMap{{Index: 0, Name: "fn1"}, {Index: 1, Name: "fn2"}},
diff --git a/internal/integration_test/vs/runtime.go b/internal/integration_test/vs/runtime.go
index 9292a30d..4c2e0187 100644
--- a/internal/integration_test/vs/runtime.go
+++ b/internal/integration_test/vs/runtime.go
@@ -2,6 +2,7 @@ package vs
 
 import (
 	"context"
+	"errors"
 	"fmt"
 
 	"github.com/tetratelabs/wazero"
@@ -167,11 +168,11 @@ func (m *wazeroModule) CallI64_I64(ctx context.Context, funcName string, param u
 
 func (m *wazeroModule) WriteMemory(ctx context.Context, offset uint32, bytes []byte) error {
 	if !m.mod.Memory().Write(ctx, offset, bytes) {
-		return fmt.Errorf("Memory.Write(%d, %d) out of range of memory size %d",
-			offset, len(bytes), m.mod.Memory().Size(ctx))
+		return errors.New("out of memory writing name")
 	}
 	return nil
 }
+
 func (m *wazeroModule) Close(ctx context.Context) (err error) {
 	if mod := m.mod; mod != nil {
 		err = mod.Close(ctx)
diff --git a/internal/testing/enginetest/enginetest.go b/internal/testing/enginetest/enginetest.go
index cbef206f..8e1c7464 100644
--- a/internal/testing/enginetest/enginetest.go
+++ b/internal/testing/enginetest/enginetest.go
@@ -329,15 +329,15 @@ func runTestModuleEngine_Call_HostFn_ModuleContext(t *testing.T, et EngineTester
 	})
 
 	m := &wasm.Module{
-		HostFunctionSection: []*reflect.Value{&host},
-		FunctionSection:     []wasm.Index{0},
-		TypeSection:         []*wasm.FunctionType{sig},
+		TypeSection:     []*wasm.FunctionType{sig},
+		FunctionSection: []wasm.Index{0},
+		CodeSection:     []*wasm.Code{{GoFunc: &host}},
 	}
 	m.BuildFunctionDefinitions()
 	err := e.CompileModule(testCtx, m)
 	require.NoError(t, err)
 
-	module := &wasm.ModuleInstance{Memory: memory}
+	module := &wasm.ModuleInstance{Memory: memory, TypeIDs: []wasm.FunctionTypeID{0}}
 	_, ns := wasm.NewStore(features, e)
 	modCtx := wasm.NewCallContext(ns, module, nil)
 
@@ -672,10 +672,10 @@ func setupCallTests(t *testing.T, e wasm.Engine) (*wasm.ModuleInstance, *wasm.Mo
 
 	hostFnVal := reflect.ValueOf(divBy)
 	hostModule := &wasm.Module{
-		HostFunctionSection: []*reflect.Value{&hostFnVal},
-		TypeSection:         []*wasm.FunctionType{ft},
-		FunctionSection:     []wasm.Index{0},
-		ExportSection:       []*wasm.Export{{Name: hostFnName, Type: wasm.ExternTypeFunc, Index: 0}},
+		TypeSection:     []*wasm.FunctionType{ft},
+		FunctionSection: []wasm.Index{0},
+		CodeSection:     []*wasm.Code{{GoFunc: &hostFnVal}},
+		ExportSection:   []*wasm.Export{{Name: hostFnName, Type: wasm.ExternTypeFunc, Index: 0}},
 		NameSection: &wasm.NameSection{
 			ModuleName:    "host",
 			FunctionNames: wasm.NameMap{{Index: wasm.Index(0), Name: hostFnName}},
diff --git a/internal/wasm/binary/encoder_test.go b/internal/wasm/binary/encoder_test.go
index 64ba4136..12e7cec3 100644
--- a/internal/wasm/binary/encoder_test.go
+++ b/internal/wasm/binary/encoder_test.go
@@ -214,7 +214,7 @@ func TestModule_Encode_HostFunctionSection_Unsupported(t *testing.T) {
 	fn := reflect.ValueOf(func(wasm.Module) {})
 
 	captured := require.CapturePanic(func() {
-		EncodeModule(&wasm.Module{HostFunctionSection: []*reflect.Value{&fn}})
+		EncodeModule(&wasm.Module{CodeSection: []*wasm.Code{{GoFunc: &fn}}})
 	})
-	require.EqualError(t, captured, "BUG: HostFunctionSection is not encodable")
+	require.EqualError(t, captured, "BUG: GoFunc is not encodable")
 }
diff --git a/internal/wasm/counts_test.go b/internal/wasm/counts_test.go
index e26b85cb..26fd5de4 100644
--- a/internal/wasm/counts_test.go
+++ b/internal/wasm/counts_test.go
@@ -1,10 +1,8 @@
 package wasm
 
 import (
-	"reflect"
 	"testing"
 
-	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/internal/testing/require"
 )
 
@@ -202,7 +200,6 @@ func TestModule_SectionElementCount(t *testing.T) {
 	i32, f32 := ValueTypeI32, ValueTypeF32
 	zero := uint32(0)
 	empty := &ConstantExpression{Opcode: OpcodeI32Const, Data: const0}
-	fn := reflect.ValueOf(func(api.Module) {})
 
 	tests := []struct {
 		name     string
@@ -219,16 +216,6 @@ func TestModule_SectionElementCount(t *testing.T) {
 			input:    &Module{NameSection: &NameSection{ModuleName: "simple"}},
 			expected: map[string]uint32{"custom": 1},
 		},
-		{
-			name:     "HostFunctionSection",
-			input:    &Module{HostFunctionSection: []*reflect.Value{&fn}},
-			expected: map[string]uint32{"host_function": 1},
-		},
-		{
-			name:     "NameSection and HostFunctionSection",
-			input:    &Module{NameSection: &NameSection{ModuleName: "simple"}, HostFunctionSection: []*reflect.Value{&fn}},
-			expected: map[string]uint32{"custom": 1, "host_function": 1},
-		},
 		{
 			name: "TypeSection",
 			input: &Module{
@@ -312,11 +299,6 @@ func TestModule_SectionElementCount(t *testing.T) {
 					actual[SectionIDName(i)] = size
 				}
 			}
-
-			// SectionIDHostFunction is intentionally not after data
-			if size := tc.input.SectionElementCount(SectionIDHostFunction); size > 0 {
-				actual[SectionIDName(SectionIDHostFunction)] = size
-			}
 			require.Equal(t, tc.expected, actual)
 		})
 	}
diff --git a/internal/wasm/function_definition_test.go b/internal/wasm/function_definition_test.go
index 9400b043..1c4bfc3e 100644
--- a/internal/wasm/function_definition_test.go
+++ b/internal/wasm/function_definition_test.go
@@ -9,7 +9,7 @@ import (
 )
 
 func TestModule_BuildFunctionDefinitions(t *testing.T) {
-	nopCode := &Code{nil, []byte{OpcodeEnd}}
+	nopCode := &Code{Body: []byte{OpcodeEnd}}
 	fnV := reflect.ValueOf(func() {})
 	tests := []struct {
 		name            string
@@ -34,9 +34,9 @@ func TestModule_BuildFunctionDefinitions(t *testing.T) {
 		{
 			name: "host func",
 			m: &Module{
-				TypeSection:         []*FunctionType{v_v},
-				FunctionSection:     []Index{0},
-				HostFunctionSection: []*reflect.Value{&fnV},
+				TypeSection:     []*FunctionType{v_v},
+				FunctionSection: []Index{0},
+				CodeSection:     []*Code{{GoFunc: &fnV}},
 			},
 			expected: []*FunctionDefinition{
 				{
@@ -59,6 +59,11 @@ func TestModule_BuildFunctionDefinitions(t *testing.T) {
 				},
 				GlobalSection:   []*Global{{}},
 				FunctionSection: []Index{1, 2, 0},
+				CodeSection: []*Code{
+					{Body: []byte{OpcodeEnd}},
+					{Body: []byte{OpcodeEnd}},
+					{Body: []byte{OpcodeEnd}},
+				},
 				TypeSection: []*FunctionType{
 					v_v,
 					{Params: []ValueType{ValueTypeF64, ValueTypeI32}, Results: []ValueType{ValueTypeV128, ValueTypeI64}},
@@ -119,6 +124,7 @@ func TestModule_BuildFunctionDefinitions(t *testing.T) {
 					{Name: "function_index=2", Type: ExternTypeFunc, Index: 2},
 				},
 				FunctionSection: []Index{1, 0},
+				CodeSection:     []*Code{{Body: []byte{OpcodeEnd}}, {Body: []byte{OpcodeEnd}}},
 				TypeSection: []*FunctionType{
 					v_v,
 					{Params: []ValueType{ValueTypeF64, ValueTypeI32}, Results: []ValueType{ValueTypeV128, ValueTypeI64}},
@@ -191,7 +197,7 @@ func TestModule_BuildFunctionDefinitions(t *testing.T) {
 					},
 				},
 				FunctionSection: []Index{0, 0, 0, 0, 0},
-				CodeSection:     []*Code{nopCode, nopCode, nopCode, nopCode},
+				CodeSection:     []*Code{nopCode, nopCode, nopCode, nopCode, nopCode},
 			},
 			expected: []*FunctionDefinition{
 				{moduleName: "module", index: 0, debugName: "module.$0", importDesc: &[2]string{"i", "f"}, funcType: v_v},
diff --git a/internal/wasm/host_test.go b/internal/wasm/host_test.go
index 7f3c5e34..941f0bab 100644
--- a/internal/wasm/host_test.go
+++ b/internal/wasm/host_test.go
@@ -68,8 +68,8 @@ func TestNewHostModule(t *testing.T) {
 					{Params: []ValueType{i32, i32}, Results: []ValueType{i32}, ParamNumInUint64: 2, ResultNumInUint64: 1},
 					{Params: []ValueType{i32, i32, i32, i32}, Results: []ValueType{i32}, ParamNumInUint64: 4, ResultNumInUint64: 1},
 				},
-				FunctionSection:     []Index{0, 1},
-				HostFunctionSection: []*reflect.Value{&fnArgsSizesGet, &fnFdWrite},
+				FunctionSection: []Index{0, 1},
+				CodeSection:     []*Code{{GoFunc: &fnArgsSizesGet}, {GoFunc: &fnFdWrite}},
 				ExportSection: []*Export{
 					{Name: "args_sizes_get", Type: ExternTypeFunc, Index: 0},
 					{Name: "fd_write", Type: ExternTypeFunc, Index: 1},
@@ -90,11 +90,11 @@ func TestNewHostModule(t *testing.T) {
 				functionSwap: swap,
 			},
 			expected: &Module{
-				TypeSection:         []*FunctionType{{Params: []ValueType{i32, i32}, Results: []ValueType{i32, i32}, ParamNumInUint64: 2, ResultNumInUint64: 2}},
-				FunctionSection:     []Index{0},
-				HostFunctionSection: []*reflect.Value{&fnSwap},
-				ExportSection:       []*Export{{Name: "swap", Type: ExternTypeFunc, Index: 0}},
-				NameSection:         &NameSection{ModuleName: "swapper", FunctionNames: NameMap{{Index: 0, Name: "swap"}}},
+				TypeSection:     []*FunctionType{{Params: []ValueType{i32, i32}, Results: []ValueType{i32, i32}, ParamNumInUint64: 2, ResultNumInUint64: 2}},
+				FunctionSection: []Index{0},
+				CodeSection:     []*Code{{GoFunc: &fnSwap}},
+				ExportSection:   []*Export{{Name: "swap", Type: ExternTypeFunc, Index: 0}},
+				NameSection:     &NameSection{ModuleName: "swapper", FunctionNames: NameMap{{Index: 0, Name: "swap"}}},
 			},
 		},
 		{
@@ -153,8 +153,8 @@ func TestNewHostModule(t *testing.T) {
 				TypeSection: []*FunctionType{
 					{Params: []ValueType{i32, i32}, Results: []ValueType{i32}, ParamNumInUint64: 2, ResultNumInUint64: 1},
 				},
-				FunctionSection:     []Index{0},
-				HostFunctionSection: []*reflect.Value{&fnArgsSizesGet},
+				FunctionSection: []Index{0},
+				CodeSection:     []*Code{{GoFunc: &fnArgsSizesGet}},
 				GlobalSection: []*Global{
 					{
 						Type: &GlobalType{ValType: i32},
@@ -199,14 +199,13 @@ func requireHostModuleEquals(t *testing.T, expected, actual *Module) {
 	require.Equal(t, expected.ExportSection, actual.ExportSection)
 	require.Equal(t, expected.StartSection, actual.StartSection)
 	require.Equal(t, expected.ElementSection, actual.ElementSection)
-	require.Nil(t, actual.CodeSection) // Host functions are implemented in Go, not Wasm!
 	require.Equal(t, expected.DataSection, actual.DataSection)
 	require.Equal(t, expected.NameSection, actual.NameSection)
 
 	// Special case because reflect.Value can't be compared with Equals
-	require.Equal(t, len(expected.HostFunctionSection), len(actual.HostFunctionSection))
-	for i := range expected.HostFunctionSection {
-		require.Equal(t, (*expected.HostFunctionSection[i]).Type(), (*actual.HostFunctionSection[i]).Type())
+	require.Equal(t, len(expected.CodeSection), len(actual.CodeSection))
+	for _, c := range expected.CodeSection {
+		require.Equal(t, c.GoFunc.Type(), c.GoFunc.Type())
 	}
 }
 
diff --git a/internal/wasm/module_test.go b/internal/wasm/module_test.go
index 43e45e8e..8984388e 100644
--- a/internal/wasm/module_test.go
+++ b/internal/wasm/module_test.go
@@ -3,7 +3,6 @@ package wasm
 import (
 	"fmt"
 	"math"
-	"reflect"
 	"testing"
 
 	"github.com/tetratelabs/wazero/api"
@@ -57,7 +56,6 @@ func TestSectionIDName(t *testing.T) {
 		{"element", SectionIDElement, "element"},
 		{"code", SectionIDCode, "code"},
 		{"data", SectionIDData, "data"},
-		{"host_function", SectionIDHostFunction, "host_function"},
 		{"unknown", 100, "unknown"},
 	}
 
@@ -327,8 +325,6 @@ func TestValidateConstExpression(t *testing.T) {
 
 func TestModule_Validate_Errors(t *testing.T) {
 	zero := Index(0)
-	fn := reflect.ValueOf(func(api.Module) {})
-
 	tests := []struct {
 		name        string
 		input       *Module
@@ -344,16 +340,6 @@ func TestModule_Validate_Errors(t *testing.T) {
 			},
 			expectedErr: "invalid start function: func[0] has an invalid type",
 		},
-		{
-			name: "CodeSection and HostFunctionSection",
-			input: &Module{
-				TypeSection:         []*FunctionType{v_v},
-				FunctionSection:     []uint32{0},
-				CodeSection:         []*Code{{Body: []byte{OpcodeEnd}}},
-				HostFunctionSection: []*reflect.Value{&fn},
-			},
-			expectedErr: "cannot mix functions and host functions in the same module",
-		},
 	}
 
 	for _, tt := range tests {
@@ -781,7 +767,7 @@ func TestModule_buildGlobals(t *testing.T) {
 }
 
 func TestModule_buildFunctions(t *testing.T) {
-	nopCode := &Code{nil, []byte{OpcodeEnd}}
+	nopCode := &Code{Body: []byte{OpcodeEnd}}
 	m := &Module{
 		TypeSection:     []*FunctionType{v_v},
 		ImportSection:   []*Import{{Type: ExternTypeFunc}},