Change dafny client to native ones in native go land

smithy-lang · Nov 4, 2024 · aab2c4f · aab2c4f
1 parent 27a8eea
commit aab2c4f
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 15 deletions.
diff --git a/...dafny-codegen/src/main/java/software/amazon/polymorph/smithygo/codegen/SymbolVisitor.java b/...dafny-codegen/src/main/java/software/amazon/polymorph/smithygo/codegen/SymbolVisitor.java
@@ -273,9 +273,7 @@ private String getDefaultShapeName(Shape shape) {
       ServiceShape.class
     );
     if (shape.hasTrait(ServiceTrait.class)) {
-      return DafnyNameResolver.getDafnyClientName(
-        shape.expectTrait(ServiceTrait.class).getSdkId()
-      );
+      return SmithyNameResolver.getAwsServiceClientName();
     }
     return StringUtils.capitalize(
       removeLeadingInvalidIdentCharacters(shape.getId().getName(serviceShape))
@@ -393,7 +391,7 @@ public Symbol mapShape(MapShape shape) {
   private Symbol.Builder symbolBuilderFor(Shape shape, String typeName) {
     final String namespace;
     if (shape.hasTrait(ServiceTrait.class)) {
-      namespace = DafnyNameResolver.dafnyTypesNamespace(shape);
+      namespace = SmithyNameResolver.smithyTypesNamespaceAws(shape.expectTrait(ServiceTrait.class), false);
     } else {
       namespace = SmithyNameResolver.smithyTypesNamespace(shape);
     }
@@ -584,11 +582,12 @@ public Symbol structureShape(StructureShape shape) {
       );
       var isResource = !shape.expectTrait(ReferenceTrait.class).isService();
       if (isResource || referredShape.hasTrait(ServiceTrait.class)) {
+        final var typeName = referredShape.hasTrait(ServiceTrait.class) ? getDefaultShapeName(referredShape) : "I".concat(getDefaultShapeName(referredShape)) ;
         builder.putProperty(
           "Referred",
           symbolBuilderFor(
             referredShape,
-            "I".concat(getDefaultShapeName(referredShape))
+            typeName
           )
             .putProperty(SymbolUtils.POINTABLE, false)
             .build()

diff --git a/...java/software/amazon/polymorph/smithygo/localservice/nameresolver/SmithyNameResolver.java b/...java/software/amazon/polymorph/smithygo/localservice/nameresolver/SmithyNameResolver.java
@@ -25,6 +25,7 @@ public class SmithyNameResolver {
       "bool" //boolean shape
     )
   );
+  private static String awsServiceClientName = "Client";
 
   public static Boolean isShapeFromAWSSDK(final Shape shape) {
     // Is there a better way to do this?
@@ -216,6 +217,10 @@ public static String getAwsServiceClient(final ServiceTrait serviceTrait) {
     return SmithyNameResolver
       .smithyTypesNamespaceAws(serviceTrait, false)
       .concat(DOT)
-      .concat("Client");
+      .concat(getAwsServiceClientName());
+  }
+
+  public static String getAwsServiceClientName() {
+    return awsServiceClientName;
   }
 }
diff --git a/...ftware/amazon/polymorph/smithygo/localservice/shapevisitor/DafnyToSmithyShapeVisitor.java b/...ftware/amazon/polymorph/smithygo/localservice/shapevisitor/DafnyToSmithyShapeVisitor.java
@@ -158,13 +158,15 @@ protected String referenceStructureShape(final StructureShape shape) {
             ),
             DafnyNameResolver.dafnyTypesNamespace(serviceShape)
           );
-          return "return %1$s.(%2$s)".formatted(
-              dataSource,
-              DafnyNameResolver.getDafnyInterfaceClient(
-                (ServiceShape) serviceShape,
-                serviceShape.expectTrait(ServiceTrait.class)
-              )
-            );
+          return """
+              shim, ok := %1$s.(*%2$swrapped.Shim)
+              if !ok {
+                  panic("Not able to convert client to native")
+              }
+              return *shim.Client
+              """.formatted(dataSource, DafnyNameResolver.dafnyNamespace(
+                resourceOrService.expectTrait(ServiceTrait.class)
+              ));
         } else {
           return "return %1$s{%2$s}".formatted(
               namespace.concat(

diff --git a/...ftware/amazon/polymorph/smithygo/localservice/shapevisitor/SmithyToDafnyShapeVisitor.java b/...ftware/amazon/polymorph/smithygo/localservice/shapevisitor/SmithyToDafnyShapeVisitor.java
@@ -13,6 +13,7 @@
 import software.amazon.polymorph.smithygo.localservice.nameresolver.SmithyNameResolver;
 import software.amazon.polymorph.traits.DafnyUtf8BytesTrait;
 import software.amazon.polymorph.traits.ReferenceTrait;
+import software.amazon.smithy.aws.traits.ServiceTrait;
 import software.amazon.smithy.codegen.core.CodegenException;
 import software.amazon.smithy.codegen.core.Symbol;
 import software.amazon.smithy.model.shapes.BlobShape;
@@ -132,8 +133,20 @@ protected String referenceStructureShape(final StructureShape shape) {
     }
 
     if (resourceOrService.asServiceShape().isPresent()) {
+      writer.addImportFromModule(
+      SmithyNameResolver.getGoModuleNameForSmithyNamespace(
+        resourceOrService.toShapeId().getNamespace()
+        ),
+        DafnyNameResolver.dafnyTypesNamespace(resourceOrService)
+      );
+      final var shim = "%swrapped.Shim".formatted(
+          DafnyNameResolver.dafnyNamespace(
+            resourceOrService.expectTrait(ServiceTrait.class)
+          )
+        );
+      final var clientConversion = "&%s{Client: &%s}".formatted(shim, dataSource);
       if (!this.isOptional) {
-        return dataSource;
+        return clientConversion;
       } else {
         final var goCodeBlock =
           """
@@ -143,7 +156,7 @@ protected String referenceStructureShape(final StructureShape shape) {
               }
               return Wrappers.Companion_Option_.Create_Some_(%s)
           }()""";
-        return goCodeBlock.formatted(dataSource, dataSource);
+        return goCodeBlock.formatted(clientConversion, clientConversion);
       }
     }