Add RenameShapes model transformer

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

@@ -18,17 +18,21 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.ServiceLoader;
 import java.util.function.BiFunction;
 import java.util.function.BiPredicate;
 import java.util.function.Function;
 import java.util.function.Predicate;
+import java.util.function.Supplier;
 import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.loader.ModelAssembler;
 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.Shape;
+import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.shapes.ShapeIndex;
 import software.amazon.smithy.model.traits.Trait;
 import software.amazon.smithy.model.traits.TraitDefinition;
@@ -132,6 +136,44 @@ public Model removeShapesIf(Model model, Predicate<Shape> predicate) {
         return filterShapes(model, FunctionalUtils.not(predicate));
     }
 
+    /**
+     *  Renames shapes using ShapeId pairs while ensuring that the
+     *  transformed model is in a consistent state.
+     *
+     *  <p>This transformer ensures that when an aggregate shape is renamed, all
+     *  members are updated in the model.
+     *
+     * @param model Model to transform.
+     * @param renamed Map of shapeIds
+     * @return Returns the transformed model.base.
+     */
+    public Model renameShapes(
+            Model model,
+            Map<ShapeId, ShapeId> renamed
+    ) {
+        return this.renameShapes(model, renamed, Model::assembler);
+    }
+
+    /**
+     *  Renames shapes using ShapeId pairs while ensuring that the
+     *  transformed model is in a consistent state.
+     *
+     *  <p>This transformer ensures that when an aggregate shape is renamed, all
+     *  members are updated in the model.
+     *
+     * @param model Model to transform.
+     * @param renamed Map of shapeIds
+     * @param modelAssemblerSupplier Supplier used to create {@link ModelAssembler}s in each transform.
+     * @return Returns the transformed model.
+     */
+    public Model renameShapes(
+            Model model,
+            Map<ShapeId, ShapeId> renamed,
+            Supplier<ModelAssembler> modelAssemblerSupplier
+    ) {
+        return new RenameShapes(renamed, modelAssemblerSupplier).transform(this, model);
+    }
+
     /**
      * Filters shapes out of the model that do not match the given predicate.
      *
@@ -141,7 +183,7 @@ public Model removeShapesIf(Model model, Predicate<Shape> predicate) {
      *
      * @param model Model to transform.
      * @param predicate Predicate that filters shapes.
-     * @return Returns the transformed model.base.
+     * @return Returns the transformed model.
      */
     public Model filterShapes(Model model, Predicate<Shape> predicate) {
         return new FilterShapes(predicate).transform(this, model);

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

@@ -0,0 +1,115 @@
+/*
+ * 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.HashMap;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.loader.ModelAssembler;
+import software.amazon.smithy.model.node.ArrayNode;
+import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.node.NodeVisitor;
+import software.amazon.smithy.model.node.ObjectNode;
+import software.amazon.smithy.model.node.StringNode;
+import software.amazon.smithy.model.shapes.ModelSerializer;
+import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.utils.Pair;
+
+/**
+ *  Renames shapes using ShapeId pairs while ensuring that the
+ *  transformed model is in a consistent state.
+ *
+ *  <p>Member shapes are updated when their containing shape is updated.
+ *
+ *  <p>Trait references to ShapeId values are also updated.
+ */
+final class RenameShapes {
+    private final Map<ShapeId, ShapeId> renamed;
+    private final ModelAssembler assembler;
+
+    RenameShapes(Map<ShapeId, ShapeId> renamed, Supplier<ModelAssembler> modelAssemblerSupplier) {
+        this.renamed = new HashMap<>(renamed);
+        this.assembler = modelAssemblerSupplier.get();
+    }
+
+    Model transform(ModelTransformer transformer, Model model) {
+        // Remove any no-op pairs to avoid renaming shapes unnecessarily.
+        renamed.keySet().removeIf(fromId -> fromId.equals(renamed.get(fromId)));
+        if (renamed.isEmpty()) {
+            return model;
+        }
+
+        // Creates a set that will be used for checking if a string value needs to be renamed or not.
+        Set<String> toRename = renamed.keySet().stream()
+                .map(ShapeId::toString)
+                .collect(Collectors.toSet());
+
+        // This transformer converts the model into an ObjectNode. This approach was chosen because the
+        // JSON AST format includes fully qualified shape ID values, making it possible rename shapes across
+        // the model by only needing to compare and replace StringNode values.
+        ModelSerializer serializer = ModelSerializer.builder().build();
+        ObjectNode node = serializer.serialize(model);
+
+        // Use visitor to traverse node and rebuild model.
+        Node newModel = node.accept(new RenameShapeVisitor(toRename, renamed));
+
+        return assembler.addDocumentNode(newModel)
+                .assemble()
+                .unwrap();
+    }
+
+    private static final class RenameShapeVisitor extends NodeVisitor.Default<Node> {
+
+        private final Set<String> toRename;
+        private final Map<ShapeId, ShapeId> shapeMapping;
+
+        RenameShapeVisitor(Set<String> toRename, Map<ShapeId, ShapeId> shapeMapping) {
+            this.toRename = toRename;
+            this.shapeMapping = shapeMapping;
+        }
+
+        @Override
+        protected Node getDefault(Node node) {
+            return node;
+        }
+
+        @Override
+        public Node arrayNode(ArrayNode node) {
+            return node.getElements().stream()
+                    .map(element -> element.accept(this))
+                    .collect(ArrayNode.collect());
+        }
+
+        @Override
+        public Node objectNode(ObjectNode node) {
+            return node.getMembers().entrySet().stream()
+                    .map(entry -> Pair.of(entry.getKey().accept(this), entry.getValue().accept(this)))
+                    .collect(ObjectNode.collect(pair -> pair.getLeft().expectStringNode(), Pair::getRight));
+        }
+
+        @Override
+        public Node stringNode(StringNode node) {
+            if (toRename.contains(node.getValue())) {
+                ShapeId nodeShapeId = node.expectShapeId();
+                return new StringNode(shapeMapping.get(nodeShapeId).toString(), node.getSourceLocation());
+            }
+            return node;
+        }
+    }
+}