From 5e9bda092a891fb26770174cbc6f5cc1351003bf Mon Sep 17 00:00:00 2001
From: Ian Wahbe <ian@wahbe.com>
Date: Thu, 8 Dec 2022 14:36:44 +0100
Subject: [PATCH 1/2] Don't use *schema.Package in nodejs codegen

---
 pkg/codegen/nodejs/doc.go                     |   4 +-
 pkg/codegen/nodejs/gen.go                     | 103 +++++++++++-------
 pkg/codegen/nodejs/gen_program.go             |  26 +++--
 pkg/codegen/nodejs/gen_program_expressions.go |   8 +-
 4 files changed, 89 insertions(+), 52 deletions(-)

diff --git a/pkg/codegen/nodejs/doc.go b/pkg/codegen/nodejs/doc.go
index 4f60bbab14d7..f346105c41a1 100644
--- a/pkg/codegen/nodejs/doc.go
+++ b/pkg/codegen/nodejs/doc.go
@@ -77,7 +77,7 @@ func (d DocLanguageHelper) GetLanguageTypeString(pkg *schema.Package, moduleName
 	}
 
 	modCtx := &modContext{
-		pkg: pkg,
+		pkg: pkg.Reference(),
 		mod: moduleName,
 	}
 	typeName := modCtx.typeString(t, input, nil)
