Extensions for options now working

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

Hide file tree

Showing 3 changed files with 23 additions and 58 deletions.
diff --git a/protoc.bat b/protoc.bat
@@ -0,0 +1,2 @@
+mkdir output_test
+for /R "output" %%f in (*.proto) do Z:\protoc-3.20.0-win64\bin\protoc.exe -I=%cd%\output --java_out=output_test\ "%%f"
diff --git a/src/main/kotlin/com/valaphee/protod/Main.kt b/src/main/kotlin/com/valaphee/protod/Main.kt
@@ -18,15 +18,12 @@ 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
 
 fun main() {
     val bytes = File("C:\\Program Files (x86)\\Battle.net\\Battle.net.13401\\battle.net.dll").readBytes()
-    val fileDescriptorProtos = mutableListOf<DescriptorProtos.FileDescriptorProto>()
+    val files = mutableListOf<DescriptorProtos.FileDescriptorProto>()
     String(bytes, Charsets.US_ASCII).occurrencesOf(".proto").forEach {
         var offset = 0
         while (true) {
@@ -40,7 +37,7 @@ fun main() {
                         while (true) {
                             val codedInputStream = CodedInputStream.newInstance(bytes, begin, offset0)
                             try {
-                                fileDescriptorProtos += DescriptorProtos.FileDescriptorProto.parseFrom(codedInputStream)
+                                files += DescriptorProtos.FileDescriptorProto.parseFrom(codedInputStream)
 
                                 break
                             } catch (_: Exception) {
@@ -56,37 +53,11 @@ fun main() {
         }
     }
 
-    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 messages = mutableMapOf<String, DescriptorProtos.DescriptorProto>()
+    files.forEach { file -> file.messageTypeList.forEach { message -> messages[".${file.`package`}.${message.name}"] = message } }
+    val messageExtensions = mutableMapOf<String, MutableMap<Int, DescriptorProtos.FieldDescriptorProto>>()
+    files.forEach { file -> file.extensionList.forEach { extension -> messages[extension.typeName]?.let { messageExtensions.getOrPut(extension.extendee) { mutableMapOf() }[extension.number] = extension } } }
 
     val outputPath = File("output")
-    fileDescriptorProtos.forEach {
-        File(outputPath, it.name).apply { parentFile.mkdirs() }.printWriter().use { printWriter ->
-            ProtoWriter(printWriter).also {
-                it.extensionRegistry = extensionRegistry
-            }.print(it)
-        }
-    }
+    files.forEach { file -> File(outputPath, file.name).apply { parentFile.mkdirs() }.printWriter().use { printWriter -> ProtoWriter(printWriter, messages, messageExtensions).print(file) } }
 }
diff --git a/src/main/kotlin/com/valaphee/protod/ProtoWriter.kt b/src/main/kotlin/com/valaphee/protod/ProtoWriter.kt
@@ -17,23 +17,24 @@
 package com.valaphee.protod
 
 import com.google.protobuf.DescriptorProtos
-import com.google.protobuf.ExtensionRegistry
 import com.google.protobuf.GeneratedMessageV3
+import com.google.protobuf.UnknownFieldSet
 import java.io.PrintWriter
 
 /**
  * @author Kevin Ludwig
  */
 class ProtoWriter(
     private val printWriter: PrintWriter,
-    private val indent: String = "  ",
+    private val messages: MutableMap<String, DescriptorProtos.DescriptorProto>,
+    private val messageExtensions: MutableMap<String, out Map<Int, DescriptorProtos.FieldDescriptorProto>>
 ) {
     private var first = false
     private var indentLevel = 0
 
-    lateinit var extensionRegistry: ExtensionRegistry
-
     fun print(file: DescriptorProtos.FileDescriptorProto) {
+        println("""syntax = "proto2";""")
+        println()
         println("""package ${file.`package`};""")
         printImports(file.dependencyList)
         printOptions(file.options)
@@ -42,12 +43,12 @@ class ProtoWriter(
     }
 
     private fun println() {
-        repeat(indentLevel) { printWriter.print(indent) }
+        repeat(indentLevel) { printWriter.print("  ") }
         printWriter.println()
     }
 
     private fun println(value: String) {
-        repeat(indentLevel) { printWriter.print(indent) }
+        repeat(indentLevel) { printWriter.print("  ") }
         printWriter.println(value)
     }
 
@@ -57,24 +58,15 @@ class ProtoWriter(
     }
 
     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"))
+        options.allFields.forEach {
+            println("""option ${it.key.name} = ${when (it.value) {
+                is String -> """"${it.value}""""
+                else -> it.value
+            }};""")
         }
+        val messageExtensions = messageExtensions[".${options.descriptorForType.fullName}"] ?: emptyMap()
+        options.unknownFields.asMap().forEach { unknownField -> messageExtensions[unknownField.key]?.let { messageExtension -> messages[messageExtension.typeName]?.fieldList?.let { messageExtensionFields -> unknownField.value.lengthDelimitedList.forEach { UnknownFieldSet.parseFrom(it).asMap().forEach { messageExtensionField -> println("""option (${messageExtension.name}).${messageExtensionFields.single { it.number == messageExtensionField.key}?.name} = ${messageExtensionField.value.varintList.firstOrNull() ?: messageExtensionField.value.fixed32List.firstOrNull() ?: messageExtensionField.value.fixed64List.firstOrNull() ?: messageExtensionField.value.lengthDelimitedList.firstOrNull()?.toStringUtf8()?.let { """"$it"""" } ?: TODO()};""") } } } } }
     }
 
     private fun printMessages(messageTypeList: List<DescriptorProtos.DescriptorProto>) {