Missing top level module declarations

smithy-lang · Sep 18, 2024 · a45277d · a45277d
1 parent f90b85b
commit a45277d
Showing 1 changed file with 16 additions and 19 deletions.
diff --git a/codegen/smithy-dafny-codegen/src/main/java/software/amazon/polymorph/CodegenEngine.java b/codegen/smithy-dafny-codegen/src/main/java/software/amazon/polymorph/CodegenEngine.java
@@ -874,10 +874,24 @@ public void patchAfterTranspiling() {
     }
   }
 
+  private static final TokenTree EXTRA_SINGLE_CRATE_DECLARATIONS = TokenTree.of(
+    """
+    mod standard_library_conversions;
+    mod standard_library_externs;
+    """
+  );
+
   private void patchRustAfterTranspiling() {
     final MergedServicesGenerator generator = new MergedServicesGenerator(model, serviceShape, namespaces, localServiceTest);
 
-    final String extraDeclarations = generator.generatorForShape(serviceShape).topLevelModuleDeclarations().toString();
+    final TokenTree extraRootServiceDeclarations = generator.generatorForShape(serviceShape).topLevelModuleDeclarations();
+    String extraDeclarations = TokenTree.of(extraRootServiceDeclarations, EXTRA_SINGLE_CRATE_DECLARATIONS).lineSeparated().toString();
+    if (!awsSdkStyle) {
+      extraDeclarations = extraDeclarations
+        + System.lineSeparator()
+        + extraTopLevelDeclarationsForLocalService(serviceShape)
+        + System.lineSeparator();
+    }
     final Path implementationFromDafnyPath = libraryRoot
       .resolve("runtimes")
       .resolve("rust")
@@ -899,15 +913,7 @@ private void patchRustAfterTranspiling() {
     }
   }
 
-  private String extraDeclarationsForSDK() {
-    return """
-    mod client;
-    mod conversions;
-    mod standard_library_conversions;
-    """;
-  }
-
-  private String extraDeclarationsForLocalService() {
+  private String extraTopLevelDeclarationsForLocalService(ServiceShape serviceShape) {
     final String configStructName = serviceShape
       .expectTrait(LocalServiceTrait.class)
       .getConfigId()
@@ -916,17 +922,8 @@ private String extraDeclarationsForLocalService() {
 
     return IOUtils.evalTemplate(
       """
-      // (extra-declarations)
-
-      $sharedTopLevelDecls:L
-
-      /// Copied from StandardLibrary
-      mod standard_library_conversions;
-      mod standard_library_externs;
-
       pub use client::Client;
       pub use types::$configSnakeCase:L::$configStructName:L;
-
       """,
       Map.of(
         "configSnakeCase",