Skip to content

Commit

Permalink
perf: async remap entries in jar (#82)
Browse files Browse the repository at this point in the history
  • Loading branch information
@SettingDust
SettingDust authored Nov 27, 2024
1 parent 6bcec62 commit 8a61eaa
Show file tree
Hide file tree
Showing 4 changed files with 133 additions and 75 deletions.
4 changes: 2 additions & 2 deletions src/main/kotlin/xyz/bluspring/kilt/loader/KiltEarlierInitializer.kt
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
package xyz.bluspring.kilt.loader

import de.florianmichael.asmfabricloader.api.event.PrePrePreLaunchEntrypoint
import kotlinx.coroutines.runBlocking

class KiltEarlierInitializer : PrePrePreLaunchEntrypoint {
override fun onLanguageAdapterLaunch() {
KiltLoader.INSTANCE.scanModJob
KiltLoader.INSTANCE.runScanMods()
runBlocking { KiltLoader.INSTANCE.scanMods() }
}
}
38 changes: 25 additions & 13 deletions src/main/kotlin/xyz/bluspring/kilt/loader/KiltLoader.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,8 @@ package xyz.bluspring.kilt.loader
import com.electronwill.nightconfig.core.CommentedConfig
import com.electronwill.nightconfig.toml.TomlParser
import com.google.gson.JsonParser
import kotlinx.coroutines.*
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.SupervisorJob
import net.fabricmc.api.EnvType
import net.fabricmc.loader.api.FabricLoader
import net.fabricmc.loader.impl.FabricLoaderImpl
Expand All @@ -14,11 +15,21 @@ import net.minecraft.SharedConstants
import net.minecraft.server.Bootstrap
import net.minecraftforge.common.ForgeStatesProvider
import net.minecraftforge.eventbus.api.Event
import net.minecraftforge.fml.*
import net.minecraftforge.fml.ModList
import net.minecraftforge.fml.ModLoader
import net.minecraftforge.fml.ModLoadingContext
import net.minecraftforge.fml.ModLoadingPhase
import net.minecraftforge.fml.ModLoadingStage
import net.minecraftforge.fml.common.Mod
import net.minecraftforge.fml.config.ConfigTracker
import net.minecraftforge.fml.config.ModConfig
import net.minecraftforge.fml.event.lifecycle.*
import net.minecraftforge.fml.event.lifecycle.FMLClientSetupEvent
import net.minecraftforge.fml.event.lifecycle.FMLCommonSetupEvent
import net.minecraftforge.fml.event.lifecycle.FMLConstructModEvent
import net.minecraftforge.fml.event.lifecycle.FMLDedicatedServerSetupEvent
import net.minecraftforge.fml.event.lifecycle.FMLLoadCompleteEvent
import net.minecraftforge.fml.event.lifecycle.InterModEnqueueEvent
import net.minecraftforge.fml.event.lifecycle.InterModProcessEvent
import net.minecraftforge.fml.loading.FMLPaths
import net.minecraftforge.fml.loading.moddiscovery.ModAnnotation
import net.minecraftforge.fml.loading.moddiscovery.ModClassVisitor
Expand Down Expand Up @@ -51,7 +62,16 @@ import java.util.function.Consumer
import java.util.jar.JarFile
import java.util.jar.Manifest
import java.util.zip.ZipFile
import kotlin.io.path.*
import kotlin.io.path.createDirectories
import kotlin.io.path.createFile
import kotlin.io.path.div
import kotlin.io.path.exists
import kotlin.io.path.forEachDirectoryEntry
import kotlin.io.path.isDirectory
import kotlin.io.path.name
import kotlin.io.path.nameWithoutExtension
import kotlin.io.path.toPath
import kotlin.io.path.writeBytes
import kotlin.system.exitProcess

class KiltLoader {
Expand All @@ -66,15 +86,7 @@ class KiltLoader {

val scope = CoroutineScope(SupervisorJob())

val scanModJob by lazy {
scope.launch(Dispatchers.IO) { scanMods() }
}

fun runScanMods() {
runBlocking { scanModJob.join() }
}

private suspend fun scanMods() {
suspend fun scanMods() {
Kilt.logger.info("Scanning the mods directory for Forge mods...")
DeltaTimeProfiler.push("scanMods")

Expand Down
140 changes: 80 additions & 60 deletions src/main/kotlin/xyz/bluspring/kilt/loader/remap/KiltRemapper.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,19 @@ import com.google.gson.JsonParser
import it.unimi.dsi.fastutil.objects.Object2ReferenceFunction
import it.unimi.dsi.fastutil.objects.Object2ReferenceMaps
import it.unimi.dsi.fastutil.objects.Object2ReferenceOpenHashMap
import kotlinx.atomicfu.locks.synchronized
import kotlinx.coroutines.CoroutineStart
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.Job
import kotlinx.coroutines.flow.asFlow
import kotlinx.coroutines.flow.filter
import kotlinx.coroutines.flow.collect
import kotlinx.coroutines.flow.flowOn
import kotlinx.coroutines.flow.launchIn
import kotlinx.coroutines.flow.toSet
import kotlinx.coroutines.joinAll
import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.stream.consumeAsFlow
import kotlinx.coroutines.withContext
import net.fabricmc.loader.api.FabricLoader
import net.fabricmc.loader.impl.game.GameProviderHelper
Expand Down Expand Up @@ -42,6 +48,10 @@ import xyz.bluspring.kilt.loader.remap.fixers.WorkaroundFixer
import xyz.bluspring.kilt.util.CaseInsensitiveStringHashSet
import xyz.bluspring.kilt.util.ClassNameHashSet
import xyz.bluspring.kilt.util.KiltHelper
import xyz.bluspring.kilt.util.filterAsync
import xyz.bluspring.kilt.util.flatMapAsync
import xyz.bluspring.kilt.util.mapAsync
import xyz.bluspring.kilt.util.onEachAsync
import java.io.ByteArrayOutputStream
import java.io.File
import java.nio.file.Path
Expand Down Expand Up @@ -114,32 +124,37 @@ object KiltRemapper {
private lateinit var remappedModsDir: Path

// SRG name -> (parent class name, intermediary/mapped name)
val srgMappedFields = srgIntermediaryMapping.classes.flatMap {
it.fields.map { f ->
f.original to
if (!forceProductionRemap)
mappingResolver.mapFieldName(
"intermediary",
it.mapped.replace("/", "."),
f.mapped,
f.mappedDescriptor
)
else
f.mapped
}
}.associateBy { it.first }
val srgMappedFields = runBlocking {
srgIntermediaryMapping.classes.asFlow().flatMapAsync {
it.fields.asFlow().mapAsync { f ->
f.original to
if (!forceProductionRemap)
mappingResolver.mapFieldName(
"intermediary",
it.mapped.replace("/", "."),
f.mapped,
f.mappedDescriptor
)
else
f.mapped
}
}.toSet().associateBy { it.first }
}

// SRG name -> (parent class name, intermediary/mapped name)
val srgMappedMethods = mutableMapOf<String, MutableMap<String, String>>()
val srgMappedMethods =
Object2ReferenceMaps.synchronize(Object2ReferenceOpenHashMap<String, MutableMap<String, String>>())

init {
srgIntermediaryMapping.classes.forEach {
srgIntermediaryMapping.classes.asFlow().flowOn(Dispatchers.IO).onEachAsync {
it.methods.forEach m@{ f ->
// otherwise FunctionalInterface methods don't get remapped properly???
if (!f.mapped.startsWith("method_") && !FabricLoader.getInstance().isDevelopmentEnvironment)
return@m

val map = srgMappedMethods.computeIfAbsent(f.original) { mutableMapOf() }
val map = srgMappedMethods.getOrPut(f.original) {
Object2ReferenceMaps.synchronize(Object2ReferenceOpenHashMap())
}
val mapped = if (!forceProductionRemap)
(mappingResolver.mapMethodName(
"intermediary",
Expand All @@ -152,7 +167,7 @@ object KiltRemapper {

map[f.parent.original] = mapped
}
}
}.launchIn(Kilt.loader.scope)
}

suspend fun remapMods(modLoadingQueue: ConcurrentLinkedQueue<ForgeMod>, remappedModsDir: Path): List<Exception> {
Expand Down Expand Up @@ -233,7 +248,7 @@ object KiltRemapper {
}.build()

val remapper = KiltEnhancedRemapper(classProvider, srgIntermediaryMapping, logConsumer)
val entryToClassNodes = Object2ReferenceOpenHashMap<JarEntry, ClassNode>()
val entryToClassNodes = Object2ReferenceMaps.synchronize(Object2ReferenceOpenHashMap<JarEntry, ClassNode>())

val mixinClasses = ClassNameHashSet()
val refmaps = CaseInsensitiveStringHashSet()
Expand All @@ -250,9 +265,11 @@ object KiltRemapper {

manifest.entries.keys.removeIf { it == "SHA-256-Digest" || it == "SHA-1-Digest" }

jarOutput.putNextEntry(manifestEntry)
jarOutput.write(ByteArrayOutputStream().also { manifest.write(it) }.toByteArray())
jarOutput.closeEntry()
synchronized(jarOutput) {
jarOutput.putNextEntry(manifestEntry)
jarOutput.write(ByteArrayOutputStream().also { manifest.write(it) }.toByteArray())
jarOutput.closeEntry()
}

return@withContext manifest
}
Expand Down Expand Up @@ -492,17 +509,19 @@ object KiltRemapper {
this.add("named:intermediary", newMappings)
})

jarOutput.putNextEntry(entry)
jarOutput.write(Kilt.gson.toJson(refmapData).toByteArray())
jarOutput.closeEntry()
synchronized(jarOutput) {
jarOutput.putNextEntry(entry)
jarOutput.write(Kilt.gson.toJson(refmapData).toByteArray())
jarOutput.closeEntry()
}

return
}

jar.entries().iterator().asFlow()
jar.stream().consumeAsFlow()
.flowOn(Dispatchers.IO)
.filter { !it.name.equals("META-INF/MANIFEST.MF", true) }
.filter {
.filterAsync { !it.name.equals("META-INF/MANIFEST.MF", true) }
.filterAsync {
val isHash = it.name.endsWith(".rsa", true) || it.name.endsWith(".sf", true)
if (isHash) {
// ignore JAR signatures.
Expand All @@ -512,19 +531,21 @@ object KiltRemapper {
}
!isHash
}
.collect { entry ->
.onEachAsync { entry ->
when {
entry.name in refmaps -> remapRefmap(jar, entry, remapper, jarOutput)

// Keep the other resources
!entry.name.endsWith(".class") -> withContext(Dispatchers.IO) {
jarOutput.putNextEntry(entry)
jarOutput.write(jar.getInputStream(entry).readAllBytes())
jarOutput.closeEntry()
!entry.name.endsWith(".class") -> {
synchronized(jarOutput) {
jarOutput.putNextEntry(entry)
jarOutput.write(jar.getInputStream(entry).readAllBytes())
jarOutput.closeEntry()
}
}

else -> {
val classReader = ClassReader(withContext(Dispatchers.IO) { jar.getInputStream(entry) })
val classReader = ClassReader(jar.getInputStream(entry))

// we need the info for this for the class writer
val classNode = ClassNode(Opcodes.ASM9)
Expand All @@ -534,14 +555,15 @@ object KiltRemapper {
}
}
}
.collect()

val classesToProcess =
entryToClassNodes
.values
.intersect(KiltHelper.getForgeClassNodes().toSet())
.toList()

suspend fun remapClass(
fun remapClass(
remapper: KiltEnhancedRemapper,
originalNode: ClassNode,
mixinClasses: ClassNameHashSet,
Expand Down Expand Up @@ -571,7 +593,7 @@ object KiltRemapper {
val classWriter = ClassWriter(0)
remappedNode.accept(classWriter)

withContext(Dispatchers.IO) {
synchronized(jarOutput) {
jarOutput.putNextEntry(entry)
jarOutput.write(classWriter.toByteArray())
jarOutput.closeEntry()
Expand All @@ -582,9 +604,10 @@ object KiltRemapper {
}
}

for ((entry, originalNode) in entryToClassNodes) {
remapClass(remapper, originalNode, mixinClasses, classesToProcess, jarOutput, entry, exceptions)
}
entryToClassNodes.asSequence().asFlow().flowOn(Dispatchers.IO)
.onEachAsync { (entry, originalNode) ->
remapClass(remapper, originalNode, mixinClasses, classesToProcess, jarOutput, entry, exceptions)
}.collect()

mod.remappedModFile = modifiedJarFile.toFile()
withContext(Dispatchers.IO) { jarOutput.close() }
Expand All @@ -595,7 +618,7 @@ object KiltRemapper {
fun remapModAsync(
mod: ForgeMod,
mods: Collection<ForgeMod>
): Job = Kilt.loader.scope.launch(Dispatchers.IO) {
): Job = Kilt.loader.scope.launch(Dispatchers.IO, start = CoroutineStart.LAZY) {
runCatching {
mod.dependencies.mapNotNull {
if (modToJob[it.modId] == null)
Expand Down Expand Up @@ -760,38 +783,35 @@ object KiltRemapper {
runCatching { srgFile.createFile() }

withContext(Dispatchers.IO) { JarOutputStream(srgFile.outputStream()) }.use { outputJar ->
for (entry in gameJar.entries()) {
if (entry.name.endsWith(".class")) {
val classReader = ClassReader(withContext(Dispatchers.IO) {
gameJar.getInputStream(entry)
})
gameJar.stream().consumeAsFlow().flowOn(Dispatchers.IO)
.onEachAsync { entry ->
if (entry.name.endsWith(".class")) {
val classReader = ClassReader(gameJar.getInputStream(entry))

val classNode = ClassNode(Opcodes.ASM9)
classReader.accept(classNode, 0)
val classNode = ClassNode(Opcodes.ASM9)
classReader.accept(classNode, 0)

val classWriter = ClassWriter(0)
val classWriter = ClassWriter(0)

val visitor =
EnhancedClassRemapper(classWriter, srgRemapper, RenamingTransformer(srgRemapper, false))
classNode.accept(visitor)
ConflictingStaticMethodFixer.fixClass(classNode)
val visitor =
EnhancedClassRemapper(classWriter, srgRemapper, RenamingTransformer(srgRemapper, false))
classNode.accept(visitor)
ConflictingStaticMethodFixer.fixClass(classNode)

// We need to remap to the SRG name, otherwise the remapper completely fails in production environments.
val srgName = intermediarySrgMapping.remapClass(entry.name.removePrefix("/").removeSuffix(".class"))
// We need to remap to the SRG name, otherwise the remapper completely fails in production environments.
val srgName =
intermediarySrgMapping.remapClass(entry.name.removePrefix("/").removeSuffix(".class"))

withContext(Dispatchers.IO) {
outputJar.putNextEntry(JarEntry("$srgName.class"))
outputJar.write(classWriter.toByteArray())
outputJar.closeEntry()
}
} else {
withContext(Dispatchers.IO) {
} else {
outputJar.putNextEntry(entry)
outputJar.write(gameJar.getInputStream(entry).readAllBytes())
outputJar.closeEntry()
}
}
}
.collect()
}

logger.info("Remapped Minecraft from Intermediary to SRG. (took ${System.currentTimeMillis() - startTime} ms)")
Expand Down
26 changes: 26 additions & 0 deletions src/main/kotlin/xyz/bluspring/kilt/util/CoroutineUtils.kt
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
package xyz.bluspring.kilt.util

import kotlinx.coroutines.async
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.channelFlow
import kotlinx.coroutines.flow.map
import kotlinx.coroutines.flow.transform

// https://github.com/Kotlin/kotlinx.coroutines/issues/1147

fun <T, R> Flow<T>.mapAsync(transform: suspend (value: T) -> R): Flow<R> =
channelFlow {
collect { e -> send(async { transform(e) }) }
}.map { it.await() }

fun <T> Flow<T>.filterAsync(predicate: suspend (value: T) -> Boolean): Flow<T> =
mapAsync { it to predicate(it) }.transform { if (it.second) emit(it.first) }

fun <T> Flow<T>.onEachAsync(action: suspend (value: T) -> Unit): Flow<T> =
mapAsync {
action(it)
it
}

fun <T, R> Flow<T>.flatMapAsync(transform: suspend (value: T) -> Flow<R>): Flow<R> =
mapAsync { transform(it) }.transform { it.collect { emit(it) } }

0 comments on commit 8a61eaa

Please sign in to comment.