Update NodeMapper to handle Trait sourceLocation

smithy-model/src/main/java/software/amazon/smithy/model/node/DefaultNodeDeserializers.java

@@ -374,7 +374,8 @@ private static boolean isBeanOrBuilderSetter(Method method, Class<?> type, Strin
             }
 
             // Must either return the target class itself (like a builder) or void.
-            if (method.getReturnType() != void.class && method.getReturnType() != type) {
+            // Ideally we should attempt to resolve any generics and make an assertion of the concrete type.
+            if (!(method.getReturnType() == void.class || method.getReturnType().isAssignableFrom(type))) {
                 return false;
             }
 
@@ -431,6 +432,7 @@ private static Class<?> resolveClassFromType(Type type) {
             } else if (type instanceof WildcardType) {
                 return resolveClassFromType(((WildcardType) type).getUpperBounds()[0]);
             } else if (type instanceof TypeVariable<?>) {
+                // TODO: implement this to enable improved builder detection
                 throw new IllegalArgumentException("TypeVariable targets are not implemented: " + type);
             } else if (type instanceof GenericArrayType) {
                 throw new IllegalArgumentException("GenericArrayType targets are not implemented: " + type);

smithy-model/src/main/java/software/amazon/smithy/model/node/DefaultNodeSerializers.java

@@ -65,6 +65,7 @@ public Node serialize(ToNode value, Set<Object> serializedObjects, NodeMapper ma
                 return null;
             }
 
+            // TODO: make sure every instance of `toNode` is setting this
             return value.toNode();
         }
     };
@@ -362,7 +363,13 @@ public Node serialize(Object value, Set<Object> serializedObjects, NodeMapper ma
             // multiple times (like in a List<T>).
             serializedObjects.remove(value);
 
-            return new ObjectNode(mappings, SourceLocation.NONE);
+            // Pass on the source location if it is present.
+            SourceLocation sourceLocation = SourceLocation.NONE;
+            if (value instanceof FromSourceLocation) {
+                sourceLocation = ((FromSourceLocation) value).getSourceLocation();
+            }
+
+            return new ObjectNode(mappings, sourceLocation);
         }
 
         private boolean canSerialize(NodeMapper mapper, Node value) {

smithy-model/src/test/java/software/amazon/smithy/model/node/NodeMapperTest.java

@@ -14,7 +14,6 @@
 import java.net.URI;
 import java.net.URL;
 import java.nio.file.Path;
-import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
@@ -34,9 +33,14 @@
 import software.amazon.smithy.model.FromSourceLocation;
 import software.amazon.smithy.model.SourceLocation;
 import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.model.traits.AbstractTrait;
+import software.amazon.smithy.model.traits.AbstractTrait.Provider;
+import software.amazon.smithy.model.traits.AbstractTraitBuilder;
 import software.amazon.smithy.model.traits.RangeTrait;
+import software.amazon.smithy.model.traits.Trait;
 import software.amazon.smithy.utils.ListUtils;
 import software.amazon.smithy.utils.SmithyBuilder;
+import software.amazon.smithy.utils.ToSmithyBuilder;
 
 public class NodeMapperTest {
     @Test
@@ -1198,15 +1202,46 @@ public void deserializesIntoValueWithSourceLocation() {
     }
 
     @Test
-    public void doesNotSerializeSourceLocation() {
+    public void deserializesIntoTraitWithSourceLocation() {
+        NodeMapper mapper = new NodeMapper();
+        Map<StringNode, Node> mappings = new HashMap<>();
+        mappings.put(StringNode.from("foo"), StringNode.from("foo"));
+        SourceLocation sourceLocation = new SourceLocation("/abc/def");
+        Node node = new ObjectNode(mappings, sourceLocation);
+        SourceLocationBearerTrait trait = mapper.deserialize(node, SourceLocationBearerTrait.class);
+
+        assertThat(trait.getSourceLocation(), equalTo(sourceLocation));
+        assertThat(trait.getFoo(), equalTo("foo"));
+    }
+
+    @Test
+    public void serializesSourceLocationFromValue() {
         NodeMapper mapper = new NodeMapper();
         HasSourceLocation hs = new HasSourceLocation();
-        hs.setSourceLocation(new SourceLocation("/foo/baz"));
+        SourceLocation sourceLocation = new SourceLocation("/foo/baz");
+        hs.setSourceLocation(sourceLocation);
         hs.setFoo("hi");
         Node result = mapper.serialize(hs);
 
         assertThat(result.expectObjectNode().getStringMap(), hasKey("foo"));
+        // The sourceLocation needs to be set on the node, not as a key.
+        assertThat(result.expectObjectNode().getStringMap(), not(hasKey("sourceLocation")));
+        assertThat(result.getSourceLocation(), equalTo(sourceLocation));
+    }
+
+    @Test
+    public void serializesSourceLocationFromTrait() {
+        SourceLocation sourceLocation = new SourceLocation("/foo/baz");
+        SourceLocationBearerTrait trait = SourceLocationBearerTrait.builder()
+                .sourceLocation(sourceLocation)
+                .foo("foo")
+                .build();
+        Node result = trait.createNode();
+
+        assertThat(result.expectObjectNode().getStringMap(), hasKey("foo"));
+        // The sourceLocation needs to be set on the node, not as a key.
         assertThat(result.expectObjectNode().getStringMap(), not(hasKey("sourceLocation")));
+        assertThat(result.getSourceLocation(), equalTo(sourceLocation));
     }
 
     private static class HasSourceLocation implements FromSourceLocation {
@@ -1231,6 +1266,50 @@ public String getFoo() {
         }
     }
 
+    private static class SourceLocationBearerTrait extends AbstractTrait implements ToSmithyBuilder<SourceLocationBearerTrait> {
+        public static final ShapeId ID = ShapeId.from("smithy.test#sourceLocationBearer");
+        private final String foo;
+
+        public SourceLocationBearerTrait(Builder builder) {
+            super(ID, builder.getSourceLocation());
+            this.foo = builder.foo;
+        }
+
+        public String getFoo() {
+            return foo;
+        }
+
+        @Override
+        protected Node createNode() {
+            NodeMapper mapper = new NodeMapper();
+            mapper.disableToNodeForClass(SourceLocationBearerTrait.class);
+            return mapper.serialize(this);
+        }
+
+        @Override
+        public SmithyBuilder<SourceLocationBearerTrait> toBuilder() {
+            return new Builder().sourceLocation(getSourceLocation());
+        }
+
+        public static Builder builder() {
+            return new Builder();
+        }
+
+        public static final class Builder extends AbstractTraitBuilder<SourceLocationBearerTrait, Builder> {
+            private String foo;
+
+            @Override
+            public SourceLocationBearerTrait build() {
+                return new SourceLocationBearerTrait(this);
+            }
+
+            public Builder foo(String foo) {
+                this.foo = foo;
+                return this;
+            }
+        }
+    }
+
     @Test
     public void deserializesEnums() {
         NodeMapper mapper = new NodeMapper();