Allow a schema.*Resource to give its own aliases

.gitignore

@@ -22,3 +22,5 @@
 examples/**/pulumi-resource-*
 
 /.vscode
+
+examples/*/*.json

infer/component.go

@@ -59,18 +59,19 @@ func Component[R ComponentResource[I, O], I any, O pulumi.ComponentResource]() I
 type derivedComponentController[R ComponentResource[I, O], I any, O pulumi.ComponentResource] struct{}
 
 func (rc *derivedComponentController[R, I, O]) GetSchema(reg schema.RegisterDerivativeType) (
-	pschema.ResourceSpec, error) {
+	pschema.ResourceSpec, schema.Associated, error) {
 	r, err := getResourceSchema[R, I, O](true)
 	if err := err.ErrorOrNil(); err != nil {
-		return pschema.ResourceSpec{}, err
+		return pschema.ResourceSpec{}, schema.Associated{}, err
 	}
 	if err := registerTypes[I](reg); err != nil {
-		return pschema.ResourceSpec{}, err
+		return pschema.ResourceSpec{}, schema.Associated{}, err
 	}
 	if err := registerTypes[O](reg); err != nil {
-		return pschema.ResourceSpec{}, err
+		return pschema.ResourceSpec{}, schema.Associated{}, err
 	}
-	return r, nil
+
+	return r, schema.Associated{}, nil
 }
 
 func (rc *derivedComponentController[R, I, O]) GetToken() (tokens.Type, error) {

infer/configuration.go

@@ -56,12 +56,12 @@ func (*config[T]) underlyingType() reflect.Type {
 }
 
 func (*config[T]) GetToken() (tokens.Type, error) { return "pulumi:providers:pkg", nil }
-func (*config[T]) GetSchema(reg schema.RegisterDerivativeType) (pschema.ResourceSpec, error) {
+func (*config[T]) GetSchema(reg schema.RegisterDerivativeType) (pschema.ResourceSpec, schema.Associated, error) {
 	if err := registerTypes[T](reg); err != nil {
-		return pschema.ResourceSpec{}, err
+		return pschema.ResourceSpec{}, schema.Associated{}, err
 	}
 	r, errs := getResourceSchema[T, T, T](false)
-	return r, errs.ErrorOrNil()
+	return r, schema.Associated{}, errs.ErrorOrNil()
 }
 
 func (c *config[T]) checkConfig(ctx p.Context, req p.CheckRequest) (p.CheckResponse, error) {
@@ -70,7 +70,7 @@ func (c *config[T]) checkConfig(ctx p.Context, req p.CheckRequest) (p.CheckRespo
 		t = reflect.New(v.Type().Elem()).Interface().(T)
 	}
 
-	r, mErr := c.GetSchema(func(tk tokens.Type, typ pschema.ComplexTypeSpec) bool { return false })
+	r, _, mErr := c.GetSchema(func(tk tokens.Type, typ pschema.ComplexTypeSpec) bool { return false })
 	if mErr != nil {
 		return p.CheckResponse{}, fmt.Errorf("could not get config secrets: %w", mErr)
 	}
@@ -208,7 +208,7 @@ func (c *config[T]) configure(ctx p.Context, req p.ConfigureRequest) error {
 	if c.t == nil {
 		c.t = new(T)
 	}
-	schema, mErr := c.GetSchema(func(tk tokens.Type, typ pschema.ComplexTypeSpec) bool { return false })
+	schema, _, mErr := c.GetSchema(func(tk tokens.Type, typ pschema.ComplexTypeSpec) bool { return false })
 	if mErr != nil {
 		return mErr
 	}
@@ -246,7 +246,7 @@ func (c *config[T]) handleConfigFailures(ctx p.Context, err mapper.MappingError)
 	}
 
 	pkgName := ctx.RuntimeInformation().PackageName
-	schema, mErr := c.GetSchema(func(tk tokens.Type, typ pschema.ComplexTypeSpec) bool { return false })
+	schema, _, mErr := c.GetSchema(func(tk tokens.Type, typ pschema.ComplexTypeSpec) bool { return false })
 	if mErr != nil {
 		return mErr
 	}

infer/function.go

@@ -74,31 +74,33 @@ func fnToken(tk tokens.Type) tokens.Type {
 	return tokens.NewTypeToken(tk.Module(), tokens.TypeName(name))
 }
 
-func (*derivedInvokeController[F, I, O]) GetSchema(reg schema.RegisterDerivativeType) (pschema.FunctionSpec, error) {
+func (*derivedInvokeController[F, I, O]) GetSchema(
+	reg schema.RegisterDerivativeType,
+) (pschema.FunctionSpec, schema.Associated, error) {
 	var f F
 	descriptions := getAnnotated(reflect.TypeOf(f))
 
 	input, err := objectSchema(reflect.TypeOf(new(I)))
 	if err != nil {
-		return pschema.FunctionSpec{}, err
+		return pschema.FunctionSpec{}, schema.Associated{}, err
 	}
 	output, err := objectSchema(reflect.TypeOf(new(O)))
 	if err != nil {
-		return pschema.FunctionSpec{}, err
+		return pschema.FunctionSpec{}, schema.Associated{}, err
 	}
 
 	if err := registerTypes[I](reg); err != nil {
-		return pschema.FunctionSpec{}, err
+		return pschema.FunctionSpec{}, schema.Associated{}, err
 	}
 	if err := registerTypes[O](reg); err != nil {
-		return pschema.FunctionSpec{}, err
+		return pschema.FunctionSpec{}, schema.Associated{}, err
 	}
 
 	return pschema.FunctionSpec{
 		Description: descriptions.Descriptions[""],
 		Inputs:      input,
 		Outputs:     output,
-	}, nil
+	}, schema.Associated{}, nil
 }
 
 func objectSchema(t reflect.Type) (*pschema.ObjectTypeSpec, error) {

infer/resource.go

@@ -396,15 +396,23 @@ type derivedResourceController[R CustomResource[I, O], I, O any] struct {
 func (*derivedResourceController[R, I, O]) isInferredResource() {}
 
 func (rc *derivedResourceController[R, I, O]) GetSchema(reg schema.RegisterDerivativeType) (
-	pschema.ResourceSpec, error) {
+	pschema.ResourceSpec, schema.Associated, error) {
 	if err := registerTypes[I](reg); err != nil {
-		return pschema.ResourceSpec{}, err
+		return pschema.ResourceSpec{}, schema.Associated{}, err
 	}
 	if err := registerTypes[O](reg); err != nil {
-		return pschema.ResourceSpec{}, err
+		return pschema.ResourceSpec{}, schema.Associated{}, err
 	}
 	r, errs := getResourceSchema[R, I, O](false)
-	return r, errs.ErrorOrNil()
+	stateTk, err := introspect.GetToken("pkg", reflect.New(reflect.TypeOf((*O)(nil)).Elem()).Interface())
+	if err != nil {
+		return pschema.ResourceSpec{}, schema.Associated{}, err
+	}
+	return r, schema.Associated{
+		Aliases: []schema.TokenAlias{
+			{Token: stateTk},
+		},
+	}, errs.ErrorOrNil()
 }
 
 func (rc *derivedResourceController[R, I, O]) GetToken() (tokens.Type, error) {
@@ -505,7 +513,7 @@ func (rc *derivedResourceController[R, I, O]) Diff(ctx p.Context, req p.DiffRequ
 	_, hasUpdate := ((interface{})(*r)).(CustomUpdate[I, O])
 	var forceReplace func(string) bool
 	if hasUpdate {
-		schema, err := rc.GetSchema(func(tk tokens.Type, typ pschema.ComplexTypeSpec) (unknown bool) { return false })
+		schema, _, err := rc.GetSchema(func(tk tokens.Type, typ pschema.ComplexTypeSpec) (unknown bool) { return false })
 		if err != nil {
 			return p.DiffResponse{}, err
 		}

middleware/schema/schema.go

@@ -46,7 +46,7 @@ type Resource interface {
 	// Return the Resource's schema definition. The passed in function should be called on
 	// types transitively referenced by the resource. See the documentation of
 	// RegisterDerivativeType for more details.
-	GetSchema(RegisterDerivativeType) (schema.ResourceSpec, error)
+	GetSchema(RegisterDerivativeType) (schema.ResourceSpec, Associated, error)
 }
 
 // A Function that can generate its own schema definition.
@@ -56,7 +56,18 @@ type Function interface {
 	// Return the Function's schema definition. The passed in function should be called on
 	// types transitively referenced by the function. See the documentation of
 	// RegisterDerivativeType for more details.
-	GetSchema(RegisterDerivativeType) (schema.FunctionSpec, error)
+	GetSchema(RegisterDerivativeType) (schema.FunctionSpec, Associated, error)
+}
+
+// Associated is a bag of associated information that effects the global generation
+// process.
+type Associated struct {
+	Aliases []TokenAlias
+}
+
+// A token that aliases to the *Resource that returned this TokenAlias.
+type TokenAlias struct {
+	Token tokens.Type
 }
 
 type cache struct {
@@ -271,13 +282,17 @@ func (s *state) generateSchema(ctx p.Context) (schema.PackageSpec, error) {
 		pkg.Types[tkString] = renamePackage(t, info.PackageName, s.ModuleMap)
 		return true
 	}
-	errs := addElements(s.Resources, pkg.Resources, info.PackageName, registerDerivative, s.ModuleMap)
-	e := addElements(s.Invokes, pkg.Functions, info.PackageName, registerDerivative, s.ModuleMap)
+	aliases, errs := addElements(s.Resources, pkg.Resources, info.PackageName, registerDerivative, s.ModuleMap)
+	_, e := addElements(s.Invokes, pkg.Functions, info.PackageName, registerDerivative, s.ModuleMap)
+
 	errs.Errors = append(errs.Errors, e.Errors...)
 
 	if s.Provider != nil {
-		_, prov, err := addElement[Resource, schema.ResourceSpec](
+		_, prov, associated, err := addElement[Resource, schema.ResourceSpec](
 			info.PackageName, registerDerivative, s.ModuleMap, s.Provider)
+		for _, alias := range associated.Aliases {
+			aliases[tokenFromType(alias.Token)] = tokens.Type("pulumi:providers:" + pkg.Name)
+		}
 		if err != nil {
 			errs.Errors = append(errs.Errors, err)
 		}
@@ -290,42 +305,90 @@ func (s *state) generateSchema(ctx p.Context) (schema.PackageSpec, error) {
 	if err := errs.ErrorOrNil(); err != nil {
 		return schema.PackageSpec{}, err
 	}
+
+	pkg = transformTypes(pkg, func(typ schema.TypeSpec) schema.TypeSpec {
+		const internalRefPrefx = "#/types/"
+		if !strings.HasPrefix(typ.Ref, internalRefPrefx) {
+			return typ
+		}
+		changeTo, ok := aliases[tokenFromType(tokens.Type(typ.Ref))]
+		if !ok {
+			return typ
+		}
+		typ.Ref = "#/resources/" + assignTo(changeTo, pkg.Name, s.ModuleMap).String()
+		return typ
+	})
+
+	for from := range aliases {
+		// We delete aliased resources and types, since they are unreferencable
+		token := from.reconstruct(tokens.PackageName(pkg.Name)).String()
+		delete(pkg.Resources, token)
+		delete(pkg.Types, token)
+	}
+
 	return pkg, nil
 }
 
 type canGetSchema[T any] interface {
 	GetToken() (tokens.Type, error)
-	GetSchema(RegisterDerivativeType) (T, error)
+	GetSchema(RegisterDerivativeType) (T, Associated, error)
+}
+
+type internalToken struct {
+	mod  tokens.ModuleName
+	name tokens.TypeName
+}
+
+func tokenFromType(tk tokens.Type) internalToken {
+	return internalToken{
+		mod:  tk.Module().Name(),
+		name: tk.Name(),
+	}
+}
+
+func (it internalToken) reconstruct(pkg tokens.PackageName) tokens.Type {
+	return tokens.NewTypeToken(
+		tokens.NewModuleToken(
+			tokens.NewPackageToken(pkg),
+			it.mod,
+		),
+		it.name,
+	)
 }
 
 func addElements[T canGetSchema[S], S any](els []T, m map[string]S,
 	pkgName string, reg RegisterDerivativeType,
-	modMap map[tokens.ModuleName]tokens.ModuleName) multierror.Error {
+	modMap map[tokens.ModuleName]tokens.ModuleName,
+) (map[internalToken]tokens.Type, multierror.Error) {
+	assoc := make(map[internalToken]tokens.Type, len(els))
 	errs := multierror.Error{}
 	for _, f := range els {
-		tk, element, err := addElement[T, S](pkgName, reg, modMap, f)
+		tk, element, associated, err := addElement[T, S](pkgName, reg, modMap, f)
+		for _, alias := range associated.Aliases {
+			assoc[tokenFromType(alias.Token)] = tk
+		}
 		if err != nil {
 			errs.Errors = append(errs.Errors, err)
 			continue
 		}
 		m[tk.String()] = element
 	}
-	return errs
+	return assoc, errs
 }
 
 func addElement[T canGetSchema[S], S any](pkgName string, reg RegisterDerivativeType,
-	modMap map[tokens.ModuleName]tokens.ModuleName, f T) (tokens.Type, S, error) {
+	modMap map[tokens.ModuleName]tokens.ModuleName, f T) (tokens.Type, S, Associated, error) {
 	var s S
 	tk, err := f.GetToken()
 	if err != nil {
-		return "", s, err
+		return "", s, Associated{}, err
 	}
 	tk = assignTo(tk, pkgName, modMap)
-	fun, err := f.GetSchema(reg)
+	fun, associated, err := f.GetSchema(reg)
 	if err != nil {
-		return "", s, fmt.Errorf("failed to get schema for '%s': %w", tk, err)
+		return "", s, Associated{}, fmt.Errorf("failed to get schema for '%s': %w", tk, err)
 	}
-	return tk, renamePackage(fun, pkgName, modMap), nil
+	return tk, renamePackage(fun, pkgName, modMap), associated, nil
 }
 
 func assignTo(tk tokens.Type, pkg string, modMap map[tokens.ModuleName]tokens.ModuleName) tokens.Type {
@@ -356,9 +419,7 @@ func fixReference(ref, pkg string, modMap map[tokens.ModuleName]tokens.ModuleNam
 	return kind + string(assignTo(tk, pkg, modMap))
 }
 
-// renamePackage sets internal package references to point to the package with the name
-// `pkg`.
-func renamePackage[T any](typ T, pkg string, modMap map[tokens.ModuleName]tokens.ModuleName) T {
+func transformTypes[T any](typ T, transform func(schema.TypeSpec) schema.TypeSpec) T {
 	var rename func(reflect.Value)
 	rename = func(v reflect.Value) {
 		switch v.Kind() {
@@ -369,9 +430,8 @@ func renamePackage[T any](typ T, pkg string, modMap map[tokens.ModuleName]tokens
 			rename(v.Elem())
 		case reflect.Struct:
 			if v.Type() == reflect.TypeOf(schema.TypeSpec{}) {
-				field := v.FieldByName("Ref")
-				rewritten := fixReference(field.String(), pkg, modMap)
-				field.SetString(rewritten)
+				rewritten := transform(v.Interface().(schema.TypeSpec))
+				v.Set(reflect.ValueOf(rewritten))
 			}
 			for _, f := range reflect.VisibleFields(v.Type()) {
 				f := v.FieldByIndex(f.Index)
@@ -399,4 +459,14 @@ func renamePackage[T any](typ T, pkg string, modMap map[tokens.ModuleName]tokens
 	v := reflect.ValueOf(t)
 	rename(v)
 	return *t
+
+}
+
+// renamePackage sets internal package references to point to the package with the name
+// `pkg`.
+func renamePackage[T any](typ T, pkg string, modMap map[tokens.ModuleName]tokens.ModuleName) T {
+	return transformTypes(typ, func(t schema.TypeSpec) schema.TypeSpec {
+		t.Ref = fixReference(t.Ref, pkg, modMap)
+		return t
+	})
 }