*squish*

smithy-lang · Sep 19, 2024 · 75ecb60 · 75ecb60
1 parent fdd8c9a
commit 75ecb60
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 27 deletions.
diff --git a/...c/main/java/software/amazon/polymorph/smithyrust/generator/AbstractRustShimGenerator.java b/...c/main/java/software/amazon/polymorph/smithyrust/generator/AbstractRustShimGenerator.java
@@ -1462,7 +1462,7 @@ protected String rustUnionName(final UnionShape unionShape) {
 
   protected String qualifiedRustUnionName(final UnionShape unionShape) {
     return "%s::%s".formatted(
-        getRustTypesModuleName(),
+        mergedGenerator.generatorForShape(unionShape).getRustTypesModuleName(),
         rustUnionName(unionShape)
       );
   }
@@ -1513,6 +1513,27 @@ protected HashMap<String, String> unionMemberVariables(
     return variables;
   }
 
+  /**
+   * Generates values for variables commonly used in structure-member-specific templates.
+   */
+  protected HashMap<String, String> structureMemberVariables(
+    final MemberShape memberShape
+  ) {
+
+    final HashMap<String, String> variables = new HashMap<>();
+    final String memberName = memberShape.getMemberName();
+    final Shape targetShape = model.expectShape(memberShape.getTarget());
+    variables.put("memberName", memberName);
+    variables.put("fieldName", toSnakeCase(memberName));
+    variables.put(
+      "fieldType",
+      mergedGeneratorRustTypeForShape(targetShape)
+    );
+    return variables;
+  }
+
+
+
 
   protected String qualifiedRustServiceErrorType() {
     return "%s::error::Error".formatted(

diff --git a/...src/main/java/software/amazon/polymorph/smithyrust/generator/RustAwsSdkShimGenerator.java b/...src/main/java/software/amazon/polymorph/smithyrust/generator/RustAwsSdkShimGenerator.java
@@ -20,6 +20,7 @@
 import software.amazon.smithy.model.traits.BoxTrait;
 import software.amazon.smithy.model.traits.DefaultTrait;
 import software.amazon.smithy.model.traits.EnumTrait;
+import software.amazon.smithy.model.traits.UnitTypeTrait;
 
 import java.nio.file.Path;
 import java.util.HashMap;
@@ -175,17 +176,25 @@ private TokenTree operationClientFunction(
       operationVariables(bindingShape, operationShape)
     );
 
-    final ShapeId outputShapeId = operationShape.getOutputShape();
-    final String outputType = outputShapeId.equals(
-        ShapeId.from("smithy.api#Unit")
-      )
+    final StructureShape inputShape = operationIndex.getInputShape(operationShape).get();
+    final StructureShape outputShape = operationIndex.getOutputShape(operationShape).get();
+    final String outputType = outputShape.hasTrait(UnitTypeTrait.class)
       ? "()"
       : evalTemplate(
         "std::rc::Rc<crate::r#$dafnyTypesModuleName:L::$operationOutputName:L>",
         variables
       );
     variables.put("outputType", outputType);
 
+    variables.put("fluentSetters", inputShape.members()
+      .stream()
+      .map(member -> evalTemplate(".set_$fieldName:L(inner_input.$fieldName:L)", structureMemberVariables(member)))
+      .collect(Collectors.joining("\n")));
+
+    variables.put("outputToDafnyMapper", outputShape.hasTrait(UnitTypeTrait.class)
+        ? "|x| ()"
+        : evalTemplate("$rustRootModuleName:L::conversions::$snakeCaseOperationName:L::_$snakeCaseOperationName:L_response::to_dafny", variables));
+
     return TokenTree.of(
       evalTemplate(
         """
@@ -195,10 +204,13 @@ private TokenTree operationClientFunction(
             std::rc::Rc<crate::r#$dafnyTypesModuleName:L::Error>
           >
         > {
-          let native_result =\s
-            dafny_tokio_runtime.block_on($rustRootModuleName:L::conversions::$snakeCaseOperationName:L::_$snakeCaseOperationName:L_request::from_dafny(input.clone()).send());
+          let inner_input = $rustRootModuleName:L::conversions::$snakeCaseOperationName:L::_$snakeCaseOperationName:L_request::from_dafny(input.clone());
+          let native_result = dafny_tokio_runtime.block_on(
+            self.inner.$snakeCaseOperationName:L()
+              $fluentSetters:L
+              .send());
           crate::standard_library_conversions::result_to_dafny(&native_result,\s
-            $rustRootModuleName:L::conversions::$snakeCaseOperationName:L::_$snakeCaseOperationName:L_response::to_dafny,
+            $outputToDafnyMapper:L,
             $rustRootModuleName:L::conversions::$snakeCaseOperationName:L::to_dafny_error)
         }
         """,

diff --git a/...rc/main/java/software/amazon/polymorph/smithyrust/generator/RustLibraryShimGenerator.java b/...rc/main/java/software/amazon/polymorph/smithyrust/generator/RustLibraryShimGenerator.java
@@ -1672,25 +1672,6 @@ protected String syntheticOperationOutputName(OperationShape operationShape) {
     return operationName(operationShape) + "Output";
   }
 
-  /**
-   * Generates values for variables commonly used in structure-member-specific templates.
-   */
-  private HashMap<String, String> structureMemberVariables(
-    final MemberShape memberShape
-  ) {
-
-    final HashMap<String, String> variables = new HashMap<>();
-    final String memberName = memberShape.getMemberName();
-    final Shape targetShape = model.expectShape(memberShape.getTarget());
-    variables.put("memberName", memberName);
-    variables.put("fieldName", toSnakeCase(memberName));
-    variables.put(
-      "fieldType",
-      mergedGeneratorRustTypeForShape(targetShape)
-    );
-    return variables;
-  }
-
   private Map<String, String> dependentServiceErrorVariables(
     final ServiceShape serviceShape,
     final ServiceShape dependentServiceShape