Move the validation into the for loop body in method to avoid two loops.

Author: joewyz

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

@@ -9,7 +9,6 @@
 import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Optional;
 import java.util.Set;
 import java.util.logging.Logger;
 import software.amazon.smithy.model.Model;
@@ -30,11 +29,11 @@
  * containing shape rather than the member.
  */
 final class RemoveShapes {
+    private static final Logger LOGGER = Logger.getLogger(RemoveShapes.class.getName());
+
     private final Collection<? extends Shape> toRemove;
     private final List<ModelTransformerPlugin> plugins;
 
-    private static final Logger LOGGER = Logger.getLogger(RemoveShapes.class.getName());
-
     RemoveShapes(Collection<? extends Shape> toRemove, List<ModelTransformerPlugin> plugins) {
         this.toRemove = toRemove;
         this.plugins = plugins;
@@ -45,8 +44,8 @@ Model transform(ModelTransformer transformer, Model model) {
 
         // Iteratively add each shape that needs to be removed from the index using multiple rounds.
         Set<Shape> removed = new HashSet<>(toRemove);
-        validateMembersFromMixin(model, removed);
         for (Shape removedShape : toRemove) {
+            validateShapeCopiedFromMixin(model, removedShape);
             builder.removeShape(removedShape.getId());
             // We don't need to remove members from the builder since
             // members are automatically removed with the container.
@@ -62,25 +61,24 @@ Model transform(ModelTransformer transformer, Model model) {
         return result;
     }
 
-    private void validateMembersFromMixin(Model model, Collection<? extends Shape> shapesToExclude) {
-        for (Shape shape : shapesToExclude) {
-            // Only validate MemberShapes
-            if (!shape.isMemberShape()) {
-                continue;
-            }
-            MemberShape memberShape = shape.asMemberShape().get();
-            Optional<Shape> container = model.getShape(memberShape.getContainer());
-            if (container.isPresent()) {
-                for (ShapeId mixinId : container.get().getMixins()) {
-                    Shape mixinShape = model.expectShape(mixinId);
-                    if (mixinShape.getMemberNames().contains(shape.getId().getMember().get())) {
-                        LOGGER.warning(format("Removing mixed in member `%s` from mixin shape `%s` "
+    private void validateShapeCopiedFromMixin(Model model, Shape shape) {
+        // Only validate MemberShapes
+        if (!shape.isMemberShape()) {
+            return;
+        }
+        MemberShape memberShape = shape.asMemberShape().get();
+        Shape container = model.getShape(memberShape.getContainer())
+                .orElseThrow(() -> new ModelTransformException(
+                        format("Cannot find the container shape for member `%s`.",
+                                memberShape.getMemberName())));
+        for (ShapeId mixinId : container.getMixins()) {
+            Shape mixinShape = model.expectShape(mixinId);
+            if (mixinShape.getMemberNames().contains(shape.getId().getMember().get())) {
+                LOGGER.warning(format("Removing mixed in member `%s` from mixin shape `%s` "
                                 + "in `%s` will result in an inconsistent model.",
-                                memberShape.getMemberName(),
-                                mixinShape.getId(),
-                                container.get().getId().getName()));
-                    }
-                }
+                        memberShape.getMemberName(),
+                        mixinShape.getId(),
+                        container.getId().getName()));
             }
         }
     }