make module optionally accept host interfaces as generic parameters

main.go

@@ -7,10 +7,11 @@ import (
 )
 
 var (
-	endian = flag.String("endian", "", "endianness of the generated code (little or big)")
-	nanbox = flag.Bool("nanbox", false, "whether to try to canonicalize NaNs")
-	nohost = flag.Bool("nohost", false, "disable generating interfaces for imports")
-	nohead = flag.Bool("nohead", false, "disable the header comment (including build tags)")
+	endian  = flag.String("endian", "", "endianness of the generated code (little or big)")
+	nanbox  = flag.Bool("nanbox", false, "whether to try to canonicalize NaNs")
+	nohost  = flag.Bool("nohost", false, "disable generating interfaces for imports")
+	nohead  = flag.Bool("nohead", false, "disable the header comment (including build tags)")
+	generic = flag.Bool("generic", false, "add generic type parameters for host module imports")
 )
 
 // https://pkg.go.dev/golang.org/x/sys/cpu#pkg-constants

module.go

@@ -5,13 +5,10 @@ import (
 	"go/token"
 	"slices"
 	"strconv"
+	"strings"
 )
 
-var modRecvList = &ast.FieldList{List: []*ast.Field{{
-	Names: []*ast.Ident{newID("m")},
-	Type:  &ast.StarExpr{X: newID("Module")}}}}
-
-func (t *translator) createModuleStruct() ast.Decl {
+func (t *translator) createModuleStruct(hostInterfaces []*ast.GenDecl) (*ast.GenDecl, *ast.TypeSpec) {
 	var fields []*ast.Field
 	// Tables: owned are []any; imported *[]any.
 	for _, tab := range t.tables {
@@ -66,13 +63,54 @@ func (t *translator) createModuleStruct() ast.Decl {
 				Type:  ast.NewIdent(exported(imp.module))})
 		}
 	}
-	return &ast.GenDecl{
-		Tok: token.TYPE,
-		Specs: []ast.Spec{
-			&ast.TypeSpec{
-				Name: newID("Module"),
-				Type: &ast.StructType{Fields: &ast.FieldList{List: fields}}}},
+
+	// If host module imports are provided,
+	// use these to generate generic type
+	// parameters for the module struct.
+	var typeParams *ast.FieldList
+	if len(hostInterfaces) > 0 {
+		seen := set[string]{}
+		typeParams = new(ast.FieldList)
+		for _, decl := range hostInterfaces {
+			typeSpec := decl.Specs[0].(*ast.TypeSpec)
+			var paramName string
+
+			// Use shortest possible prefix
+			// of the type parameter for
+			// the generic parameter name.
+			typeName := typeSpec.Name.Name
+			typeName = strings.TrimPrefix(typeName, "X")
+			for i := 1; i < len(typeName); i++ {
+				if seen.has(typeName[:i]) {
+					continue
+				}
+				paramName = typeName[:i]
+				seen.add(paramName)
+				break
+			}
+
+			if paramName == "" {
+				panic("could not determine type param prefix for: " + typeName)
+			}
+
+			// Append type parameter with shorted name to list.
+			typeParams.List = append(typeParams.List, &ast.Field{
+				Names: []*ast.Ident{ast.NewIdent(strings.ToUpper(paramName))},
+				Type:  newID(typeSpec.Name.Name),
+			})
+		}
+	}
+
+	typeSpec := &ast.TypeSpec{
+		Name:       newID("Module"),
+		Type:       &ast.StructType{Fields: &ast.FieldList{List: fields}},
+		TypeParams: typeParams,
 	}
+
+	return &ast.GenDecl{
+		Tok:   token.TYPE,
+		Specs: []ast.Spec{typeSpec},
+	}, typeSpec
 }
 
 func (t *translator) createNewFunc() ast.Decl {
@@ -84,7 +122,7 @@ func (t *translator) createNewFunc() ast.Decl {
 				Lhs: []ast.Expr{newID("m")},
 				Rhs: []ast.Expr{&ast.UnaryExpr{
 					Op: token.AND,
-					X:  &ast.CompositeLit{Type: newID("Module")}}}}},
+					X:  &ast.CompositeLit{Type: t.modRecv.Type.(*ast.StarExpr).X}}}}},
 	}
 
 	// Create the params, and initialize fields.
@@ -284,14 +322,16 @@ func (t *translator) createNewFunc() ast.Decl {
 	body.List = append(body.List, &ast.ReturnStmt{Results: []ast.Expr{newID("m")}})
 
 	return &ast.FuncDecl{
-		Name: newID("New"),
-		Type: &ast.FuncType{
-			Params:  &ast.FieldList{List: params},
-			Results: &ast.FieldList{List: []*ast.Field{{Type: &ast.StarExpr{X: newID("Module")}}}}},
-		Body: body}
+		Name: newID("New"), Type: &ast.FuncType{
+			Params:     &ast.FieldList{List: params},
+			TypeParams: t.modType.TypeParams,
+			Results:    &ast.FieldList{List: []*ast.Field{{Type: t.modRecv.Type}}},
+		},
+		Body: body,
+	}
 }
 
-func (t *translator) createHostInterfaces() []ast.Decl {
+func (t *translator) createHostInterfaces() []*ast.GenDecl {
 	ifaces := map[string][]*ast.Field{}
 
 	for _, imp := range t.imports {
@@ -312,7 +352,7 @@ func (t *translator) createHostInterfaces() []ast.Decl {
 			Type:  typ})
 	}
 
-	decls := make([]ast.Decl, 0, len(ifaces))
+	decls := make([]*ast.GenDecl, 0, len(ifaces))
 	for name, methods := range ifaces {
 		decls = append(decls, &ast.GenDecl{
 			Tok: token.TYPE,
@@ -321,6 +361,7 @@ func (t *translator) createHostInterfaces() []ast.Decl {
 				Name:   ast.NewIdent(exported(name)),
 				Type:   &ast.InterfaceType{Methods: &ast.FieldList{List: methods}}}}})
 	}
