Add Stream and Iterable as a repeated container

SuduIDE · Sep 14, 2023 · 79ce310 · 79ce310
1 parent 8612781
commit 79ce310
Show file tree

Hide file tree

Showing 9 changed files with 146 additions and 50 deletions.
diff --git a/generator/src/main/java/org/sudu/protogen/descriptors/Method.java b/generator/src/main/java/org/sudu/protogen/descriptors/Method.java
@@ -80,9 +80,10 @@ protected Optional<Boolean> getNullableOption() {
         return Options.wrapExtension(methodDescriptor.getOptions(), protogen.Options.nullable);
     }
 
-    public Optional<RepeatedContainer> getStreamToContainer() {
+    public RepeatedContainer getStreamToContainer() {
         return Options.wrapExtension(methodDescriptor.getOptions(), protogen.Options.streamToContainer)
-                .map(RepeatedContainer::fromGrpc);
+                .map(RepeatedContainer::fromGrpc)
+                .orElse(RepeatedContainer.ITERATOR);
     }
 
     protected Optional<String> getNameOption() {

diff --git a/generator/src/main/java/org/sudu/protogen/descriptors/RepeatedContainer.java b/generator/src/main/java/org/sudu/protogen/descriptors/RepeatedContainer.java
@@ -4,33 +4,134 @@
 import com.squareup.javapoet.CodeBlock;
 import protogen.Options;
 
+import java.util.Iterator;
+import java.util.List;
+import java.util.Set;
+import java.util.Spliterators;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
 
 public enum RepeatedContainer {
-    LIST(ClassName.get("java.util", "List"), CodeBlock.of("$T.toList()", Collectors.class)),
-    SET(ClassName.get("java.util", "Set"), CodeBlock.of("$T.toSet()", Collectors.class));
+
+    LIST(ClassName.get(List.class)) {
+        @Override
+        public CodeBlock getCollectorExpr() {
+            return CodeBlock.of(".collect($T.toList())", Collectors.class);
+        }
+
+        @Override
+        public CodeBlock getToStreamExpr(CodeBlock thisRef) {
+            return CodeBlock.of("$L.stream()", thisRef);
+        }
+    },
+
+    SET(ClassName.get(Set.class)) {
+        @Override
+        public CodeBlock getCollectorExpr() {
+            return CodeBlock.of(".collect($T.toSet())", Collectors.class);
+        }
+
+        @Override
+        public CodeBlock getToStreamExpr(CodeBlock thisRef) {
+            return CodeBlock.of("$L.stream()", thisRef);
+        }
+    },
+
+    ITERATOR(ClassName.get(Iterator.class)) {
+        @Override
+        public CodeBlock getCollectorExpr() {
+            return CodeBlock.of(".iterator()");
+        }
+
+        @Override
+        public CodeBlock getToStreamExpr(CodeBlock thisRef) {
+            return CodeBlock.of(
+                    "$T.stream($T.spliteratorUnknownSize($L, 0), false)",
+                    StreamSupport.class,
+                    Spliterators.class,
+                    thisRef
+            );
+        }
+
+        @Override
+        public CodeBlock convertListToInstance(CodeBlock thisRef) {
+            return CodeBlock.of("$L.iterator()", thisRef);
+        }
+
+        @Override
+        public CodeBlock convertInstanceToIterable(CodeBlock thisRef) {
+            return CodeBlock.of("() -> $L", thisRef);
+        }
+    },
+
+    STREAM(ClassName.get(Stream.class)) {
+        @Override
+        public CodeBlock getCollectorExpr() {
+            return CodeBlock.of("");
+        }
+
+        @Override
+        public CodeBlock getToStreamExpr(CodeBlock thisRef) {
+            return thisRef;
+        }
+
+        @Override
+        public CodeBlock convertListToInstance(CodeBlock thisRef) {
+            return CodeBlock.of("$L.stream()", thisRef);
+        }
+
+        @Override
+        public CodeBlock convertInstanceToIterable(CodeBlock thisRef) {
+            return CodeBlock.of("$L.toList()", thisRef);
+        }
+    };
 
     private final ClassName typeName;
-    private final CodeBlock collectorExpr;
 
-    RepeatedContainer(ClassName typeName, CodeBlock collectorExpr) {
+    RepeatedContainer(ClassName typeName) {
         this.typeName = typeName;
-        this.collectorExpr = collectorExpr;
     }
 
     public static RepeatedContainer fromGrpc(Options.RepeatedContainer proto) {
         return switch (proto) {
             case UNRECOGNIZED -> throw new IllegalArgumentException();
             case LIST -> RepeatedContainer.LIST;
             case SET -> RepeatedContainer.SET;
+            case ITERATOR -> RepeatedContainer.ITERATOR;
+            case STREAM -> RepeatedContainer.STREAM;
         };
     }
 
-    public ClassName getTypeName() {
-        return typeName;
+    /**
+     * Builds an expressions transforming an instance of the container into stream
+     */
+    public abstract CodeBlock getToStreamExpr(CodeBlock thisRef);
+
+    /**
+     * Builds an expressions collecting a stream into the container
+     */
+    public abstract CodeBlock getCollectorExpr();
+
+    /**
+     * Builds an expression converting a java.Util.List instance into the container
+     * Used to build .getSomeList().... expressions
+     * Necessary, because getToStreamExpr consumes container instance, but protobuf-message getter returns List
+     */
+    public CodeBlock convertListToInstance(CodeBlock thisRef) {
+        return thisRef;
+    }
+
+    /*
+     * Builds an expression converting the container to Iterable<>
+     * Used to build .addSome(..) calls in protobuf-generated builders
+     * It's necessary, because there is no way to directly pass stream or iterator as a parameter
+     */
+    public CodeBlock convertInstanceToIterable(CodeBlock thisRef) {
+        return thisRef;
     }
 
-    public CodeBlock getCollectorExpr() {
-        return collectorExpr;
+    public ClassName getTypeName() {
+        return typeName;
     }
 }
diff --git a/generator/src/main/java/org/sudu/protogen/generator/client/MethodGeneratorBase.java b/generator/src/main/java/org/sudu/protogen/generator/client/MethodGeneratorBase.java
@@ -5,7 +5,6 @@
 import org.sudu.protogen.descriptors.Method;
 import org.sudu.protogen.generator.GenerationContext;
 import org.sudu.protogen.generator.field.FieldGenerator;
-import org.sudu.protogen.generator.type.IteratorType;
 import org.sudu.protogen.generator.type.RepeatedType;
 import org.sudu.protogen.generator.type.TypeModel;
 
@@ -88,12 +87,7 @@ protected TypeModel getReturnType() {
             if (type.getTypeName().toString().equalsIgnoreCase("void")) {
                 return type;
             }
-            var containerO = method.getStreamToContainer();
-            if (containerO.isPresent()) {
-                return new RepeatedType(type, containerO.get());
-            } else {
-                return new IteratorType(type);
-            }
+            return new RepeatedType(type, method.getStreamToContainer());
         }
         return type;
     }

diff --git a/...or/src/main/java/org/sudu/protogen/generator/client/StreamingGrpcCallMethodGenerator.java b/...or/src/main/java/org/sudu/protogen/generator/client/StreamingGrpcCallMethodGenerator.java
@@ -3,9 +3,7 @@
 import com.squareup.javapoet.*;
 import org.apache.commons.lang3.Validate;
 import org.sudu.protogen.descriptors.Method;
-import org.sudu.protogen.descriptors.RepeatedContainer;
 import org.sudu.protogen.generator.GenerationContext;
-import org.sudu.protogen.generator.type.IteratorType;
 import org.sudu.protogen.generator.type.RepeatedType;
 import org.sudu.protogen.generator.type.TypeModel;
 import org.sudu.protogen.utils.Name;
@@ -36,35 +34,22 @@ protected CodeBlock body(TypeModel returnType, List<ParameterSpec> params) {
             return CodeBlock.of("");
         }
         CodeBlock call = commonGen.onlyCallReturnExpression(returnType, params);
-        if (method.getStreamToContainer().isPresent()) {
-            Validate.validState(returnType instanceof RepeatedType);
-            RepeatedType repType = (RepeatedType) returnType;
-            return CodeBlock.builder()
-                    .addStatement("var response = $L", call)
-                    .addStatement(new IteratorWrapper().wrapToIterable("response", repType.getElementModel()))
-                    .build();
-        } else {
-            return CodeBlock.builder()
-                    .addStatement("var response = $L", call)
-                    .build();
-        }
+        Validate.validState(returnType instanceof RepeatedType);
+        RepeatedType repType = (RepeatedType) returnType;
+        return CodeBlock.builder()
+                .addStatement("var response = $L", call)
+                .addStatement(new IteratorWrapper().wrapToIterable("response", repType.getElementModel()))
+                .build();
     }
 
     @Override
     protected CodeBlock returnExpression(TypeModel returnType, List<ParameterSpec> params) {
         if (returnType.getTypeName().toString().equalsIgnoreCase("void")) {
             return commonGen.onlyCallReturnExpression(returnType, params);
         }
-        var containerO = method.getStreamToContainer();
-        if (containerO.isPresent()) {
-            RepeatedContainer container = containerO.get();
-            return CodeBlock.of("$T.stream(iterable.spliterator(), false).collect($L)",
-                    StreamSupport.class, container.getCollectorExpr());
-        } else {
-            Validate.validState(returnType instanceof IteratorType);
-            IteratorType itType = (IteratorType) returnType;
-            return new IteratorWrapper().wrapIterator("response", itType.getIteratedType());
-        }
+        var container = method.getStreamToContainer();
+        return CodeBlock.of("$T.stream(iterable.spliterator(), false)$L",
+                StreamSupport.class, container.getCollectorExpr());
     }
 
     @Override

diff --git a/...c/main/java/org/sudu/protogen/generator/field/processors/PrimitiveFieldTypeProcessor.java b/...c/main/java/org/sudu/protogen/generator/field/processors/PrimitiveFieldTypeProcessor.java
@@ -18,14 +18,14 @@ class PrimitiveFieldTypeProcessor extends FieldTypeProcessor.Chain {
             case FLOAT -> boxIfNullable(field, TypeName.FLOAT);
             case DOUBLE -> boxIfNullable(field, TypeName.DOUBLE);
             case BOOLEAN -> boxIfNullable(field, TypeName.BOOLEAN);
-            case STRING -> new TypeModel(ClassName.get(String.class));
+            case STRING -> boxIfNullable(field, ClassName.get(String.class));
             case BYTE_STRING -> new TypeModel(ClassName.get("com.google.protobuf", "ByteString"));
             default -> next(field, context);
         };
     }
 
     private TypeModel boxIfNullable(@NotNull Field field, TypeName primitive) {
-        return field.isNullable() || field.isList()
+        return field.isNullable()
                 ? new TypeModel(primitive.box())
                 : new PrimitiveTypeModel(primitive);
     }

diff --git a/generator/src/main/java/org/sudu/protogen/generator/type/RepeatedType.java b/generator/src/main/java/org/sudu/protogen/generator/type/RepeatedType.java
@@ -32,15 +32,15 @@ public CodeBlock toGrpcTransformer(CodeBlock expr) {
             CodeBlock mappingLambda = CodeBlock.builder()
                     .add("$L -> $L", lambdaParameter, elementModel.toGrpcTransformer(lambdaParameter))
                     .build();
-            return listMapper(expr, mappingLambda);
+            expr = listMapper(expr, mappingLambda);
         }
-        return expr;
+        return repeatedType.convertInstanceToIterable(expr);
     }
 
     @Override
     public CodeBlock fromGrpcTransformer(CodeBlock expr) {
-        // todo why don't we need to do this with lists? comment it
-        if (!(elementModel instanceof PrimitiveTypeModel) || repeatedType != RepeatedContainer.LIST) {
+        expr = repeatedType.convertListToInstance(expr);
+        if (!(elementModel instanceof PrimitiveTypeModel)) {
             CodeBlock lambdaParameter = CodeBlock.builder().add("i").build();
             CodeBlock mappingLambda = CodeBlock.builder()
                     .add("$L -> $L", lambdaParameter, elementModel.fromGrpcTransformer(lambdaParameter))
@@ -53,10 +53,10 @@ public CodeBlock fromGrpcTransformer(CodeBlock expr) {
     @NotNull
     private CodeBlock listMapper(CodeBlock expr, CodeBlock mapper) {
         return CodeBlock.builder()
-                .add("$L.stream()\n", expr)
+                .add("$L\n", repeatedType.getToStreamExpr(expr))
                 .indent()
                 .add(".map($L)\n", mapper)
-                .add(".collect($L)", repeatedType.getCollectorExpr())
+                .add("$L", repeatedType.getCollectorExpr())
                 .unindent()
                 .build();
     }

diff --git a/generator/src/main/proto/protogen/options.proto b/generator/src/main/proto/protogen/options.proto
@@ -12,6 +12,10 @@ enum RepeatedContainer {
   LIST = 0;
   // java.util.Set; collects using Collectors.toSet()
   SET = 1;
+  //
+  ITERATOR = 2;
+  //
+  STREAM = 3;
 }
 
 enum AccessModifier {

diff --git a/tests/src/test/proto/client.proto b/tests/src/test/proto/client.proto
@@ -76,6 +76,12 @@ service SomeService {
   rpc streamDomainList(google.protobuf.Empty) returns (stream GrpcDomainResponse) {
     option (.protogen.stream_to_container) = LIST;
   }
+  rpc streamDomainIterator(google.protobuf.Empty) returns (stream GrpcDomainResponse) {
+    option (.protogen.stream_to_container) = ITERATOR;
+  }
+  rpc streamDomainStream(google.protobuf.Empty) returns (stream GrpcDomainResponse) {
+    option (.protogen.stream_to_container) = STREAM;
+  }
   rpc streamMulti(google.protobuf.Empty) returns (stream GrpcMultiFieldResponse);
 }
 

diff --git a/tests/src/test/proto/general.proto b/tests/src/test/proto/general.proto
@@ -48,4 +48,9 @@ message GrpcTasksBatch {
 message GrpcRegisteredMapList {
   map<string, google.protobuf.Timestamp> m = 1;
   repeated google.protobuf.Timestamp l = 2;
+}
+
+message GrpcRepeatedTypes {
+  repeated string iterator = 1 [(.protogen.repeated_container) = ITERATOR];
+  repeated string stream = 2 [(.protogen.repeated_container) = STREAM];
 }