fix: Misc Rust bug fixes, mostly around @positional (#581)

smithy-lang · Sep 12, 2024 · 779d38d · 779d38d
1 parent 3e098eb
commit 779d38d
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 67 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
@@ -14,9 +14,11 @@
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
+import software.amazon.polymorph.smithydafny.DafnyNameResolver;
 import software.amazon.polymorph.traits.DafnyUtf8BytesTrait;
 import software.amazon.polymorph.traits.PositionalTrait;
 import software.amazon.polymorph.traits.ReferenceTrait;
+import software.amazon.polymorph.utils.DafnyNameResolverHelpers;
 import software.amazon.polymorph.utils.IOUtils;
 import software.amazon.polymorph.utils.MapUtils;
 import software.amazon.polymorph.utils.ModelUtils;
@@ -505,14 +507,19 @@ protected TokenTree fromDafny(
         }
       }
       case INTEGER -> {
-        if (isRustOption) {
+        if (isDafnyOption) {
           yield TokenTree.of(
             "crate::standard_library_conversions::oint_from_dafny(%s.clone())".formatted(
                 dafnyValue
               )
           );
         } else {
-          yield TokenTree.of(dafnyValue, ".clone()");
+          TokenTree result = TokenTree.of(dafnyValue, ".clone()");
+          if (isRustOption) {
+            result =
+              TokenTree.of(TokenTree.of("Some("), result, TokenTree.of(")"));
+          }
+          yield result;
         }
       }
       case LONG -> {
@@ -908,23 +915,17 @@ protected RustFile enumConversionModule(final EnumShape enumShape) {
   protected HashMap<String, String> serviceVariables() {
     final HashMap<String, String> variables = new HashMap<>();
     variables.put("serviceName", service.getId().getName(service));
-    variables.put("dafnyModuleName", getDafnyModuleName());
     variables.put("dafnyInternalModuleName", getDafnyInternalModuleName());
     variables.put("dafnyTypesModuleName", getDafnyTypesModuleName());
     variables.put("rustTypesModuleName", getRustTypesModuleName());
     return variables;
   }
 
-  protected String getDafnyModuleName() {
-    return service
-      .getId()
-      .getNamespace()
-      .replace(".", "::")
-      .toLowerCase(Locale.ROOT);
-  }
-
   protected String getDafnyInternalModuleName() {
-    return "%s::internaldafny".formatted(getDafnyModuleName());
+    String dafnyExternName = DafnyNameResolver.dafnyExternNamespace(
+      service.getId().getNamespace()
+    );
+    return dafnyExternName.replace(".", "::").toLowerCase(Locale.ROOT);
   }
 
   protected String getDafnyTypesModuleName() {

diff --git a/...src/main/java/software/amazon/polymorph/smithyrust/generator/RustAwsSdkShimGenerator.java b/...src/main/java/software/amazon/polymorph/smithyrust/generator/RustAwsSdkShimGenerator.java
@@ -511,8 +511,8 @@ pub fn to_dafny(
   }
 
   @Override
-  protected String getDafnyModuleName() {
-    return "software::amazon::cryptography::services::%s".formatted(
+  protected String getDafnyInternalModuleName() {
+    return "software::amazon::cryptography::services::%s::internaldafny".formatted(
         getSdkId().toLowerCase()
       );
   }

diff --git a/...rc/main/java/software/amazon/polymorph/smithyrust/generator/RustLibraryShimGenerator.java b/...rc/main/java/software/amazon/polymorph/smithyrust/generator/RustLibraryShimGenerator.java
@@ -657,7 +657,7 @@ private RustFile operationOuterModule(final OperationShape operationShape) {
       StructureShape outputShape = operationIndex
         .getOutputShape(operationShape)
         .get();
-      if (inputShape.hasTrait(PositionalTrait.class)) {
+      if (outputShape.hasTrait(PositionalTrait.class)) {
         variables.put(
           "outputFromDafny",
           fromDafny(outputShape, "inner_result.value()", false, false)
@@ -1175,7 +1175,7 @@ protected Set<RustFile> operationConversionModules(
       operationName(operationShape)
     );
 
-    Optional<StructureShape> inputStructure = operationIndex.getOutputShape(
+    Optional<StructureShape> inputStructure = operationIndex.getInputShape(
       operationShape
     );
     final boolean hasInputStructure =
@@ -1307,7 +1307,6 @@ private TokenTree operationStructureFromDafnyFunction(
       structureId,
       StructureShape.class
     );
-    final boolean isPositional = structureShape.hasTrait(PositionalTrait.class);
     final Map<String, String> variables = MapUtils.merge(
       serviceVariables(),
       operationVariables(operationShape),
@@ -1318,46 +1317,24 @@ private TokenTree operationStructureFromDafnyFunction(
       fluentMemberSettersForStructure(structureShape).toString()
     );
 
-    if (isPositional) {
-      return TokenTree.of(
-        evalTemplate(
-          """
-          #[allow(dead_code)]
-          pub fn from_dafny(
-              dafny_value: ::std::rc::Rc<
-                  crate::r#$dafnyTypesModuleName:L::$structureName:L,
-              >,
-          ) -> crate::operation::$snakeCaseOperationName:L::$rustStructureName:L {
-              crate::operation::$snakeCaseOperationName:L::$rustStructureName:L::builder()
-                  $fluentMemberSetters:L
-                  .build()
-                  .unwrap()
-          }
-          """,
-          variables
-        )
-      );
-    } else {
-      // unwrap() is safe as long as the builder is infallible
-      return TokenTree.of(
-        evalTemplate(
-          """
-          #[allow(dead_code)]
-          pub fn from_dafny(
-              dafny_value: ::std::rc::Rc<
-                  crate::r#$dafnyTypesModuleName:L::$structureName:L,
-              >,
-          ) -> crate::operation::$snakeCaseOperationName:L::$rustStructureName:L {
-              crate::operation::$snakeCaseOperationName:L::$rustStructureName:L::builder()
-                  $fluentMemberSetters:L
-                  .build()
-                  .unwrap()
-          }
-          """,
-          variables
-        )
-      );
-    }
+    return TokenTree.of(
+      evalTemplate(
+        """
+        #[allow(dead_code)]
+        pub fn from_dafny(
+            dafny_value: ::std::rc::Rc<
+                crate::r#$dafnyTypesModuleName:L::$structureName:L,
+            >,
+        ) -> crate::operation::$snakeCaseOperationName:L::$rustStructureName:L {
+            crate::operation::$snakeCaseOperationName:L::$rustStructureName:L::builder()
+                $fluentMemberSetters:L
+                .build()
+                .unwrap()
+        }
+        """,
+        variables
+      )
+    );
   }
 
   private RustFile wrappedModule() {
@@ -1434,7 +1411,7 @@ private String wrappedClientOperationImpl(
     StructureShape outputShape = operationIndex
       .getOutputShape(operationShape)
       .get();
-    if (inputShape.hasTrait(PositionalTrait.class)) {
+    if (outputShape.hasTrait(PositionalTrait.class)) {
       variables.put(
         "outputToDafny",
         toDafny(outputShape, "inner_result", false, false).toString()
@@ -1641,14 +1618,26 @@ protected TokenTree toDafny(
         }
       }
       case BOOLEAN -> {
-        if (isRustOption) {
-          yield TokenTree.of(
-            "crate::standard_library_conversions::obool_to_dafny(&%s)".formatted(
-                rustValue
-              )
-          );
+        if (isDafnyOption) {
+          if (isRustOption) {
+            yield TokenTree.of(
+              "crate::standard_library_conversions::obool_to_dafny(&%s)".formatted(
+                  rustValue
+                )
+            );
+          } else {
+            yield TokenTree.of(
+              "crate::standard_library_conversions::obool_to_dafny(Some(%s))".formatted(
+                  rustValue
+                )
+            );
+          }
         } else {
-          yield TokenTree.of("%s.clone()".formatted(rustValue));
+          if (isRustOption) {
+            yield TokenTree.of("%s.clone().unwrap()".formatted(rustValue));
+          } else {
+            yield TokenTree.of("%s.clone()".formatted(rustValue));
+          }
         }
       }
       case INTEGER -> {
@@ -1667,7 +1656,11 @@ protected TokenTree toDafny(
             );
           }
         } else {
-          yield TokenTree.of("%s.clone()".formatted(rustValue));
+          if (isRustOption) {
+            yield TokenTree.of("%s.clone().unwrap()".formatted(rustValue));
+          } else {
+            yield TokenTree.of("%s.clone()".formatted(rustValue));
+          }
         }
       }
       case LONG -> {
@@ -1830,7 +1823,7 @@ protected TokenTree toDafny(
           if (isRustOption) {
             yield TokenTree.of(
               """
-              %s::conversions::%s::to_dafny(&%s.clone().unwrap())
+              %s::conversions::%s::to_dafny(%s.clone().unwrap())
               """.formatted(prefix, structureShapeName, rustValue)
             );
           } else {