Working on extension registry

valaphee · Apr 4, 2022 · 58ec40b · 58ec40b
1 parent bf9051f
commit 58ec40b
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 16 deletions.
diff --git a/src/main/kotlin/com/valaphee/protod/Main.kt b/src/main/kotlin/com/valaphee/protod/Main.kt
@@ -18,6 +18,9 @@ package com.valaphee.protod
 
 import com.google.protobuf.CodedInputStream
 import com.google.protobuf.DescriptorProtos
+import com.google.protobuf.Descriptors
+import com.google.protobuf.DynamicMessage
+import com.google.protobuf.ExtensionRegistry
 import com.valaphee.protod.util.occurrencesOf
 import java.io.File
 
@@ -53,13 +56,37 @@ fun main() {
         }
     }
 
-    val extendedMessages = mutableMapOf<String, MutableMap<Int, String>>()
-    fileDescriptorProtos.forEach {
-        it.extensionList.forEach {
-            extendedMessages.getOrPut(it.extendee) { mutableMapOf() }[it.number] = it.name
+    val fileDescriptorProtos0 = fileDescriptorProtos.toMutableList()
+    val fileDescriptors = mutableMapOf<String, Descriptors.FileDescriptor>()
+    var changed = true
+    while (fileDescriptorProtos0.isNotEmpty() && changed) {
+        changed = false
+
+        val fileDescriptorProtoIterator = fileDescriptorProtos0.iterator()
+        while (fileDescriptorProtoIterator.hasNext()) {
+            val fileDescriptorProto = fileDescriptorProtoIterator.next()
+            if (fileDescriptorProto.dependencyList.all { fileDescriptors.contains(it) }) {
+                try {
+                    val fileDescriptor = Descriptors.FileDescriptor.buildFrom(fileDescriptorProto, fileDescriptorProto.dependencyList.map { fileDescriptors[it] }.toTypedArray())
+                    fileDescriptors[fileDescriptor.name] = fileDescriptor
+                    fileDescriptorProtoIterator.remove()
+                    changed = true
+                } catch (ex: Exception) {
+                    ex.printStackTrace()
+                }
+            }
         }
     }
 
+    val extensionRegistry = ExtensionRegistry.newInstance()
+    fileDescriptors.values.forEach { it.extensions.forEach {if (it.javaType == Descriptors.FieldDescriptor.JavaType.MESSAGE) extensionRegistry.add(it, DynamicMessage.newBuilder(it.messageType).apply { it.messageType.fields.forEach { setField(it, it.defaultValue) } }.build()) else extensionRegistry.add(it) } }
+
     val outputPath = File("output")
-    fileDescriptorProtos.forEach { File(outputPath, it.name).apply { parentFile.mkdirs() }.printWriter().use { printWriter -> ProtoWriter(printWriter).print(it) } }
+    fileDescriptorProtos.forEach {
+        File(outputPath, it.name).apply { parentFile.mkdirs() }.printWriter().use { printWriter ->
+            ProtoWriter(printWriter).also {
+                it.extensionRegistry = extensionRegistry
+            }.print(it)
+        }
+    }
 }
diff --git a/src/main/kotlin/com/valaphee/protod/ProtoWriter.kt b/src/main/kotlin/com/valaphee/protod/ProtoWriter.kt
@@ -17,25 +17,28 @@
 package com.valaphee.protod
 
 import com.google.protobuf.DescriptorProtos
-import com.google.protobuf.Message
+import com.google.protobuf.ExtensionRegistry
+import com.google.protobuf.GeneratedMessageV3
 import java.io.PrintWriter
 
 /**
  * @author Kevin Ludwig
  */
 class ProtoWriter(
     private val printWriter: PrintWriter,
-    private val indent: String = "  "
+    private val indent: String = "  ",
 ) {
     private var first = false
     private var indentLevel = 0
 
-    fun print(fileDescriptorProto: DescriptorProtos.FileDescriptorProto) {
-        println("""package ${fileDescriptorProto.`package`};""")
-        printImports(fileDescriptorProto.dependencyList)
-        printOptions(fileDescriptorProto.options)
-        printMessages(fileDescriptorProto.messageTypeList)
-        printServices(fileDescriptorProto.serviceList)
+    lateinit var extensionRegistry: ExtensionRegistry
+
+    fun print(file: DescriptorProtos.FileDescriptorProto) {
+        println("""package ${file.`package`};""")
+        printImports(file.dependencyList)
+        printOptions(file.options)
+        printMessages(file.messageTypeList)
+        printServices(file.serviceList)
     }
 
     private fun println() {
@@ -53,9 +56,25 @@ class ProtoWriter(
         dependencyList.forEach { println("""import "$it";""") }
     }
 
-    private fun printOptions(fileOptions: Message) {
-        if (fileOptions.allFields.isNotEmpty() || fileOptions.unknownFields.asMap().isNotEmpty()) if (first) first = false else println()
-        fileOptions.allFields.forEach { println("""option ${it.key} = ${it.value};""") }
+    private fun printOptions(options: GeneratedMessageV3.ExtendableMessage<*>) {
+        options.toByteString().newCodedInput()
+
+
+        //println(extensionRegistry.findImmutableExtensionByName("${options.descriptorForType.fullName}"))
+
+        if (options.allFields.isNotEmpty() || options.unknownFields.asMap().isNotEmpty()) if (first) first = false else println()
+        options.allFields.forEach { println("""option ${it.key} = ${it.value};""") }
+    }
+
+    private fun printOptions(options: DescriptorProtos.ServiceOptions) {
+        val options0 = DescriptorProtos.ServiceOptions.parseFrom(options.toByteString().newCodedInput(), extensionRegistry)
+        if (options0.allFields.isNotEmpty() || options.unknownFields.asMap().isNotEmpty()) if (first) first = false else println()
+        options0.allFields.forEach { println("""option ${it.key} = ${it.value};""") }
+        if (options0.unknownFields.asMap().isNotEmpty()) {
+            System.out.println("MISSING ")
+
+            println(extensionRegistry.findImmutableExtensionByName("google.protobuf.ServiceOptions"))
+        }
     }
 
     private fun printMessages(messageTypeList: List<DescriptorProtos.DescriptorProto>) {
@@ -84,6 +103,7 @@ class ProtoWriter(
             println("""service ${it.name} {""")
             first = true
             indentLevel++
+            printOptions(it.options)
             printMethods(it.methodList)
             first = first0
             indentLevel--