Merge pull request #85 from salesforce/bugfix/jprotoc-nested

Correctly handle nested types

Author: rmichela

jprotoc-test/src/test/proto/nested.proto

@@ -0,0 +1,30 @@
+syntax = "proto3";
+
+package nested;
+
+message Outer {  // Level 0
+    enum FooEnum {
+        FOO = 0;
+        BAR = 1;
+        CHEESE = 2;
+    }
+    message MiddleAA {  // Level 1
+
+        message Inner {   // Level 2
+            int64 ival = 1;
+            bool  booly = 2;
+            Outer.FooEnum enum = 3;
+        }
+    }
+    message MiddleBB {  // Level 1
+        message Inner {   // Level 2
+            int32 ival = 1;
+            bool  booly = 2;
+            Outer.FooEnum enum = 3;
+        }
+    }
+}
+
+service Nested {
+    rpc doNested (Outer.MiddleAA.Inner) returns (Outer.MiddleBB.Inner) {}
+}

jprotoc/src/main/java/com/salesforce/jprotoc/Generator.java

@@ -18,6 +18,8 @@
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.io.StringWriter;
+import java.util.Collections;
+import java.util.List;
 import java.util.stream.Stream;
 
 /**
@@ -28,10 +30,25 @@ public abstract class Generator {
 
     /**
      * Processes a generator request into a set of files to output.
+     *
+     * @deprecated use {@link #generateFiles(PluginProtos.CodeGeneratorRequest)} and return a List instead of a Stream.
      * @param request The raw generator request from protoc.
      * @return The completed files to write out.
      */
