From a45277dd8c613eaf07171ee7fb0aaba04f8f4704 Mon Sep 17 00:00:00 2001
From: Robin Salkeld <salkeldr@amazon.com>
Date: Wed, 18 Sep 2024 08:52:12 -0700
Subject: [PATCH] Missing top level module declarations

---
 .../amazon/polymorph/CodegenEngine.java       | 35 +++++++++----------
 1 file changed, 16 insertions(+), 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
index 13be3d1d4..7f412bb5b 100644
--- 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",