Add codegenerated scope macro

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ServerCodegenVisitor.kt

@@ -63,6 +63,7 @@ import software.amazon.smithy.rust.codegen.server.smithy.generators.ConstrainedT
 import software.amazon.smithy.rust.codegen.server.smithy.generators.MapConstraintViolationGenerator
 import software.amazon.smithy.rust.codegen.server.smithy.generators.PubCrateConstrainedCollectionGenerator
 import software.amazon.smithy.rust.codegen.server.smithy.generators.PubCrateConstrainedMapGenerator
+import software.amazon.smithy.rust.codegen.server.smithy.generators.ScopeMacroGenerator
 import software.amazon.smithy.rust.codegen.server.smithy.generators.ServerBuilderGenerator
 import software.amazon.smithy.rust.codegen.server.smithy.generators.ServerBuilderGeneratorWithoutPublicConstrainedTypes
 import software.amazon.smithy.rust.codegen.server.smithy.generators.ServerEnumGenerator
@@ -610,6 +611,8 @@ open class ServerCodegenVisitor(
                 codegenContext,
                 serverProtocol,
             ).render(this)
+
+            ScopeMacroGenerator(codegenContext).render(this)
         }
     }
 

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ScopeMacroGenerator.kt

@@ -0,0 +1,137 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rust.codegen.server.smithy.generators
+
+import software.amazon.smithy.model.knowledge.TopDownIndex
+import software.amazon.smithy.rust.codegen.core.rustlang.RustWriter
+import software.amazon.smithy.rust.codegen.core.rustlang.Writable
+import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
+import software.amazon.smithy.rust.codegen.core.rustlang.writable
+import software.amazon.smithy.rust.codegen.server.smithy.ServerCargoDependency
+import software.amazon.smithy.rust.codegen.server.smithy.ServerCodegenContext
+
+class ScopeMacroGenerator(
+    private val codegenContext: ServerCodegenContext,
+) {
+    private val runtimeConfig = codegenContext.runtimeConfig
+    private val codegenScope =
+        arrayOf(
+            "SmithyHttpServer" to ServerCargoDependency.smithyHttpServer(runtimeConfig).toType(),
+        )
+
+    /** Calculate all `operationShape`s contained within the `ServiceShape`. */
+    private val index = TopDownIndex.of(codegenContext.model)
+    private val operations = index.getContainedOperations(codegenContext.serviceShape).toSortedSet(compareBy { it.id })
+
+    fun macro(): Writable = writable {
+        val firstOperationName = codegenContext.symbolProvider.toSymbol(operations.first()).name
+        val operationNames = operations
+            .map { codegenContext.symbolProvider.toSymbol(it).name }
+            .joinToString(" ")
+        val operationBranches = operations
+            .map { codegenContext.symbolProvider.toSymbol(it).name }
+            .map {
+                """
+                // $it Match found, pop from both `member` and `not_member`
+                (@ $ name: ident, $ predicate: ident ($it $($ member: ident)*) ($($ temp: ident)*) ($it $($ not_member: ident)*)) => {
+                    scope! { @ $ name, $ predicate ($($ member)*) ($($ temp)*) ($($ not_member)*) }
+                };
+                // $it match not found, pop from `not_member` into `temp` stack
+                (@ $ name: ident, $ predicate: ident ($it $($ member: ident)*) ($($ temp: ident)*) ($ other: ident $($ not_member: ident)*)) => {
+                    scope! { @ $ name, $ predicate ($it $($ member)*) ($ other $($ temp)*) ($($ not_member)*) }
+                };
+                """
+            }.joinToString("")
+
+        rustTemplate(
+            """
+            /// A macro to help with scoping plugins to a subset of all operations.
+            ///
+            /// In contrast to [`aws_smithy_http_server::scope`](#{SmithyHttpServer}::scope), this macro has knowledge
+            /// of the service and any operations _not_ specified will be placed in the opposing group.
+            ///
+            /// ## Example
+            ///
+            /// ```rust
+            /// scope! {
+            ///     /// Includes [`$firstOperationName`], excluding all other operations
+            ///     struct ScopeA {
+            ///         includes: [$firstOperationName],
+            ///     }
+            /// }
+            ///
+            /// scope! {
+            ///     /// Excludes [`$firstOperationName`], excluding all other operations
+            ///     struct ScopeB {
+            ///         excludes: [$firstOperationName]
+            ///     }
+            /// }
+            /// ```
+            ##[macro_export]
+            macro_rules! scope {
+                // Completed, render impls
+                (@ $ name: ident, $ predicate: ident () ($($ temp: ident)*) ($($ not_member: ident)*)) => {
+                    $(
+                        impl #{SmithyHttpServer}::plugin::scoped::Membership<$ temp> for $ name {
+                            type Contains = #{SmithyHttpServer}::plugin::scoped::$ predicate;
+                        }
+                    )*
+                    $(
+                        impl #{SmithyHttpServer}::plugin::scoped::Membership<$ not_member> for $ name {
+                            type Contains = #{SmithyHttpServer}::plugin::scoped::$ predicate;
+                        }
+                    )*
+                };
+                // All `not_member`s exhausted, move `temp` into `not_member`
+                (@ $ name: ident, $ predicate: ident ($($ member: ident)*) ($($ temp: ident)*) ()) => {
+                    scope! { @ $ name, $ predicate ($($ member)*) () ($($ temp)*) }
+                };
+                $operationBranches
+                (
+                    $(##[$ attrs:meta])*
+                    $ vis:vis struct $ name:ident {
+                        includes: [$($ include:ident),*]
+                    }
+                ) => {
+                    use $ crate::operation_shape::*;
+                    use #{SmithyHttpServer}::scope as scope_runtime;
+                    scope_runtime! {
+                        $(##[$ attrs])*
+                        $ vis struct $ name {
+                            includes: [$($ include),*],
+                            excludes: []
+                        }
+                    }
+                    scope! { @ $ name, False ($($ include)*) () ($operationNames) }
+                };
+                (
+                    $(##[$ attrs:meta])*
+                    $ vis:vis struct $ name:ident {
+                        excludes: [$($ exclude:ident),*]
+                    }
+                ) => {
+                    use $ crate::operation_shape::*;
+                    use #{SmithyHttpServer}::scope as scope_runtime;
+
+                    scope_runtime! {
+                        $(##[$ attrs])*
+                        $ vis struct $ name {
+                            includes: [],
+                            excludes: [$($ exclude),*]
+                        }
+                    }
+                    scope! { @ $ name, True ($($ exclude)*) () ($operationNames) }
+                };
+            }
+            """,
+            *codegenScope,
+        )
+    }
+
+    fun render(writer: RustWriter) {
+        macro()(writer)
+    }
+}

