Add custom IDL serialization ordering support

A new `SmithyIdlSerializationOrder` interface extracts the comparator
methods from `SmithyIdlComponentOrder` into a public contract that
the enum now implements. Users can provide custom implementations to
control shape, trait, and metadata ordering during IDL serialization.

Author: kstich

.changes/next-release/feature-d774f63260bae2ee8d11ff1898618d4096b6daac.json

@@ -0,0 +1,7 @@
+{
+  "type": "feature",
+  "description": "Add support for custom comparators for sorting IDL serialization output",
+  "pull_requests": [
+    "[#3058](https://github.com/smithy-lang/smithy/pull/3058)"
+  ]
+}

smithy-model/src/main/java/software/amazon/smithy/model/shapes/SmithyIdlComponentOrder.java

@@ -10,13 +10,18 @@
 import software.amazon.smithy.model.FromSourceLocation;
 import software.amazon.smithy.model.SourceLocation;
 import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.traits.Trait;
 import software.amazon.smithy.model.traits.TraitDefinition;
 import software.amazon.smithy.utils.MapUtils;
 
 /**
  * Defines how shapes, traits, and metadata are sorted when serializing a model with {@link SmithyIdlModelSerializer}.
+ *
+ * <p>This enum provides the built-in orderings. For custom ordering logic, implement
+ * {@link SmithyIdlSerializationOrder} directly and pass it to
+ * {@link SmithyIdlModelSerializer.Builder#componentOrder(SmithyIdlSerializationOrder)}.
  */