@@ -114,7 +114,7 @@ func (d DocLanguageHelper) GetMethodResultName(pkg *schema.Package, modName stri
 	if info, ok := pkg.Language["nodejs"].(NodePackageInfo); ok {
 		if info.LiftSingleValueMethodReturns && m.Function.Outputs != nil && len(m.Function.Outputs.Properties) == 1 {
 			modCtx := &modContext{
-				pkg: pkg,
+				pkg: pkg.Reference(),
 				mod: modName,
 			}
 			return modCtx.typeString(m.Function.Outputs.Properties[0].Type, false, nil)
diff --git a/pkg/codegen/nodejs/gen.go b/pkg/codegen/nodejs/gen.go
index 40241014db09..23adc5882478 100644
--- a/pkg/codegen/nodejs/gen.go
+++ b/pkg/codegen/nodejs/gen.go
@@ -128,7 +128,7 @@ func externalModuleName(s string) string {
 }
 
 type modContext struct {
-	pkg              *schema.Package
+	pkg              schema.PackageReference
 	mod              string
 	types            []*schema.ObjectType
 	enums            []*schema.EnumType
@@ -180,15 +180,17 @@ func (mod *modContext) tokenToModName(tok string) string {
 	return modName
 }
 
-func (mod *modContext) namingContext(pkg *schema.Package) (namingCtx *modContext, pkgName string, external bool) {
+func (mod *modContext) namingContext(pkg schema.PackageReference) (namingCtx *modContext, pkgName string, external bool) {
 	namingCtx = mod
-	if pkg != nil && pkg != mod.pkg {
+	if pkg != nil && !codegen.PkgEquals(pkg, mod.pkg) {
 		external = true
-		pkgName = pkg.Name + "."
+		pkgName = pkg.Name() + "."
 
 		var info NodePackageInfo
-		contract.AssertNoError(pkg.ImportLanguages(map[string]schema.Language{"nodejs": Importer}))
-		if v, ok := pkg.Language["nodejs"].(NodePackageInfo); ok {
+		def, err := pkg.Definition()
+		contract.AssertNoError(err)
+		contract.AssertNoError(def.ImportLanguages(map[string]schema.Language{"nodejs": Importer}))
+		if v, ok := def.Language["nodejs"].(NodePackageInfo); ok {
 			info = v
 		}
 		namingCtx = &modContext{
@@ -200,7 +202,7 @@ func (mod *modContext) namingContext(pkg *schema.Package) (namingCtx *modContext
 	return
 }
 
-func (mod *modContext) objectType(pkg *schema.Package, details *typeDetails, tok string, input, args, enum bool) string {
+func (mod *modContext) objectType(pkg schema.PackageReference, details *typeDetails, tok string, input, args, enum bool) string {
 
 	root := "outputs."
 	if input {
@@ -238,7 +240,7 @@ func (mod *modContext) objectType(pkg *schema.Package, details *typeDetails, tok
 func (mod *modContext) resourceType(r *schema.ResourceType) string {
 	if strings.HasPrefix(r.Token, "pulumi:providers:") {
 		pkgName := strings.TrimPrefix(r.Token, "pulumi:providers:")
-		if pkgName != mod.pkg.Name {
+		if pkgName != mod.pkg.Name() {
 			pkgName = externalModuleName(pkgName)
 		}
 
@@ -247,7 +249,7 @@ func (mod *modContext) resourceType(r *schema.ResourceType) string {
 
 	pkg := mod.pkg
 	if r.Resource != nil {
-		pkg = r.Resource.Package
+		pkg = r.Resource.PackageReference
 	}
 	namingCtx, pkgName, external := mod.namingContext(pkg)
 	if external {
@@ -298,14 +300,14 @@ func (mod *modContext) typeAst(t schema.Type, input bool, constValue interface{}
 		}
 		return tstypes.Identifier(fmt.Sprintf("pulumi.Input<%s>", typ))
 	case *schema.EnumType:
-		return tstypes.Identifier(mod.objectType(t.Package, nil, t.Token, input, false, true))
+		return tstypes.Identifier(mod.objectType(t.PackageReference, nil, t.Token, input, false, true))
 	case *schema.ArrayType:
 		return tstypes.Array(mod.typeAst(t.ElementType, input, constValue))
 	case *schema.MapType:
 		return tstypes.StringMap(mod.typeAst(t.ElementType, input, constValue))
 	case *schema.ObjectType:
 		details := mod.details(t)
-		return tstypes.Identifier(mod.objectType(t.Package, details, t.Token, input, t.IsInputShape(), false))
+		return tstypes.Identifier(mod.objectType(t.PackageReference, details, t.Token, input, t.IsInputShape(), false))
 	case *schema.ResourceType:
 		return tstypes.Identifier(mod.resourceType(t))
 	case *schema.TokenType:
@@ -670,7 +672,7 @@ func (mod *modContext) genResource(w io.Writer, r *schema.Resource) (resourceFil
 
 	pulumiType := r.Token
 	if r.IsProvider {
-		pulumiType = mod.pkg.Name
+		pulumiType = mod.pkg.Name()
 	}
 
 	fmt.Fprintf(w, "    /** @internal */\n")
@@ -1314,7 +1316,9 @@ func (mod *modContext) getTypeImportsForResource(t schema.Type, recurse bool, ex
 	}
 
 	var nodePackageInfo NodePackageInfo
-	if languageInfo, hasLanguageInfo := mod.pkg.Language["nodejs"]; hasLanguageInfo {
+	def, err := mod.pkg.Definition()
+	contract.AssertNoError(err)
+	if languageInfo, hasLanguageInfo := def.Language["nodejs"]; hasLanguageInfo {
 		nodePackageInfo = languageInfo.(NodePackageInfo)
 	}
 
@@ -1337,16 +1341,16 @@ func (mod *modContext) getTypeImportsForResource(t schema.Type, recurse bool, ex
 		return mod.getTypeImports(t.ElementType, recurse, externalImports, imports, seen)
 	case *schema.EnumType:
 		// If the enum is from another package, add an import for the external package.
-		if t.Package != nil && t.Package != mod.pkg {
-			pkg := t.Package.Name
+		if t.PackageReference != nil && !codegen.PkgEquals(t.PackageReference, mod.pkg) {
+			pkg := t.PackageReference.Name()
 			writeImports(pkg)
 			return false
 		}
 		return true
 	case *schema.ObjectType:
 		// If it's from another package, add an import for the external package.
-		if t.Package != nil && t.Package != mod.pkg {
-			pkg := t.Package.Name
+		if t.PackageReference != nil && !codegen.PkgEquals(t.PackageReference, mod.pkg) {
+			pkg := t.PackageReference.Name()
 			writeImports(pkg)
 			return false
 		}
@@ -1357,8 +1361,8 @@ func (mod *modContext) getTypeImportsForResource(t schema.Type, recurse bool, ex
 		return true
 	case *schema.ResourceType:
 		// If it's from another package, add an import for the external package.
-		if t.Resource != nil && t.Resource.Package != mod.pkg {
-			pkg := t.Resource.Package.Name
+		if t.Resource != nil && !codegen.PkgEquals(t.Resource.PackageReference, mod.pkg) {
+			pkg := t.Resource.PackageReference.Name()
 			writeImports(pkg)
 			return false
 		}
@@ -1510,7 +1514,7 @@ func (mod *modContext) genConfig(w io.Writer, variables []*schema.Property) erro
 	fmt.Fprintf(w, "declare var exports: any;\n")
 
 	// Create a config bag for the variables to pull from.
-	fmt.Fprintf(w, "const __config = new pulumi.Config(\"%v\");\n", mod.pkg.Name)
+	fmt.Fprintf(w, "const __config = new pulumi.Config(\"%v\");\n", mod.pkg.Name())
 	fmt.Fprintf(w, "\n")
 
 	// Emit an entry for all config variables.
@@ -1561,9 +1565,11 @@ func (mod *modContext) sdkImports(nested, utilities bool) []string {
 			fmt.Sprintf(`import * as outputs from "%s/types/output";`, relRoot),
 		}...)
 
-		if mod.pkg.Language["nodejs"].(NodePackageInfo).ContainsEnums {
+		def, err := mod.pkg.Definition()
+		contract.AssertNoError(err)
+		if def.Language["nodejs"].(NodePackageInfo).ContainsEnums {
 			code := `import * as enums from "%s/types/enums";`
-			if lookupNodePackageInfo(mod.pkg).UseTypeOnlyReferences {
+			if lookupNodePackageInfo(def).UseTypeOnlyReferences {
 				code = `import type * as enums from "%s/types/enums";`
 			}
 			imports = append(imports, fmt.Sprintf(code, relRoot))
@@ -1754,7 +1760,9 @@ func (mod *modContext) isReservedSourceFileName(name string) bool {
 	case "utilities.ts":
 		return mod.mod == ""
 	case "vars.ts":
-		return len(mod.pkg.Config) > 0
+		config, err := mod.pkg.Config()
+		contract.AssertNoError(err)
+		return len(config) > 0
 	default:
 		return false
 	}
@@ -1800,26 +1808,34 @@ func (mod *modContext) gen(fs codegen.Fs) error {
 		fs.Add(p, []byte(contents))
 	}
 
+	def, err := mod.pkg.Definition()
+	if err != nil {
+		return err
+	}
+
 	// Utilities, config, readme
 	switch mod.mod {
 	case "":
 		buffer := &bytes.Buffer{}
 		mod.genHeader(buffer, nil, nil, nil)
-		mod.genUtilitiesFile(buffer)
+		err := mod.genUtilitiesFile(buffer)
+		if err != nil {
+			return err
+		}
 		fs.Add(path.Join(modDir, "utilities.ts"), buffer.Bytes())
 
 		// Ensure that the top-level (provider) module directory contains a README.md file.
-		readme := mod.pkg.Language["nodejs"].(NodePackageInfo).Readme
+		readme := def.Language["nodejs"].(NodePackageInfo).Readme
 		if readme == "" {
-			readme = mod.pkg.Description
+			readme = def.Description
 			if readme != "" && readme[len(readme)-1] != '\n' {
 				readme += "\n"
 			}
-			if mod.pkg.Attribution != "" {
+			if def.Attribution != "" {
 				if len(readme) != 0 {
 					readme += "\n"
 				}
-				readme += mod.pkg.Attribution
+				readme += def.Attribution
 			}
 		}
 		if readme != "" && readme[len(readme)-1] != '\n' {
@@ -1827,9 +1843,9 @@ func (mod *modContext) gen(fs codegen.Fs) error {
 		}
 		fs.Add(path.Join(modDir, "README.md"), []byte(readme))
 	case "config":
-		if len(mod.pkg.Config) > 0 {
+		if len(def.Config) > 0 {
 			buffer := &bytes.Buffer{}
-			if err := mod.genConfig(buffer, mod.pkg.Config); err != nil {
+			if err := mod.genConfig(buffer, def.Config); err != nil {
 				return err
 			}
 			addFile(otherFileType, "vars.ts", buffer.String())
@@ -1904,7 +1920,7 @@ func (mod *modContext) gen(fs codegen.Fs) error {
 
 	// Nested types
 	// Importing enums always imports inputs and outputs, so if we have enums we generate inputs and outputs
-	if len(mod.types) > 0 || (mod.pkg.Language["nodejs"].(NodePackageInfo).ContainsEnums && mod.mod == "types") {
+	if len(mod.types) > 0 || (def.Language["nodejs"].(NodePackageInfo).ContainsEnums && mod.mod == "types") {
 		input, output, err := mod.genTypes()
 		if err != nil {
 			return err
@@ -1975,7 +1991,10 @@ func (mod *modContext) genIndex(exports []fileInfo) string {
 		}
 	}
 
-	info, _ := mod.pkg.Language["nodejs"].(NodePackageInfo)
+	def, err := mod.pkg.Definition()
+	contract.AssertNoError(err)
+
+	info, _ := def.Language["nodejs"].(NodePackageInfo)
 	if info.ContainsEnums {
 		if mod.mod == "types" {
 			children.Add("enums")
@@ -2050,7 +2069,7 @@ func (mod *modContext) genResourceModule(w io.Writer) {
 				continue
 			}
 
-			registrations.Add(mod.pkg.TokenToRuntimeModule(r.Token))
+			registrations.Add(schema.TokenToRuntimeModule(r.Token))
 		}
 
 		fmt.Fprintf(w, "\nconst _module = {\n")
@@ -2078,12 +2097,12 @@ func (mod *modContext) genResourceModule(w io.Writer) {
 		fmt.Fprintf(w, "    },\n")
 		fmt.Fprintf(w, "};\n")
 		for _, name := range registrations.SortedValues() {
-			fmt.Fprintf(w, "pulumi.runtime.registerResourceModule(\"%v\", \"%v\", _module)\n", mod.pkg.Name, name)
+			fmt.Fprintf(w, "pulumi.runtime.registerResourceModule(\"%v\", \"%v\", _module)\n", mod.pkg.Name(), name)
 		}
 	}
 
 	if provider != nil {
-		fmt.Fprintf(w, "pulumi.runtime.registerResourcePackage(\"%v\", {\n", mod.pkg.Name)
+		fmt.Fprintf(w, "pulumi.runtime.registerResourcePackage(\"%v\", {\n", mod.pkg.Name())
 		fmt.Fprintf(w, "    version: utilities.getVersion(),\n")
 		fmt.Fprintf(w, "    constructProvider: (name: string, type: string, urn: string): pulumi.ProviderResource => {\n")
 		fmt.Fprintf(w, "        if (type !== \"%v\") {\n", provider.Token)
@@ -2352,7 +2371,7 @@ func generateModuleContextMap(tool string, pkg *schema.Package, extraFiles map[s
 		mod, ok := modules[modName]
 		if !ok {
 			mod = &modContext{
-				pkg:                          pkg,
+				pkg:                          pkg.Reference(),
 				mod:                          modName,
 				tool:                         tool,
 				compatibility:                info.Compatibility,
@@ -2571,7 +2590,7 @@ func GeneratePackage(tool string, pkg *schema.Package, extraFiles map[string][]b
 	return files, nil
 }
 
-func (mod *modContext) genUtilitiesFile(w io.Writer) {
+func (mod *modContext) genUtilitiesFile(w io.Writer) error {
 	const body = `
 export function getEnv(...vars: string[]): string | undefined {
     for (const v of vars) {
@@ -2636,12 +2655,16 @@ export function lazyLoad(exports: any, props: string[], loadModule: any) {
     }
 }
 `
+	def, err := mod.pkg.Definition()
+	if err != nil {
+		return err
+	}
 	var pluginDownloadURL string
-	if url := mod.pkg.PluginDownloadURL; url != "" {
+	if url := def.PluginDownloadURL; url != "" {
 		pluginDownloadURL = fmt.Sprintf(", pluginDownloadURL: %q", url)
 	}
-	_, err := fmt.Fprintf(w, body, pluginDownloadURL)
-	contract.AssertNoError(err)
+	_, err = fmt.Fprintf(w, body, pluginDownloadURL)
+	return err
 }
 
 func genInstallScript(pluginDownloadURL string) string {
diff --git a/pkg/codegen/nodejs/gen_program.go b/pkg/codegen/nodejs/gen_program.go
index e8caac5626ee..ec2dae39ff2b 100644
--- a/pkg/codegen/nodejs/gen_program.go
+++ b/pkg/codegen/nodejs/gen_program.go
@@ -73,7 +73,10 @@ func GenerateProgram(program *pcl.Program) (map[string][]byte, hcl.Diagnostics,
 	}
 
 	var index bytes.Buffer
-	g.genPreamble(&index, program, preambleHelperMethods)
+	err = g.genPreamble(&index, program, preambleHelperMethods)
+	if err != nil {
+		return nil, nil, err
+	}
 	for _, n := range nodes {
 		if g.asyncMain {
 			break
@@ -281,7 +284,7 @@ func (g *generator) genComment(w io.Writer, comment syntax.Comment) {
 	}
 }
 
-func (g *generator) genPreamble(w io.Writer, program *pcl.Program, preambleHelperMethods codegen.StringSet) {
+func (g *generator) genPreamble(w io.Writer, program *pcl.Program, preambleHelperMethods codegen.StringSet) error {
 	// Print the @pulumi/pulumi import at the top.
 	g.Fprintln(w, `import * as pulumi from "@pulumi/pulumi";`)
 
@@ -296,8 +299,12 @@ func (g *generator) genPreamble(w io.Writer, program *pcl.Program, preambleHelpe
 				continue
 			}
 			pkgName := "@pulumi/" + pkg
-			if r.Schema != nil && r.Schema.Package != nil {
-				if info, ok := r.Schema.Package.Language["nodejs"].(NodePackageInfo); ok && info.PackageName != "" {
+			if r.Schema != nil && r.Schema.PackageReference != nil {
+				def, err := r.Schema.PackageReference.Definition()
+				if err != nil {
+					return err
+				}
+				if info, ok := def.Language["nodejs"].(NodePackageInfo); ok && info.PackageName != "" {
 					pkgName = info.PackageName
 				}
 				npmToPuPkgName[pkgName] = pkg
@@ -345,6 +352,7 @@ func (g *generator) genPreamble(w io.Writer, program *pcl.Program, preambleHelpe
 	for _, preambleHelperMethodBody := range preambleHelperMethods.SortedValues() {
 		g.Fprintf(w, "%s\n\n", preambleHelperMethodBody)
 	}
+	return nil
 }
 
 func (g *generator) genNode(w io.Writer, n pcl.Node) {
@@ -384,18 +392,20 @@ func resourceTypeName(r *pcl.Resource) (string, string, string, hcl.Diagnostics)
 	pkg, module, member, diagnostics := r.DecomposeToken()
 
 	if r.Schema != nil {
-		module = moduleName(module, r.Schema.Package)
+		module = moduleName(module, r.Schema.PackageReference)
 	}
 
 	return makeValidIdentifier(pkg), module, title(member), diagnostics
 }
 
-func moduleName(module string, pkg *schema.Package) string {
+func moduleName(module string, pkg schema.PackageReference) string {
 	// Normalize module.
 	if pkg != nil {
-		err := pkg.ImportLanguages(map[string]schema.Language{"nodejs": Importer})
+		def, err := pkg.Definition()
+		contract.AssertNoError(err)
+		err = def.ImportLanguages(map[string]schema.Language{"nodejs": Importer})
 		contract.AssertNoError(err)
-		if lang, ok := pkg.Language["nodejs"]; ok {
+		if lang, ok := def.Language["nodejs"]; ok {
 			pkgInfo := lang.(NodePackageInfo)
 			if m, ok := pkgInfo.ModuleToPackage[module]; ok {
 				module = m
diff --git a/pkg/codegen/nodejs/gen_program_expressions.go b/pkg/codegen/nodejs/gen_program_expressions.go
index c4b142f6c195..c5daeb8bd5eb 100644
--- a/pkg/codegen/nodejs/gen_program_expressions.go
+++ b/pkg/codegen/nodejs/gen_program_expressions.go
@@ -309,11 +309,15 @@ func enumName(enum *model.EnumType) (string, error) {
 	if !ok {
 		return "", fmt.Errorf("Could not get associated enum")
 	}
-	if name := e.(*schema.EnumType).Package.Language["nodejs"].(NodePackageInfo).PackageName; name != "" {
+	def, err := e.(*schema.EnumType).PackageReference.Definition()
+	if err != nil {
+		return "", err
+	}
+	if name := def.Language["nodejs"].(NodePackageInfo).PackageName; name != "" {
 		pkg = name
 	}
 	if mod := components[1]; mod != "" && mod != "index" {
-		if pkg := e.(*schema.EnumType).Package; pkg != nil {
+		if pkg := e.(*schema.EnumType).PackageReference; pkg != nil {
 			mod = moduleName(mod, pkg)
 		}
 		pkg += "." + mod

From 7ae9c181d05f255a8c06490844bdcdb66ea623e4 Mon Sep 17 00:00:00 2001
From: Ian Wahbe <ian@wahbe.com>
Date: Thu, 8 Dec 2022 11:45:46 +0100
Subject: [PATCH 2/2] Don't use *schema.Package in go codegen

---
 pkg/codegen/go/doc.go                     |   5 +-
 pkg/codegen/go/gen.go                     | 214 ++++++++++++++--------
 pkg/codegen/go/gen_crd2pulumi.go          |   5 +-
 pkg/codegen/go/gen_program.go             |   3 +-
 pkg/codegen/go/gen_program_expressions.go |  11 +-
 pkg/codegen/go/gen_test.go                |  17 +-
 6 files changed, 165 insertions(+), 90 deletions(-)

diff --git a/pkg/codegen/go/doc.go b/pkg/codegen/go/doc.go
index ca641a1dc74a..8f4dda888aa3 100644
--- a/pkg/codegen/go/doc.go
+++ b/pkg/codegen/go/doc.go
@@ -26,6 +26,7 @@ import (
 	"github.com/golang/glog"
 	"github.com/pulumi/pulumi/pkg/v3/codegen"
 	"github.com/pulumi/pulumi/pkg/v3/codegen/schema"
+	"github.com/pulumi/pulumi/sdk/v3/go/common/util/contract"
 )
 
 const pulumiSDKVersion = "v3"
@@ -98,7 +99,9 @@ func (d DocLanguageHelper) GetLanguageTypeString(pkg *schema.Package, moduleName
 
 // GeneratePackagesMap generates a map of Go packages for resources, functions and types.
 func (d *DocLanguageHelper) GeneratePackagesMap(pkg *schema.Package, tool string, goInfo GoPackageInfo) {
-	d.packages = generatePackageContextMap(tool, pkg, goInfo, nil)
+	var err error
+	d.packages, err = generatePackageContextMap(tool, pkg.Reference(), goInfo, nil)
+	contract.AssertNoError(err)
 }
 
 // GetPropertyName returns the property name specific to Go.
diff --git a/pkg/codegen/go/gen.go b/pkg/codegen/go/gen.go
index 298b0b8a82c1..f8e26cc9fdfd 100644
--- a/pkg/codegen/go/gen.go
+++ b/pkg/codegen/go/gen.go
@@ -103,7 +103,7 @@ func Title(s string) string {
 	return s
 }
 
-func tokenToPackage(pkg *schema.Package, overrides map[string]string, tok string) string {
+func tokenToPackage(pkg schema.PackageReference, overrides map[string]string, tok string) string {
 	mod := pkg.TokenToModule(tok)
 	if override, ok := overrides[mod]; ok {
 		mod = override
@@ -140,7 +140,7 @@ func (c *Cache) setContextMap(pkg *schema.Package, m map[string]*pkgContext) {
 }
 
 type pkgContext struct {
-	pkg             *schema.Package
+	pkg             schema.PackageReference
 	mod             string
 	importBasePath  string
 	rootPackageName string
@@ -226,7 +226,9 @@ func (pkg *pkgContext) tokenToType(tok string) string {
 		return name
 	}
 	if mod == "" {
-		mod = packageRoot(pkg.pkg)
+		var err error
+		mod, err = packageRoot(pkg.pkg)
+		contract.AssertNoError(err)
 	}
 
 	var importPath string
@@ -624,30 +626,26 @@ func (pkg *pkgContext) isExternalReference(t schema.Type) bool {
 
 // Return if `t` is external to `pkg`. If so, the associated foreign schema.Package is returned.
 func (pkg *pkgContext) isExternalReferenceWithPackage(t schema.Type) (
-	isExternal bool, extPkg *schema.Package, token string) {
-	var err error
+	isExternal bool, extPkg schema.PackageReference, token string) {
 	switch typ := t.(type) {
 	case *schema.ObjectType:
-		isExternal = typ.Package != nil && pkg.pkg != nil && typ.Package != pkg.pkg
+		isExternal = typ.PackageReference != nil && !codegen.PkgEquals(typ.PackageReference, pkg.pkg)
 		if isExternal {
-			extPkg, err = typ.PackageReference.Definition()
-			contract.AssertNoError(err)
+			extPkg = typ.PackageReference
 			token = typ.Token
 		}
 		return
 	case *schema.ResourceType:
-		isExternal = typ.Resource != nil && pkg.pkg != nil && typ.Resource.Package != pkg.pkg
+		isExternal = typ.Resource != nil && pkg.pkg != nil && !codegen.PkgEquals(typ.Resource.PackageReference, pkg.pkg)
 		if isExternal {
-			extPkg, err = typ.Resource.PackageReference.Definition()
-			contract.AssertNoError(err)
+			extPkg = typ.Resource.PackageReference
 			token = typ.Token
 		}
 		return
 	case *schema.EnumType:
-		isExternal = pkg.pkg != nil && typ.Package != pkg.pkg
+		isExternal = pkg.pkg != nil && !codegen.PkgEquals(typ.PackageReference, pkg.pkg)
 		if isExternal {
-			extPkg, err = typ.PackageReference.Definition()
-			contract.AssertNoError(err)
+			extPkg = typ.PackageReference
 			token = typ.Token
 		}
 		return
@@ -666,7 +664,7 @@ func (pkg *pkgContext) resolveResourceType(t *schema.ResourceType) string {
 	extPkgCtx, _ := pkg.contextForExternalReference(t)
 	resType := extPkgCtx.tokenToResource(t.Token)
 	if !strings.Contains(resType, ".") {
-		resType = fmt.Sprintf("%s.%s", extPkgCtx.pkg.Name, resType)
+		resType = fmt.Sprintf("%s.%s", extPkgCtx.pkg.Name(), resType)
 	}
 	return resType
 }
@@ -694,8 +692,10 @@ func (pkg *pkgContext) contextForExternalReference(t schema.Type) (*pkgContext,
 	contract.Assert(isExternal)
 
 	var goInfo GoPackageInfo
-	contract.AssertNoError(extPkg.ImportLanguages(map[string]schema.Language{"go": Importer}))
-	if info, ok := extPkg.Language["go"].(GoPackageInfo); ok {
+	extDef, err := extPkg.Definition()
+	contract.AssertNoError(err)
+	contract.AssertNoError(extDef.ImportLanguages(map[string]schema.Language{"go": Importer}))
+	if info, ok := extDef.Language["go"].(GoPackageInfo); ok {
 		goInfo = info
 	} else {
 		goInfo.ImportBasePath = extractImportBasePath(extPkg)
@@ -705,7 +705,9 @@ func (pkg *pkgContext) contextForExternalReference(t schema.Type) (*pkgContext,
 
 	// Ensure that any package import aliases we have specified locally take precedence over those
 	// specified in the remote package.
-	if ourPkgGoInfoI, has := pkg.pkg.Language["go"]; has {
+	def, err := pkg.pkg.Definition()
+	contract.AssertNoError(err)
+	if ourPkgGoInfoI, has := def.Language["go"]; has {
 		ourPkgGoInfo := ourPkgGoInfoI.(GoPackageInfo)
 		if len(ourPkgGoInfo.PackageImportAliases) > 0 {
 			pkgImportAliases = make(map[string]string)
@@ -722,11 +724,12 @@ func (pkg *pkgContext) contextForExternalReference(t schema.Type) (*pkgContext,
 
 	var maps map[string]*pkgContext
 
-	if extMap, ok := pkg.externalPackages.lookupContextMap(extPkg); ok {
+	if extMap, ok := pkg.externalPackages.lookupContextMap(extDef); ok {
 		maps = extMap
 	} else {
-		maps = generatePackageContextMap(pkg.tool, extPkg, goInfo, pkg.externalPackages)
-		pkg.externalPackages.setContextMap(extPkg, maps)
+		maps, err = generatePackageContextMap(pkg.tool, extPkg, goInfo, pkg.externalPackages)
+		contract.AssertNoError(err)
+		pkg.externalPackages.setContextMap(extDef, maps)
 	}
 	extPkgCtx := maps[""]
 	extPkgCtx.pkgImportAliases = pkgImportAliases
@@ -1792,7 +1795,10 @@ func (pkg *pkgContext) genResource(w io.Writer, r *schema.Resource, generateReso
 		fmt.Fprint(w, "\topts = append(opts, replaceOnChanges)\n")
 	}
 
-	pkg.GenPkgDefaultsOptsCall(w, false /*invoke*/)
+	err := pkg.GenPkgDefaultsOptsCall(w, false /*invoke*/)
+	if err != nil {
+		return err
+	}
 
 	// Finally make the call to registration.
 	fmt.Fprintf(w, "\tvar resource %s\n", name)
@@ -2129,7 +2135,10 @@ func (pkg *pkgContext) genFunction(w io.Writer, f *schema.Function) error {
 		outputsType = name + "Result"
 	}
 
-	pkg.GenPkgDefaultsOptsCall(w, true /*invoke*/)
+	err := pkg.GenPkgDefaultsOptsCall(w, true /*invoke*/)
+	if err != nil {
+		return err
+	}
 
 	fmt.Fprintf(w, "\tvar rv %s\n", outputsType)
 	fmt.Fprintf(w, "\terr := ctx.Invoke(\"%s\", %s, &rv, opts...)\n", f.Token, inputsVar)
@@ -2735,13 +2744,13 @@ func (pkg *pkgContext) getTypeImports(t schema.Type, recurse bool, importsAndAli
 	}
 }
 
-func extractImportBasePath(extPkg *schema.Package) string {
-	version := extPkg.Version.Major
+func extractImportBasePath(extPkg schema.PackageReference) string {
+	version := extPkg.Version().Major
 	var vPath string
 	if version > 1 {
 		vPath = fmt.Sprintf("/v%d", version)
 	}
-	return fmt.Sprintf("github.com/pulumi/pulumi-%s/sdk%s/go/%s", extPkg.Name, vPath, extPkg.Name)
+	return fmt.Sprintf("github.com/pulumi/pulumi-%s/sdk%s/go/%s", extPkg.Name(), vPath, extPkg.Name())
 }
 
 func (pkg *pkgContext) getImports(member interface{}, importsAndAliases map[string]string) {
@@ -2799,7 +2808,9 @@ func (pkg *pkgContext) genHeader(w io.Writer, goImports []string, importsAndAlia
 
 	var pkgName string
 	if pkg.mod == "" {
-		pkgName = packageName(pkg.pkg)
+		def, err := pkg.pkg.Definition()
+		contract.AssertNoError(err)
+		pkgName = packageName(def)
 	} else {
 		pkgName = path.Base(pkg.mod)
 	}
@@ -2869,7 +2880,7 @@ func (pkg *pkgContext) genConfig(w io.Writer, variables []*schema.Property) erro
 		}
 
 		printCommentWithDeprecationMessage(w, p.Comment, p.DeprecationMessage, false)
-		configKey := fmt.Sprintf("\"%s:%s\"", pkg.pkg.Name, cgstrings.Camel(p.Name))
+		configKey := fmt.Sprintf("\"%s:%s\"", pkg.pkg.Name(), cgstrings.Camel(p.Name))
 
 		fmt.Fprintf(w, "func Get%s(ctx *pulumi.Context) %s {\n", Title(p.Name), getType)
 		if p.DefaultValue != nil {
@@ -2896,7 +2907,7 @@ func (pkg *pkgContext) genConfig(w io.Writer, variables []*schema.Property) erro
 // Pulumi runtime. The generated ResourceModule supports the deserialization of resource references into fully-
 // hydrated Resource instances. If this is the root module, this function also generates a ResourcePackage
 // definition and its registration to support rehydrating providers.
-func (pkg *pkgContext) genResourceModule(w io.Writer) {
+func (pkg *pkgContext) genResourceModule(w io.Writer) error {
 	contract.Assert(len(pkg.resources) != 0)
 	allResourcesAreOverlays := true
 	for _, r := range pkg.resources {
@@ -2907,7 +2918,7 @@ func (pkg *pkgContext) genResourceModule(w io.Writer) {
 	}
 	if allResourcesAreOverlays {
 		// If all resources in this module are overlays, skip further code generation.
-		return
+		return nil
 	}
 
 	basePath := pkg.importBasePath
@@ -2928,7 +2939,11 @@ func (pkg *pkgContext) genResourceModule(w io.Writer) {
 
 	// If there are any internal dependencies, include them as blank imports.
 	if topLevelModule {
-		if goInfo, ok := pkg.pkg.Language["go"].(GoPackageInfo); ok {
+		def, err := pkg.pkg.Definition()
+		if err != nil {
+			return err
+		}
+		if goInfo, ok := def.Language["go"].(GoPackageInfo); ok {
 			for _, dep := range goInfo.InternalDependencies {
 				imports[dep] = "_"
 			}
@@ -2985,7 +3000,7 @@ func (pkg *pkgContext) genResourceModule(w io.Writer) {
 		fmt.Fprintf(w, "}\n\n")
 
 		fmt.Fprintf(w, "func (p *pkg) ConstructProvider(ctx *pulumi.Context, name, typ, urn string) (pulumi.ProviderResource, error) {\n")
-		fmt.Fprintf(w, "\tif typ != \"pulumi:providers:%s\" {\n", pkg.pkg.Name)
+		fmt.Fprintf(w, "\tif typ != \"pulumi:providers:%s\" {\n", pkg.pkg.Name())
 		fmt.Fprintf(w, "\t\treturn nil, fmt.Errorf(\"unknown provider type: %%s\", typ)\n")
 		fmt.Fprintf(w, "\t}\n\n")
 		fmt.Fprintf(w, "\tr := &Provider{}\n")
@@ -3017,7 +3032,7 @@ func (pkg *pkgContext) genResourceModule(w io.Writer) {
 	if len(registrations) > 0 {
 		for _, mod := range registrations.SortedValues() {
 			fmt.Fprintf(w, "\tpulumi.RegisterResourceModule(\n")
-			fmt.Fprintf(w, "\t\t%q,\n", pkg.pkg.Name)
+			fmt.Fprintf(w, "\t\t%q,\n", pkg.pkg.Name())
 			fmt.Fprintf(w, "\t\t%q,\n", mod)
 			fmt.Fprintf(w, "\t\t&module{version},\n")
 			fmt.Fprintf(w, "\t)\n")
@@ -3025,15 +3040,16 @@ func (pkg *pkgContext) genResourceModule(w io.Writer) {
 	}
 	if provider != nil {
 		fmt.Fprintf(w, "\tpulumi.RegisterResourcePackage(\n")
-		fmt.Fprintf(w, "\t\t%q,\n", pkg.pkg.Name)
+		fmt.Fprintf(w, "\t\t%q,\n", pkg.pkg.Name())
 		fmt.Fprintf(w, "\t\t&pkg{version},\n")
 		fmt.Fprintf(w, "\t)\n")
 	}
-	fmt.Fprintf(w, "}\n")
+	_, err := fmt.Fprintf(w, "}\n")
+	return err
 }
 
 // generatePackageContextMap groups resources, types, and functions into Go packages.
-func generatePackageContextMap(tool string, pkg *schema.Package, goInfo GoPackageInfo, externalPkgs *Cache) map[string]*pkgContext {
+func generatePackageContextMap(tool string, pkg schema.PackageReference, goInfo GoPackageInfo, externalPkgs *Cache) (map[string]*pkgContext, error) {
 	packages := map[string]*pkgContext{}
 
 	// Share the cache
@@ -3085,7 +3101,11 @@ func generatePackageContextMap(tool string, pkg *schema.Package, goInfo GoPackag
 		}
 	}
 
-	if len(pkg.Config) > 0 {
+	config, err := pkg.Config()
+	if err != nil {
+		return nil, err
+	}
+	if len(config) > 0 {
 		_ = getPkg("config")
 	}
 
@@ -3176,13 +3196,17 @@ func generatePackageContextMap(tool string, pkg *schema.Package, goInfo GoPackag
 	}
 
 	// Rewrite cyclic types. See the docs on rewriteCyclicFields for the motivation.
-	rewriteCyclicObjectFields(pkg)
+	def, err := pkg.Definition()
+	if err != nil {
+		return nil, err
+	}
+	rewriteCyclicObjectFields(def)
 
 	// Use a string set to track object types that have already been processed.
 	// This avoids recursively processing the same type. For example, in the
 	// Kubernetes package, JSONSchemaProps have properties whose type is itself.
 	seenMap := codegen.NewStringSet()
-	for _, t := range pkg.Types {
+	for _, t := range def.Types {
 		switch typ := t.(type) {
 		case *schema.ArrayType:
 			details := getPkgFromType(typ.ElementType).detailsForType(codegen.UnwrapType(typ.ElementType))
@@ -3293,8 +3317,8 @@ func generatePackageContextMap(tool string, pkg *schema.Package, goInfo GoPackag
 		}
 	}
 
-	scanResource(pkg.Provider)
-	for _, r := range pkg.Resources {
+	scanResource(def.Provider)
+	for _, r := range def.Resources {
 		scanResource(r)
 	}
 
@@ -3402,7 +3426,7 @@ func generatePackageContextMap(tool string, pkg *schema.Package, goInfo GoPackag
 		}
 	}
 
-	for _, t := range pkg.Types {
+	for _, t := range def.Types {
 		scanType(t)
 	}
 
@@ -3410,7 +3434,7 @@ func generatePackageContextMap(tool string, pkg *schema.Package, goInfo GoPackag
 	// input or output property type metadata, in case they have
 	// types used in array or pointer element positions.
 	if !goInfo.DisableFunctionOutputVersions || goInfo.GenerateExtraInputTypes {
-		for _, f := range pkg.Functions {
+		for _, f := range def.Functions {
 			if f.NeedsOutputVersion() || goInfo.GenerateExtraInputTypes {
 				optional := false
 				if f.Inputs != nil {
@@ -3423,7 +3447,7 @@ func generatePackageContextMap(tool string, pkg *schema.Package, goInfo GoPackag
 		}
 	}
 
-	for _, f := range pkg.Functions {
+	for _, f := range def.Functions {
 		if f.IsMethod {
 			continue
 		}
@@ -3454,7 +3478,7 @@ func generatePackageContextMap(tool string, pkg *schema.Package, goInfo GoPackag
 		}
 	}
 
-	return packages
+	return packages, nil
 }
 
 // LanguageResource is derived from the schema and can be used by downstream codegen.
@@ -3479,7 +3503,10 @@ func LanguageResources(tool string, pkg *schema.Package) (map[string]LanguageRes
 	if goInfo, ok := pkg.Language["go"].(GoPackageInfo); ok {
 		goPkgInfo = goInfo
 	}
-	packages := generatePackageContextMap(tool, pkg, goPkgInfo, globalCache)
+	packages, err := generatePackageContextMap(tool, pkg.Reference(), goPkgInfo, globalCache)
+	if err != nil {
+		return nil, err
+	}
 
 	// emit each package
 	var pkgMods []string
@@ -3517,19 +3544,23 @@ func LanguageResources(tool string, pkg *schema.Package) (map[string]LanguageRes
 // source file should be under this root. For example:
 //
 // root = aws => sdk/go/aws/*.go
-func packageRoot(pkg *schema.Package) string {
+func packageRoot(pkg schema.PackageReference) (string, error) {
+	def, err := pkg.Definition()
+	if err != nil {
+		return "", err
+	}
 	var info GoPackageInfo
-	if goInfo, ok := pkg.Language["go"].(GoPackageInfo); ok {
+	if goInfo, ok := def.Language["go"].(GoPackageInfo); ok {
 		info = goInfo
 	}
 	if info.RootPackageName != "" {
 		// package structure is flat
-		return ""
+		return "", nil
 	}
 	if info.ImportBasePath != "" {
-		return path.Base(info.ImportBasePath)
+		return path.Base(info.ImportBasePath), nil
 	}
-	return goPackage(pkg.Name)
+	return goPackage(pkg.Name()), nil
 }
 
 // packageName is the go package name for the generated package.
@@ -3541,7 +3572,9 @@ func packageName(pkg *schema.Package) string {
 	if info.RootPackageName != "" {
 		return info.RootPackageName
 	}
-	return goPackage(packageRoot(pkg))
+	root, err := packageRoot(pkg.Reference())
+	contract.AssertNoErrorf(err, "We generated the ref from a pkg, so we know its a valid ref")
+	return goPackage(root)
 }
 
 func GeneratePackage(tool string, pkg *schema.Package) (map[string][]byte, error) {
@@ -3553,7 +3586,10 @@ func GeneratePackage(tool string, pkg *schema.Package) (map[string][]byte, error
 	if goInfo, ok := pkg.Language["go"].(GoPackageInfo); ok {
 		goPkgInfo = goInfo
 	}
-	packages := generatePackageContextMap(tool, pkg, goPkgInfo, NewCache())
+	packages, err := generatePackageContextMap(tool, pkg.Reference(), goPkgInfo, NewCache())
+	if err != nil {
+		return nil, err
+	}
 
 	// emit each package
 	var pkgMods []string
@@ -3563,7 +3599,10 @@ func GeneratePackage(tool string, pkg *schema.Package) (map[string][]byte, error
 	sort.Strings(pkgMods)
 
 	name := packageName(pkg)
-	pathPrefix := packageRoot(pkg)
+	pathPrefix, err := packageRoot(pkg.Reference())
+	if err != nil {
+		return nil, err
+	}
 
 	files := codegen.Fs{}
 
@@ -3602,8 +3641,8 @@ func GeneratePackage(tool string, pkg *schema.Package) (map[string][]byte, error
 		switch mod {
 		case "":
 			buffer := &bytes.Buffer{}
-			if pkg.pkg.Description != "" {
-				printComment(buffer, pkg.pkg.Description, false)
+			if pkg.pkg.Description() != "" {
+				printComment(buffer, pkg.pkg.Description(), false)
 			} else {
 				fmt.Fprintf(buffer, "// Package %[1]s exports types, functions, subpackages for provisioning %[1]s resources.\n", name)
 			}
@@ -3612,9 +3651,13 @@ func GeneratePackage(tool string, pkg *schema.Package) (map[string][]byte, error
 			setFile(path.Join(mod, "doc.go"), buffer.String())
 
 		case "config":
-			if len(pkg.pkg.Config) > 0 {
+			config, err := pkg.pkg.Config()
+			if err != nil {
+				return nil, err
+			}
+			if len(config) > 0 {
 				buffer := &bytes.Buffer{}
-				if err := pkg.genConfig(buffer, pkg.pkg.Config); err != nil {
+				if err := pkg.genConfig(buffer, config); err != nil {
 					return nil, err
 				}
 
@@ -3729,12 +3772,15 @@ func GeneratePackage(tool string, pkg *schema.Package) (map[string][]byte, error
 			}
 			pkg.genHeader(buffer, []string{"fmt", "os", "reflect", "regexp", "strconv", "strings"}, importsAndAliases)
 
-			packageRegex := fmt.Sprintf("^.*/pulumi-%s/sdk(/v\\d+)?", pkg.pkg.Name)
+			packageRegex := fmt.Sprintf("^.*/pulumi-%s/sdk(/v\\d+)?", pkg.pkg.Name())
 			if pkg.rootPackageName != "" {
 				packageRegex = fmt.Sprintf("^%s(/v\\d+)?", pkg.importBasePath)
 			}
 
-			pkg.GenUtilitiesFile(buffer, packageRegex)
+			err := pkg.GenUtilitiesFile(buffer, packageRegex)
+			if err != nil {
+				return nil, err
+			}
 
 			setFile(path.Join(mod, "pulumiUtilities.go"), buffer.String())
 		}
@@ -3742,7 +3788,10 @@ func GeneratePackage(tool string, pkg *schema.Package) (map[string][]byte, error
 		// If there are resources in this module, register the module with the runtime.
 		if len(pkg.resources) != 0 && !allResourcesAreOverlays(pkg.resources) {
 			buffer := &bytes.Buffer{}
-			pkg.genResourceModule(buffer)
+			err := pkg.genResourceModule(buffer)
+			if err != nil {
+				return nil, err
+			}
 
 			setFile(path.Join(mod, "init.go"), buffer.String())
 		}
@@ -3802,7 +3851,7 @@ func goPackage(name string) string {
 	return strings.ReplaceAll(name, "-", "")
 }
 
-func (pkg *pkgContext) GenUtilitiesFile(w io.Writer, packageRegex string) {
+func (pkg *pkgContext) GenUtilitiesFile(w io.Writer, packageRegex string) error {
 	const utilitiesFile = `
 type envParser func(v string) interface{}
 
@@ -3876,14 +3925,20 @@ func isZero(v interface{}) bool {
 }
 `
 	_, err := fmt.Fprintf(w, utilitiesFile, packageRegex)
-	contract.AssertNoError(err)
-	pkg.GenPkgDefaultOpts(w)
+	if err != nil {
+		return err
+	}
+	return pkg.GenPkgDefaultOpts(w)
 }
 
-func (pkg *pkgContext) GenPkgDefaultOpts(w io.Writer) {
-	url := pkg.pkg.PluginDownloadURL
+func (pkg *pkgContext) GenPkgDefaultOpts(w io.Writer) error {
+	p, err := pkg.pkg.Definition()
+	if err != nil {
+		return err
+	}
+	url := p.PluginDownloadURL
 	if url == "" {
-		return
+		return nil
 	}
 	const template string = `
 // pkg%[1]sDefaultOpts provides package level defaults to pulumi.Option%[1]s.
@@ -3895,28 +3950,35 @@ func pkg%[1]sDefaultOpts(opts []pulumi.%[1]sOption) []pulumi.%[1]sOption {
 `
 	pluginDownloadURL := fmt.Sprintf("pulumi.PluginDownloadURL(%q)", url)
 	version := ""
-	if info := pkg.pkg.Language["go"]; info != nil {
-		if info.(GoPackageInfo).RespectSchemaVersion && pkg.pkg.Version != nil {
-			version = fmt.Sprintf(", pulumi.Version(%q)", pkg.pkg.Version.String())
+	if info := p.Language["go"]; info != nil {
+		if info.(GoPackageInfo).RespectSchemaVersion && pkg.pkg.Version() != nil {
+			version = fmt.Sprintf(", pulumi.Version(%q)", p.Version.String())
 		}
 	}
 	for _, typ := range []string{"Resource", "Invoke"} {
 		_, err := fmt.Fprintf(w, template, typ, pluginDownloadURL, version)
-		contract.AssertNoError(err)
+		if err != nil {
+			return err
+		}
 	}
+	return nil
 }
 
 // GenPkgDefaultsOptsCall generates a call to Pkg{TYPE}DefaultsOpts.
-func (pkg *pkgContext) GenPkgDefaultsOptsCall(w io.Writer, invoke bool) {
+func (pkg *pkgContext) GenPkgDefaultsOptsCall(w io.Writer, invoke bool) error {
 	// The `pkg%sDefaultOpts` call won't do anything, so we don't insert it.
-	if pkg.pkg.PluginDownloadURL == "" {
-		return
+	p, err := pkg.pkg.Definition()
+	if err != nil {
+		return err
+	}
+	if p.PluginDownloadURL == "" {
+		return nil
 	}
 	pkg.needsUtils = true
 	typ := "Resource"
 	if invoke {
 		typ = "Invoke"
 	}
-	_, err := fmt.Fprintf(w, "\topts = pkg%sDefaultOpts(opts)\n", typ)
-	contract.AssertNoError(err)
+	_, err = fmt.Fprintf(w, "\topts = pkg%sDefaultOpts(opts)\n", typ)
+	return err
 }
diff --git a/pkg/codegen/go/gen_crd2pulumi.go b/pkg/codegen/go/gen_crd2pulumi.go
index cd4c86e90df4..28808c867556 100644
--- a/pkg/codegen/go/gen_crd2pulumi.go
+++ b/pkg/codegen/go/gen_crd2pulumi.go
@@ -18,7 +18,10 @@ func CRDTypes(tool string, pkg *schema.Package) (map[string]*bytes.Buffer, error
 	if goInfo, ok := pkg.Language["go"].(GoPackageInfo); ok {
 		goPkgInfo = goInfo
 	}
-	packages := generatePackageContextMap(tool, pkg, goPkgInfo, nil)
+	packages, err := generatePackageContextMap(tool, pkg.Reference(), goPkgInfo, nil)
+	if err != nil {
+		return nil, err
+	}
 
 	var pkgMods []string
 	for mod := range packages {
diff --git a/pkg/codegen/go/gen_program.go b/pkg/codegen/go/gen_program.go
index edbbec1d1cf8..fb6fd53b738b 100644
--- a/pkg/codegen/go/gen_program.go
+++ b/pkg/codegen/go/gen_program.go
@@ -272,7 +272,8 @@ func getPackages(tool string, pkg *schema.Package, cache *Cache) map[string]*pkg
 	if goInfo, ok := pkg.Language["go"].(GoPackageInfo); ok {
 		goPkgInfo = goInfo
 	}
-	v := generatePackageContextMap(tool, pkg, goPkgInfo, cache)
+	v, err := generatePackageContextMap(tool, pkg.Reference(), goPkgInfo, cache)
+	contract.AssertNoError(err)
 	packageContexts.Store(pkg, v)
 	return v
 }
diff --git a/pkg/codegen/go/gen_program_expressions.go b/pkg/codegen/go/gen_program_expressions.go
index 6660137df39d..03ade930cdb1 100644
--- a/pkg/codegen/go/gen_program_expressions.go
+++ b/pkg/codegen/go/gen_program_expressions.go
@@ -371,7 +371,10 @@ func outputVersionFunctionArgTypeName(t model.Type, cache *Cache) (string, error
 		return "", fmt.Errorf("Expected a schema.ObjectType, got %s", schemaType.String())
 	}
 
-	pkg := &pkgContext{pkg: &schema.Package{Name: "main"}, externalPackages: cache}
+	pkg := &pkgContext{
+		pkg:              (&schema.Package{Name: "main"}).Reference(),
+		externalPackages: cache,
+	}
 
 	var ty string
 	if pkg.isExternalReference(objType) {
@@ -781,8 +784,10 @@ func (g *generator) argumentTypeName(expr model.Expression, destType model.Type,
 	}
 
 	if schemaType, ok := pcl.GetSchemaForType(destType); ok {
-		pkg := &pkgContext{pkg: &schema.Package{Name: "main"}, externalPackages: g.externalCache}
-		return pkg.argsType(schemaType)
+		return (&pkgContext{
+			pkg:              (&schema.Package{Name: "main"}).Reference(),
+			externalPackages: g.externalCache,
+		}).argsType(schemaType)
 	}
 
 	switch destType := destType.(type) {
diff --git a/pkg/codegen/go/gen_test.go b/pkg/codegen/go/gen_test.go
index 3f94022ab2d1..2acecde4d190 100644
--- a/pkg/codegen/go/gen_test.go
+++ b/pkg/codegen/go/gen_test.go
@@ -135,7 +135,8 @@ func TestGenerateTypeNames(t *testing.T) {
 		if goInfo, ok := pkg.Language["go"].(GoPackageInfo); ok {
 			goPkgInfo = goInfo
 		}
-		packages := generatePackageContextMap("test", pkg, goPkgInfo, nil)
+		packages, err := generatePackageContextMap("test", pkg.Reference(), goPkgInfo, nil)
+		require.NoError(t, err)
 
 		root, ok := packages[""]
 		require.True(t, ok)
@@ -252,7 +253,7 @@ func TestTokenToType(t *testing.T) {
 	}{
 		{
 			pkg: &pkgContext{
-				pkg:            importSpec(t, awsSpec),
+				pkg:            importSpec(t, awsSpec).Reference(),
 				importBasePath: awsImportBasePath,
 			},
 			token:    "aws:s3/BucketWebsite:BucketWebsite",
@@ -260,7 +261,7 @@ func TestTokenToType(t *testing.T) {
 		},
 		{
 			pkg: &pkgContext{
-				pkg:            importSpec(t, awsSpec),
+				pkg:            importSpec(t, awsSpec).Reference(),
 				importBasePath: awsImportBasePath,
 				pkgImportAliases: map[string]string{
 					"github.com/pulumi/pulumi-aws/sdk/v4/go/aws/s3": "awss3",
@@ -271,7 +272,7 @@ func TestTokenToType(t *testing.T) {
 		},
 		{
 			pkg: &pkgContext{
-				pkg:            importSpec(t, googleNativeSpec),
+				pkg:            importSpec(t, googleNativeSpec).Reference(),
 				importBasePath: googleNativeImportBasePath,
 				pkgImportAliases: map[string]string{
 					"github.com/pulumi/pulumi-google-native/sdk/go/google/dns/v1": "dns",
@@ -316,7 +317,7 @@ func TestTokenToResource(t *testing.T) {
 	}{
 		{
 			pkg: &pkgContext{
-				pkg:            importSpec(t, awsSpec),
+				pkg:            importSpec(t, awsSpec).Reference(),
 				importBasePath: awsImportBasePath,
 			},
 			token:    "aws:s3/Bucket:Bucket",
@@ -324,7 +325,7 @@ func TestTokenToResource(t *testing.T) {
 		},
 		{
 			pkg: &pkgContext{
-				pkg:            importSpec(t, awsSpec),
+				pkg:            importSpec(t, awsSpec).Reference(),
 				importBasePath: awsImportBasePath,
 				pkgImportAliases: map[string]string{
 					"github.com/pulumi/pulumi-aws/sdk/v4/go/aws/s3": "awss3",
@@ -335,7 +336,7 @@ func TestTokenToResource(t *testing.T) {
 		},
 		{
 			pkg: &pkgContext{
-				pkg:            importSpec(t, googleNativeSpec),
+				pkg:            importSpec(t, googleNativeSpec).Reference(),
 				importBasePath: googleNativeImportBasePath,
 				pkgImportAliases: map[string]string{
 					"github.com/pulumi/pulumi-google-native/sdk/go/google/dns/v1": "dns",
@@ -368,7 +369,7 @@ func TestGenHeader(t *testing.T) {
 
 	pkg := &pkgContext{
 		tool: "a tool",
-		pkg:  &schema.Package{Name: "test-pkg"},
+		pkg:  (&schema.Package{Name: "test-pkg"}).Reference(),
 	}
 
 	s := func() string {