Avoid allocating a separate map

smithy-lang · Feb 27, 2023 · 847fff2 · 847fff2
1 parent d237159
commit 847fff2
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 43 deletions.
diff --git a/...aits/src/main/java/software/amazon/smithy/openapi/traits/SpecificationExtensionTrait.java b/...aits/src/main/java/software/amazon/smithy/openapi/traits/SpecificationExtensionTrait.java
@@ -5,7 +5,6 @@
 
 package software.amazon.smithy.openapi.traits;
 
-import java.util.Optional;
 import software.amazon.smithy.model.node.Node;
 import software.amazon.smithy.model.node.ObjectNode;
 import software.amazon.smithy.model.shapes.ShapeId;
@@ -47,12 +46,16 @@ public Trait createTrait(ShapeId target, Node value) {
     }
 
     /**
-     * Get the explicit name for the target specification extension.
+     * Gets the extension name for a given trait shape.
+     * Either an explicitly configured extension name, or a default transformation of the shape ID.
      *
-     * @return Explicit name for the target specification extension.
+     * @param traitShapeId Trait shape to get the extension name.
+     * @return Extension name for the given trait shape.
      */
-    public Optional<String> as() {
-        return Optional.ofNullable(as);
+    public String extensionNameFor(ShapeId traitShapeId) {
+        return as != null
+                ? as
+                : "x-" + traitShapeId.toString().replaceAll("[.#]", "-");
     }
 
     public static SpecificationExtensionTrait.Builder builder() {

diff --git a/...napi/src/main/java/software/amazon/smithy/openapi/fromsmithy/OpenApiJsonSchemaMapper.java b/...napi/src/main/java/software/amazon/smithy/openapi/fromsmithy/OpenApiJsonSchemaMapper.java
@@ -93,7 +93,7 @@ public Schema.Builder updateSchema(Shape shape, Schema.Builder builder, JsonSche
             }
         }
 
-        SpecificationExtensionsMapper.getSpecificationExtensions(model, shape).forEach(builder::putExtension);
+        SpecificationExtensionsMapper.findSpecificationExtensions(model, shape, builder::putExtension);
 
         // Remove unsupported JSON Schema keywords.
         UNSUPPORTED_KEYWORD_DIRECTIVES.forEach(builder::disableProperty);

diff --git a/...java/software/amazon/smithy/openapi/fromsmithy/mappers/SpecificationExtensionsMapper.java b/...java/software/amazon/smithy/openapi/fromsmithy/mappers/SpecificationExtensionsMapper.java
@@ -5,17 +5,16 @@
 
 package software.amazon.smithy.openapi.fromsmithy.mappers;
 
-import java.util.HashMap;
 import java.util.Map;
-import java.util.Optional;
+import java.util.function.BiConsumer;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.node.Node;
 import software.amazon.smithy.model.shapes.OperationShape;
 import software.amazon.smithy.model.shapes.Shape;
-import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.traits.Trait;
 import software.amazon.smithy.openapi.fromsmithy.Context;
 import software.amazon.smithy.openapi.fromsmithy.OpenApiMapper;
+import software.amazon.smithy.openapi.model.Component;
 import software.amazon.smithy.openapi.model.OpenApi;
 import software.amazon.smithy.openapi.model.OperationObject;
 import software.amazon.smithy.openapi.traits.SpecificationExtensionTrait;
@@ -25,13 +24,8 @@
  */
 public class SpecificationExtensionsMapper implements OpenApiMapper {
     @Override
-    public void before(Context<? extends Trait> context, OpenApi.Builder openapi) {
-        Map<String, Node> specificationExtensions = getSpecificationExtensions(context.getModel(), context.getService());
-        if (specificationExtensions.isEmpty()) {
-            return;
-        }
-
-        specificationExtensions.forEach(openapi::putExtension);
+    public OpenApi after(Context<? extends Trait> context, OpenApi openapi) {
+        return attachAllExtensionsFromShape(openapi, context.getModel(), context.getService());
     }
 
     @Override
@@ -42,39 +36,27 @@ public OperationObject updateOperation(
             String httpMethodName,
             String path
     ) {
-        Map<String, Node> specificationExtensions = getSpecificationExtensions(context.getModel(), shape);
-        if (specificationExtensions.isEmpty()) {
-            return operation;
-        }
+        return attachAllExtensionsFromShape(operation, context.getModel(), shape);
+    }
 
-        OperationObject.Builder builder = operation.toBuilder();
-        specificationExtensions.forEach(builder::putExtension);
-        return builder.build();
+    private static <T extends Component> T attachAllExtensionsFromShape(T component, Model model, Shape shape) {
+        Map<String, Node> extensions = component.getExtensions();
+        findSpecificationExtensions(model, shape, extensions::put);
+        return component;
     }
 
     /**
-     * Get all specification extension trait values attached to the given shape.
+     * Find all specification extension trait values attached to the given shape.
      *
-     * @param model Model the shape belongs to.
-     * @param shape Shape to get extensions for.
-     * @return All specification extension values.
+     * @param model    Model the shape belongs to.
+     * @param shape    Shape to get extensions for.
+     * @param consumer Consumer called for each specification extension key-value pair found.
      */
-    public static Map<String, Node> getSpecificationExtensions(Model model, Shape shape) {
-        Map<String, Node> specificationExtensions = new HashMap<>();
-
+    public static void findSpecificationExtensions(Model model, Shape shape, BiConsumer<String, Node> consumer) {
         shape.getAllTraits().forEach((traitShapeId, trait) ->
-                determineExtensionName(traitShapeId, model)
-                        .ifPresent(name -> specificationExtensions.put(name, trait.toNode())));
-
-        return specificationExtensions;
-    }
-
-    private static Optional<String> determineExtensionName(ShapeId traitShapeId, Model model) {
-        return model
-                .getShape(traitShapeId)
-                .flatMap(shape -> shape.getTrait(SpecificationExtensionTrait.class))
-                .map(specificationExtensionTrait -> specificationExtensionTrait
-                        .as()
-                        .orElseGet(() -> "x-" + traitShapeId.toString().replaceAll("[.#]", "-")));
+                model.getShape(traitShapeId)
+                        .flatMap(traitShape -> traitShape.getTrait(SpecificationExtensionTrait.class))
+                        .map(specificationExtensionTrait -> specificationExtensionTrait.extensionNameFor(traitShapeId))
+                        .ifPresent(name -> consumer.accept(name, trait.toNode())));
     }
 }