Support @default

Add support for the @default trait. Issue: #1860 Reference for the trait: https://awslabs.github.io/smithy/2.0/spec/type-refinement-traits.html Signed-off-by: Daniele Ahmed <ahmeddan@amazon.de>
smithy-lang · Oct 20, 2022 · 6e45344 · 6e45344
1 parent e74ecf3
commit 6e45344
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 2 deletions.
diff --git a/codegen-core/common-test-models/simple.smithy b/codegen-core/common-test-models/simple.smithy
@@ -1,5 +1,4 @@
-$version: "1.0"
-
+$version: "2"
 namespace com.amazonaws.simple
 
 use aws.protocols#restJson1
@@ -17,6 +16,7 @@ service SimpleService {
     operations: [
         Healthcheck,
         StoreServiceBlob,
+        DefaultValue,
     ],
 }
 
@@ -132,3 +132,22 @@ structure StoreServiceBlobInput {
 structure StoreServiceBlobOutput {
 
 }
+
+@http(method: "POST", uri: "/default-value")
+operation DefaultValue {
+    input: DefaultValueInput,
+}
+
+structure DefaultValueInput {
+    s: Language = "en"
+    s1: String = "en"
+    i: Integer = 0
+    i1: LanguageLevel = 1
+}
+
+enum Language {
+    EN = "en"
+}
+
+@default(1)
+integer LanguageLevel
diff --git a/...ain/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/BuilderGenerator.kt b/...ain/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/BuilderGenerator.kt
@@ -7,8 +7,23 @@ package software.amazon.smithy.rust.codegen.core.smithy.generators
 
 import software.amazon.smithy.codegen.core.Symbol
 import software.amazon.smithy.model.Model
+import software.amazon.smithy.model.node.BooleanNode
+import software.amazon.smithy.model.node.NullNode
+import software.amazon.smithy.model.node.NumberNode
+import software.amazon.smithy.model.node.StringNode
+import software.amazon.smithy.model.shapes.BooleanShape
+import software.amazon.smithy.model.shapes.DocumentShape
+import software.amazon.smithy.model.shapes.EnumShape
+import software.amazon.smithy.model.shapes.FloatShape
+import software.amazon.smithy.model.shapes.IntEnumShape
+import software.amazon.smithy.model.shapes.IntegerShape
+import software.amazon.smithy.model.shapes.ListShape
+import software.amazon.smithy.model.shapes.MapShape
 import software.amazon.smithy.model.shapes.MemberShape
+import software.amazon.smithy.model.shapes.StringShape
 import software.amazon.smithy.model.shapes.StructureShape
+import software.amazon.smithy.model.traits.DefaultTrait
+import software.amazon.smithy.model.traits.EnumDefinition
 import software.amazon.smithy.rust.codegen.core.rustlang.RustModule
 import software.amazon.smithy.rust.codegen.core.rustlang.RustReservedWords
 import software.amazon.smithy.rust.codegen.core.rustlang.RustType
@@ -34,6 +49,8 @@ import software.amazon.smithy.rust.codegen.core.smithy.isOptional
 import software.amazon.smithy.rust.codegen.core.smithy.makeOptional
 import software.amazon.smithy.rust.codegen.core.smithy.rustType
 import software.amazon.smithy.rust.codegen.core.util.dq
+import software.amazon.smithy.rust.codegen.core.util.expectTrait
+import software.amazon.smithy.rust.codegen.core.util.hasTrait
 import software.amazon.smithy.rust.codegen.core.util.toSnakeCase
 
 fun StructureShape.builderSymbol(symbolProvider: RustSymbolProvider): Symbol {
@@ -258,6 +275,42 @@ class BuilderGenerator(
                 withBlock("$memberName: self.$memberName", ",") {
                     // Write the modifier
                     when {
+                        !memberSymbol.isOptional() && member.hasTrait<DefaultTrait>() -> {
+                            val node = member.expectTrait<DefaultTrait>().toNode()!!
+                            when (val target = model.expectShape(member.target)) {
+                                is EnumShape, is IntEnumShape -> {
+                                    val value = when (target) {
+                                        is IntEnumShape -> node.expectNumberNode().value
+                                        else -> node.expectStringNode().value
+                                    }
+                                    val enumValues = when (target) {
+                                        is IntEnumShape -> target.enumValues
+                                        is EnumShape -> target.enumValues
+                                        else -> mapOf<String, String>()
+                                    }
+                                    val variant = enumValues
+                                        .entries
+                                        .filter { entry -> entry.value == value }
+                                        .map { entry -> symbolProvider.toEnumVariantName(EnumDefinition.builder().name(entry.key).value(entry.value.toString()).build())!! }
+                                        .first()
+                                    val symbol = symbolProvider.toSymbol(target)
+                                    rust(".unwrap_or($symbol::${variant.name})")
+                                }
+                                is IntegerShape, is FloatShape -> rust(".unwrap_or(${node.expectNumberNode().value})")
+                                is BooleanShape -> rust(".unwrap_or(${node.expectBooleanNode().value})")
+                                is StringShape -> rust(".unwrap_or(String::from(${node.expectStringNode().value.dq()}))")
+                                is ListShape, is MapShape -> rust(".unwrap_or_default()")
+                                is DocumentShape -> {
+                                    when (node) {
+                                        is NullNode -> rust(".unwrap_or_default()")
+                                        is BooleanNode -> rust(".unwrap_or(${node.value})")
+                                        is StringNode -> rust(".unwrap_or(String::from(${node.value.dq()}))")
+                                        is NumberNode -> rust(".unwrap_or(${node.value})")
+                                        else -> rust(".unwrap_or_default()")
+                                    }
+                                }
+                            }
+                        }
                         !memberSymbol.isOptional() && default == Default.RustDefault -> rust(".unwrap_or_default()")
                         !memberSymbol.isOptional() -> withBlock(
                             ".ok_or(",