Add pagination flattening transform (#2454)

Adds a transform that flattens service-level pagination info into pagination traits on operations.

smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/CodegenDirector.java

@@ -302,6 +302,19 @@ public void changeStringEnumsToEnumShapes(boolean synthesizeEnumNames) {
         });
     }
 
+    /**
+     * Flattens service-level pagination information into operation pagination traits.
+     *
+     * @see ModelTransformer#flattenPaginationInfoIntoOperations(Model, ServiceShape)
+     */
+    public void flattenPaginationInfoIntoOperations() {
+        transforms.add((model, transformer) -> {
+            LOGGER.finest("Flattening pagination info into operation traits for directed codegen");
+            return transformer.flattenPaginationInfoIntoOperations(model,
+                    model.expectShape(service, ServiceShape.class));
+        });
+    }
+
     /**
      * Sets the shapes order for code generation.
      *

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

@@ -0,0 +1,52 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.model.transform;
+
+import java.util.HashSet;
+import java.util.Optional;
+import java.util.Set;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.knowledge.PaginatedIndex;
+import software.amazon.smithy.model.knowledge.PaginationInfo;
+import software.amazon.smithy.model.shapes.OperationShape;
+import software.amazon.smithy.model.shapes.ServiceShape;
+import software.amazon.smithy.model.shapes.Shape;
+import software.amazon.smithy.model.traits.PaginatedTrait;
+
+/**
+ * Flattens pagination info from service shapes into operation-level pagination traits.
+ */
+final class FlattenPaginationInfo {
+
+    private final ServiceShape service;
+
+    FlattenPaginationInfo(ServiceShape service) {
+        this.service = service;
+    }
+
+    public Model transform(ModelTransformer transformer, Model model) {
+        Optional<PaginatedTrait> serviceLevelPagination = service.getTrait(PaginatedTrait.class);
+        if (!serviceLevelPagination.isPresent()) {
+            return model;
+        }
+        PaginatedIndex paginatedIndex = PaginatedIndex.of(model);
+
+        // Merge service-level information into each operation's pagination trait.
+        Set<Shape> updatedShapes = new HashSet<>();
+        for (OperationShape operationShape : model.getOperationShapesWithTrait(PaginatedTrait.class)) {
+            PaginationInfo paginationInfo = paginatedIndex.getPaginationInfo(service, operationShape).get();
+            OperationShape updatedShape = operationShape.toBuilder()
+                    .addTrait(paginationInfo.getPaginatedTrait())
+                    .build();
+            updatedShapes.add(updatedShape);
+        }
+
+        // Remove the paginated trait from the service as it's info has been flattened into the operations
+        updatedShapes.add(service.toBuilder().removeTrait(PaginatedTrait.ID).build());
+
+        return transformer.replaceShapes(model, updatedShapes);
+    }
+}

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

@@ -677,4 +677,14 @@ public Model removeInvalidDefaults(Model model) {
     public Model deconflictErrorsWithSharedStatusCode(Model model, ServiceShape forService) {
         return new DeconflictErrorsWithSharedStatusCode(forService).transform(this, model);
     }
+
+    /**
+     * Flattens all service-level pagination information into operation-level pagination traits.
+     *
+     * @param model Model to transform.
+     * @return Returns the transformed model.
+     */
+    public Model flattenPaginationInfoIntoOperations(Model model, ServiceShape forService) {
+        return new FlattenPaginationInfo(forService).transform(this, model);
+    }
 }

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

@@ -0,0 +1,42 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.model.transform;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+
+import org.junit.jupiter.api.Test;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.shapes.ServiceShape;
+import software.amazon.smithy.model.shapes.Shape;
+import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.model.traits.PaginatedTrait;
+
+public class FlattenPaginationInfoTest {
+    private static final ShapeId serviceId = ShapeId.from("smithy.example#PaginatedService");
+    private static final ShapeId operationId = ShapeId.from("smithy.example#PaginatedOperation");
+
+    @Test
+    void compareTransform() {
+        Model before = Model.assembler()
+                .addImport(FlattenPaginationInfoTest.class.getResource("flatten-pagination-before.smithy"))
+                .assemble()
+                .unwrap();
+        ServiceShape service = before.expectShape(serviceId).asServiceShape().get();
+        Model result = ModelTransformer.create().flattenPaginationInfoIntoOperations(before, service);
+
+        Shape resultService = result.expectShape(serviceId);
+        assertFalse(resultService.hasTrait(PaginatedTrait.class));
+
+
+        Shape resultOperation = result.expectShape(operationId);
+        PaginatedTrait resultTrait = resultOperation.expectTrait(PaginatedTrait.class);
+        assertEquals(resultTrait.getInputToken().get(), "nextToken");
+        assertEquals(resultTrait.getOutputToken().get(), "nextToken");
+        assertEquals(resultTrait.getPageSize().get(), "maxResults");
+        assertEquals(resultTrait.getItems().get(), "foos");
+    }
+}

smithy-model/src/test/resources/software/amazon/smithy/model/transform/flatten-pagination-before.smithy

@@ -0,0 +1,29 @@
+$version: "2.0"
+
+namespace smithy.example
+
+@paginated(inputToken: "nextToken", outputToken: "nextToken")
+service PaginatedService {
+    operations: [
+        PaginatedOperation
+    ]
+}
+
+@paginated(pageSize: "maxResults", items: "foos")
+operation PaginatedOperation {
+    input := {
+        maxResults: Integer
+        nextToken: String
+    }
+    output := {
+        nextToken: String
+
+        @required
+        foos: StringList
+    }
+}
+
+@private
+list StringList {
+    member: String
+}