From 786b8bb6ddb54abb8ecf956dbca7e5775014957e Mon Sep 17 00:00:00 2001
From: Dmitrii Duzhinskii <w1vvv1w@ya.ru>
Date: Fri, 29 Sep 2023 13:09:48 +0300
Subject: [PATCH] Make body generation more clearly

by applying the strategy pattern, which fixed a bug with streaming void clients
---
 .../generator/client/ClientGenerator.java     |  5 +-
 .../client/StubCallMethodGenerator.java       | 50 ++++++++++++-------
 .../protogen/generator/type/VoidType.java     |  5 ++
 3 files changed, 39 insertions(+), 21 deletions(-)

diff --git a/generator/src/main/java/org/sudu/protogen/generator/client/ClientGenerator.java b/generator/src/main/java/org/sudu/protogen/generator/client/ClientGenerator.java
index 3d61fc4..39b3807 100644
--- a/generator/src/main/java/org/sudu/protogen/generator/client/ClientGenerator.java
+++ b/generator/src/main/java/org/sudu/protogen/generator/client/ClientGenerator.java
@@ -79,15 +79,14 @@ private Stream<MethodSpec> generateRpcMethod(Method method) {
     }
 
     protected TypeModel getMethodReturnType(Method method) {
-        if (method.getOutputType().getFields().isEmpty()) {
-            return new VoidType();
-        }
         var responseType = context.processType(method.getOutputType());
         TypeModel type;
         if (responseType != null) {
             type = responseType;
         } else if (method.doUnfoldResponse(responseType)) {
             type = new UnfoldedType(context.processType(method.unfoldedResponseField()), method.getOutputType());
+        } else if (method.getOutputType().getFields().isEmpty()) {
+            type = new VoidType();
         } else {
             throw new IllegalStateException(("Unable to create a method returning %s because request consist of more than " +
                     "1 field and doesn't have a domain object.").formatted(method.getOutputType().getFullName()));
diff --git a/generator/src/main/java/org/sudu/protogen/generator/client/StubCallMethodGenerator.java b/generator/src/main/java/org/sudu/protogen/generator/client/StubCallMethodGenerator.java
index 776ae37..6febbef 100644
--- a/generator/src/main/java/org/sudu/protogen/generator/client/StubCallMethodGenerator.java
+++ b/generator/src/main/java/org/sudu/protogen/generator/client/StubCallMethodGenerator.java
@@ -50,28 +50,42 @@ private class BodyGenerator implements Supplier<CodeBlock> {
 
         @Override
         public CodeBlock get() {
-            CodeBlock returnExpr = CodeBlock.of("$N.$L(request)", stubField, method.getName());
-            CodeBlock body = CodeBlock.of("");
-
-            if (returnType instanceof RepeatedType repType) { // i.e. method is streaming non-void
-                body = CodeBlock.builder()
-                        .addStatement(
-                                "var iterator = $L",
-                                returnExpr
-                        )
-                        .build();
+            if (method.isOutputStreaming()) {
+                return new StreamingBodyGenerator().get();
+            } else {
+                return new CommonBodyGenerator().get();
+            }
+        }
+
+        private class CommonBodyGenerator extends BodyGenerator {
+            @Override
+            public CodeBlock get() {
+                CodeBlock returnExpr = CodeBlock.of("$N.$L(request)", stubField, method.getName());
+                if (returnType.getTypeName() != TypeName.VOID) {
+                    returnExpr = CodeBlock.of("return $L", returnType.fromGrpcTransformer(returnExpr));
+                }
+                return CodeBlock.builder().addStatement(returnExpr).build();
+            }
+        }
+
+        private class StreamingBodyGenerator extends BodyGenerator {
+
+            @Override
+            public CodeBlock get() {
+                if (!(returnType instanceof RepeatedType repType)) throw new IllegalArgumentException();
+                CodeBlock body = CodeBlock.of("var iterator = $N.$L(request);\n", stubField, method.getName());
                 // I write mapping here manually because input is always an Iterator<Grpc..> and output is specified by the RepeatedContainer option
                 // So RepeatedType.fromGrpcTransformer is not suitable because it does only T<U> <--> T<V> mappings
                 CodeBlock mappingExpr = CodeBlock.of("i -> $L", repType.getElementModel().fromGrpcTransformer(CodeBlock.of("i")));
-                returnExpr = CodeBlock.of("$L\n.map($L)$L", RepeatedContainer.ITERATOR.getToStreamExpr(CodeBlock.of("iterator")), mappingExpr, repType.getRepeatedType().getCollectorExpr());
-            } else {
-                // Don't transform if returnType is streaming (RepeatedType), see the comment above
-                returnExpr = returnType.fromGrpcTransformer(returnExpr);
-            }
-            if (returnType.getTypeName() != TypeName.VOID) {
-                returnExpr = CodeBlock.of("return $L", returnExpr);
+                return CodeBlock.builder()
+                        .add(body)
+                        .addStatement("return $L\n.map($L)$L",
+                                RepeatedContainer.ITERATOR.getToStreamExpr(CodeBlock.of("iterator")),
+                                mappingExpr,
+                                repType.getRepeatedType().getCollectorExpr()
+                        )
+                        .build();
             }
-            return CodeBlock.builder().add(body).addStatement(returnExpr).build();
         }
 
         private BodyGenerator withIfNotFound() {
diff --git a/generator/src/main/java/org/sudu/protogen/generator/type/VoidType.java b/generator/src/main/java/org/sudu/protogen/generator/type/VoidType.java
index 4b012b5..eb276bc 100644
--- a/generator/src/main/java/org/sudu/protogen/generator/type/VoidType.java
+++ b/generator/src/main/java/org/sudu/protogen/generator/type/VoidType.java
@@ -13,4 +13,9 @@ public VoidType() {
     public CodeBlock toGrpcTransformer(CodeBlock expr) {
         return CodeBlock.of("null");
     }
+
+    @Override
+    public CodeBlock fromGrpcTransformer(CodeBlock expr) {
+        return CodeBlock.of("(Void)null");
+    }
 }