fix(smithy-dafny): local Service config in another namespace (#207)

Author: texastony

codegen/smithy-dafny-codegen/src/main/java/software/amazon/polymorph/smithydafny/DafnyApiCodegen.java

@@ -119,19 +119,23 @@ public Map<Path, TokenTree> generate() {
             DafnyNameResolverHelpers.dafnyExternNamespaceForShapeId(serviceShape.getId()),
             typesModuleName
           ));
-
+        final TokenTree abstractServiceModule = generateAbstractServiceModule(serviceShape);
+        final TokenTree abstractOperationsModule = generateAbstractOperationsModule(serviceShape);
+        final TokenTree modeledErrorDataType = generateModeledErrorDataType();
         // A smithy model may reference a model in a different package.
         // In which case we need to import it.
+        // Everything MUST BE GENERATED BEFORE `typesModulePrelude` is calculated.
+        // Otherwise, a foreign shape maybe referenced WITHOUT an import.
         final TokenTree typesModulePrelude = TokenTree
           .of(Stream
             .concat(
-              nameResolver.modulePreludeStandardImports(),
+              DafnyNameResolver.modulePreludeStandardImports(),
               nameResolver
                 .dependentModels()
                 .stream()
                 .map(d ->
-                  "import " + nameResolver.dafnyTypesModuleForNamespace(d.namespace())))
-            .map(i -> Token.of(i))
+                  "import " + DafnyNameResolver.dafnyTypesModuleForNamespace(d.namespace())))
+            .map(Token::of)
           )
           .lineSeparated();
 
@@ -152,7 +156,7 @@ public Map<Path, TokenTree> generate() {
               // Error types are generated *after*
               // all other types to account
               // for any dependent modules
-              generateModeledErrorDataType()
+              modeledErrorDataType
             )
             .lineSeparated()
             .braced();
@@ -163,8 +167,8 @@ public Map<Path, TokenTree> generate() {
             includeDirectives,
             typesModuleHeader,
             typesModuleBody,
-            generateAbstractServiceModule(serviceShape),
-            generateAbstractOperationsModule(serviceShape)
+            abstractServiceModule,
+            abstractOperationsModule
           )
           .lineSeparated();
         return Map.of(path, fullCode);
@@ -425,7 +429,7 @@ public TokenTree generateReferenceTraitDefinition(final ShapeId shapeWithReferen
           // the module needs to be included
           // because we are obviously using it!
           final String sideEffect = nameResolver
-            .dafnyModulePrefixForShapeId(model.expectShape(referenceTrait.getReferentId()));
+            .dafnyModulePrefixForShape(model.expectShape(referenceTrait.getReferentId()));
           return null;
         }
 
@@ -1350,6 +1354,7 @@ abstractModulePrelude, generateAbstractLocalService(serviceShape)
       }
     }
 
