Update NodeMapper to handle Trait sourceLocation (#865)

* Update NodeMapper to handle Trait sourceLocation

This updates NodeMapper so that it can properly set source location
on traits that use the AbstractTraitBuilder. Unfortunately the fix
implemented here is a hack that looks for that specific case, and
the same root cause may bite us elsewhere.

The true root cause is that our default deserializer is unable to
resolve generics. AbstractTraitBuilder generically defines the
return value of the setter `sourceLocation` to `B extends
AbstractTraitBuilder`. When we use reflection to look at the return
type, all we see is `AbstractTraitBuilder`. We *can* also get the
type variable, but actually resolving it will take a decent amount
of effort as that functionality isn't built in.

* Make NodeMapper pass along source location

This updates the serializer for NodeMapper to pass along the
source location.

* Check if builder return type is self-assignable

This updates the builder detection in the node deserializer to check
to see if the builder type is assignable to a given method's return
type. This removes the need for a hack enabling source location, but
it's still not perfect. Ideally we should attempt to resolve any
generic return types.

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();