Fix trait codegen for structure's string member with @idRef

smithy-trait-codegen/src/it/java/software/amazon/smithy/traitcodegen/test/CreatesTraitTest.java

@@ -41,6 +41,7 @@
 import com.example.traits.structures.BasicAnnotationTrait;
 import com.example.traits.structures.NestedA;
 import com.example.traits.structures.NestedB;
+import com.example.traits.structures.StructWithIdrefMemberTrait;
 import com.example.traits.structures.StructWithListOfMapTrait;
 import com.example.traits.structures.StructWithUniqueItemsListTrait;
 import com.example.traits.structures.StructureTrait;
@@ -200,6 +201,12 @@ static Stream<Arguments> createTraitTests() {
                                 .addItems(SetUtils.of("a", "b", "c"))
                                 .build()
                                 .toNode()),
+                Arguments.of(StructWithIdrefMemberTrait.ID,
+                        StructWithIdrefMemberTrait.builder()
+                                .idRefMemberA(ShapeId.from("test.smithy.traitcodegen#a"))
+                                .idRefMemberB(ShapeId.from("test.smithy.traitcodegen#b"))
+                                .build()
+                                .toNode()),
                 // Timestamps
                 Arguments.of(TimestampTrait.ID, Node.from("1985-04-12T23:20:50.52Z")),
                 Arguments.of(DateTimeTimestampTrait.ID, Node.from("1985-04-12T23:20:50.52Z")),

smithy-trait-codegen/src/it/java/software/amazon/smithy/traitcodegen/test/LoadsFromModelTest.java

@@ -48,6 +48,7 @@
 import com.example.traits.structures.BasicAnnotationTrait;
 import com.example.traits.structures.NestedA;
 import com.example.traits.structures.NestedB;
+import com.example.traits.structures.StructWithIdrefMemberTrait;
 import com.example.traits.structures.StructWithListOfMapTrait;
 import com.example.traits.structures.StructWithUniqueItemsListTrait;
 import com.example.traits.structures.StructureTrait;
@@ -324,6 +325,13 @@ static Stream<Arguments> loadsModelTests() {
                                 Optional.of("a"),
                                 "getItems",
                                 Optional.of(SetUtils.of(SetUtils.of("b", "c"), SetUtils.of("d", "e"))))),
+                Arguments.of("structures/struct-with-idref-member-trait.smithy",
+                        StructWithIdrefMemberTrait.class,
+                        MapUtils.of(
+                                "getIdRefMemberA",
+                                Optional.of(ShapeId.from("test.smithy.traitcodegen#a")),
+                                "getIdRefMemberB",
+                                Optional.of(ShapeId.from("test.smithy.traitcodegen#b")))),
                 // Timestamps
                 Arguments.of("timestamps/struct-with-nested-timestamps.smithy",
                         StructWithNestedTimestampsTrait.class,

smithy-trait-codegen/src/it/resources/software/amazon/smithy/traitcodegen/test/structures/struct-with-idref-member-trait.smithy

@@ -0,0 +1,12 @@
+$version: "2.0"
+
+namespace test.smithy.traitcodegen
+
+use test.smithy.traitcodegen.structures#StructWithIdrefMember
+
+@StructWithIdrefMember(idRefMemberA: "test.smithy.traitcodegen#a", idRefMemberB: "test.smithy.traitcodegen#b")
+structure myStruct {}
+
+structure a {}
+
+structure b {}

smithy-trait-codegen/src/main/java/software/amazon/smithy/traitcodegen/generators/BuilderGenerator.java

@@ -37,6 +37,7 @@
 import software.amazon.smithy.model.shapes.UnionShape;
 import software.amazon.smithy.model.traits.AbstractTraitBuilder;
 import software.amazon.smithy.model.traits.DefaultTrait;
+import software.amazon.smithy.model.traits.IdRefTrait;
 import software.amazon.smithy.model.traits.StringListTrait;
 import software.amazon.smithy.model.traits.TimestampFormatTrait;
 import software.amazon.smithy.model.traits.TraitDefinition;
@@ -371,6 +372,10 @@ public Void mapShape(MapShape shape) {
 
         @Override
         public Void memberShape(MemberShape shape) {
+            if (shape.hasTrait(IdRefTrait.ID)) {
+                this.getDefault(shape);
+                return null;
+            }
             return model.expectShape(shape.getTarget()).accept(this);
         }
     }

smithy-trait-codegen/src/main/java/software/amazon/smithy/traitcodegen/generators/FromNodeGenerator.java

@@ -35,6 +35,7 @@
 import software.amazon.smithy.model.shapes.StructureShape;
 import software.amazon.smithy.model.shapes.TimestampShape;
 import software.amazon.smithy.model.shapes.UnionShape;
+import software.amazon.smithy.model.traits.IdRefTrait;
 import software.amazon.smithy.model.traits.TimestampFormatTrait;
 import software.amazon.smithy.model.traits.UniqueItemsTrait;
 import software.amazon.smithy.traitcodegen.SymbolProperties;
@@ -242,6 +243,14 @@ private MemberGenerator(MemberShape member) {
 
         @Override
         public Void memberShape(MemberShape shape) {
+            if (shape.hasTrait(IdRefTrait.ID)) {
+                writer.writeInline(memberPrefix + "Member($1S, n -> $3C, builder::$2L)",
+                        fieldName,
+                        memberName,
+                        (Runnable) () -> shape
+                                .accept(new FromNodeMapperVisitor(writer, model, "n", 1, symbolProvider)));
+                return null;
+            }
             return model.expectShape(shape.getTarget()).accept(this);
         }
 

smithy-trait-codegen/src/test/java/software/amazon/smithy/traitcodegen/TraitCodegenPluginTest.java

@@ -26,7 +26,7 @@
 import software.amazon.smithy.model.node.ObjectNode;
 
 public class TraitCodegenPluginTest {
-    private static final int EXPECTED_NUMBER_OF_FILES = 67;
+    private static final int EXPECTED_NUMBER_OF_FILES = 68;
 
     private MockManifest manifest;
     private Model model;

smithy-trait-codegen/src/test/resources/META-INF/smithy/manifest

@@ -42,6 +42,7 @@ structures/annotation-trait.smithy
 structures/structure-trait.smithy
 structures/struct-with-listofmap-trait.smithy
 structures/struct-with-uniqueitems-list-trait.smithy
+structures/struct-with-idref-member-trait.smithy
 timestamps/date-time-format-timestamp-trait.smithy
 timestamps/epoch-seconds-format-timestamp-trait.smithy
 timestamps/http-date-format-timestamp-trait.smithy

smithy-trait-codegen/src/test/resources/META-INF/smithy/structures/struct-with-idref-member-trait.smithy

@@ -0,0 +1,13 @@
+$version: "2.0"
+
+namespace test.smithy.traitcodegen.structures
+
+@trait
+structure StructWithIdrefMember {
+    @idRef(selector: "structure")
+    idRefMemberA: String
+    idRefMemberB: IdRefMemberB
+}
+
+@idRef(selector: "structure")
+string IdRefMemberB