fix(codegen): use TopDownIndex::getContainedOperations() for operation iterations (#1109)

Author: Steven Yuan

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/RuntimeConfigGenerator.java

@@ -28,6 +28,7 @@
 import software.amazon.smithy.codegen.core.SymbolProvider;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.knowledge.ServiceIndex;
+import software.amazon.smithy.model.knowledge.TopDownIndex;
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.typescript.codegen.auth.AuthUtils;
@@ -229,8 +230,9 @@ private void generateHttpAuthSchemeConfig(
 
         // feat(experimentalIdentityAndAuth): gather HttpAuthSchemes to generate
         ServiceIndex serviceIndex = ServiceIndex.of(model);
+        TopDownIndex topDownIndex = TopDownIndex.of(model);
         Map<ShapeId, HttpAuthScheme> allEffectiveHttpAuthSchemes =
-            AuthUtils.getAllEffectiveNoAuthAwareAuthSchemes(service, serviceIndex, authIndex);
+            AuthUtils.getAllEffectiveNoAuthAwareAuthSchemes(service, serviceIndex, authIndex, topDownIndex);
         List<HttpAuthSchemeTarget> targetAuthSchemes = getHttpAuthSchemeTargets(target, allEffectiveHttpAuthSchemes);
 
         // Generate only if the "inherited" target is different than the current target

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/auth/AuthUtils.java

@@ -17,6 +17,7 @@
 import software.amazon.smithy.codegen.core.SymbolDependency;
 import software.amazon.smithy.model.knowledge.ServiceIndex;
 import software.amazon.smithy.model.knowledge.ServiceIndex.AuthSchemeMode;
+import software.amazon.smithy.model.knowledge.TopDownIndex;
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.traits.Trait;
@@ -75,15 +76,16 @@ private AuthUtils() {}
     public static Map<ShapeId, HttpAuthScheme> getAllEffectiveNoAuthAwareAuthSchemes(
         ServiceShape serviceShape,
         ServiceIndex serviceIndex,
-        SupportedHttpAuthSchemesIndex authIndex
+        SupportedHttpAuthSchemesIndex authIndex,
+        TopDownIndex topDownIndex
     ) {
         Map<ShapeId, HttpAuthScheme> effectiveAuthSchemes = new TreeMap<>();
         var serviceEffectiveAuthSchemes =
             serviceIndex.getEffectiveAuthSchemes(serviceShape, AuthSchemeMode.NO_AUTH_AWARE);
         for (ShapeId shapeId : serviceEffectiveAuthSchemes.keySet()) {
             effectiveAuthSchemes.put(shapeId, authIndex.getHttpAuthScheme(shapeId));
         }
-        for (var operation : serviceShape.getAllOperations()) {
+        for (var operation : topDownIndex.getContainedOperations(serviceShape)) {
             var operationEffectiveAuthSchemes =
                 serviceIndex.getEffectiveAuthSchemes(serviceShape, operation, AuthSchemeMode.NO_AUTH_AWARE);
             for (ShapeId shapeId : operationEffectiveAuthSchemes.keySet()) {

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/auth/http/HttpAuthSchemeProviderGenerator.java

@@ -16,6 +16,8 @@
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.knowledge.ServiceIndex;
 import software.amazon.smithy.model.knowledge.ServiceIndex.AuthSchemeMode;
+import software.amazon.smithy.model.knowledge.TopDownIndex;
+import software.amazon.smithy.model.shapes.OperationShape;
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.traits.Trait;
@@ -59,6 +61,7 @@ public class HttpAuthSchemeProviderGenerator implements Runnable {
 
     private final SupportedHttpAuthSchemesIndex authIndex;
     private final ServiceIndex serviceIndex;
+    private final TopDownIndex topDownIndex;
     private final ServiceShape serviceShape;
     private final Symbol serviceSymbol;
     private final String serviceName;
@@ -88,11 +91,12 @@ public HttpAuthSchemeProviderGenerator(
 
         this.authIndex = new SupportedHttpAuthSchemesIndex(integrations, model, settings);
         this.serviceIndex = ServiceIndex.of(model);
+        this.topDownIndex = TopDownIndex.of(model);
         this.serviceShape = settings.getService(model);
         this.serviceSymbol = symbolProvider.toSymbol(serviceShape);
         this.serviceName = CodegenUtils.getServiceName(settings, model, symbolProvider);
         this.effectiveHttpAuthSchemes =
-            AuthUtils.getAllEffectiveNoAuthAwareAuthSchemes(serviceShape, serviceIndex, authIndex);
+            AuthUtils.getAllEffectiveNoAuthAwareAuthSchemes(serviceShape, serviceIndex, authIndex, topDownIndex);
         this.httpAuthSchemeParameters =
             AuthUtils.collectHttpAuthSchemeParameters(effectiveHttpAuthSchemes.values());
     }
@@ -395,7 +399,8 @@ private void generateDefaultHttpAuthSchemeProviderFunction() {
                 w.openBlock("switch (authParameters.operation) {", "};", () -> {
                     var serviceAuthSchemes = serviceIndex.getEffectiveAuthSchemes(
                         serviceShape, AuthSchemeMode.NO_AUTH_AWARE);
-                    for (ShapeId operationShapeId : serviceShape.getAllOperations()) {
+                    for (OperationShape operationShape : topDownIndex.getContainedOperations(serviceShape)) {
+                        ShapeId operationShapeId = operationShape.getId();
                         var operationAuthSchemes = serviceIndex.getEffectiveAuthSchemes(
                             serviceShape, operationShapeId, AuthSchemeMode.NO_AUTH_AWARE);
                         // Skip operation generation if operation auth schemes are equivalent to the default service

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/auth/http/integration/AddHttpAuthSchemePlugin.java

@@ -11,6 +11,7 @@
 import java.util.stream.Collectors;
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.model.knowledge.ServiceIndex;
+import software.amazon.smithy.model.knowledge.TopDownIndex;
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.rulesengine.traits.EndpointRuleSetTrait;
@@ -145,8 +146,9 @@ private getIdentityProviderConfigProvider() {
                         s.getModel(),
                         s.getSettings());
                     ServiceIndex serviceIndex = ServiceIndex.of(s.getModel());
-                    Map<ShapeId, HttpAuthScheme> httpAuthSchemes =
-                        AuthUtils.getAllEffectiveNoAuthAwareAuthSchemes(s.getService(), serviceIndex, authIndex);
+                    TopDownIndex topDownIndex = TopDownIndex.of(s.getModel());
+                    Map<ShapeId, HttpAuthScheme> httpAuthSchemes = AuthUtils.getAllEffectiveNoAuthAwareAuthSchemes(
+                        s.getService(), serviceIndex, authIndex, topDownIndex);
                     for (HttpAuthScheme scheme : httpAuthSchemes.values()) {
                         if (scheme == null) {
                             continue;
@@ -177,8 +179,9 @@ public void customize(TypeScriptCodegenContext codegenContext) {
                 codegenContext.settings());
             ServiceShape serviceShape = codegenContext.settings().getService(codegenContext.model());
             ServiceIndex serviceIndex = ServiceIndex.of(codegenContext.model());
+            TopDownIndex topDownIndex = TopDownIndex.of(codegenContext.model());
             Map<ShapeId, HttpAuthScheme> httpAuthSchemes =
-                AuthUtils.getAllEffectiveNoAuthAwareAuthSchemes(serviceShape, serviceIndex, authIndex);
+                AuthUtils.getAllEffectiveNoAuthAwareAuthSchemes(serviceShape, serviceIndex, authIndex, topDownIndex);
             Map<String, ConfigField> configFields =
                 AuthUtils.collectConfigFields(httpAuthSchemes.values());
             Map<Symbol, ResolveConfigFunction> resolveConfigFunctions =

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/auth/http/integration/HttpAuthRuntimeExtensionIntegration.java

@@ -9,6 +9,7 @@
 import java.util.Map;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.knowledge.ServiceIndex;
+import software.amazon.smithy.model.knowledge.TopDownIndex;
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.typescript.codegen.CodegenUtils;
@@ -60,11 +61,12 @@ public void customize(TypeScriptCodegenContext codegenContext) {
             codegenContext.model(),
             codegenContext.settings());
         ServiceIndex serviceIndex = ServiceIndex.of(codegenContext.model());
+        TopDownIndex topDownIndex = TopDownIndex.of(codegenContext.model());
         String serviceName = CodegenUtils.getServiceName(
             codegenContext.settings(), codegenContext.model(), codegenContext.symbolProvider());
         ServiceShape serviceShape = codegenContext.settings().getService(codegenContext.model());
         Map<ShapeId, HttpAuthScheme> effectiveAuthSchemes =
-            AuthUtils.getAllEffectiveNoAuthAwareAuthSchemes(serviceShape, serviceIndex, authIndex);
+            AuthUtils.getAllEffectiveNoAuthAwareAuthSchemes(serviceShape, serviceIndex, authIndex, topDownIndex);
         Map<String, ConfigField> configFields = AuthUtils.collectConfigFields(effectiveAuthSchemes.values());
 
         generateHttpAuthExtensionConfigurationInterface(

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/AddHttpApiKeyAuthPlugin.java

@@ -13,9 +13,9 @@
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.knowledge.ServiceIndex;
+import software.amazon.smithy.model.knowledge.TopDownIndex;
 import software.amazon.smithy.model.shapes.OperationShape;
 import software.amazon.smithy.model.shapes.ServiceShape;
-import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.traits.HttpApiKeyAuthTrait;
 import software.amazon.smithy.model.traits.OptionalAuthTrait;
 import software.amazon.smithy.typescript.codegen.CodegenUtils;
@@ -155,8 +155,8 @@ private static boolean hasEffectiveHttpApiKeyAuthTrait(
         ServiceShape service
     ) {
         ServiceIndex serviceIndex = ServiceIndex.of(model);
-        for (ShapeId id : service.getAllOperations()) {
-            OperationShape operation = model.expectShape(id, OperationShape.class);
+        TopDownIndex topDownIndex = TopDownIndex.of(model);
+        for (OperationShape operation : topDownIndex.getContainedOperations(service)) {
             if (operation.hasTrait(OptionalAuthTrait.ID)) {
                 continue;
             }