Decouple the CUE reference resolver logic from the main parser

cmd/cli/loaders/cue.go

@@ -21,6 +21,11 @@ func cueLoader(opts Options) ([]*ast.Schema, error) {
 		return nil, err
 	}
 
+	libraries, err := opts.cueIncludeImports()
+	if err != nil {
+		return nil, err
+	}
+
 	allSchemas := make([]*ast.Schema, 0, len(opts.CueEntrypoints))
 	for _, entrypoint := range opts.CueEntrypoints {
 		pkg := filepath.Base(entrypoint)
@@ -42,6 +47,7 @@ func cueLoader(opts Options) ([]*ast.Schema, error) {
 			SchemaMetadata: ast.SchemaMeta{
 				// TODO: extract these from somewhere
 			},
+			Libraries: libraries,
 		})
 		if err != nil {
 			return nil, err
@@ -75,12 +81,12 @@ func buildBaseFSWithLibraries(opts Options) (fs.FS, error) {
 
 	var librariesFS []fs.FS
 	for _, importDefinition := range importDefinitions {
-		absPath, err := filepath.Abs(importDefinition.fsPath)
+		absPath, err := filepath.Abs(importDefinition.FSPath)
 		if err != nil {
 			return nil, err
 		}
 
-		libraryFS, err := dirToPrefixedFS(absPath, "cue.mod/pkg/"+importDefinition.importPath)
+		libraryFS, err := dirToPrefixedFS(absPath, "cue.mod/pkg/"+importDefinition.ImportPath)
 		if err != nil {
 			return nil, err
 		}

cmd/cli/loaders/kindsyscomposable.go

@@ -15,6 +15,11 @@ import (
 func kindsysComposableLoader(opts Options) ([]*ast.Schema, error) {
 	themaRuntime := thema.NewRuntime(cuecontext.New())
 
+	libraries, err := opts.cueIncludeImports()
+	if err != nil {
+		return nil, err
+	}
+
 	allSchemas := make([]*ast.Schema, 0, len(opts.KindsysComposableEntrypoints))
 	for _, entrypoint := range opts.KindsysComposableEntrypoints {
 		pkg := filepath.Base(entrypoint)
@@ -57,6 +62,7 @@ func kindsysComposableLoader(opts Options) ([]*ast.Schema, error) {
 				Variant:    variant,
 				Identifier: inferComposableKindIdentifier(props),
 			},
+			Libraries: libraries,
 		})
 		if err != nil {
 			return nil, err

cmd/cli/loaders/kindsyscore.go

@@ -17,6 +17,11 @@ import (
 func kindsysCoreLoader(opts Options) ([]*ast.Schema, error) {
 	themaRuntime := thema.NewRuntime(cuecontext.New())
 
+	libraries, err := opts.cueIncludeImports()
+	if err != nil {
+		return nil, err
+	}
+
 	allSchemas := make([]*ast.Schema, 0, len(opts.KindsysCoreEntrypoints))
 	for _, entrypoint := range opts.KindsysCoreEntrypoints {
 		pkg := filepath.Base(entrypoint)
@@ -52,6 +57,7 @@ func kindsysCoreLoader(opts Options) ([]*ast.Schema, error) {
 				Kind:       ast.SchemaKindCore,
 				Identifier: pkg, // TODO: maybe even core kinds could have one explicitly set in their schema?
 			},
+			Libraries: libraries,
 		})
 		if err != nil {
 			return nil, err

cmd/cli/loaders/kindsyscustom.go

@@ -14,6 +14,11 @@ import (
 func kindsysCustomLoader(opts Options) ([]*ast.Schema, error) {
 	themaRuntime := thema.NewRuntime(cuecontext.New())
 
+	libraries, err := opts.cueIncludeImports()
+	if err != nil {
+		return nil, err
+	}
+
 	allSchemas := make([]*ast.Schema, 0, len(opts.KindsysCustomEntrypoints))
 	for _, entrypoint := range opts.KindsysCustomEntrypoints {
 		pkg := filepath.Base(entrypoint)
@@ -49,6 +54,7 @@ func kindsysCustomLoader(opts Options) ([]*ast.Schema, error) {
 				Kind:       ast.SchemaKindCore, // TODO: is there any need for a "SchemaKindCustom"?
 				Identifier: pkg,                // TODO: maybe even core kinds could have one explicitly set in their schema?
 			},
+			Libraries: libraries,
 		})
 		if err != nil {
 			return nil, err

cmd/cli/loaders/options.go

@@ -5,6 +5,7 @@ import (
 	"strings"
 
 	"github.com/grafana/cog/internal/ast"
+	"github.com/grafana/cog/internal/simplecue"
 )
 
 type LoaderRef string
@@ -31,11 +32,6 @@ func loadersMap() map[LoaderRef]Loader {
 	}
 }
 
-type cueIncludeImport struct {
-	fsPath     string // path of the library on the filesystem
-	importPath string // path used in CUE files to import that library
-}
-
 type Loader func(opts Options) ([]*ast.Schema, error)
 
 type Options struct {
@@ -54,20 +50,22 @@ type Options struct {
 	KindRegistryVersion string
 }
 
-func (opts Options) cueIncludeImports() ([]cueIncludeImport, error) {
+func (opts Options) cueIncludeImports() ([]simplecue.LibraryInclude, error) {
 	if len(opts.CueImports) == 0 {
 		return nil, nil
 	}
 
-	imports := make([]cueIncludeImport, len(opts.CueImports))
+	imports := make([]simplecue.LibraryInclude, len(opts.CueImports))
 	for i, importDefinition := range opts.CueImports {
 		parts := strings.Split(importDefinition, ":")
 		if len(parts) != 2 {
 			return nil, fmt.Errorf("'%s' is not a valid import definition", importDefinition)
 		}
 
-		imports[i].fsPath = parts[0]
-		imports[i].importPath = parts[1]
+		imports[i] = simplecue.LibraryInclude{
+			FSPath:     parts[0],
+			ImportPath: parts[1],
+		}
 	}
 
 	return imports, nil

internal/simplecue/generator.go

@@ -5,7 +5,6 @@ import (
 	"strings"
 
 	"cuelang.org/go/cue"
-	cueast "cuelang.org/go/cue/ast"
 	"cuelang.org/go/cue/format"
 	"cuelang.org/go/pkg/strconv"
 	"github.com/grafana/cog/internal/ast"
@@ -17,6 +16,11 @@ const hintKindEnum = "enum"
 const annotationKindFieldName = "kind"
 const enumMembersAttr = "memberNames"
 
+type LibraryInclude struct {
+	FSPath     string // path of the library on the filesystem
+	ImportPath string // path used in CUE files to import that library
+}
+
 type Config struct {
 	// Package name used to generate code into.
 	Package string
@@ -25,11 +29,13 @@ type Config struct {
 	ForceVariantEnvelope bool
 
 	SchemaMetadata ast.SchemaMeta
+
+	Libraries []LibraryInclude
 }
 
 type generator struct {
-	schema          *ast.Schema
-	importsAliasMap map[string]string
+	schema      *ast.Schema
+	refResolver *referenceResolver
 }
 
 func GenerateAST(val cue.Value, c Config) (*ast.Schema, error) {
@@ -38,7 +44,9 @@ func GenerateAST(val cue.Value, c Config) (*ast.Schema, error) {
 			Package:  c.Package,
 			Metadata: c.SchemaMetadata,
 		},
-		importsAliasMap: buildImportsAliasMap(val),
+		refResolver: newReferenceResolver(val, referenceResolverConfig{
+			Libraries: c.Libraries,
+		}),
 	}
 
 	if c.ForceVariantEnvelope {
@@ -253,77 +261,6 @@ func (g *generator) structFields(v cue.Value) ([]ast.StructField, error) {
 	return fields, nil
 }
 
-// FIXME: this is probably very brittle and not always correct :|
-func (g *generator) referencePackage(source cueast.Node) (string, error) {
-	switch source.(type) { //nolint: gocritic
-	case *cueast.SelectorExpr:
-		selector := source.(*cueast.SelectorExpr)
-
-		x := selector.X.(*cueast.Ident)
-
-		return g.resolveImportAlias(x.Name), nil
-	case *cueast.Field:
-		field := source.(*cueast.Field)
-
-		if _, ok := field.Value.(*cueast.SelectorExpr); ok {
-			return g.referencePackage(field.Value)
-		}
-
-		ident := field.Value.(*cueast.Ident)
-
-		if ident.Scope == nil {
-			return g.schema.Package, nil
-		}
-
-		if _, ok := ident.Scope.(*cueast.File); !ok {
-			return g.schema.Package, nil
-		}
-
-		scope := ident.Scope.(*cueast.File)
-		if len(scope.Decls) == 0 {
-			return g.schema.Package, nil
-		}
-
-		if _, ok := scope.Decls[0].(*cueast.Package); !ok {
-			return g.schema.Package, nil
-		}
-
-		referredTypePkg := scope.Decls[0].(*cueast.Package).Name
-
-		return g.resolveImportAlias(referredTypePkg.Name), nil
-	case *cueast.Ident:
-		ident := source.(*cueast.Ident)
-		if ident.Scope == nil {
-			return g.schema.Package, nil
-		}
-
-		if _, ok := ident.Scope.(*cueast.File); !ok {
-			return g.schema.Package, nil
-		}
-
-		scope := ident.Scope.(*cueast.File)
-		if len(scope.Decls) == 0 {
-			return g.schema.Package, nil
-		}
-
-		referredTypePkg := ident.Scope.(*cueast.File).Decls[0].(*cueast.Package).Name
-
-		return g.resolveImportAlias(referredTypePkg.Name), nil
-	case *cueast.Ellipsis:
-		if source.(*cueast.Ellipsis).Type == nil {
-			return g.schema.Package, nil
-		}
-
-		if _, ok := source.(*cueast.Ellipsis).Type.(*cueast.SelectorExpr); ok {
-			return g.referencePackage(source.(*cueast.Ellipsis).Type)
-		}
-
-		return g.schema.Package, nil
-	default:
-		return "", fmt.Errorf("can't extract reference package")
-	}
-}
-
 func (g *generator) declareNode(v cue.Value) (ast.Type, error) {
 	v = g.removeTautologicalUnification(v)
 
@@ -332,7 +269,7 @@ func (g *generator) declareNode(v cue.Value) (ast.Type, error) {
 	if path.String() != "" {
 		selectors := path.Selectors()
 
-		refPkg, err := g.referencePackage(v.Source())
+		refPkg, err := g.refResolver.PackageForNode(v.Source(), g.schema.Package)
 		if err != nil {
 			return ast.Type{}, errorWithCueRef(v, err.Error())
 		}
@@ -770,11 +707,3 @@ func (g *generator) removeTautologicalUnification(v cue.Value) cue.Value {
 
 	return v
 }
-
-func (g *generator) resolveImportAlias(alias string) string {
-	if resolved, found := g.importsAliasMap[alias]; found {
-		return resolved
-	}
-
-	return alias
-}