Use @Export.refs in codegen

part of #629

pkg/codegen/java/gen.go

@@ -965,18 +965,12 @@ func (mod *modContext) genResource(ctx *classFileContext, r *schema.Resource, ar
 			fprintf(w, "     */\n")
 		}
 
-		outputExportParameters := strings.Join(
-			propertyType.ParameterTypesTransformed(func(ts TypeShape) string {
-				return ts.ToCodeClassLiteral(ctx.imports)
-			}),
-			", ",
-		)
-		outputExportType := propertyType.ToCodeClassLiteral(ctx.imports)
+		outputExportRefs, outputExportTree := propertyType.ToTree(ctx.imports)
 		outputParameterType := propertyType.ToCodeCommentedAnnotations(ctx.imports)
 		printObsoleteAttribute(ctx, prop.DeprecationMessage, "    ")
 		fprintf(w,
-			"    @%s(name=\"%s\", type=%s, parameters={%s})\n",
-			ctx.ref(names.Export), wireName, outputExportType, outputExportParameters)
+			"    @%s(name=\"%s\", refs={%s}, tree=\"%s\")\n",
+			ctx.ref(names.Export), wireName, outputExportRefs, outputExportTree)
 		fprintf(w,
 			"    private %s<%s> %s;\n", ctx.imports.Ref(names.Output), outputParameterType, propertyName)
 		fprintf(w, "\n")

pkg/codegen/java/typeshape.go

