protected Collection generateProtoMarshaller()

in kogito-codegen-modules/kogito-codegen-processes/src/main/java/org/kie/kogito/codegen/process/persistence/PersistenceGenerator.java [152:273]


    protected Collection<GeneratedFile> generateProtoMarshaller() {
        if (!hasProtoMarshaller(context())) {
            // TODO implement a validation check to verify that data classes implement Serializable
            LOGGER.debug("Proto marshaller generation is skipped because " + KOGITO_PERSISTENCE_PROTO_MARSHALLER + "=false");
            return Collections.emptyList();
        }
        Proto proto = protoGenerator.protoOfDataClasses(context().getPackageName(), "import \"kogito-types.proto\";");

        List<String> variableMarshallers = new ArrayList<>();

        String protoContent = proto.serialize();

        List<CompilationUnit> marshallers;
        try {
            marshallers = marshallerGenerator.generate(protoContent);
        } catch (IOException e) {
            throw new UncheckedIOException("Impossible to obtain marshaller CompilationUnits", e);
        }

        Collection<GeneratedFile> protoFiles = new ArrayList<>();
        try {
            String typesURI = "META-INF/kogito-types.proto";
            protoFiles.add(new GeneratedFile(GeneratedFileType.INTERNAL_RESOURCE,
                    typesURI,
                    new String(context().getClassLoader().getResourceAsStream(typesURI).readAllBytes(), StandardCharsets.UTF_8)));
        } catch (IOException e) {
            throw new UncheckedIOException("Cannot find kogito types protobuf!", e);
        }
        // generate proto files leads to problems as it has a reverse dependency of kogito-index
        String typesURI = "META-INF/application-types.proto";
        protoFiles.add(new GeneratedFile(GeneratedFileType.INTERNAL_RESOURCE,
                typesURI,
                protoContent));

        Collection<GeneratedFile> generatedFiles = new ArrayList<>(protoFiles);

        if (!marshallers.isEmpty()) {

            List<CompilationUnit> files = new ArrayList<>(marshallers);

            variableMarshallers.add("org.jbpm.flow.serialization.marshaller.StringProtostreamBaseMarshaller");
            variableMarshallers.add("org.jbpm.flow.serialization.marshaller.BooleanProtostreamBaseMarshaller");
            variableMarshallers.add("org.jbpm.flow.serialization.marshaller.DateProtostreamBaseMarshaller");
            variableMarshallers.add("org.jbpm.flow.serialization.marshaller.DoubleProtostreamBaseMarshaller");
            variableMarshallers.add("org.jbpm.flow.serialization.marshaller.FloatProtostreamBaseMarshaller");
            variableMarshallers.add("org.jbpm.flow.serialization.marshaller.IntegerProtostreamBaseMarshaller");
            variableMarshallers.add("org.jbpm.flow.serialization.marshaller.LongProtostreamBaseMarshaller");
            variableMarshallers.add("org.jbpm.flow.serialization.marshaller.InstantProtostreamBaseMarshaller");
            variableMarshallers.add("org.jbpm.flow.serialization.marshaller.SerializableProtostreamBaseMarshaller");

            for (CompilationUnit unit : files) {
                String packageName = unit.getPackageDeclaration().map(pd -> pd.getName().toString()).orElse("");
                Optional<ClassOrInterfaceDeclaration> clazz = unit.findFirst(ClassOrInterfaceDeclaration.class);
                clazz.ifPresent(c -> {
                    String clazzName = packageName + "." + c.getName().toString();
                    variableMarshallers.add(clazzName);
                    generatedFiles.add(new GeneratedFile(GeneratedFileType.SOURCE,
                            clazzName.replace('.', '/') + JAVA,
                            unit.toString()));
                });
            }

            // we build the marshaller for protostream
            TemplatedGenerator generatorProtostreamSerialization = TemplatedGenerator.builder().withTemplateBasePath(CLASS_TEMPLATES_PERSISTENCE)
                    .withFallbackContext(JavaKogitoBuildContext.CONTEXT_NAME)
                    .withPackageName(KOGITO_PROCESS_INSTANCE_PACKAGE)
                    .build(context(), "ProtostreamObjectMarshaller");
            CompilationUnit parsedClazzFile = generatorProtostreamSerialization.compilationUnitOrThrow();
            String packageName = parsedClazzFile.getPackageDeclaration().map(pd -> pd.getName().toString()).orElse("");
            ClassOrInterfaceDeclaration clazz = parsedClazzFile.findFirst(ClassOrInterfaceDeclaration.class)
                    .orElseThrow(() -> new InvalidTemplateException(generatorProtostreamSerialization, "Failed to find template for ProtostreamObjectMarshaller"));

            ConstructorDeclaration constructor = clazz.getDefaultConstructor()
                    .orElseThrow(() -> new InvalidTemplateException(generatorProtostreamSerialization, "Failed to find default constructor in template for ProtostreamObjectMarshaller"));

            // register protofiles and marshallers
            BlockStmt body = new BlockStmt();
            Expression newFileDescriptorSource = new ObjectCreationExpr(null, new ClassOrInterfaceType(null, FileDescriptorSource.class.getCanonicalName()), NodeList.nodeList());
            Expression getClassLoader = new MethodCallExpr(new MethodCallExpr(null, "getClass", NodeList.nodeList()), "getClassLoader", NodeList.nodeList());

            Expression chainExpression = newFileDescriptorSource;
            for (GeneratedFile generatedFile : protoFiles) {
                String path = generatedFile.relativePath();
                String name = generatedFile.path().getFileName().toString();
                if (!name.endsWith(".proto")) {
                    continue;
                }
                Expression getISKogito = new MethodCallExpr(getClassLoader, "getResourceAsStream", NodeList.nodeList(new StringLiteralExpr(path)));
                chainExpression =
                        new MethodCallExpr(new EnclosedExpr(chainExpression), "addProtoFile",
                                NodeList.nodeList(new StringLiteralExpr(name), getISKogito));
            }

            body.addStatement(new MethodCallExpr(new NameExpr("context"), "registerProtoFiles", NodeList.nodeList(chainExpression)));
            for (String baseMarshallers : variableMarshallers) {
                Expression newMarshallerExpr = new ObjectCreationExpr(null, new ClassOrInterfaceType(null, baseMarshallers), NodeList.nodeList());
                body.addStatement(new MethodCallExpr(new NameExpr("context"), "registerMarshaller", NodeList.nodeList(newMarshallerExpr)));
            }
            CatchClause catchClause = new CatchClause(new Parameter().setType(IOException.class).setName("e"), new BlockStmt());
            TryStmt tryStmt = new TryStmt(body, NodeList.nodeList(catchClause), null);
            constructor.getBody().addStatement(tryStmt);
            String fqnProtoStreamMarshaller = packageName + "." + clazz.getName().toString();
            generatedFiles.add(new GeneratedFile(GeneratedFileType.SOURCE,
                    fqnProtoStreamMarshaller.replace('.', '/') + JAVA,
                    parsedClazzFile.toString()));

            String objectMarshallerStrategyServiceDescriptor = "";
            try {
                //try to find an existing ObjectMarshallerStrategy descriptor in the classpath to be appended to the ProtoStream generated one
                objectMarshallerStrategyServiceDescriptor =
                        new String(getClass().getResourceAsStream("/META-INF/services/org.jbpm.flow.serialization.ObjectMarshallerStrategy").readAllBytes(), StandardCharsets.UTF_8);
            } catch (Exception e) {
                LOGGER.warn("No existing ObjectMarshallerStrategy found the the classpath to be included with the ProtoS generated one for SPI.");
            }
            objectMarshallerStrategyServiceDescriptor += "\n" + fqnProtoStreamMarshaller + "\n";

            generatedFiles.add(new GeneratedFile(GeneratedFileType.INTERNAL_RESOURCE,
                    "META-INF/services/org.jbpm.flow.serialization.ObjectMarshallerStrategy",
                    objectMarshallerStrategyServiceDescriptor));
        }
        return generatedFiles;
    }