Add method to get non-trait shapes

smithy-build/src/main/java/software/amazon/smithy/build/PluginContext.java

@@ -28,7 +28,9 @@
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.node.Node;
 import software.amazon.smithy.model.node.ObjectNode;
+import software.amazon.smithy.model.shapes.ShapeIndex;
 import software.amazon.smithy.model.shapes.ToShapeId;
+import software.amazon.smithy.model.transform.ModelTransformer;
 import software.amazon.smithy.model.validation.ValidationEvent;
 import software.amazon.smithy.utils.SetUtils;
 import software.amazon.smithy.utils.SmithyBuilder;
@@ -46,6 +48,7 @@ public final class PluginContext {
     private final FileManifest fileManifest;
     private final ClassLoader pluginClassLoader;
     private final Set<Path> sources;
+    private ShapeIndex nonTraitsIndex;
 
     private PluginContext(Builder builder) {
         model = SmithyBuilder.requiredState("model", builder.model);
@@ -144,6 +147,30 @@ public Optional<ClassLoader> getPluginClassLoader() {
         return Optional.ofNullable(pluginClassLoader);
     }
 
+    /**
+     * Gets all shapes from a model as a {@code ShapeIndex} where shapes that
+     * define traits or shapes that are only used as part of a trait
+     * definition have been removed.
+     *
+     * <p>This is typically functionality used by code generators when
+     * generating data structures from a model. It's useful because it only
+     * provides shapes that are used to describe data structures rather than
+     * shapes used to describe metadata about the data structures.
+     *
+     * <p>Note: this method just calls {@link ModelTransformer#getNonTraitShapes}.
+     * It's added to {@code PluginContext} to make it more easily available
+     * to code generators.
+     *
+     * @return Returns a ShapeIndex containing matching shapes.
+     */
+    public synchronized ShapeIndex getNonTraitShapes() {
+        if (nonTraitsIndex == null) {
+            nonTraitsIndex = ModelTransformer.create().getNonTraitShapes(model);
+        }
+
+        return nonTraitsIndex;
+    }
+
     /**
      * Gets the source models, or models that are considered the subject
      * of the build.

smithy-build/src/test/java/software/amazon/smithy/build/PluginContextTest.java

@@ -8,6 +8,8 @@
 import org.junit.jupiter.api.Test;
 import software.amazon.smithy.build.model.ProjectionConfig;
 import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.shapes.ShapeIndex;
+import software.amazon.smithy.model.transform.ModelTransformer;
 import software.amazon.smithy.utils.ListUtils;
 
 public class PluginContextTest {
@@ -42,4 +44,21 @@ public void hasSources() {
 
         assertThat(context.getSources(), contains(Paths.get("/foo/baz")));
     }
+
+    @Test
+    public void createsNonTraitShapeIndex() {
+        Model model = Model.assembler()
+                .addImport(getClass().getResource("simple-model.json"))
+                .assemble()
+                .unwrap();
+        ShapeIndex scrubbed = ModelTransformer.create().getNonTraitShapes(model);
+        PluginContext context = PluginContext.builder()
+                .fileManifest(new MockManifest())
+                .model(model)
+                .sources(ListUtils.of(Paths.get("/foo/baz")))
+                .build();
+
+        assertThat(context.getNonTraitShapes(), equalTo(scrubbed));
+        assertThat(context.getNonTraitShapes(), equalTo(scrubbed)); // trigger loading from cache
+    }
 }

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

@@ -18,6 +18,7 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.Optional;
 import java.util.ServiceLoader;
 import java.util.function.BiFunction;
 import java.util.function.BiPredicate;
@@ -28,6 +29,7 @@
 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.ShapeIndex;
 import software.amazon.smithy.model.traits.Trait;
 import software.amazon.smithy.model.traits.TraitDefinition;
 import software.amazon.smithy.utils.FunctionalUtils;
@@ -335,4 +337,31 @@ public Model removeUnreferencedTraitDefinitions(Model model, Predicate<Shape> ke
     public Model scrubTraitDefinitions(Model model) {
         return new ScrubTraitDefinitions().transform(this, model);
     }
+
+    /**
+     * Gets all shapes from a model as a {@code ShapeIndex} where shapes that
+     * define traits or shapes that are only used as part of a trait
+     * definition have been removed.
+     *
+     * @param model Model that contains shapes.
+     * @return Returns a ShapeIndex containing matching shapes.
+     */
+    public ShapeIndex getNonTraitShapes(Model model) {
+        ShapeIndex currentIndex = model.getShapeIndex();
+        ShapeIndex.Builder indexBuilder = ShapeIndex.builder();
+
+        // ScrubTraitDefinitions is used to removed traits and trait shapes.
+        // However, the returned model can't be returned directly because
+        // as traits are removed, uses of that trait are removed. Instead,
+        // a ShapeIndex is created by getting all shape IDs from the modified
+        // model, grabbing shapes from the original model, and building a new
+        // ShapeIndex.
+        scrubTraitDefinitions(model).getShapeIndex().shapes()
+                .map(Shape::getId)
+                .map(currentIndex::getShape)
+                .map(Optional::get)
+                .forEach(indexBuilder::addShape);
+
+        return indexBuilder.build();
+    }
 }

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

@@ -26,6 +26,8 @@
 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.EnumTrait;
+import software.amazon.smithy.model.traits.ReadonlyTrait;
 
 public class ModelTransformerTest {
 
@@ -46,6 +48,18 @@ public void discoversOnRemoveClassesWithSpi() {
                           Matchers.equalTo(Optional.of(Collections.emptyList())));
     }
 
+    @Test
+    public void removesTraitShapesButNotTraitUsage() {
+        ModelTransformer transformer = ModelTransformer.create();
+        Model model = createTestModel();
+        ShapeIndex index = transformer.getNonTraitShapes(model);
+        ShapeId operation = ShapeId.from("ns.foo#MyOperation");
+
+        assertThat(index.getShape(operation), Matchers.not(Optional.empty()));
+        assertThat(index.getShape(operation).get().getTrait(ReadonlyTrait.class), Matchers.not(Optional.empty()));
+        assertThat(index.getShape(EnumTrait.ID), Matchers.equalTo(Optional.empty()));
+    }
+
     private Model createTestModel() {
         return Model.assembler()
                 .addImport(ModelTransformerTest.class.getResource("test-model.json"))