Add SortMembers transform

smithy-model/src/main/java/software/amazon/smithy/model/transform/ModelTransformer.java

@@ -17,6 +17,7 @@
 
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -31,6 +32,7 @@
 import software.amazon.smithy.model.neighbor.UnreferencedShapes;
 import software.amazon.smithy.model.neighbor.UnreferencedTraitDefinitions;
 import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.shapes.MemberShape;
 import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.traits.Trait;
@@ -436,4 +438,21 @@ public Model getModelWithoutTraitShapes(Model model, Predicate<Shape> keepFilter
 
         return builder.build();
     }
+
+    /**
+     * Reorders the members of structure and union shapes using the given
+     * {@link Comparator}.
+     *
+     * <p>Note that by default, Smithy models retain the order in which
+     * members are defined in the model. However, in programming languages
+     * where this isn't important, it may be desirable to order members
+     * alphabetically or using some other kind of order.
+     *
+     * @param model Model that contains shapes.
+     * @param comparator Comparator used to order members of unions and structures.
+     * @return Returns a model that contains matching shapes.
+     */
+    public Model sortMembers(Model model, Comparator<MemberShape> comparator) {
+        return new SortMembers(comparator).transform(this, model);
+    }
 }

smithy-model/src/main/java/software/amazon/smithy/model/transform/SortMembers.java

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.smithy.model.transform;
+
+import java.util.Comparator;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.TreeSet;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.shapes.MemberShape;
+import software.amazon.smithy.model.shapes.Shape;
+import software.amazon.smithy.model.shapes.StructureShape;
+import software.amazon.smithy.model.shapes.UnionShape;
+
+/**
+ *  Reorders members of structure and union shapes using a {@link Comparator}.
+ */
+final class SortMembers {
+
+    private final Comparator<MemberShape> comparator;
+
+    SortMembers(Comparator<MemberShape> comparator) {
+        this.comparator = comparator;
+    }
+
+    Model transform(ModelTransformer transformer, Model model) {
+        Set<Shape> toReplace = new HashSet<>();
+
+        model.shapes(StructureShape.class).forEach(structure -> {
+            if (!structure.getAllMembers().isEmpty()) {
+                Set<MemberShape> members = new TreeSet<>(comparator);
+                members.addAll(structure.getAllMembers().values());
+                toReplace.add(structure.toBuilder().members(members).build());
+            }
+        });
+
+        model.shapes(UnionShape.class).forEach(structure -> {
+            if (!structure.getAllMembers().isEmpty()) {
+                Set<MemberShape> members = new TreeSet<>(comparator);
+                members.addAll(structure.getAllMembers().values());
+                toReplace.add(structure.toBuilder().members(members).build());
+            }
+        });
+
+        return transformer.replaceShapes(model, toReplace);
+    }
+}

smithy-model/src/test/java/software/amazon/smithy/model/transform/SortMembersTest.java

@@ -0,0 +1,44 @@
+package software.amazon.smithy.model.transform;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+
+import java.util.Comparator;
+import org.hamcrest.Matchers;
+import org.junit.jupiter.api.Test;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.loader.ModelAssembler;
+import software.amazon.smithy.model.shapes.MemberShape;
+import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.model.shapes.StructureShape;
+import software.amazon.smithy.model.shapes.UnionShape;
+
+public class SortMembersTest {
+    @Test
+    public void sortsModelMembers() {
+        ModelAssembler assembler = Model.assembler();
+        UnionShape u = UnionShape.builder()
+                .id("com.foo#U")
+                .addMember("zoo", ShapeId.from("smithy.api#String"))
+                .addMember("abc", ShapeId.from("smithy.api#String"))
+                .build();
+        StructureShape s = StructureShape.builder()
+                .id("com.foo#S")
+                .addMember("zoo", ShapeId.from("smithy.api#String"))
+                .addMember("abc", ShapeId.from("smithy.api#String"))
+                .build();
+        assembler.addShapes(u, s);
+        Model model = assembler.assemble().unwrap();
+        ModelTransformer transformer = ModelTransformer.create();
+        Model sortedModel = transformer.sortMembers(model, Comparator.comparing(MemberShape::getMemberName));
+
+        // Members use given order by default.
+        assertThat(u.getMemberNames(), Matchers.contains("zoo", "abc"));
+        assertThat(s.getMemberNames(), Matchers.contains("zoo", "abc"));
+
+        // Members in the new model use the sorted order.
+        assertThat(sortedModel.expectShape(u.getId(), UnionShape.class).getMemberNames(),
+                   Matchers.contains("abc", "zoo"));
+        assertThat(sortedModel.expectShape(s.getId(), StructureShape.class).getMemberNames(),
+                   Matchers.contains("abc", "zoo"));
+    }
+}