+
 	return decls
 }
 
@@ -429,7 +470,7 @@ func (t *translator) createExportMethods() []ast.Decl {
 		}
 
 		decls = append(decls, &ast.FuncDecl{
-			Recv: modRecvList,
+			Recv: &ast.FieldList{List: []*ast.Field{t.modRecv}},
 			Name: methodName,
 			Type: &ast.FuncType{Results: results},
 			Body: &ast.BlockStmt{List: body}})

translate.go

@@ -44,9 +44,15 @@ var stdlib = map[string]string{
 type translator struct {
 	in  *bufio.Reader
 	out ast.File
+
+	// WASM module Go type.
+	modType *ast.TypeSpec
+	modRecv *ast.Field
+
 	// Dependencies.
 	packages set[string]
 	helpers  set[string]
+
 	// Sections.
 	types     []funcType
 	imports   []importDef
@@ -89,25 +95,66 @@ func translate(r io.Reader, w io.Writer) error {
 			break
 		}
 	}
+
+	var hostInterfaces []*ast.GenDecl
+
 	if t.memory != nil && (exported || t.memory.imported) {
 		t.out.Decls = append(t.createMemoryTypes(), t.out.Decls...)
 	}
+
 	if !*nohost && len(t.imports) > 0 {
-		t.out.Decls = append(t.createHostInterfaces(), t.out.Decls...)
+		hostInterfaces = t.createHostInterfaces()
+		t.out.Decls = append(toDecl(hostInterfaces...), t.out.Decls...)
+
+		if !*generic {
+			// Drop host interface slice here, to
+			// prevent generic type parameters being
+			// generated in module creation below.
+			hostInterfaces = nil
+		}
+	}
+
+	// Next define module, with host interface generic param(s).
+	modDecl, modType := t.createModuleStruct(hostInterfaces)
+	t.modRecv = &ast.Field{Names: []*ast.Ident{newID("m")}}
+	t.modType = modType
+
+	// Set module receiver type depending
+	// on if generic type params do exist.
+	if modType.TypeParams != nil &&
+		len(modType.TypeParams.List) > 0 {
+		var paramNames []ast.Expr
+		for _, field := range modType.TypeParams.List {
+			paramNames = append(paramNames, field.Names[0])
+		}
+		t.modRecv.Type = &ast.StarExpr{X: &ast.IndexListExpr{X: newID("Module"), Indices: paramNames}}
+	} else {
+		t.modRecv.Type = &ast.StarExpr{X: newID("Module")}
 	}
 
 	t.out.Decls = append([]ast.Decl{
-		t.createModuleStruct(),
-		t.createNewFunc()},
-		t.out.Decls...)
+		modDecl,
+		t.createNewFunc(),
+	}, t.out.Decls...)
 
 	// Fill in missing names.
 	if t.out.Name == nil {
 		t.out.Name = newID("wasm2go")
 	}
 	for i, fn := range t.functions {
-		if fn.decl != nil && fn.decl.Name.Name == "" {
-			fn.decl.Name.Name = "f" + strconv.Itoa(i)
+		if fn.decl != nil {
+			if fn.decl.Name.Name == "" {
+				// If no function name, generate one.
+				fn.decl.Name.Name = "f" + strconv.Itoa(i)
+			}
+
+			// Now module type information is known,
+			// update all module receiver methods with it.
+			for _, field := range fn.decl.Recv.List {
+				if len(field.Names) > 0 && field.Names[0].Name == "m" {
+					field.Type = t.modRecv.Type
+				}
+			}
 		}
 	}
 	if t.memory != nil && t.memory.id.Name == "" {
@@ -382,7 +429,16 @@ func (t *translator) readImportSection() error {
 				typ: typ,
 				decl: &ast.FuncDecl{
 					Name: &ast.Ident{},
-					Recv: modRecvList,
+					Recv: &ast.FieldList{
+						List: []*ast.Field{
+							{
+								Names: []*ast.Ident{newID("m")},
+								// Don't set receiving field type
+								// yet, this gets set after module
+								// generic type params are known.
+							},
+						},
+					},
 					Type: typ.toAST(),
 					Body: &ast.BlockStmt{List: []ast.Stmt{stmt}}},
 			}
@@ -487,7 +543,16 @@ func (t *translator) readFunctionSection() error {
 		fn.typ = t.types[index]
 		fn.decl = &ast.FuncDecl{
 			Name: &ast.Ident{},
-			Recv: modRecvList,
+			Recv: &ast.FieldList{
+				List: []*ast.Field{
+					{
+						Names: []*ast.Ident{newID("m")},
+						// Don't set receiving field type
+						// yet, this gets set after module
+						// generic type params are known.
+					},
+				},
+			},
 			Type: fn.typ.toAST(),
 		}
 		t.out.Decls = append(t.out.Decls, fn.decl)

util.go

@@ -0,0 +1,11 @@
+package main
+
+import "go/ast"
+
+func toDecl[D ast.Decl](decl ...D) []ast.Decl {
+	decls := make([]ast.Decl, 0, len(decl))
+	for _, d := range decl {
+		decls = append(decls, d)
+	}
+	return decls
+}