-public enum SmithyIdlComponentOrder {
+public enum SmithyIdlComponentOrder implements SmithyIdlSerializationOrder {
     /**
      * Sort shapes, traits, and metadata alphabetically. Member order, however, is not sorted.
      */
@@ -46,29 +51,36 @@ public enum SmithyIdlComponentOrder {
      */
     PREFERRED;
 
-    Comparator<Shape> shapeComparator() {
+    @Override
+    public Comparator<Shape> shapeComparator() {
         return this == PREFERRED ? new PreferredShapeComparator() : toShapeIdComparator();
     }
 
-    <T extends FromSourceLocation & ToShapeId> Comparator<T> toShapeIdComparator() {
+    @Override
+    public Comparator<Trait> traitComparator() {
+        return toShapeIdComparator();
+    }
+
+    @Override
+    public Comparator<Map.Entry<String, Node>> metadataComparator() {
         switch (this) {
-            case PREFERRED:
             case ALPHA_NUMERIC:
-                return Comparator.comparing(ToShapeId::toShapeId);
+            case PREFERRED:
+                return Map.Entry.comparingByKey();
             case SOURCE_LOCATION:
             default:
-                return new SourceComparator<>();
+                return new MetadataComparator();
         }
     }
 
-    Comparator<Map.Entry<String, Node>> metadataComparator() {
+    private <T extends FromSourceLocation & ToShapeId> Comparator<T> toShapeIdComparator() {
         switch (this) {
-            case ALPHA_NUMERIC:
             case PREFERRED:
-                return Map.Entry.comparingByKey();
+            case ALPHA_NUMERIC:
+                return Comparator.comparing(ToShapeId::toShapeId);
             case SOURCE_LOCATION:
             default:
-                return new MetadataComparator();
+                return new SourceComparator<>();
         }
     }
 

smithy-model/src/main/java/software/amazon/smithy/model/shapes/SmithyIdlModelSerializer.java

@@ -60,7 +60,7 @@ public final class SmithyIdlModelSerializer {
     private final Predicate<Trait> traitFilter;
     private final Function<Shape, Path> shapePlacer;
     private final Path basePath;
-    private final SmithyIdlComponentOrder componentOrder;
+    private final SmithyIdlSerializationOrder componentOrder;
     private final String inlineInputSuffix;
     private final String inlineOutputSuffix;
     private final boolean inferInlineIoSuffixes;
@@ -382,7 +382,7 @@ public static final class Builder implements SmithyBuilder<SmithyIdlModelSeriali
         private Function<Shape, Path> shapePlacer = SmithyIdlModelSerializer::placeShapesByNamespace;
         private Path basePath = null;
         private boolean serializePrelude = false;
-        private SmithyIdlComponentOrder componentOrder = SmithyIdlComponentOrder.PREFERRED;
+        private SmithyIdlSerializationOrder componentOrder = SmithyIdlComponentOrder.PREFERRED;
         private String inlineInputSuffix = DEFAULT_INLINE_INPUT_SUFFIX;
         private String inlineOutputSuffix = DEFAULT_INLINE_OUTPUT_SUFFIX;
         private boolean inferInlineIoSuffixes = false;
@@ -477,6 +477,21 @@ public Builder componentOrder(SmithyIdlComponentOrder componentOrder) {
             return this;
         }
 
+        /**
+         * Defines how components are sorted in the model using a custom ordering configuration.
+         *
+         * <p>Implement {@link SmithyIdlSerializationOrder} to provide custom comparators for shapes,
+         * traits, and metadata. The built-in orderings are available as constants on
+         * {@link SmithyIdlComponentOrder}.
+         *
+         * @param componentOrder Custom component ordering configuration.
+         * @return Returns the builder.
+         */
+        public Builder componentOrder(SmithyIdlSerializationOrder componentOrder) {
+            this.componentOrder = Objects.requireNonNull(componentOrder);
+            return this;
+        }
+
         /**
          * Defines what suffixes are checked on operation input shapes to determine whether
          * inline syntax should be used.
@@ -556,15 +571,15 @@ private static final class ShapeSerializer extends ShapeVisitor.Default<Void> {
         private final Predicate<Trait> traitFilter;
         private final Model model;
         private final Set<ShapeId> inlineableShapes;
-        private final SmithyIdlComponentOrder componentOrder;
+        private final SmithyIdlSerializationOrder componentOrder;
 
         ShapeSerializer(
                 SmithyCodeWriter codeWriter,
                 NodeSerializer nodeSerializer,
                 Predicate<Trait> traitFilter,
                 Model model,
                 Set<ShapeId> inlineableShapes,
-                SmithyIdlComponentOrder componentOrder
+                SmithyIdlSerializationOrder componentOrder
         ) {
             this.codeWriter = codeWriter;
             this.nodeSerializer = nodeSerializer;
@@ -703,7 +718,7 @@ private void serializeTraits(Map<ShapeId, Trait> traits, TraitFeature... traitFe
                 }
             }
 
-            Comparator<Trait> traitComparator = componentOrder.toShapeIdComparator();
+            Comparator<Trait> traitComparator = componentOrder.traitComparator();
 
             traits.values()
                     .stream()

smithy-model/src/main/java/software/amazon/smithy/model/shapes/SmithyIdlSerializationOrder.java

@@ -0,0 +1,63 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.model.shapes;
+
+import java.util.Comparator;
+import java.util.Map;
+import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.traits.Trait;
+
+/**
+ * Configures how shapes, traits, and metadata are ordered when serializing a model
+ * with {@link SmithyIdlModelSerializer}.
+ *
+ * <p>The built-in orderings are available as constants on {@link SmithyIdlComponentOrder}.
+ * Implement this interface directly to provide custom ordering logic.
+ *
+ * <p>Only {@link #shapeComparator()} is required. The default implementations of
+ * {@link #traitComparator()} and {@link #metadataComparator()} sort alphabetically,
+ * which is suitable for most use cases.
+ *
+ * @see SmithyIdlComponentOrder
+ * @see SmithyIdlModelSerializer.Builder#componentOrder(SmithyIdlSerializationOrder)
+ */
+public interface SmithyIdlSerializationOrder {
+
+    /**
+     * A shape comparator that sorts alphabetically by shape ID.
+     */
+    Comparator<Shape> ALPHABETICAL = Comparator.comparing(Shape::toShapeId);
+
+    /**
+     * Returns a comparator used to sort shapes within a namespace file.
+     *
+     * @return Comparator for ordering shapes.
+     */
+    Comparator<Shape> shapeComparator();
+
+    /**
+     * Returns a comparator used to sort traits applied to a shape.
+     *
+     * <p>The default implementation sorts traits alphabetically by their shape ID. Custom
+     * implementations that use source-location-based ordering should consider overriding
+     * this method to also sort traits by source location for consistency.
+     *
+     * @return Comparator for ordering traits.
+     */
+    default Comparator<Trait> traitComparator() {
+        return Comparator.comparing(Trait::toShapeId);
+    }
+
+    /**
+     * Returns a comparator used to sort metadata entries.
+     *
+     * <p>The default implementation sorts metadata alphabetically by key.
+     *
+     * @return Comparator for ordering metadata entries.
+     */
+    default Comparator<Map.Entry<String, Node>> metadataComparator() {
+        return Map.Entry.comparingByKey();
+    }
+}

smithy-model/src/test/java/software/amazon/smithy/model/shapes/SmithyIdlModelSerializerTest.java

@@ -22,6 +22,7 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.Comparator;
 import java.util.Map;
 import java.util.Objects;
 import java.util.stream.Stream;
@@ -37,6 +38,8 @@
 import software.amazon.smithy.model.traits.DocumentationTrait;
 import software.amazon.smithy.model.traits.DynamicTrait;
 import software.amazon.smithy.model.traits.RequiredTrait;
+import software.amazon.smithy.model.traits.SensitiveTrait;
+import software.amazon.smithy.model.traits.Trait;
 import software.amazon.smithy.model.traits.synthetic.OriginalShapeIdTrait;
 import software.amazon.smithy.utils.IoUtils;
 import software.amazon.smithy.utils.MapUtils;
@@ -392,4 +395,108 @@ public void unresolvableIdRefSerializedAsQuotedString() {
         Model model2 = Model.assembler().addUnparsedModel("test.smithy", modelResult).assemble().unwrap();
         assertThat(model2, equalTo(model));
     }
+
+    @Test
+    public void usesCustomComponentOrderConfig() {
+        // Build a simple model with multiple shapes to verify custom ordering.
+        Model model = Model.builder()
+                .addShape(StringShape.builder().id("com.example#Zebra").build())
+                .addShape(StringShape.builder().id("com.example#Alpha").build())
+                .addShape(StructureShape.builder().id("com.example#Middle").build())
+                .build();
+
+        // Custom config that reverses the default alphabetical shape order.
+        SmithyIdlSerializationOrder reverseOrder = new SmithyIdlSerializationOrder() {
+            @Override
+            public Comparator<Shape> shapeComparator() {
+                return Comparator.comparing(Shape::toShapeId).reversed();
+            }
+        };
+
+        Map<Path, String> serialized = SmithyIdlModelSerializer.builder()
+                .componentOrder(reverseOrder)
+                .build()
+                .serialize(model);
+        String result = serialized.values().iterator().next();
+
+        // Verify shapes appear in reverse alphabetical order: Zebra, Middle, Alpha.
+        int zebraPos = result.indexOf("Zebra");
+        int middlePos = result.indexOf("Middle");
+        int alphaPos = result.indexOf("Alpha");
+        assertTrue(zebraPos < middlePos, "Zebra should appear before Middle in reverse order");
+        assertTrue(middlePos < alphaPos, "Middle should appear before Alpha in reverse order");
+    }
+
+    @Test
+    public void usesCustomTraitComparator() {
+        // Use two annotation traits (not documentation) whose relative order is observable.
+        Model model = Model.builder()
+                .addShape(StructureShape.builder()
+                        .id("com.example#MyStruct")
+                        .addTrait(new RequiredTrait())
+                        .addTrait(new SensitiveTrait())
+                        .build())
+                .build();
+
+        // Reversed trait comparator: sensitive (s) should appear before required (r) alphabetically,
+        // but reversed means required comes first.
+        SmithyIdlSerializationOrder customOrder = new SmithyIdlSerializationOrder() {
+            @Override
+            public Comparator<Shape> shapeComparator() {
+                return SmithyIdlSerializationOrder.ALPHABETICAL;
+            }
+
+            @Override
+            public Comparator<Trait> traitComparator() {
+                return Comparator.comparing(Trait::toShapeId).reversed();
+            }
+        };
+
+        Map<Path, String> serialized = SmithyIdlModelSerializer.builder()
+                .componentOrder(customOrder)
+                .build()
+                .serialize(model);
+        String result = serialized.get(Paths.get("com.example.smithy"));
+
+        // In reversed order, sensitive (s) comes before required (r).
+        int sensitivePos = result.indexOf("@sensitive");
+        int requiredPos = result.indexOf("@required");
+        assertTrue(sensitivePos < requiredPos,
+                "@sensitive should appear before @required in reversed trait order");
+    }
+
+    @Test
+    public void usesCustomMetadataComparator() {
+        // Use two metadata keys whose relative order is observable.
+        Model model = Model.builder()
+                .addShape(StringShape.builder().id("com.example#Placeholder").build())
+                .putMetadataProperty("zeta", Node.from("z"))
+                .putMetadataProperty("alpha", Node.from("a"))
+                .build();
+
+        // Reversed metadata comparator: alpha (a) should appear before zeta (z) alphabetically,
+        // but reversed means zeta comes first.
+        SmithyIdlSerializationOrder customOrder = new SmithyIdlSerializationOrder() {
+            @Override
+            public Comparator<Shape> shapeComparator() {
+                return SmithyIdlSerializationOrder.ALPHABETICAL;
+            }
+
+            @Override
+            public Comparator<Map.Entry<String, Node>> metadataComparator() {
+                return Map.Entry.<String, Node>comparingByKey().reversed();
+            }
+        };
+
+        Map<Path, String> serialized = SmithyIdlModelSerializer.builder()
+                .componentOrder(customOrder)
+                .build()
+                .serialize(model);
+        String result = serialized.get(Paths.get("metadata.smithy"));
+
+        // In reversed order, zeta (z) comes before alpha (a).
+        int zetaPos = result.indexOf("zeta");
+        int alphaPos = result.indexOf("alpha");
+        assertTrue(zetaPos < alphaPos, "zeta should appear before alpha in reversed metadata order");
+    }
 }