@@ -3,6 +3,7 @@
 package java
 
 import (
+	"encoding/json"
 	"fmt"
 	"strings"
 
@@ -84,6 +85,10 @@ func (ts TypeShape) ToCodeCommentedAnnotations(imports *names.Imports) string {
 }
 
 func (ts TypeShape) ToCodeWithOptions(imports *names.Imports, opts TypeShapeStringOptions) string {
+	// guard against unexpected output, should not happen outside of tests
+	if ts.Type.String() == "" {
+		return ""
+	}
 	var annotationsString string
 	if !opts.SkipAnnotations {
 		annotationsString = strings.Join(ts.Annotations, " ")
@@ -120,6 +125,65 @@ func (ts TypeShape) ToCodeWithOptions(imports *names.Imports, opts TypeShapeStri
 	return fmt.Sprintf("%s%s%s%s", annotationsString, imports.Ref(ts.Type), parametersString, classLiteral)
 }
 
+func (ts TypeShape) ToTree(imports *names.Imports) (string, string) {
+	refs := []string{}
+	refs, tree := ts.toTreeInternal(refs, imports)
+	treeJSON, err := json.Marshal(tree)
+	if err != nil {
+		panic(err)
+	}
+	return strings.Join(refs, ","), string(treeJSON)
+}
+
+func (ts TypeShape) toTreeInternal(refs []string, imports *names.Imports) ([]string, []interface{}) {
+	refIndex := func(classLiteral string) (int, bool) {
+		return sliceIndex(len(refs), func(i int) bool {
+			return refs[i] == classLiteral
+		})
+	}
+	requireRefIndex := func(classLiteral string) int {
+		classRef, ok := refIndex(classLiteral)
+		if !ok {
+			panic(fmt.Errorf(
+				"expected a class reference index for a class literal: '%s'", classLiteral,
+			))
+		}
+		return classRef
+	}
+	rootClassLiteral := ts.ToCodeClassLiteral(imports)
+	// guard against unexpected output, should not happen outside of tests
+	if rootClassLiteral == "" {
+		return []string{}, []interface{}{}
+	}
+	_, ok := refIndex(rootClassLiteral)
+	if !ok {
+		refs = append(refs, rootClassLiteral)
+	}
+	classRef := requireRefIndex(rootClassLiteral)
+	tree := []interface{}{classRef}
+
+	for _, param := range ts.Parameters {
+		subRefs, subTree := param.toTreeInternal(refs, imports)
+		refs = subRefs
+		if len(subTree) == 1 {
+			tree = append(tree, subTree[0])
+		} else {
+			tree = append(tree, subTree)
+		}
+	}
+
+	return refs, tree
+}
+
+func sliceIndex(limit int, predicate func(i int) bool) (int, bool) {
+	for i := 0; i < limit; i++ {
+		if predicate(i) {
+			return i, true
+		}
+	}
+	return -1, false
+}
+
 func (ts TypeShape) ParameterTypes(imports *names.Imports) []string {
 	return ts.ParameterTypesTransformed(func(ts TypeShape) string {
 		return imports.Ref(ts.Type)

pkg/codegen/java/typeshape_test.go

@@ -0,0 +1,105 @@
+package java
+
+import (
+	"github.com/pulumi/pulumi-java/pkg/codegen/java/names"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestTypeShape_ToTree(t *testing.T) {
+	tests := []struct {
+		name    string
+		ts      TypeShape
+		refs    string
+		tree    string
+		imports string
+	}{
+		{
+			name:    "Empty",
+			ts:      TypeShape{},
+			refs:    "",
+			tree:    "[]",
+			imports: "",
+		},
+		{
+			name: "String",
+			ts: TypeShape{
+				Type: names.String,
+			},
+			refs:    "String.class",
+			tree:    "[0]",
+			imports: "import java.lang.String;",
+		},
+		{
+			name: "List<String>",
+			ts: TypeShape{
+				Type:       names.List,
+				Parameters: []TypeShape{{Type: names.String}},
+			},
+			refs:    "List.class,String.class",
+			tree:    "[0,1]",
+			imports: "import java.lang.String;\nimport java.util.List;",
+		},
+		{
+			name: "Either<Integer,String>",
+			ts: TypeShape{
+				Type: names.Either,
+				Parameters: []TypeShape{
+					{Type: names.Integer},
+					{Type: names.String},
+				},
+			},
+			refs:    "Either.class,Integer.class,String.class",
+			tree:    "[0,1,2]",
+			imports: "import com.pulumi.core.Either;\nimport java.lang.Integer;\nimport java.lang.String;",
+		},
+		{
+			name: "List<List<List<String>>>",
+			ts: TypeShape{
+				Type: names.List,
+				Parameters: []TypeShape{
+					{Type: names.List, Parameters: []TypeShape{
+						{Type: names.List, Parameters: []TypeShape{
+							{Type: names.String},
+						},
+						},
+					}},
+				},
+			},
+			refs:    "List.class,String.class",
+			tree:    "[0,[0,[0,1]]]",
+			imports: "import java.lang.String;\nimport java.util.List;",
+		},
+		{
+			name: "Either<Map<String,List<String>>,Map<String,List<String>>>",
+			ts: TypeShape{
+				Type: names.Either,
+				Parameters: []TypeShape{
+					{Type: names.Map, Parameters: []TypeShape{
+						{Type: names.String},
+						{Type: names.List, Parameters: []TypeShape{{Type: names.String}}},
+					}},
+					{Type: names.Map, Parameters: []TypeShape{
+						{Type: names.List, Parameters: []TypeShape{{Type: names.String}}},
+						{Type: names.String},
+					}},
+				},
+			},
+			refs:    "Either.class,Map.class,String.class,List.class",
+			tree:    "[0,[1,2,[3,2]],[1,[3,2],2]]",
+			imports: "import com.pulumi.core.Either;\nimport java.lang.String;\nimport java.util.List;\nimport java.util.Map;",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			pkg := names.Ident("com").FQN().Dot("pulumi").Dot("tests")
+			cls := names.Ident("TypeShapeTest")
+			imports := names.NewImports(pkg, cls)
+			gotRefs, gotTree := tt.ts.ToTree(imports)
+			gotImports := imports.ImportCode()
+			assert.Equal(t, tt.refs, gotRefs)
+			assert.Equal(t, tt.tree, gotTree)
+			assert.Equal(t, tt.imports, gotImports)
+		})
+	}
+}