+    // This method needs to be called before typesModulePrelude is calculated
     public TokenTree generateAbstractServiceModule(ServiceShape serviceShape) {
       final TokenTree abstractModulePrelude = TokenTree
         .of(DafnyNameResolver.abstractModulePrelude(serviceShape))
@@ -1358,11 +1363,12 @@ public TokenTree generateAbstractServiceModule(ServiceShape serviceShape) {
         .formatted(nameResolver.abstractServiceModuleName(serviceShape)));
 
       if (serviceShape.hasTrait(ServiceTrait.class)) {
+          TokenTree body = generateAbstractAwsServiceClass(serviceShape);
         return moduleHeader
           .append(Token
             .of(
               abstractModulePrelude,
-              generateAbstractAwsServiceClass(serviceShape)
+              body
             )
             .lineSeparated()
             .braced()
@@ -1829,8 +1835,8 @@ public TokenTree generateAbstractLocalService(ServiceShape serviceShape)  {
         if (!serviceShape.hasTrait(LocalServiceTrait.class)) throw new IllegalStateException("MUST be an LocalService");
         final LocalServiceTrait localServiceTrait = serviceShape.expectTrait(LocalServiceTrait.class);
 
-        final String configTypeName = localServiceTrait.getConfigId().getName();
-        final String defaultFunctionMethodName = "Default%s".formatted(configTypeName);
+        final String configTypeName = nameResolver.baseTypeForShape(localServiceTrait.getConfigId());
+        final String defaultFunctionMethodName = "Default%s".formatted(localServiceTrait.getConfigId().getName());
 
         final TokenTree defaultConfig = TokenTree
             .of("function method %s(): %s".formatted(defaultFunctionMethodName, configTypeName));
@@ -1940,8 +1946,8 @@ public TokenTree generateAbstractWrappedLocalService(ServiceShape serviceShape)
                 .of(DafnyNameResolver.wrappedAbstractModulePrelude(serviceShape))
                 .lineSeparated();
 
-        final String configTypeName = localServiceTrait.getConfigId().getName();
-        final String defaultFunctionMethodName = "Default%s".formatted(configTypeName);
+        final String configTypeName = nameResolver.baseTypeForShape(localServiceTrait.getConfigId());
+        final String defaultFunctionMethodName = "Default%s".formatted(localServiceTrait.getConfigId().getName());
 
         final TokenTree defaultConfig = TokenTree
                 .of("function method Wrapped%s(): %s".formatted(defaultFunctionMethodName, configTypeName));

codegen/smithy-dafny-codegen/src/main/java/software/amazon/polymorph/smithydafny/DafnyNameResolver.java

@@ -68,7 +68,7 @@ public String baseTypeForShape(final ShapeId shapeId) {
                     // currently unused in model and unsupported in StandardLibrary.UInt
                     // BYTE, SHORT
                     INTEGER, LONG, DOUBLE,
-                    LIST, MAP -> dafnyModulePrefixForShapeId(shape) + shapeName;
+                    LIST, MAP -> dafnyModulePrefixForShape(shape) + shapeName;
             case STRUCTURE -> {
                 if (shape.hasTrait(ReferenceTrait.class)) {
                     yield baseTypeForShape(shape.expectTrait(ReferenceTrait.class).getReferentId());
@@ -108,7 +108,7 @@ private String baseTypeForMember(final MemberShape memberShape) {
     }
 
     private String dafnyTypeNameShape(final Shape shape) {
-        return dafnyModulePrefixForShapeId(shape) + shape.getId().getName();
+        return dafnyModulePrefixForShape(shape) + shape.getId().getName();
     }
 
     public String generatedTypeForShape(final ShapeId shapeId) {
@@ -120,7 +120,7 @@ public static String traitNameForServiceClient(final ServiceShape serviceShape)
     }
 
     public String traitForServiceClient(final ServiceShape serviceShape) {
-        return dafnyModulePrefixForShapeId(serviceShape) + traitNameForServiceClient(serviceShape);
+        return dafnyModulePrefixForShape(serviceShape) + traitNameForServiceClient(serviceShape);
     }
 
     public static String classNameForServiceClient(ServiceShape shape) {
@@ -133,7 +133,7 @@ public static String classNameForServiceClient(ServiceShape shape) {
     }
 
     public String traitForResource(final ResourceShape resourceShape) {
-        return dafnyModulePrefixForShapeId(resourceShape) + traitNameForResource(resourceShape);
+        return dafnyModulePrefixForShape(resourceShape) + traitNameForResource(resourceShape);
     }
 
     public static String traitNameForResource(final ResourceShape shape) {
@@ -260,7 +260,7 @@ public static String dafnyInternalConfigModuleForNamespace(final String namespac
     }
 
 
-    public String dafnyModulePrefixForShapeId(final Shape shape) {
+    public String dafnyModulePrefixForShape(final Shape shape) {
         final String shapeNamespace = shape.getId().getNamespace();
         if (!namespace.equals(shapeNamespace)) {