feat: add type registry in Go module contexts (#1396)

fixes #1386

buildengine/build_go_test.go

@@ -245,3 +245,74 @@ func TestGoModVersion(t *testing.T) {
 	}
 	testBuild(t, bctx, fmt.Sprintf("go version %q is not recent enough for this module, needs minimum version \"9000.1.1\"", runtime.Version()[2:]), []assertion{})
 }
+
+func TestGeneratedTypeRegistry(t *testing.T) {
+	sch := &schema.Schema{
+		Modules: []*schema.Module{
+			{Name: "another", Decls: []schema.Decl{
+				&schema.Enum{
+					Name:   "TypeEnum",
+					Export: true,
+					Variants: []*schema.EnumVariant{
+						{Name: "A", Value: &schema.TypeValue{Value: &schema.Int{}}},
+						{Name: "B", Value: &schema.TypeValue{Value: &schema.String{}}},
+					},
+				},
+			}},
+		},
+	}
+	expected := `// Code generated by FTL. DO NOT EDIT.
+package main
+
+import (
+	"context"
+	"reflect"
+
+	"github.com/TBD54566975/ftl/common/plugin"
+	"github.com/TBD54566975/ftl/go-runtime/server"
+	"github.com/TBD54566975/ftl/backend/protos/xyz/block/ftl/v1/ftlv1connect"
+	"github.com/TBD54566975/ftl/go-runtime/ftl"
+
+	"ftl/other"
+	"ftl/another"
+)
+
+func main() {
+	verbConstructor := server.NewUserVerbServer("other",
+		server.HandleCall(other.Echo),
+	)
+	ctx := context.Background()
+	typeRegistry := ftl.NewTypeRegistry()
+	typeRegistry.RegisterSumType(reflect.TypeFor[other.D](), map[string]reflect.Type{
+		"Bool": reflect.TypeFor[other.Bool](),
+		"Bytes": reflect.TypeFor[other.Bytes](),
+		"Float": reflect.TypeFor[other.Float](),
+		"Int": reflect.TypeFor[other.Int](),
+		"U": reflect.TypeFor[other.U](),
+		"List": reflect.TypeFor[other.List](),
+		"Map": reflect.TypeFor[other.Map](),
+		"String": reflect.TypeFor[other.String](),
+		"Struct": reflect.TypeFor[other.Struct](),
+	})
+	typeRegistry.RegisterSumType(reflect.TypeFor[other.SecondSumType](), map[string]reflect.Type{
+		"A": reflect.TypeFor[other.A](),
+		"B": reflect.TypeFor[other.B](),
+	})
+	typeRegistry.RegisterSumType(reflect.TypeFor[another.TypeEnum](), map[string]reflect.Type{
+		"A": reflect.TypeFor[another.A](),
+		"B": reflect.TypeFor[another.B](),
+	})
+	ctx = context.WithValue(ctx, "typeRegistry", typeRegistry)
+
+	plugin.Start(ctx, "other", verbConstructor, ftlv1connect.VerbServiceName, ftlv1connect.NewVerbServiceHandler)
+}
+`
+	bctx := buildContext{
+		moduleDir: "testdata/projects/other",
+		buildDir:  "_ftl",
+		sch:       sch,
+	}
+	testBuild(t, bctx, "", []assertion{
+		assertGeneratedMain(expected),
+	})
+}

buildengine/build_test.go

@@ -77,6 +77,17 @@ func assertGeneratedModule(generatedModulePath string, expectedContent string) a
 	}
 }
 
