improve fill missing speed

wagyourtail · wagyourtail · commit 84e7e05ee2c6 · 2025-03-18T20:46:55.000-05:00
diff --git a/cli/src/main/kotlin/xyz/wagyourtail/unimined/mapping/cli/Main.kt b/cli/src/main/kotlin/xyz/wagyourtail/unimined/mapping/cli/Main.kt
@@ -18,7 +18,7 @@ import xyz.wagyourtail.unimined.mapping.cli.Main.Companion.LOGGER
 import xyz.wagyourtail.unimined.mapping.formats.FormatRegistry
 import xyz.wagyourtail.unimined.mapping.propogator.Propagator
 import xyz.wagyourtail.unimined.mapping.tree.MemoryMappingTree
-import xyz.wagyourtail.unimined.mapping.visitor.delegate.copyTo
+import xyz.wagyourtail.unimined.mapping.visitor.delegate.copyNames
 import kotlin.time.measureTime
 
 fun main(vararg args: String) {
@@ -72,7 +72,7 @@ class Main: CliktCommand(printHelpOnEmptyArgs = true) {
         for ((from, to) in copyMissingMap) {
             LOGGER.info { "Copying missing names from $from to ${to.joinToString(", ")}" }
             val t = measureTime {
-                mappings.accept(mappings.copyTo(Namespace(from), to.map { Namespace(it) }.toSet()))
+                mappings.accept(mappings.copyNames(Namespace(from), to.map { Namespace(it) }.toSet()))
             }
             LOGGER.info { "Copied missing names in ${t.inWholeMilliseconds}ms" }
         }
diff --git a/src/commonMain/kotlin/xyz/wagyourtail/unimined/mapping/resolver/MappingResolver.kt b/src/commonMain/kotlin/xyz/wagyourtail/unimined/mapping/resolver/MappingResolver.kt
@@ -1,9 +1,6 @@
 package xyz.wagyourtail.unimined.mapping.resolver
 
 import io.github.oshai.kotlinlogging.KotlinLogging
-import kotlinx.coroutines.async
-import kotlinx.coroutines.awaitAll
-import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
 import okio.Buffer
@@ -25,7 +22,8 @@ import xyz.wagyourtail.unimined.mapping.formats.zip.ZipFS
 import xyz.wagyourtail.unimined.mapping.tree.MemoryMappingTree
 import xyz.wagyourtail.unimined.mapping.util.*
 import xyz.wagyourtail.unimined.mapping.visitor.MappingVisitor
-import xyz.wagyourtail.unimined.mapping.visitor.delegate.copyTo
+import xyz.wagyourtail.unimined.mapping.visitor.delegate.DelegateClassVisitor
+import xyz.wagyourtail.unimined.mapping.visitor.delegate.NameCopyDelegate
 import xyz.wagyourtail.unimined.mapping.visitor.delegate.nsFiltered
 import kotlin.jvm.JvmOverloads
 
@@ -229,13 +227,15 @@ abstract class MappingResolver<T : MappingResolver<T>>(val name: String) {
 
                 // fill in missing names from dependent namespaces
                 val filled = unmappedNs.toMutableSet()
+                val toFill = mutableListOf<Pair<Namespace, Set<Namespace>>>()
                 for (entry in sorted) {
-                    val toFill = entry.provides.map { it.first }.toSet() - filled
-                    if (toFill.isNotEmpty()) {
-                        resolved.accept(resolved.copyTo(entry.requires, toFill))
-                        filled.addAll(toFill)
+                    val targets = entry.provides.map { it.first }.toSet() - filled
+                    if (targets.isNotEmpty()) {
+                        filled.addAll(targets)
+                        toFill.add(entry.requires to targets)
                     }
                 }
+                resolved.fillMissingNames(*toFill.toTypedArray())
 
                 LOGGER.info { "Re-resolving fields and methods..." }
 
diff --git a/src/commonMain/kotlin/xyz/wagyourtail/unimined/mapping/tree/MemoryMappingTree.kt b/src/commonMain/kotlin/xyz/wagyourtail/unimined/mapping/tree/MemoryMappingTree.kt
@@ -12,9 +12,11 @@ import xyz.wagyourtail.unimined.mapping.jvms.four.two.one.PackageName
 import xyz.wagyourtail.unimined.mapping.tree.node._class.ClassNode
 import xyz.wagyourtail.unimined.mapping.tree.node._constant.ConstantGroupNode
 import xyz.wagyourtail.unimined.mapping.tree.node._package.PackageNode
+import xyz.wagyourtail.unimined.mapping.visitor.ClassVisitor
 import xyz.wagyourtail.unimined.mapping.visitor.ConstantGroupVisitor
 import xyz.wagyourtail.unimined.mapping.visitor.MappingVisitor
 import xyz.wagyourtail.unimined.mapping.visitor.PackageVisitor
+import xyz.wagyourtail.unimined.mapping.visitor.delegate.*
 
 class MemoryMappingTree : AbstractMappingTree() {
     private val _namespaces = mutableListOf<Namespace>()
@@ -157,4 +159,32 @@ class MemoryMappingTree : AbstractMappingTree() {
         }
     }
 
+    /**
+     * function to fill missing names
+     */
+    suspend fun fillMissingNames(vararg toFill: Pair<Namespace, Set<Namespace>>) {
+        coroutineScope {
+            listOf(
+                async {
+                    packages.parallelMap {
+                        val nameMap = it.names.toMutableMap()
+                        NameCopyDelegate.fillNames(toFill, nameMap)
+                        it.setNames(nameMap)
+                    }
+                },
+                async {
+                    classes.parallelMap {
+                        val nameMap = it.names.toMutableMap()
+                        NameCopyDelegate.fillNames(toFill, nameMap)
+                        it.setNames(nameMap)
+                        it.acceptInner(
+                            DelegateClassVisitor(it, NameCopyDelegate(*toFill)),
+                            namespaces
+                        )
+                    }
+                }
+            ).awaitAll()
+        }
+    }
+
 }
diff --git a/src/commonMain/kotlin/xyz/wagyourtail/unimined/mapping/visitor/delegate/NameCopyVisitor.kt b/src/commonMain/kotlin/xyz/wagyourtail/unimined/mapping/visitor/delegate/NameCopyVisitor.kt
@@ -7,31 +7,41 @@ import xyz.wagyourtail.unimined.mapping.jvms.four.two.one.InternalName
 import xyz.wagyourtail.unimined.mapping.jvms.four.two.one.PackageName
 import xyz.wagyourtail.unimined.mapping.visitor.*
 
-fun MappingVisitor.copyTo(from: Namespace, to: Set<Namespace>, onlyMissing: Boolean = true): MappingVisitor {
-    return DelegateMappingVisitor(this, NameCopyDelegate(from, to, onlyMissing))
+fun MappingVisitor.copyNames(from: Namespace, to: Set<Namespace>, onlyMissing: Boolean = true): MappingVisitor {
+    return DelegateMappingVisitor(this, NameCopyDelegate(from to to, onlyMissing = onlyMissing))
 }
 
-private class NameCopyDelegate(val from: Namespace, val to: Set<Namespace>, val onlyMissing: Boolean) : NullDelegator() {
+fun MappingVisitor.copyNames(vararg from: Pair<Namespace, Set<Namespace>>, onlyMissing: Boolean = true): MappingVisitor {
+    return DelegateMappingVisitor(this, NameCopyDelegate(*from, onlyMissing = onlyMissing))
+}
+
+class NameCopyDelegate(vararg val from: Pair<Namespace, Set<Namespace>>, val onlyMissing: Boolean = true) : NullDelegator() {
+
+    companion object {
 
-    fun Set<Namespace>.ifOnlyMissing(): Set<Namespace> {
-        return if (onlyMissing) to - this else to
+        fun Set<Namespace>.ifOnlyMissing(to: Set<Namespace>, onlyMissing: Boolean): Set<Namespace> {
+            return if (onlyMissing) to - this else to
+        }
+
+        fun <T> fillNames(toFill: Array<out Pair<Namespace, Set<Namespace>>>, names: MutableMap<Namespace, T>, onlyMissing: Boolean = true, fillWith: (T) -> T = { it }) {
+            for ((from, to) in toFill) {
+                val name = names[from] ?: continue
+                for (namespace in names.keys.ifOnlyMissing(to, onlyMissing)) {
+                    names[namespace] = fillWith(name)
+                }
+            }
+        }
     }
 
     override fun visitClass(delegate: MappingVisitor, names: Map<Namespace, InternalName>): ClassVisitor? {
         val nameMap = names.toMutableMap()
-        val name = nameMap[from] ?: return default.visitClass(delegate, nameMap)
-        for (namespace in nameMap.keys.ifOnlyMissing()) {
-            nameMap[namespace] = name
-        }
+        fillNames(from, nameMap, onlyMissing)
         return default.visitClass(delegate, nameMap)
     }
 
     override fun visitPackage(delegate: MappingVisitor, names: Map<Namespace, PackageName>): PackageVisitor? {
         val nameMap = names.toMutableMap()
-        val name = nameMap[from] ?: return default.visitPackage(delegate, nameMap)
-        for (namespace in nameMap.keys.ifOnlyMissing()) {
-            nameMap[namespace] = name
-        }
+        fillNames(from, nameMap, onlyMissing)
         return default.visitPackage(delegate, nameMap)
     }
 
@@ -40,10 +50,7 @@ private class NameCopyDelegate(val from: Namespace, val to: Set<Namespace>, val
         names: Map<Namespace, Pair<String, FieldDescriptor?>>
     ): FieldVisitor? {
         val nameMap = names.toMutableMap()
-        val name = nameMap[from] ?: return default.visitField(delegate, nameMap)
-        for (namespace in nameMap.keys.ifOnlyMissing()) {
-            nameMap[namespace] = name.first to null
-        }
+        fillNames(from, nameMap, onlyMissing) { it.first to null }
         return default.visitField(delegate, nameMap)
     }
 
@@ -52,10 +59,7 @@ private class NameCopyDelegate(val from: Namespace, val to: Set<Namespace>, val
         names: Map<Namespace, Pair<String, MethodDescriptor?>>
     ): MethodVisitor? {
         val nameMap = names.toMutableMap()
-        val name = nameMap[from] ?: return default.visitMethod(delegate, nameMap)
-        for (namespace in nameMap.keys.ifOnlyMissing()) {
-            nameMap[namespace] = name.first to null
-        }
+        fillNames(from, nameMap, onlyMissing) { it.first to null }
         return default.visitMethod(delegate, nameMap)
     }
 
@@ -66,10 +70,7 @@ private class NameCopyDelegate(val from: Namespace, val to: Set<Namespace>, val
         names: Map<Namespace, String>
     ): ParameterVisitor? {
         val nameMap = names.toMutableMap()
-        val name = nameMap[from] ?: return default.visitParameter(delegate, index, lvOrd, nameMap)
-        for (namespace in nameMap.keys.ifOnlyMissing()) {
-            nameMap[namespace] = name
-        }
+        fillNames(from, nameMap, onlyMissing)
         return default.visitParameter(delegate, index, lvOrd, nameMap)
     }
 
@@ -80,10 +81,7 @@ private class NameCopyDelegate(val from: Namespace, val to: Set<Namespace>, val
         names: Map<Namespace, String>
     ): LocalVariableVisitor? {
         val nameMap = names.toMutableMap()
-        val name = nameMap[from] ?: return default.visitLocalVariable(delegate, lvOrd, startOp, nameMap)
-        for (namespace in nameMap.keys.ifOnlyMissing()) {
-            nameMap[namespace] = name
-        }
+        fillNames(from, nameMap, onlyMissing)
         return default.visitLocalVariable(delegate, lvOrd, startOp, nameMap)
     }
 
diff --git a/src/commonTest/kotlin/xyz/wagyourtail/unimined/mapping/test/delegate/NameCopyTest.kt b/src/commonTest/kotlin/xyz/wagyourtail/unimined/mapping/test/delegate/NameCopyTest.kt
@@ -0,0 +1,68 @@
+package xyz.wagyourtail.unimined.mapping.test.delegate
+
+import kotlinx.coroutines.test.runTest
+import okio.Buffer
+import okio.use
+import xyz.wagyourtail.unimined.mapping.Namespace
+import xyz.wagyourtail.unimined.mapping.formats.umf.UMFReader
+import xyz.wagyourtail.unimined.mapping.formats.umf.UMFWriter
+import xyz.wagyourtail.unimined.mapping.visitor.delegate.copyNames
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+class NameCopyTest {
+
+    val ORIGINAL = """
+        umf 1 0
+        intermediary named extra
+        c a b
+         f a;Z b
+         m a;()V b
+        c d
+         f d;Z
+        """.trimIndent()
+
+    val FILLED_EXTRA = """
+        umf 1 0
+        intermediary named extra
+        c a b b
+         f a;Z b b
+         m a;()V b b
+        c d d d
+         f d;Z d d
+        """.trimIndent()
+
+    @Test
+    fun testNameCopy() = runTest{
+        val mappings = Buffer().use { input ->
+            input.writeUtf8(ORIGINAL)
+            UMFReader.read(input)
+        }
+
+        val output = Buffer().use { output ->
+            mappings.accept(UMFWriter.write(output, true).copyNames(Namespace("intermediary") to setOf(Namespace("named")), Namespace("named") to setOf(Namespace("extra"))))
+            output.readUtf8()
+        }
+
+        assertEquals(FILLED_EXTRA.trimEnd(), output.trimEnd().replace('\t', ' '))
+
+    }
+
+    @Test
+    fun testParallelTreeVersion() = runTest {
+        val mappings = Buffer().use { input ->
+            input.writeUtf8(ORIGINAL)
+            UMFReader.read(input)
+        }
+
+        mappings.fillMissingNames(Namespace("intermediary") to setOf(Namespace("named")), Namespace("named") to setOf(Namespace("extra")))
+
+        val output = Buffer().use { output ->
+            mappings.accept(UMFWriter.write(output, true))
+            output.readUtf8()
+        }
+
+        assertEquals(FILLED_EXTRA.trimEnd(), output.trimEnd().replace('\t', ' '))
+    }
+
+}
diff --git a/src/jvmMain/kotlin/xyz/wagyourtail/unimined/mapping/propogator/Propagator.kt b/src/jvmMain/kotlin/xyz/wagyourtail/unimined/mapping/propogator/Propagator.kt
@@ -19,11 +19,9 @@ import kotlin.streams.asStream
 class Propagator(tree: AbstractMappingTree, override val fns: Namespace, jars: Set<Path>): InheritanceTree(tree) {
 
     override val classes: Map<InternalName, ClassInfo> by lazy {
-        runBlocking {
-            jars.parallelStream().flatMap {
-                scanJar(it).asStream()
-            }.asSequence().toMap()
-        }
+        jars.parallelStream().flatMap {
+            scanJar(it).asStream()
+        }.asSequence().toMap()
     }
 
     private fun scanJar(jar: Path): Sequence<Pair<InternalName, ClassInfo>> {
