Fix sourceLocation propagation in Traits

smithy-aws-apigateway-traits/src/main/java/software/amazon/smithy/aws/apigateway/traits/AuthorizersTrait.java

@@ -101,14 +101,15 @@ public Map<String, AuthorizerDefinition> getAuthorizers() {
 
     @Override
     public Builder toBuilder() {
-        return builder().authorizers(authorizers);
+        return builder().sourceLocation(getSourceLocation()).authorizers(authorizers);
     }
 
     @Override
     protected Node createNode() {
         return authorizers.entrySet().stream()
                 .sorted(Comparator.comparing(Map.Entry::getKey))
-                .collect(ObjectNode.collectStringKeys(Map.Entry::getKey, Map.Entry::getValue));
+                .collect(ObjectNode.collectStringKeys(Map.Entry::getKey, Map.Entry::getValue))
+                .toBuilder().sourceLocation(getSourceLocation()).build();
     }
 
     /**

smithy-aws-apigateway-traits/src/main/java/software/amazon/smithy/aws/apigateway/traits/IntegrationTrait.java

@@ -342,6 +342,7 @@ protected ObjectNode createNode() {
     @Override
     public Builder toBuilder() {
         return builder()
+                .sourceLocation(getSourceLocation())
                 .type(type)
                 .uri(uri)
                 .credentials(credentials)

smithy-aws-apigateway-traits/src/main/java/software/amazon/smithy/aws/apigateway/traits/MockIntegrationTrait.java

@@ -163,6 +163,7 @@ protected ObjectNode createNode() {
     @Override
     public Builder toBuilder() {
         return builder()
+                .sourceLocation(getSourceLocation())
                 .passThroughBehavior(passThroughBehavior)
                 .contentHandling(contentHandling)
                 .requestParameters(requestParameters)

smithy-aws-cloudformation-traits/src/main/java/software/amazon/smithy/aws/cloudformation/traits/CfnResourceTrait.java

@@ -76,7 +76,7 @@ protected Node createNode() {
 
     @Override
     public SmithyBuilder<CfnResourceTrait> toBuilder() {
-        return builder().name(name).additionalSchemas(additionalSchemas);
+        return builder().sourceLocation(getSourceLocation()).name(name).additionalSchemas(additionalSchemas);
     }
 
     public static final class Provider extends AbstractTrait.Provider {

smithy-aws-iam-traits/src/main/java/software/amazon/smithy/aws/iam/traits/DefineConditionKeysTrait.java

@@ -53,7 +53,7 @@ public Provider() {
 
         @Override
         public Trait createTrait(ShapeId target, Node value) {
-            Builder builder = builder();
+            Builder builder = builder().sourceLocation(value);
             for (Map.Entry<StringNode, Node> entry : value.expectObjectNode().getMembers().entrySet()) {
                 ConditionKeyDefinition definition = ConditionKeyDefinition.fromNode(
                         entry.getValue().expectObjectNode());
@@ -87,7 +87,8 @@ protected Node createNode() {
         return conditionKeys.entrySet().stream()
                 .map(entry -> new AbstractMap.SimpleImmutableEntry<>(
                         Node.from(entry.getKey()), entry.getValue().toNode()))
-                .collect(ObjectNode.collect(Map.Entry::getKey, Map.Entry::getValue));
+                .collect(ObjectNode.collect(Map.Entry::getKey, Map.Entry::getValue))
+                .toBuilder().sourceLocation(getSourceLocation()).build();
     }
 
     @Override

smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/ArnReferenceTrait.java

@@ -54,7 +54,7 @@ public Provider() {
         @Override
         public Trait createTrait(ShapeId target, Node value) {
             ObjectNode objectNode = value.expectObjectNode();
-            Builder builder = builder();
+            Builder builder = builder().sourceLocation(value);
             objectNode.getStringMember(TYPE)
                     .map(StringNode::getValue)
                     .ifPresent(builder::type);
@@ -110,10 +110,12 @@ public Builder toBuilder() {
 
     @Override
     protected Node createNode() {
-        return Node.objectNode()
+        return Node.objectNodeBuilder()
+                .sourceLocation(getSourceLocation())
                 .withOptionalMember(TYPE, getType().map(Node::from))
                 .withOptionalMember(SERVICE, getService().map(ShapeId::toString).map(Node::from))
-                .withOptionalMember(RESOURCE, getResource().map(ShapeId::toString).map(Node::from));
+                .withOptionalMember(RESOURCE, getResource().map(ShapeId::toString).map(Node::from))
+                .build();
     }
 
     /** Builder for {@link ArnReferenceTrait}. */

smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/ArnTrait.java

@@ -70,8 +70,8 @@ public Provider() {
 
         @Override
         public Trait createTrait(ShapeId target, Node value) {
-            Builder builder = builder();
             ObjectNode objectNode = value.expectObjectNode();
+            Builder builder = builder().sourceLocation(value);
             builder.template(objectNode.expectStringMember(TEMPLATE).getValue());
             builder.absolute(objectNode.getBooleanMemberOrDefault(ABSOLUTE));
             builder.noRegion(objectNode.getBooleanMemberOrDefault(NO_REGION));
@@ -133,11 +133,13 @@ public List<String> getLabels() {
 
     @Override
     protected Node createNode() {
-        return Node.objectNode()
+        return Node.objectNodeBuilder()
+                .sourceLocation(getSourceLocation())
                 .withMember(TEMPLATE, Node.from(getTemplate()))
                 .withMember(ABSOLUTE, Node.from(isAbsolute()))
                 .withMember(NO_ACCOUNT, Node.from(isNoAccount()))
-                .withMember(NO_REGION, Node.from(isNoRegion()));
+                .withMember(NO_REGION, Node.from(isNoRegion()))
+                .build();
     }
 
     @Override

smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/ServiceTrait.java

@@ -61,7 +61,7 @@ public Provider() {
         @Override
         public Trait createTrait(ShapeId target, Node value) {
             ObjectNode objectNode = value.expectObjectNode();
-            Builder builder = builder();
+            Builder builder = builder().sourceLocation(value);
             String sdkId = objectNode.getStringMember("sdkId")
                     .map(StringNode::getValue)
                     .orElseThrow(() -> new SourceException(String.format(
@@ -169,12 +169,14 @@ public Builder toBuilder() {
 
     @Override
     protected Node createNode() {
-        return Node.objectNode()
+        return Node.objectNodeBuilder()
+                .sourceLocation(getSourceLocation())
                 .withMember("sdkId", Node.from(sdkId))
                 .withMember("arnNamespace", Node.from(getArnNamespace()))
                 .withMember("cloudFormationName", Node.from(getCloudFormationName()))
                 .withMember("cloudTrailEventSource", Node.from(getCloudTrailEventSource()))
-                .withMember("endpointPrefix", Node.from(getEndpointPrefix()));
+                .withMember("endpointPrefix", Node.from(getEndpointPrefix()))
+                .build();
     }
 
     /** Builder for {@link ServiceTrait}. */

smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/auth/CognitoUserPoolsTrait.java

@@ -53,6 +53,7 @@ public Provider() {
         public Trait createTrait(ShapeId target, Node value) {
             ObjectNode objectNode = value.expectObjectNode();
             return builder()
+                    .sourceLocation(value)
                     .providerArns(objectNode.expectArrayMember(PROVIDER_ARNS).getElementsAs(StringNode::getValue))
                     .build();
         }
@@ -76,13 +77,15 @@ public List<String> getProviderArns() {
 
     @Override
     public Builder toBuilder() {
-        return builder().providerArns(providerArns);
+        return builder().sourceLocation(getSourceLocation()).providerArns(providerArns);
     }
 
     @Override
     protected Node createNode() {
-        return Node.objectNode()
-                .withMember(PROVIDER_ARNS, providerArns.stream().map(Node::from).collect(ArrayNode.collect()));
+        return Node.objectNodeBuilder()
+                .sourceLocation(getSourceLocation())
+                .withMember(PROVIDER_ARNS, providerArns.stream().map(Node::from).collect(ArrayNode.collect()))
+                .build();
     }
 
     /** Builder for {@link CognitoUserPoolsTrait}. */

smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/auth/SigV4Trait.java

@@ -59,7 +59,7 @@ public Builder toBuilder() {
 
     @Override
     protected Node createNode() {
-        return Node.objectNode().withMember(NAME, getName());
+        return Node.objectNodeBuilder().sourceLocation(getSourceLocation()).withMember(NAME, getName()).build();
     }
 
     public static final class Builder extends AbstractTraitBuilder<SigV4Trait, Builder> {

smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/clientendpointdiscovery/ClientDiscoveredEndpointTrait.java

@@ -56,12 +56,15 @@ public boolean isRequired() {
 
     @Override
     protected Node createNode() {
-        return Node.objectNode().withMember(REQUIRED, Node.from(isRequired()));
+        return Node.objectNodeBuilder()
+                .sourceLocation(getSourceLocation())
+                .withMember(REQUIRED, Node.from(isRequired()))
+                .build();
     }
 
     @Override
     public SmithyBuilder<ClientDiscoveredEndpointTrait> toBuilder() {
-        return builder().required(required);
+        return builder().sourceLocation(getSourceLocation()).required(required);
     }
 
     /** Builder for {@link ClientDiscoveredEndpointTrait}. */
@@ -90,6 +93,7 @@ public Provider() {
         public ClientDiscoveredEndpointTrait createTrait(ShapeId target, Node value) {
             ObjectNode objectNode = value.expectObjectNode();
             return builder()
+                    .sourceLocation(value)
                     .required(objectNode.getBooleanMemberOrDefault(REQUIRED, true))
                     .build();
         }

smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/clientendpointdiscovery/ClientEndpointDiscoveryTrait.java

@@ -83,15 +83,18 @@ public ShapeId getError() {
 
     @Override
     protected Node createNode() {
-        return Node.objectNode()
+        return Node.objectNodeBuilder()
+                .sourceLocation(getSourceLocation())
                 .withMember(OPERATION, Node.from(getOperation().toString()))
-                .withOptionalMember(ERROR, getOptionalError().map(error -> Node.from(error.toString())));
+                .withOptionalMember(ERROR, getOptionalError().map(error -> Node.from(error.toString())))
+                .build();
     }
 
     @Override
     public Builder toBuilder() {
         return new Builder()
-                .operation(getOperation())
+                .sourceLocation(getSourceLocation())
+                .operation(operation)
                 .error(error);
     }
 
@@ -144,7 +147,9 @@ public Provider() {
         @Override
         public ClientEndpointDiscoveryTrait createTrait(ShapeId target, Node value) {
             ObjectNode objectNode = value.expectObjectNode();
-            Builder builder = builder().operation(objectNode.expectStringMember(OPERATION).expectShapeId());
+            Builder builder = builder()
+                    .sourceLocation(value)
+                    .operation(objectNode.expectStringMember(OPERATION).expectShapeId());
             objectNode.getStringMember(ERROR).ifPresent(error -> builder.error(error.expectShapeId()));
             return builder.build();
         }

smithy-aws-traits/src/main/java/software/amazon/smithy/aws/traits/protocols/AwsProtocolTrait.java

@@ -68,6 +68,7 @@ public List<String> getEventStreamHttp() {
     @Override
     protected Node createNode() {
         ObjectNode.Builder builder = Node.objectNodeBuilder();
+        builder.sourceLocation(getSourceLocation());
 
         if (!getHttp().isEmpty()) {
             builder.withMember(HTTP, Node.fromStrings(getHttp()));

smithy-model/src/main/java/software/amazon/smithy/model/traits/AuthDefinitionTrait.java

@@ -53,16 +53,16 @@ public static Builder builder() {
 
     @Override
     protected Node createNode() {
-        if (traits.isEmpty()) {
-            return Node.objectNode();
-        }
-
-        ArrayNode ids = traits.stream()
+        ObjectNode.Builder builder = Node.objectNodeBuilder();
+        builder.sourceLocation(getSourceLocation());
+        if (!traits.isEmpty()) {
+            ArrayNode ids = traits.stream()
                 .map(ShapeId::toString)
                 .map(Node::from)
                 .collect(ArrayNode.collect());
-
-        return Node.objectNode().withMember("traits", ids);
+            builder.withMember("traits", ids);
+        }
+        return builder.build();
     }
 
     @Override

smithy-model/src/main/java/software/amazon/smithy/model/traits/AuthTrait.java

@@ -70,6 +70,7 @@ public Trait createTrait(ShapeId target, Node value) {
 
     @Override
     protected Node createNode() {
-        return getValues().stream().map(ShapeId::toString).map(Node::from).collect(ArrayNode.collect());
+        return getValues().stream().map(ShapeId::toString).map(Node::from)
+                .collect(ArrayNode.collect(getSourceLocation()));
     }
 }

smithy-model/src/main/java/software/amazon/smithy/model/traits/EndpointTrait.java

@@ -101,7 +101,7 @@ public static Builder builder() {
 
     @Override
     public Builder toBuilder() {
-        return new Builder().hostPrefix(hostPrefix.toString());
+        return new Builder().sourceLocation(getSourceLocation()).hostPrefix(hostPrefix.toString());
     }
 
     @Override

smithy-model/src/main/java/software/amazon/smithy/model/traits/EnumTrait.java

@@ -74,7 +74,7 @@ public boolean hasNames() {
 
     @Override
     protected Node createNode() {
-        return definitions.stream().map(EnumDefinition::toNode).collect(ArrayNode.collect());
+        return definitions.stream().map(EnumDefinition::toNode).collect(ArrayNode.collect(getSourceLocation()));
     }
 
     @Override

smithy-model/src/main/java/software/amazon/smithy/model/traits/ExternalDocumentationTrait.java

@@ -66,7 +66,7 @@ public Map<String, String> getUrls() {
 
     @Override
     protected Node createNode() {
-        return ObjectNode.fromStringMap(urls);
+        return ObjectNode.fromStringMap(urls).toBuilder().sourceLocation(getSourceLocation()).build();
     }
 
     @Override

smithy-model/src/main/java/software/amazon/smithy/model/traits/HttpApiKeyAuthTrait.java

@@ -57,6 +57,7 @@ public Builder toBuilder() {
     @Override
     protected Node createNode() {
         return Node.objectNodeBuilder()
+                .sourceLocation(getSourceLocation())
                 .withMember("name", getName())
                 .withMember("in", getIn().toString())
                 .build();

smithy-model/src/main/java/software/amazon/smithy/model/traits/HttpTrait.java

@@ -89,7 +89,7 @@ public static HttpTrait.Builder builder() {
 
     @Override
     public HttpTrait.Builder toBuilder() {
-        return new Builder().method(method).uri(uri).code(code);
+        return new Builder().sourceLocation(getSourceLocation()).method(method).uri(uri).code(code);
     }
 
     /**

smithy-model/src/main/java/software/amazon/smithy/model/traits/IdRefTrait.java

@@ -78,15 +78,16 @@ public Builder toBuilder() {
 
     @Override
     protected Node createNode() {
-        ObjectNode result = Node.objectNode()
+        ObjectNode.Builder builder = Node.objectNodeBuilder()
+                .sourceLocation(getSourceLocation())
                 .withOptionalMember(
                         SELECTOR_MEMBER_ID,
                         Optional.ofNullable(selector).map(Selector::toString).map(Node::from))
                 .withOptionalMember(ERROR_MESSAGE, getErrorMessage().map(Node::from));
         if (failWhenMissing) {
-            result = result.withMember(FAIL_WHEN_MISSING_MEMBER, Node.from(true));
+            builder = builder.withMember(FAIL_WHEN_MISSING_MEMBER, Node.from(true));
         }
-        return result;
+        return builder.build();
     }
 
     public static Builder builder() {

smithy-model/src/main/java/software/amazon/smithy/model/traits/ProtocolDefinitionTrait.java

@@ -67,22 +67,19 @@ public static Builder builder() {
 
     @Override
     protected Node createNode() {
-        if (traits.isEmpty()) {
-            return Node.objectNode();
-        }
-
         ObjectNode.Builder builder = Node.objectNodeBuilder();
-
-        ArrayNode ids = traits.stream()
-                .map(ShapeId::toString)
-                .map(Node::from)
-                .collect(ArrayNode.collect());
-        builder.withMember(TRAITS, ids);
-
-        if (noInlineDocumentSupport) {
-            builder.withMember(NO_INLINE_DOCUMENT_SUPPORT, true);
+        builder.sourceLocation(getSourceLocation());
+        if (!traits.isEmpty()) {
+            ArrayNode ids = traits.stream()
+                    .map(ShapeId::toString)
+                    .map(Node::from)
+                    .collect(ArrayNode.collect());
+            builder.withMember(TRAITS, ids);
+
+            if (noInlineDocumentSupport) {
+                builder.withMember(NO_INLINE_DOCUMENT_SUPPORT, true);
+            }
         }
-
         return builder.build();
     }
 

smithy-model/src/main/java/software/amazon/smithy/model/traits/RangeTrait.java

@@ -80,7 +80,7 @@ public static Builder builder() {
     }
 
     /**
-     * Builder used to create a LongTrait.
+     * Builder used to create a RangeTrait.
      */
     public static final class Builder extends AbstractTraitBuilder<RangeTrait, Builder> {
         private BigDecimal min;

smithy-model/src/main/java/software/amazon/smithy/model/traits/ReferencesTrait.java

@@ -75,7 +75,7 @@ public List<Reference> getResourceReferences(ShapeId shapeId) {
 
     @Override
     protected Node createNode() {
-        return references.stream().map(Reference::toNode).collect(ArrayNode.collect());
+        return references.stream().map(Reference::toNode).collect(ArrayNode.collect(getSourceLocation()));
     }
 
     @Override