-    public abstract Stream<PluginProtos.CodeGeneratorResponse.File> generate(PluginProtos.CodeGeneratorRequest request) throws GeneratorException;
+    @Deprecated()
+    public Stream<PluginProtos.CodeGeneratorResponse.File> generate(PluginProtos.CodeGeneratorRequest request) throws GeneratorException {
+        return Stream.empty();
+    }
+
+    /**
+     * Processes a generator request into a set of files to output.
+     *
+     * @param request The raw generator request from protoc.
+     * @return The completed files to write out.
+     */
+    public List<PluginProtos.CodeGeneratorResponse.File> generateFiles(PluginProtos.CodeGeneratorRequest request) throws GeneratorException {
+        return Collections.emptyList();
+    }
 
     /**
      * Executes a mustache template against a generatorContext object to generate an output string.

jprotoc/src/main/java/com/salesforce/jprotoc/ProtoTypeMap.java

@@ -57,21 +57,45 @@ public static ProtoTypeMap of(@Nonnull Collection<DescriptorProtos.FileDescripto
                             getJavaOuterClassname(fileDescriptor, fileOptions);
 
 
-
+            // Identify top-level enums
             fileDescriptor.getEnumTypeList().forEach(
                 e -> types.put(
                         protoPackage + "." + e.getName(),
                         toJavaTypeName(e.getName(), enclosingClassName, javaPackage)));
 
+            // Identify top-level messages, and nested types
             fileDescriptor.getMessageTypeList().forEach(
-                m -> types.put(
-                        protoPackage + "." + m.getName(),
-                        toJavaTypeName(m.getName(), enclosingClassName, javaPackage)));
+                m -> recursivelyAddTypes(types, m, protoPackage, enclosingClassName, javaPackage)
+            );
         }
 
         return new ProtoTypeMap(types.build());
     }
 
+    private static void recursivelyAddTypes(ImmutableMap.Builder<String, String> types, DescriptorProtos.DescriptorProto m, String protoPackage, String enclosingClassName, String javaPackage) {
+        // Identify current type
+        types.put(
+            protoPackage + "." + m.getName(),
+            toJavaTypeName(m.getName(), enclosingClassName, javaPackage));
+
+        // Identify any nested Enums
+        m.getEnumTypeList().forEach(
+            e -> types.put(
+                protoPackage + "." + e.getName(),
+                toJavaTypeName(e.getName(),
+                enclosingClassName,
+                javaPackage)));
+
+        // Recursively identify any nested types
+        m.getNestedTypeList().forEach(
+            n -> recursivelyAddTypes(
+                types,
+                n,
+                protoPackage + "." + m.getName(),
+                enclosingClassName + "." + m.getName(),
+                javaPackage));
+    }
+
     /**
      * Returns the full Java type name for the given proto type.
      *

jprotoc/src/main/java/com/salesforce/jprotoc/ProtocPlugin.java

@@ -16,6 +16,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import javax.annotation.Nonnull;
 
 /**
@@ -71,15 +72,18 @@ public static void generate(
                     generatorRequestBytes, extensionRegistry);
 
             // Run each file generator, collecting the output
-            List<PluginProtos.CodeGeneratorResponse.File> outputFiles = generators
+            Stream<PluginProtos.CodeGeneratorResponse.File> oldWay = generators
                     .stream()
-                    .flatMap(gen -> gen.generate(request))
-                    .collect(Collectors.toList());
+                    .flatMap(gen -> gen.generate(request));
+
+            Stream<PluginProtos.CodeGeneratorResponse.File> newWay = generators
+                    .stream()
+                    .flatMap(gen -> gen.generateFiles(request).stream());
 
             // Send the files back to protoc
             PluginProtos.CodeGeneratorResponse response = PluginProtos.CodeGeneratorResponse
                     .newBuilder()
-                    .addAllFile(outputFiles)
+                    .addAllFile(Stream.concat(oldWay, newWay).collect(Collectors.toList()))
                     .build();
             response.writeTo(System.out);
 

jprotoc/src/main/java/com/salesforce/jprotoc/jdk8/Jdk8Generator.java

@@ -17,7 +17,6 @@
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.stream.Stream;
 
 /**
  * Generates a set of gRPC stubs that support JDK8 {@link java.util.concurrent.CompletableFuture}.
@@ -30,24 +29,32 @@ public static void main(String[] args) {
     private static final String CLASS_SUFFIX = "Grpc8";
 
     @Override
-    public Stream<PluginProtos.CodeGeneratorResponse.File> generate(PluginProtos.CodeGeneratorRequest request) throws GeneratorException {
+    public List<PluginProtos.CodeGeneratorResponse.File> generateFiles(PluginProtos.CodeGeneratorRequest request) throws GeneratorException {
         final ProtoTypeMap protoTypeMap = ProtoTypeMap.of(request.getProtoFileList());
+        List<PluginProtos.CodeGeneratorResponse.File> files = new ArrayList<>();
 
-        return request.getProtoFileList().stream()
-                .filter(protoFile -> request.getFileToGenerateList().contains(protoFile.getName()))
-                .flatMap(f -> extractContext(protoTypeMap, f))
-                .map(this::buildFile);
+        for (DescriptorProtos.FileDescriptorProto protoFile : request.getProtoFileList()) {
+            if (request.getFileToGenerateList().contains(protoFile.getName())) {
+                for (Context ctx : extractContext(protoTypeMap, protoFile)) {
+                    files.add(buildFile(ctx));
+                }
+            }
+        }
+
+        return files;
     }
 
-    private Stream<Context> extractContext(ProtoTypeMap protoTypeMap, DescriptorProtos.FileDescriptorProto proto) {
-        return proto.getServiceList().stream()
-                .map(s -> extractServiceContext(protoTypeMap, s))
-                .map(ctx -> {
-                    ctx.packageName = extractPackageName(proto); return ctx;
-                })
-                .map(ctx -> {
-                    ctx.protoName = proto.getName(); return ctx;
-                });
+    private List<Context> extractContext(ProtoTypeMap protoTypeMap, DescriptorProtos.FileDescriptorProto proto) {
+        List<Context> contexts = new ArrayList<>();
+
+        for (DescriptorProtos.ServiceDescriptorProto service : proto.getServiceList()) {
+            Context ctx = extractServiceContext(protoTypeMap, service);
+            ctx.packageName = extractPackageName(proto);
+            ctx.protoName = proto.getName();
+            contexts.add(ctx);
+        }
+
+        return contexts;
     }
 
     private String extractPackageName(DescriptorProtos.FileDescriptorProto proto) {