+func assertGeneratedMain(expectedContent string) assertion {
+	return func(t testing.TB, bctx buildContext) error {
+		t.Helper()
+		output := filepath.Join(bctx.moduleDir, bctx.buildDir, "go/main/main.go")
+		fileContent, err := os.ReadFile(output)
+		assert.NoError(t, err)
+		assert.Equal(t, expectedContent, string(fileContent))
+		return nil
+	}
+}
+
 func assertBuildProtoErrors(msgs ...string) assertion {
 	return func(t testing.TB, bctx buildContext) error {
 		t.Helper()

buildengine/engine_test.go

@@ -49,7 +49,7 @@ func TestEngine(t *testing.T) {
 	expected := map[string][]string{
 		"alpha":   {"another", "other", "builtin"},
 		"another": {"builtin"},
-		"other":   {"builtin"},
+		"other":   {"another", "builtin"},
 		"builtin": {},
 	}
 	graph, err := engine.Graph()

buildengine/testdata/projects/another/another.go

@@ -7,6 +7,19 @@ import (
 	"github.com/TBD54566975/ftl/go-runtime/ftl" // Import the FTL SDK.
 )
 
+//ftl:enum export
+type TypeEnum interface {
+	tag()
+}
+
+type A int
+
+func (A) tag() {}
+
+type B string
+
+func (B) tag() {}
+
 type EchoRequest struct {
 	Name ftl.Option[string] `json:"name"`
 }

buildengine/testdata/projects/other/other.go

@@ -5,10 +5,71 @@ import (
 	"fmt"
 
 	"github.com/TBD54566975/ftl/go-runtime/ftl" // Import the FTL SDK.
+
+	"ftl/another"
 )
 
+//ftl:enum
+type D interface {
+	tag()
+}
+
+type Bool bool
+
+func (Bool) tag() {}
+
+type Bytes []byte
+
+func (Bytes) tag() {}
+
+type Float float64
+
+func (Float) tag() {}
+
+type Int int
+
+func (Int) tag() {}
+
+//type Time time.Time
+//
+//func (Time) tag() {}
+
+type U ftl.Unit
+
+func (U) tag() {}
+
+type List []string
+
+func (List) tag() {}
+
+type Map map[string]string
+
+func (Map) tag() {}
+
+type String string
+
+func (String) tag() {}
+
+type Struct struct{}
+
+func (Struct) tag() {}
+
+//ftl:enum
+type SecondSumType interface {
+	tag2()
+}
+
+type A string
+
+func (A) tag2() {}
+
+type B EchoRequest
+
+func (B) tag2() {}
+
 type EchoRequest struct {
-	Name ftl.Option[string] `json:"name"`
+	Name            ftl.Option[string] `json:"name"`
+	ExternalSumType another.TypeEnum
 }
 
 type EchoResponse struct {

go-runtime/compile/build.go

@@ -15,6 +15,8 @@ import (
 	"unicode"
 
 	"github.com/TBD54566975/scaffolder"
+	sets "github.com/deckarep/golang-set/v2"
+	gomaps "golang.org/x/exp/maps"
 	"golang.org/x/mod/modfile"
 	"golang.org/x/mod/semver"
 	"golang.org/x/sync/errgroup"
@@ -50,6 +52,17 @@ type mainModuleContext struct {
 	Name         string
 	Verbs        []goVerb
 	Replacements []*modfile.Replace
+	SumTypes     []goSumType
+}
+
+type goSumType struct {
+	Discriminator string
+	Variants      []goSumTypeVariant
+}
+
+type goSumTypeVariant struct {
+	Name string
+	Type string
 }
 
 type ModifyFilesTransaction interface {
@@ -123,11 +136,13 @@ func Build(ctx context.Context, moduleDir string, sch *schema.Schema, filesTrans
 
 	buildDir := buildDir(moduleDir)
 	logger.Debugf("Extracting schema")
-	nativeNames, main, schemaErrs, err := ExtractModuleSchema(moduleDir)
+	parseResult, err := ExtractModuleSchema(moduleDir)
 	if err != nil {
 		return fmt.Errorf("failed to extract module schema: %w", err)
 	}
-	if len(schemaErrs) > 0 {
+	pr := parseResult.MustGet()
+	main := pr.Module
+	if schemaErrs := pr.Errors; len(schemaErrs) > 0 {
 		if err := writeSchemaErrors(config, schemaErrs); err != nil {
 			return fmt.Errorf("failed to write errors: %w", err)
 		}
@@ -146,7 +161,7 @@ func Build(ctx context.Context, moduleDir string, sch *schema.Schema, filesTrans
 	goVerbs := make([]goVerb, 0, len(main.Decls))
 	for _, decl := range main.Decls {
 		if verb, ok := decl.(*schema.Verb); ok {
-			nativeName, ok := nativeNames[verb]
+			nativeName, ok := pr.NativeNames[verb]
 			if !ok {
 				return fmt.Errorf("missing native name for verb %s", verb.Name)
 			}
@@ -167,6 +182,7 @@ func Build(ctx context.Context, moduleDir string, sch *schema.Schema, filesTrans
 		Name:         main.Name,
 		Verbs:        goVerbs,
 		Replacements: replacements,
+		SumTypes:     getSumTypes(pr.EnumRefs, main, sch, pr.NativeNames),
 	}, scaffolder.Exclude("^go.mod$"), scaffolder.Functions(funcs)); err != nil {
 		return err
 	}
@@ -219,7 +235,6 @@ func GenerateStubsForExternalLibrary(ctx context.Context, dir string, schema *sc
 		Schema:       schema,
 		Replacements: replacements,
 	})
-
 }
 
 func generateExternalModules(context ExternalModuleContext) error {
@@ -314,6 +329,23 @@ var scaffoldFuncs = scaffolder.FuncMap{
 		}
 		return false
 	},
+	"mainImports": func(ctx mainModuleContext) []string {
+		imports := sets.NewSet[string]()
+		if len(ctx.Verbs) > 0 {
+			imports.Add(ctx.Name)
+		}
+		for _, st := range ctx.SumTypes {
+			if i := strings.LastIndex(st.Discriminator, "."); i != -1 {
+				imports.Add(st.Discriminator[:i])
+			}
+			for _, v := range st.Variants {
+				if i := strings.LastIndex(v.Type, "."); i != -1 {
+					imports.Add(v.Type[:i])
+				}
+			}
+		}
+		return imports.ToSlice()
+	},
 }
 
 func genType(module *schema.Module, t schema.Type) string {
@@ -444,3 +476,46 @@ func writeSchemaErrors(config moduleconfig.ModuleConfig, errors []*schema.Error)
 	}
 	return os.WriteFile(filepath.Join(config.AbsDeployDir(), config.Errors), elBytes, 0600)
 }
+
+func getSumTypes(enumRefs []*schema.Ref, module *schema.Module, sch *schema.Schema, nativeNames NativeNames) []goSumType {
+	sumTypes := make(map[string]goSumType)
+	for _, d := range module.Decls {
+		if e, ok := d.(*schema.Enum); ok && !e.IsValueEnum() {
+			variants := make([]goSumTypeVariant, 0, len(e.Variants))
+			for _, v := range e.Variants {
+				variants = append(variants, goSumTypeVariant{
+					Name: v.Name,
+					Type: nativeNames[v],
+				})
+			}
+			stFqName := nativeNames[d]
+			sumTypes[stFqName] = goSumType{
+				Discriminator: nativeNames[d],
+				Variants:      variants,
+			}
+		}
+	}
+
+	// register sum types from other modules
+	for _, ref := range enumRefs {
+		if ref.Module == module.Name {
+			continue
+		}
+		resolved := sch.ResolveRef(ref)
+		if e, ok := resolved.(*schema.Enum); ok && !e.IsValueEnum() {
+			variants := make([]goSumTypeVariant, 0, len(e.Variants))
+			for _, v := range e.Variants {
+				variants = append(variants, goSumTypeVariant{
+					Name: v.Name,
+					Type: ref.Module + "." + v.Name,
+				})
+			}
+			stFqName := ref.Module + "." + e.Name
+			sumTypes[stFqName] = goSumType{
+				Discriminator: stFqName,
+				Variants:      variants,
+			}
+		}
+	}
+	return gomaps.Values(sumTypes)
+}