Symbol provider to inject custom symbols (#3209)

## Description
Add a SymbolProvider that uses a custom trait to return the symbol.

## Testing
Unit tests.



----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._

---------

Signed-off-by: Daniele Ahmed <ahmeddan@amazon.de>

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

@@ -0,0 +1,26 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rust.codegen.server.smithy
+
+import software.amazon.smithy.codegen.core.Symbol
+import software.amazon.smithy.model.node.Node
+import software.amazon.smithy.model.shapes.Shape
+import software.amazon.smithy.model.shapes.ShapeId
+import software.amazon.smithy.model.traits.AnnotationTrait
+import software.amazon.smithy.rust.codegen.core.smithy.RustSymbolProvider
+import software.amazon.smithy.rust.codegen.core.smithy.WrappingSymbolProvider
+import software.amazon.smithy.rust.codegen.core.util.getTrait
+
+class SyntheticCustomShapeTrait(val ID: ShapeId, val symbol: Symbol) : AnnotationTrait(ID, Node.objectNode())
+
+/**
+ * This SymbolProvider honors [`SyntheticCustomShapeTrait`] and returns this trait's symbol when available.
+ */
+class CustomShapeSymbolProvider(private val base: RustSymbolProvider) : WrappingSymbolProvider(base) {
+ override fun toSymbol(shape: Shape): Symbol {
+ return shape.getTrait<SyntheticCustomShapeTrait>()?.symbol ?: base.toSymbol(shape)
+ }
+}

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

@@ -93,5 +93,7 @@ class RustServerCodegenPlugin : ServerDecoratableBuildPlugin() {
  .let { RustReservedWordSymbolProvider(it, ServerReservedWords) }
  // Allows decorators to inject a custom symbol provider
  .let { codegenDecorator.symbolProvider(it) }
+ // Inject custom symbols.
+ .let { CustomShapeSymbolProvider(it) }
  }
 }

codegen-server/src/test/kotlin/software/amazon/smithy/rust/codegen/server/smithy/CustomShapeSymbolProviderTest.kt

@@ -0,0 +1,63 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rust.codegen.server.smithy
+
+import io.kotest.matchers.shouldBe
+import org.junit.jupiter.api.Test
+import software.amazon.smithy.codegen.core.Symbol
+import software.amazon.smithy.model.shapes.MemberShape
+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.transform.ModelTransformer
+import software.amazon.smithy.rust.codegen.core.rustlang.RustType
+import software.amazon.smithy.rust.codegen.core.smithy.rustType
+import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel
+import software.amazon.smithy.rust.codegen.core.util.lookup
+import software.amazon.smithy.rust.codegen.server.smithy.testutil.serverTestSymbolProvider
+
+class CustomShapeSymbolProviderTest {
+ private val baseModel =
+ """
+ namespace test
+
+ service TestService {
+ version: "1"
+ operations: [TestOperation]
+ }
+
+ operation TestOperation {
+ input: TestInputOutput
+ output: TestInputOutput
+ }
+
+ structure TestInputOutput {
+ myString: String,
+ }
+ """.asSmithyModel(smithyVersion = "2.0")
+ private val serviceShape = baseModel.lookup<ServiceShape>("test#TestService")
+ private val rustType = RustType.Opaque("fake-type")
+ private val symbol = Symbol.builder()
+ .name("fake-symbol")
+ .rustType(rustType)
+ .build()
+ private val model = ModelTransformer.create()
+ .mapShapes(baseModel) {
+ if (it is MemberShape) {
+ it.toBuilder().addTrait(SyntheticCustomShapeTrait(ShapeId.from("some#id"), symbol)).build()
+ } else {
+ it
+ }
+ }
+ private val symbolProvider = serverTestSymbolProvider(baseModel, serviceShape)
+ .let { CustomShapeSymbolProvider(it) }
+
+ @Test
+ fun `override with custom symbol`() {
+ val shape = model.lookup<Shape>("test#TestInputOutput\$myString")
+ symbolProvider.toSymbol(shape) shouldBe symbol
+ }
+}