design/src/server/middleware.md

@@ -208,15 +208,24 @@ A "HTTP layer" can be applied to specific operations.
 # extern crate aws_smithy_http_server;
 # use tower::{util::service_fn, Layer};
 # use std::time::Duration;
-# use pokemon_service_server_sdk::{operation_shape::GetPokemonSpecies, PokemonService, input::*, output::*, error::*};
+# use pokemon_service_server_sdk::{operation_shape::GetPokemonSpecies, input::*, output::*, error::*};
 # use aws_smithy_http_server::{operation::OperationShapeExt, plugin::*, operation::*};
 # let handler = |req: GetPokemonSpeciesInput| async { Result::<GetPokemonSpeciesOutput, GetPokemonSpeciesError>::Ok(todo!()) };
 # struct LoggingLayer;
 # impl LoggingLayer { pub fn new() -> Self { Self } }
 # impl<S> Layer<S> for LoggingLayer { type Service = S; fn layer(&self, svc: S) -> Self::Service { svc } }
+use pokemon_service_server_sdk::{PokemonService, scope};
+
+scope! {
+    /// Only log on `GetPokemonSpecies` and `GetStorage`
+    struct LoggingScope {
+        includes: [GetPokemonSpecies, GetStorage]
+    }
+}
+
 // Construct `LoggingLayer`.
 let logging_plugin = LayerPlugin(LoggingLayer::new());
