Make annotation traits lossless

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

@@ -15,7 +15,8 @@
 
 package software.amazon.smithy.aws.iam.traits;
 
-import software.amazon.smithy.model.SourceLocation;
+import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.node.ObjectNode;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.traits.AnnotationTrait;
 
@@ -25,12 +26,12 @@
 public final class DisableConditionKeyInferenceTrait extends AnnotationTrait {
     public static final ShapeId ID = ShapeId.from("aws.iam#disableConditionKeyInference");
 
-    public DisableConditionKeyInferenceTrait(SourceLocation sourceLocation) {
-        super(ID, sourceLocation);
+    public DisableConditionKeyInferenceTrait(ObjectNode node) {
+        super(ID, node);
     }
 
     public DisableConditionKeyInferenceTrait() {
-        this(SourceLocation.NONE);
+        this(Node.objectNode());
     }
 
     public static final class Provider extends AnnotationTrait.Provider<DisableConditionKeyInferenceTrait> {

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

@@ -15,19 +15,20 @@
 
 package software.amazon.smithy.aws.traits;
 
-import software.amazon.smithy.model.SourceLocation;
+import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.node.ObjectNode;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.traits.AnnotationTrait;
 
 public final class ControlPlaneTrait extends AnnotationTrait {
     public static final ShapeId ID = ShapeId.from("aws.api#controlPlane");
 
-    public ControlPlaneTrait(SourceLocation sourceLocation) {
-        super(ID, sourceLocation);
+    public ControlPlaneTrait(ObjectNode node) {
+        super(ID, node);
     }
 
     public ControlPlaneTrait() {
-        this(SourceLocation.NONE);
+        this(Node.objectNode());
     }
 
     public static final class Provider extends AnnotationTrait.Provider<ControlPlaneTrait> {

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

@@ -15,19 +15,20 @@
 
 package software.amazon.smithy.aws.traits;
 
-import software.amazon.smithy.model.SourceLocation;
+import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.node.ObjectNode;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.traits.AnnotationTrait;
 
 public final class DataPlaneTrait extends AnnotationTrait {
     public static final ShapeId ID = ShapeId.from("aws.api#dataPlane");
 
-    public DataPlaneTrait(SourceLocation sourceLocation) {
-        super(ID, sourceLocation);
+    public DataPlaneTrait(ObjectNode node) {
+        super(ID, node);
     }
 
     public DataPlaneTrait() {
-        this(SourceLocation.NONE);
+        this(Node.objectNode());
     }
 
     public static final class Provider extends AnnotationTrait.Provider<DataPlaneTrait> {

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

@@ -15,8 +15,8 @@
 
 package software.amazon.smithy.aws.traits.auth;
 
-import software.amazon.smithy.model.FromSourceLocation;
-import software.amazon.smithy.model.SourceLocation;
+import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.node.ObjectNode;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.traits.AnnotationTrait;
 
@@ -26,12 +26,12 @@
 public final class UnsignedPayloadTrait extends AnnotationTrait {
     public static final ShapeId ID = ShapeId.from("aws.auth#unsignedPayload");
 
-    public UnsignedPayloadTrait(FromSourceLocation sourceLocation) {
-        super(ID, sourceLocation.getSourceLocation());
+    public UnsignedPayloadTrait(ObjectNode node) {
+        super(ID, node);
     }
 
     public UnsignedPayloadTrait() {
-        this(SourceLocation.NONE);
+        this(Node.objectNode());
     }
 
     public static final class Provider extends AnnotationTrait.Provider<UnsignedPayloadTrait> {

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

@@ -15,7 +15,8 @@
 
 package software.amazon.smithy.aws.traits.clientendpointdiscovery;
 
-import software.amazon.smithy.model.SourceLocation;
+import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.node.ObjectNode;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.traits.AnnotationTrait;
 
@@ -26,12 +27,12 @@
 public final class ClientEndpointDiscoveryIdTrait extends AnnotationTrait {
     public static final ShapeId ID = ShapeId.from("aws.api#clientEndpointDiscoveryId");
 
-    public ClientEndpointDiscoveryIdTrait(SourceLocation sourceLocation) {
-        super(ID, sourceLocation);
+    public ClientEndpointDiscoveryIdTrait(ObjectNode node) {
+        super(ID, node);
     }
 
     public ClientEndpointDiscoveryIdTrait() {
-        this(SourceLocation.NONE);
+        this(Node.objectNode());
     }
 
     public static final class Provider extends AnnotationTrait.Provider<ClientEndpointDiscoveryIdTrait> {

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

@@ -15,7 +15,8 @@
 
 package software.amazon.smithy.aws.traits.protocols;
 
-import software.amazon.smithy.model.SourceLocation;
+import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.node.ObjectNode;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.traits.AnnotationTrait;
 
@@ -29,12 +30,12 @@ public final class AwsQueryTrait extends AnnotationTrait {
 
     public static final ShapeId ID = ShapeId.from("aws.protocols#awsQuery");
 
-    public AwsQueryTrait(SourceLocation sourceLocation) {
-        super(ID, sourceLocation);
+    public AwsQueryTrait(ObjectNode node) {
+        super(ID, node);
     }
 
     public AwsQueryTrait() {
-        this(SourceLocation.NONE);
+        this(Node.objectNode());
     }
 
     public static final class Provider extends AnnotationTrait.Provider<AwsQueryTrait> {

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

@@ -15,7 +15,8 @@
 
 package software.amazon.smithy.aws.traits.protocols;
 
-import software.amazon.smithy.model.SourceLocation;
+import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.node.ObjectNode;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.traits.AnnotationTrait;
 
@@ -30,12 +31,12 @@ public final class Ec2QueryTrait extends AnnotationTrait {
 
     public static final ShapeId ID = ShapeId.from("aws.protocols#ec2Query");
 
-    public Ec2QueryTrait(SourceLocation sourceLocation) {
-        super(ID, sourceLocation);
+    public Ec2QueryTrait(ObjectNode node) {
+        super(ID, node);
     }
 
     public Ec2QueryTrait() {
-        this(SourceLocation.NONE);
+        this(Node.objectNode());
     }
 
     public static final class Provider extends AnnotationTrait.Provider<Ec2QueryTrait> {

smithy-build/src/test/java/software/amazon/smithy/build/transforms/ExcludeTraitsTest.java

@@ -37,7 +37,7 @@ public class ExcludeTraitsTest {
     public void removesTraitsInList() {
         StringShape stringShape = StringShape.builder()
                 .id("ns.foo#baz")
-                .addTrait(new SensitiveTrait(SourceLocation.NONE))
+                .addTrait(new SensitiveTrait())
                 .addTrait(new DocumentationTrait("docs", SourceLocation.NONE))
                 .build();
         Model model = Model.assembler()

smithy-build/src/test/java/software/amazon/smithy/build/transforms/IncludeTraitsTest.java

@@ -64,7 +64,7 @@ public void removesTraitsNotInList() {
     public void includesBuiltinTraits() {
         StringShape stringShape = StringShape.builder()
                 .id("ns.foo#baz")
-                .addTrait(new SensitiveTrait(SourceLocation.NONE))
+                .addTrait(new SensitiveTrait())
                 .addTrait(new DocumentationTrait("docs", SourceLocation.NONE))
                 .build();
         Model model = Model.assembler()

smithy-diff/src/test/java/software/amazon/smithy/diff/DifferencesTest.java

@@ -85,7 +85,7 @@ public void detectsChangedShapes() {
         Shape shape1 = StringShape.builder().id("foo.bar#Baz").build();
         Shape shape2 = StringShape.builder()
                 .id("foo.bar#Baz")
-                .addTrait(new SensitiveTrait(SourceLocation.NONE))
+                .addTrait(new SensitiveTrait())
                 .build();
         Model previous = Model.assembler().addShape(shape1).assemble().unwrap();
         Model current = Model.assembler().addShape(shape2).assemble().unwrap();

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

@@ -15,7 +15,7 @@
 
 package software.amazon.smithy.model.traits;
 
-import java.util.Collections;
+import java.util.Objects;
 import java.util.function.Function;
 import software.amazon.smithy.model.SourceLocation;
 import software.amazon.smithy.model.node.ExpectationNotMetException;
@@ -26,35 +26,73 @@
 /**
  * Trait implementation for traits that are an empty object.
  */
-public abstract class AnnotationTrait extends AbstractTrait {
-    public AnnotationTrait(ShapeId id, SourceLocation sourceLocation) {
-        super(id, sourceLocation);
+public abstract class AnnotationTrait implements Trait {
+
+    private final ShapeId id;
+    private final ObjectNode node;
+
+    public AnnotationTrait(ShapeId id, ObjectNode node) {
+        this.id = Objects.requireNonNull(id);
+        this.node = Objects.requireNonNull(node);
+    }
+
+    @Override
+    public final ShapeId toShapeId() {
+        return id;
+    }
+
+    @Override
+    public final Node toNode() {
+        return node;
+    }
+
+    @Override
+    public final SourceLocation getSourceLocation() {
+        return node.getSourceLocation();
     }
 
     @Override
-    protected final Node createNode() {
-        return new ObjectNode(Collections.emptyMap(), getSourceLocation());
+    public int hashCode() {
+        return toShapeId().hashCode() * 17 + node.hashCode();
+    }
+
+    @Override
+    public boolean equals(Object other) {
+        if (other == null || other.getClass() != getClass()) {
+            return false;
+        }
+
+        if (this == other) {
+            return true;
+        }
+
+        Trait b = (Trait) other;
+        if (!toShapeId().equals(b.toShapeId())) {
+            return false;
+        }
+
+        return node.equals(b.toNode());
     }
 
     /**
      * Trait provider that expects a boolean value of true.
      */
     public static class Provider<T extends AnnotationTrait> extends AbstractTrait.Provider {
-        private final Function<SourceLocation, T> traitFactory;
+        private final Function<ObjectNode, T> traitFactory;
 
         /**
          * @param id ID of the trait being created.
          * @param traitFactory Factory function used to create the trait.
          */
-        public Provider(ShapeId id, Function<SourceLocation, T> traitFactory) {
+        public Provider(ShapeId id, Function<ObjectNode, T> traitFactory) {
             super(id);
             this.traitFactory = traitFactory;
         }
 
         @Override
         public T createTrait(ShapeId id, Node value) {
             if (value.isObjectNode()) {
-                return traitFactory.apply(value.getSourceLocation());
+                return traitFactory.apply(value.expectObjectNode());
             }
 
             throw new ExpectationNotMetException(String.format(

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

@@ -15,7 +15,8 @@
 
 package software.amazon.smithy.model.traits;
 
-import software.amazon.smithy.model.SourceLocation;
+import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.node.ObjectNode;
 import software.amazon.smithy.model.shapes.ShapeId;
 
 /**
@@ -24,12 +25,12 @@
 public final class BoxTrait extends AnnotationTrait {
     public static final ShapeId ID = ShapeId.from("smithy.api#box");
 
-    public BoxTrait(SourceLocation sourceLocation) {
-        super(ID, sourceLocation);
+    public BoxTrait(ObjectNode node) {
+        super(ID, node);
     }
 
     public BoxTrait() {
-        this(SourceLocation.NONE);
+        this(Node.objectNode());
     }
 
     public static final class Provider extends AnnotationTrait.Provider<BoxTrait> {

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

@@ -15,7 +15,8 @@
 
 package software.amazon.smithy.model.traits;
 
-import software.amazon.smithy.model.SourceLocation;
+import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.node.ObjectNode;
 import software.amazon.smithy.model.shapes.ShapeId;
 
 /**
@@ -27,12 +28,12 @@
 public final class EventHeaderTrait extends AnnotationTrait {
     public static final ShapeId ID = ShapeId.from("smithy.api#eventHeader");
 
-    public EventHeaderTrait(SourceLocation sourceLocation) {
-        super(ID, sourceLocation);
+    public EventHeaderTrait(ObjectNode node) {
+        super(ID, node);
     }
 
     public EventHeaderTrait() {
-        this(SourceLocation.NONE);
+        this(Node.objectNode());
     }
 
     public static final class Provider extends AnnotationTrait.Provider<EventHeaderTrait> {

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

@@ -15,7 +15,8 @@
 
 package software.amazon.smithy.model.traits;
 
-import software.amazon.smithy.model.SourceLocation;
+import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.node.ObjectNode;
 import software.amazon.smithy.model.shapes.ShapeId;
 
 /**
@@ -27,12 +28,12 @@
 public final class EventPayloadTrait extends AnnotationTrait {
     public static final ShapeId ID = ShapeId.from("smithy.api#eventPayload");
 
-    public EventPayloadTrait(SourceLocation sourceLocation) {
-        super(ID, sourceLocation);
+    public EventPayloadTrait(ObjectNode node) {
+        super(ID, node);
     }
 
     public EventPayloadTrait() {
-        this(SourceLocation.NONE);
+        this(Node.objectNode());
     }
 
     public static final class Provider extends AnnotationTrait.Provider<EventPayloadTrait> {