-let logging_plugin = filter_by_operation_id(logging_plugin, |name| name == GetPokemonSpecies::ID);
+let logging_plugin = Scoped::new::<LoggingScope>(logging_plugin);
 let http_plugins = PluginPipeline::new().push(logging_plugin);
 
 let app /* : PokemonService<Route<B>> */ = PokemonService::builder_with_plugins(http_plugins, IdentityPlugin)
@@ -244,11 +253,18 @@ A "model layer" can be applied to specific operations.
 # struct BufferLayer;
 # impl BufferLayer { pub fn new(size: usize) -> Self { Self } }
 # impl<S> Layer<S> for BufferLayer { type Service = S; fn layer(&self, svc: S) -> Self::Service { svc } }
-use pokemon_service_server_sdk::PokemonService;
+use pokemon_service_server_sdk::{PokemonService, scope};
+
+scope! {
+    /// Only buffer on `GetPokemonSpecies` and `GetStorage`
+    struct BufferScope {
+        includes: [GetPokemonSpecies, GetStorage]
+    }
+}
 
 // Construct `BufferLayer`.
 let buffer_plugin = LayerPlugin(BufferLayer::new(3));
-let buffer_plugin = filter_by_operation_id(buffer_plugin, |name| name != GetPokemonSpecies::ID);
+let buffer_plugin = Scoped::new::<BufferScope>(buffer_plugin);
 let model_plugins = PluginPipeline::new().push(buffer_plugin);
 
 let app /* : PokemonService<Route<B>> */ = PokemonService::builder_with_plugins(IdentityPlugin, model_plugins)

examples/pokemon-service/src/main.rs

@@ -10,7 +10,7 @@ use std::{net::SocketAddr, sync::Arc};
 use aws_smithy_http_server::{
     extension::OperationExtensionExt,
     instrumentation::InstrumentExt,
-    plugin::{alb_health_check::AlbHealthCheckLayer, IdentityPlugin, PluginPipeline},
+    plugin::{alb_health_check::AlbHealthCheckLayer, IdentityPlugin, PluginPipeline, Scoped},
     request::request_id::ServerRequestIdProviderLayer,
     AddExtensionLayer,
 };
@@ -26,7 +26,7 @@ use pokemon_service_common::{
     capture_pokemon, check_health, get_pokemon_species, get_server_statistics, setup_tracing,
     stream_pokemon_radio, State,
 };
-use pokemon_service_server_sdk::PokemonService;
+use pokemon_service_server_sdk::{scope, PokemonService};
 
 #[derive(Parser, Debug)]
 #[clap(author, version, about, long_about = None)]
@@ -44,9 +44,18 @@ pub async fn main() {
     let args = Args::parse();
     setup_tracing();
 
+    scope! {
+        /// A scope containing `GetPokemonSpecies` and `GetStorage`
+        struct PrintScope {
+            includes: [GetPokemonSpecies, GetStorage]
+        }
+    }
+    // Scope the `PrintPlugin`, defined in `plugin.rs`, to `PrintScope`
+    let print_plugin = Scoped::new::<PrintScope>(PluginPipeline::new().print());
+
     let plugins = PluginPipeline::new()
-        // Apply the `PrintPlugin` defined in `plugin.rs`
-        .print()
+        // Apply the scoped `PrintPlugin`
+        .push(print_plugin)
         // Apply the `OperationExtensionPlugin` defined in `aws_smithy_http_server::extension`. This allows other
         // plugins or tests to access a `aws_smithy_http_server::extension::OperationExtension` from
         // `Response::extensions`, or infer routing failure when it's missing.

rust-runtime/aws-smithy-http-server/src/plugin/scoped.rs

@@ -103,7 +103,11 @@ where
 
 /// A macro to help with scoping plugins to a subset of all operations.
 ///
-/// The scope must partition _all_ operations, that is, each and every operation must be included or excluded, but not both.
+/// The scope must partition _all_ operations, that is, each and every operation must be included or excluded, but not
+/// both.
+///
+/// The generated server SDK exports a similar `scope` macro which is aware of a services operations and can complete
+/// underspecified scopes automatically.
 ///
 /// # Example
 ///