[WIP] adopt model replacement from Forge vesion

+ bunch of renames to bring the 2 version closer + at-least-not-crashing levels of Optifine support
2021-05-13 21:13:10 +02:00
parent 9dacdde761
commit 40fd46b278
51 changed files with 1082 additions and 749 deletions
@@ -0,0 +1,28 @@
+package mods.betterfoliage.resource
+
+import net.fabricmc.fabric.api.resource.IdentifiableResourceReloadListener
+import net.minecraft.resource.ResourceManager
+import net.minecraft.resource.ResourceReloadListener
+import net.minecraft.util.profiler.Profiler
+import java.util.concurrent.CompletableFuture
+import java.util.concurrent.Executor
+
+/**
+ * Catch resource reload extremely early, before any of the reloaders
+ * have started working.
+ */
+interface VeryEarlyReloadListener : ResourceReloadListener, IdentifiableResourceReloadListener {
+    override fun reload(
+        synchronizer: ResourceReloadListener.Synchronizer,
+        resourceManager: ResourceManager,
+        preparationsProfiler: Profiler,
+        reloadProfiler: Profiler,
+        backgroundExecutor: Executor,
+        gameExecutor: Executor
+    ): CompletableFuture<Void> {
+        onReloadStarted(resourceManager)
+        return synchronizer.whenPrepared(null)
+    }
+
+    fun onReloadStarted(resourceManager: ResourceManager) {}
+}
@@ -1,102 +0,0 @@
-package mods.betterfoliage.resource.discovery
-
-import mods.betterfoliage.BetterFoliage
-import mods.betterfoliage.BlockModelsReloadCallback
-import mods.betterfoliage.ModelLoadingCallback
-import mods.betterfoliage.util.HasLogger
-import mods.betterfoliage.util.Invalidator
-import mods.betterfoliage.util.YarnHelper
-import mods.betterfoliage.util.get
-import net.fabricmc.fabric.api.event.client.ClientSpriteRegistryCallback
-import net.minecraft.block.BlockState
-import net.minecraft.client.MinecraftClient
-import net.minecraft.client.render.block.BlockModels
-import net.minecraft.client.render.model.BakedModel
-import net.minecraft.client.render.model.ModelLoader
-import net.minecraft.client.texture.SpriteAtlasTexture
-import net.minecraft.resource.ResourceManager
-import net.minecraft.util.Identifier
-import org.apache.logging.log4j.Level
-import org.apache.logging.log4j.Level.DEBUG
-import org.apache.logging.log4j.Level.WARN
-import java.lang.ref.WeakReference
-import java.util.*
-import java.util.concurrent.CompletableFuture
-import java.util.function.Consumer
-import java.util.function.Supplier
-
-// net.minecraft.client.render.block.BlockModels.models
-val BlockModels_models = YarnHelper.requiredField<Map<BlockState, BakedModel>>("net.minecraft.class_773", "field_4162", "Ljava/util/Map;")
-
-class BakedModelReplacer : ModelLoadingCallback, ClientSpriteRegistryCallback, BlockModelsReloadCallback, Invalidator, HasLogger {
-    override val logger get() = BetterFoliage.logDetail
-
-    val discoverers = mutableListOf<ModelDiscovery>()
-    override val callbacks = mutableListOf<WeakReference<()->Unit>>()
-
-    protected var keys = emptyMap<BlockState, BlockRenderKey>()
-
-    operator fun get(state: BlockState) = keys[state]
-    inline fun <reified T> getTyped(state: BlockState) = get(state) as? T
-
-    var currentLoader: ModelLoader? = null
-
-    override fun beginLoadModels(loader: ModelLoader, manager: ResourceManager) {
-        // Step 1: get a hold of the ModelLoader instance when model reloading starts
-        currentLoader = loader
-        log("reloading block discovery configuration")
-        BetterFoliage.blockConfig.reloadConfig(manager)
-        invalidate()
-    }
-
-    override fun registerSprites(atlasTexture: SpriteAtlasTexture, registry: ClientSpriteRegistryCallback.Registry) {
-        // Step 2: ModelLoader is finished with the unbaked models by now, we can inspect them
-        log("discovering blocks")
-        val idSet = Collections.synchronizedSet(mutableSetOf<Identifier>())
-        val allKeys = discoverers.map {
-            // run model discoverers in parallel
-            CompletableFuture.supplyAsync(Supplier {
-                it.discover(currentLoader!!, Consumer { idSet.add(it) })
-            }, MinecraftClient.getInstance())
-        }.map { it.join() }
-        idSet.forEach { registry.register(it) }
-
-        val result = mutableMapOf<BlockState, BlockRenderKey>()
-        allKeys.forEach { keys ->
-            keys.entries.forEach { (state, key) ->
-                val oldKey = result[state]
-                if (oldKey != null) log("Replacing $oldKey with $key for state $state")
-                else log(DEBUG, "Adding replacement $key for state $state")
-                result[state] = key
-            }
-        }
-
-        keys = result
-    }
-
-    override fun reloadBlockModels(blockModels: BlockModels) {
-        // Step 3: replace the baked models with our own
-        log("block model baking finished")
-        val modelMap = blockModels[BlockModels_models] as MutableMap<BlockState, BakedModel>
-        keys.forEach { (state, key) ->
-            val oldModel = modelMap[state]
-            if (oldModel == null) log(WARN, "Cannot find model for state $state, ignoring")
-            else {
-                try {
-                    val newModel = key.replace(oldModel, state)
-                    modelMap[state] = newModel
-                    log(DEBUG, "Replaced model for state $state with $key")
-                } catch (e: Exception) {
-                    log(WARN, "Error creating model for state $state with $key", e)
-                }
-            }
-        }
-    }
-
-    init {
-        ModelLoadingCallback.EVENT.register(this)
-        ClientSpriteRegistryCallback.event(SpriteAtlasTexture.BLOCK_ATLAS_TEX).register(this)
-        BlockModelsReloadCallback.EVENT.register(this)
-    }
-}
-
@@ -0,0 +1,124 @@
+package mods.betterfoliage.resource.discovery
+
+import mods.betterfoliage.BetterFoliage
+import mods.betterfoliage.BlockModelsReloadCallback
+import mods.betterfoliage.ModelLoadingCallback
+import mods.betterfoliage.util.HasLogger
+import mods.betterfoliage.util.Invalidator
+import net.fabricmc.fabric.api.event.client.ClientSpriteRegistryCallback
+import net.minecraft.block.BlockState
+import net.minecraft.client.render.block.BlockModels
+import net.minecraft.client.render.model.BakedModel
+import net.minecraft.client.render.model.ModelBakeSettings
+import net.minecraft.client.render.model.ModelLoader
+import net.minecraft.client.render.model.UnbakedModel
+import net.minecraft.client.texture.Sprite
+import net.minecraft.client.texture.SpriteAtlasTexture
+import net.minecraft.client.util.SpriteIdentifier
+import net.minecraft.resource.ResourceManager
+import net.minecraft.util.Identifier
+import org.apache.logging.log4j.Level.INFO
+import org.apache.logging.log4j.Level.WARN
+import org.apache.logging.log4j.Logger
+import java.lang.ref.WeakReference
+import java.util.function.Function
+
+data class ModelDiscoveryContext(
+    val bakery: ModelLoader,
+    val blockState: BlockState,
+    val modelLocation: Identifier,
+    val sprites: MutableSet<Identifier>,
+    val replacements: MutableMap<Identifier, ModelBakingKey>,
+    val logger: Logger
+) {
+    fun getUnbaked(location: Identifier = modelLocation) = bakery.getOrLoadModel(location)
+    fun addReplacement(key: ModelBakingKey, addToStateKeys: Boolean = true) {
+        replacements[modelLocation] = key
+        if (addToStateKeys) BetterFoliage.blockTypes.stateKeys[blockState] = key
+        logger.log(INFO, "Adding model replacement $modelLocation -> $key")
+    }
+}
+
+interface ModelDiscovery {
+    fun onModelsLoaded(
+        bakery: ModelLoader,
+        sprites: MutableSet<Identifier>,
+        replacements: MutableMap<Identifier, ModelBakingKey>
+    )
+}
+
+data class ModelBakingContext(
+    val bakery: ModelLoader,
+    val spriteGetter: Function<SpriteIdentifier, Sprite>,
+    val location: Identifier,
+    val transform: ModelBakeSettings,
+    val logger: Logger
+) {
+    fun getUnbaked() = bakery.getOrLoadModel(location)
+    fun getBaked() = bakery.bake(location, transform)
+}
+
+interface ModelBakingKey {
+    fun bake(ctx: ModelBakingContext): BakedModel? =
+        ctx.getUnbaked().bake(ctx.bakery, ctx.spriteGetter, ctx.transform, ctx.location)
+}
+
+object BakeWrapperManager : HasLogger(), Invalidator, ModelLoadingCallback, ClientSpriteRegistryCallback, BlockModelsReloadCallback {
+    init {
+        ModelLoadingCallback.EVENT.register(this)
+        ClientSpriteRegistryCallback.event(SpriteAtlasTexture.BLOCK_ATLAS_TEX).register(this)
+    }
+    val discoverers = mutableListOf<ModelDiscovery>()
+    override val callbacks = mutableListOf<WeakReference<()->Unit>>()
+
+    private val replacements = mutableMapOf<Identifier, ModelBakingKey>()
+    private val sprites = mutableSetOf<Identifier>()
+
+    override fun beginLoadModels(loader: ModelLoader, manager: ResourceManager) {
+        val startTime = System.currentTimeMillis()
+        replacements.clear()
+        sprites.clear()
+        invalidate()
+        BetterFoliage.blockTypes = BlockTypeCache()
+
+        logger.log(INFO, "starting model discovery (${discoverers.size} listeners)")
+        discoverers.forEach { listener ->
+            val replacementsLocal = mutableMapOf<Identifier, ModelBakingKey>()
+            listener.onModelsLoaded(loader, sprites, replacements)
+        }
+
+        val elapsed = System.currentTimeMillis() - startTime
+        logger.log(INFO, "finished model discovery in $elapsed ms, ${replacements.size} top-level replacements")
+    }
+
+    override fun registerSprites(atlas: SpriteAtlasTexture, registry: ClientSpriteRegistryCallback.Registry) {
+        logger.log(INFO, "Adding ${sprites.size} sprites to block atlas")
+        sprites.forEach { registry.register(it) }
+        sprites.clear()
+    }
+
+    override fun reloadBlockModels(blockModels: BlockModels) {
+        replacements.clear()
+    }
+
+    fun onBake(
+        unbaked: UnbakedModel,
+        bakery: ModelLoader,
+        spriteGetter: Function<SpriteIdentifier, Sprite>,
+        transform: ModelBakeSettings,
+        location: Identifier
+    ): BakedModel? {
+        val ctx = ModelBakingContext(bakery, spriteGetter, location, transform, detailLogger)
+        // bake replacement if available
+        replacements[location]?.let { replacement ->
+            detailLogger.log(INFO, "Baking replacement for [${unbaked::class.java.simpleName}] $location -> $replacement")
+            try {
+                return replacement.bake(ctx)
+            } catch (e: Exception) {
+                detailLogger.log(WARN, "Error while baking $replacement", e)
+                logger.log(WARN, "Error while baking $replacement", e)
+            }
+        }
+        return unbaked.bake(bakery, spriteGetter, transform, location)
+    }
+}
@@ -0,0 +1,14 @@
+package mods.betterfoliage.resource.discovery
+
+import net.minecraft.block.BlockState
+
+class BlockTypeCache {
+    val leaf = mutableSetOf<BlockState>()
+    val grass = mutableSetOf<BlockState>()
+    val dirt = mutableSetOf<BlockState>()
+
+    val stateKeys = mutableMapOf<BlockState, ModelBakingKey>()
+
+    inline fun <reified T> getTyped(state: BlockState) = stateKeys[state] as? T
+    inline fun <reified T> hasTyped(state: BlockState) = stateKeys[state] is T
+}
@@ -1,69 +0,0 @@
-package mods.betterfoliage.resource.discovery
-
-import com.google.common.base.Joiner
-import mods.betterfoliage.util.YarnHelper
-import mods.betterfoliage.util.get
-import mods.betterfoliage.util.stripStart
-import net.minecraft.block.BlockState
-import net.minecraft.client.render.model.ModelLoader
-import net.minecraft.client.render.model.json.JsonUnbakedModel
-import net.minecraft.client.texture.MissingSprite
-import net.minecraft.util.Identifier
-import org.apache.logging.log4j.Level
-import org.apache.logging.log4j.Level.DEBUG
-import org.apache.logging.log4j.Level.INFO
-import java.util.function.Consumer
-
-// net.minecraft.client.render.model.json.JsonUnbakedModel.parent
-val JsonUnbakedModel_parent = YarnHelper.requiredField<JsonUnbakedModel>("net.minecraft.class_793", "field_4253", "Lnet/minecraft/class_793;")
-// net.minecraft.client.render.model.json.JsonUnbakedModel.parentId
-val JsonUnbakedModel_parentId = YarnHelper.requiredField<Identifier>("net.minecraft.class_793", "field_4247", "Lnet/minecraft/class_2960;")
-
-fun Pair<JsonUnbakedModel, Identifier>.derivesFrom(targetLocation: Identifier): Boolean {
-    if (second.stripStart("models/") == targetLocation) return true
-    if (first[JsonUnbakedModel_parent] != null && first[JsonUnbakedModel_parentId] != null)
-        return Pair(first[JsonUnbakedModel_parent]!!, first[JsonUnbakedModel_parentId]!!).derivesFrom(targetLocation)
-    return false
-}
-
-abstract class ConfigurableModelDiscovery : ModelDiscoveryBase() {
-
-    abstract val matchClasses: IBlockMatcher
-    abstract val modelTextures: List<ModelTextureList>
-
-    override fun discover(loader: ModelLoader, atlas: Consumer<Identifier>): Map<BlockState, BlockRenderKey> {
-        log(INFO, "Starting model discovery: ${this::class.java.canonicalName}")
-        matchClasses.describe(this)
-        modelTextures.forEach { modelTex ->
-            log(DEBUG, "    model: ${modelTex.modelLocation} textures: ${modelTex.textureNames.joinToString(", ")}")
-        }
-        return super.discover(loader, atlas)
-    }
-
-    abstract fun processModel(state: BlockState, textures: List<Identifier>, atlas: Consumer<Identifier>): BlockRenderKey?
-
-    override fun processModel(ctx: ModelDiscoveryContext, atlas: Consumer<Identifier>): BlockRenderKey? {
-        val matchClass = matchClasses.matchingClass(ctx.state.block) ?: return null
-        log(DEBUG, "block state ${ctx.state.toString()}")
-        log(DEBUG, "      class ${ctx.state.block.javaClass.name} matches ${matchClass.name}")
-
-        (ctx.models.filter { it.first is JsonUnbakedModel } as List<Pair<JsonUnbakedModel, Identifier>>).forEach { (model, location) ->
-            val modelMatch = modelTextures.firstOrNull { (model to location).derivesFrom(it.modelLocation) }
-            if (modelMatch != null) {
-                log(DEBUG, "      model ${model} matches ${modelMatch.modelLocation}")
-
-                val textures = modelMatch.textureNames.map { it to model.resolveSprite(it).textureId }
-                val texMapString = Joiner.on(", ").join(textures.map { "${it.first}=${it.second}" })
-                log(DEBUG, "    sprites [$texMapString]")
-
-                if (textures.all { it.second != MissingSprite.getMissingSpriteId() }) {
-                    // found a valid model (all required textures exist)
-                    return processModel(ctx.state, textures.map { it.second }, atlas).also {
-                        log(DEBUG, "    valid model discovered: $it")
-                    }
-                }
-            }
-        }
-        return null
-    }
-}
@@ -9,13 +9,12 @@ import net.minecraft.block.Block
 import net.minecraft.resource.ResourceManager
 import net.minecraft.util.Identifier
 import org.apache.logging.log4j.Level
+import org.apache.logging.log4j.Level.INFO
 import org.apache.logging.log4j.Logger

 interface IBlockMatcher {
    fun matchesClass(block: Block): Boolean
    fun matchingClass(block: Block): Class<*>?
-
-    fun describe(logger: HasLogger)
 }

 class SimpleBlockMatcher(vararg val classes: Class<*>) : IBlockMatcher {
@@ -26,15 +25,9 @@ class SimpleBlockMatcher(vararg val classes: Class<*>) : IBlockMatcher {
        classes.forEach { if (it.isAssignableFrom(blockClass)) return it }
        return null
    }
-
-    override fun describe(logger: HasLogger) {
-        classes.forEach { klass ->
-            logger.log(Level.DEBUG, "    class whitelist: ${klass.name}")
-        }
-    }
 }

-class ConfigurableBlockMatcher(val logger: Logger, val location: Identifier) : IBlockMatcher {
+class ConfigurableBlockMatcher(val location: Identifier) : HasLogger(), IBlockMatcher {

    val blackList = mutableListOf<Class<*>>()
    val whiteList = mutableListOf<Class<*>>()
@@ -57,7 +50,7 @@ class ConfigurableBlockMatcher(val logger: Logger, val location: Identifier) : I
        blackList.clear()
        whiteList.clear()
        manager.getAllResources(location).forEach { resource ->
-            logger.info("Reading class list $location from pack ${resource.resourcePackName}")
+            detailLogger.log(INFO, "Reading class list $location from pack ${resource.resourcePackName}")
            resource.getLines().map{ it.trim() }.filter { !it.startsWith("//") && it.isNotEmpty() }.forEach { line ->
                val name = if (line.startsWith("-")) line.substring(1) else line
                val mappedName = FabricLoader.getInstance().mappingResolver.mapClassName(INTERMEDIARY, name)
@@ -75,26 +68,17 @@ class ConfigurableBlockMatcher(val logger: Logger, val location: Identifier) : I
            }
        }
    }
-
-    override fun describe(logger: HasLogger) {
-        whiteList.forEach { klass ->
-            logger.log(Level.DEBUG, "    class whitelist: ${klass.name}")
-        }
-        blackList.forEach { klass ->
-            logger.log(Level.DEBUG, "    class blacklist: ${klass.name}")
-        }
-    }
 }

 data class ModelTextureList(val modelLocation: Identifier, val textureNames: List<String>) {
    constructor(vararg args: String) : this(Identifier(args[0]), listOf(*args).drop(1))
 }

-class ModelTextureListConfiguration(val logger: Logger, val location: Identifier) {
+class ModelTextureListConfiguration(val location: Identifier) : HasLogger() {
    val modelList = mutableListOf<ModelTextureList>()
    fun readDefaults(manager: ResourceManager) {
        manager.getAllResources(location).forEach { resource ->
-            logger.info("Reading model configuration $location from pack ${resource.resourcePackName}")
+            detailLogger.log(INFO, "Reading model configuration $location from pack ${resource.resourcePackName}")
            resource.getLines().map{ it.trim() }.filter { !it.startsWith("//") && it.isNotEmpty() }.forEach { line ->
                val elements = line.split(",")
                modelList.add(ModelTextureList(Identifier(elements.first()), elements.drop(1)))
@@ -0,0 +1,98 @@
+package mods.betterfoliage.resource.discovery
+
+import com.google.common.base.Joiner
+import mods.betterfoliage.util.HasLogger
+import mods.betterfoliage.util.YarnHelper
+import mods.betterfoliage.util.get
+import net.minecraft.client.render.block.BlockModels
+import net.minecraft.client.render.model.ModelLoader
+import net.minecraft.client.render.model.json.JsonUnbakedModel
+import net.minecraft.client.render.model.json.WeightedUnbakedModel
+import net.minecraft.client.texture.MissingSprite
+import net.minecraft.util.Identifier
+import net.minecraft.util.registry.Registry
+import org.apache.logging.log4j.Level
+
+abstract class AbstractModelDiscovery : HasLogger(), ModelDiscovery {
+    override fun onModelsLoaded(
+        bakery: ModelLoader,
+        sprites: MutableSet<Identifier>,
+        replacements: MutableMap<Identifier, ModelBakingKey>
+    ) {
+        Registry.BLOCK
+            .flatMap { block -> block.stateManager.states }
+            .forEach { state ->
+                val location = BlockModels.getModelId(state)
+                val ctx = ModelDiscoveryContext(bakery, state, location, sprites, replacements, detailLogger)
+                try {
+                    processModel(ctx)
+                } catch (e: Exception) {
+                    logger.log(Level.WARN, "Discovery error in $location", e)
+                }
+            }
+    }
+
+    open fun processModel(ctx: ModelDiscoveryContext) {
+        val model = ctx.getUnbaked()
+
+        // built-in support for container models
+        if (model is WeightedUnbakedModel) {
+            // per-location replacements need to be scoped to the variant list, as replacement models
+            // may need information from the BlockState which is not available at baking time
+            val scopedReplacements = mutableMapOf<Identifier, ModelBakingKey>()
+            model.variants.forEach { variant ->
+                processModel(ctx.copy(modelLocation = variant.location, replacements = scopedReplacements))
+            }
+            if (scopedReplacements.isNotEmpty()) {
+                ctx.addReplacement(WeightedUnbakedKey(scopedReplacements), addToStateKeys = false)
+            }
+        }
+    }
+}
+
+abstract class ConfigurableModelDiscovery : AbstractModelDiscovery() {
+    abstract val matchClasses: IBlockMatcher
+    abstract val modelTextures: List<ModelTextureList>
+
+    abstract fun processModel(
+        ctx: ModelDiscoveryContext,
+        textureMatch: List<Identifier>
+    )
+
+    override fun processModel(ctx: ModelDiscoveryContext) {
+        val model = ctx.getUnbaked()
+        if (model is JsonUnbakedModel) {
+            val matchClass = matchClasses.matchingClass(ctx.blockState.block) ?: return
+
+            detailLogger.log(Level.INFO, "block state ${ctx.blockState}")
+            detailLogger.log(Level.INFO, "      model ${ctx.modelLocation}")
+            detailLogger.log(Level.INFO, "      class ${ctx.blockState.block.javaClass.name} matches ${matchClass.name}")
+
+            modelTextures
+                .filter { matcher -> ctx.bakery.modelDerivesFrom(model, ctx.modelLocation, matcher.modelLocation) }
+                .forEach { match ->
+                    detailLogger.log(Level.INFO, "      model $model matches ${match.modelLocation}")
+
+                    val materials = match.textureNames.map { it to model.resolveSprite(it) }
+                    val texMapString = Joiner.on(", ").join(materials.map { "${it.first}=${it.second.textureId}" })
+                    detailLogger.log(Level.INFO, "    sprites [$texMapString]")
+
+                    if (materials.all { it.second.textureId != MissingSprite.getMissingSpriteId() }) {
+                        // found a valid model (all required textures exist)
+                        processModel(ctx, materials.map { it.second.textureId })
+                    }
+                }
+        }
+        return super.processModel(ctx)
+    }
+}
+
+// net.minecraft.client.render.model.json.JsonUnbakedModel.parentId
+val JsonUnbakedModel_parentId = YarnHelper.requiredField<Identifier>("net.minecraft.class_793", "field_4247", "Lnet/minecraft/class_2960;")
+
+fun ModelLoader.modelDerivesFrom(model: JsonUnbakedModel, location: Identifier, target: Identifier): Boolean =
+    if (location == target) true
+    else model[JsonUnbakedModel_parentId]
+        ?.let { getOrLoadModel(it) as? JsonUnbakedModel }
+        ?.let { parent -> modelDerivesFrom(parent, model[JsonUnbakedModel_parentId]!!, target) }
+        ?: false
@@ -1,74 +0,0 @@
-package mods.betterfoliage.resource.discovery
-
-import mods.betterfoliage.util.HasLogger
-import net.minecraft.block.BlockState
-import net.minecraft.client.render.block.BlockModels
-import net.minecraft.client.render.model.BakedModel
-import net.minecraft.client.render.model.ModelLoader
-import net.minecraft.client.render.model.UnbakedModel
-import net.minecraft.client.render.model.json.JsonUnbakedModel
-import net.minecraft.client.render.model.json.ModelVariant
-import net.minecraft.client.render.model.json.WeightedUnbakedModel
-import net.minecraft.client.texture.SpriteAtlasTexture
-import net.minecraft.client.util.ModelIdentifier
-import net.minecraft.util.Identifier
-import net.minecraft.util.registry.Registry
-import java.util.function.Consumer
-
-interface BlockRenderKey {
-    fun replace(model: BakedModel, state: BlockState): BakedModel = model
-}
-
-fun ModelLoader.iterateModels(func: (ModelDiscoveryContext)->Unit) {
-    Registry.BLOCK.flatMap { block ->
-        block.stateManager.states.map { state -> state to BlockModels.getModelId(state) }
-    }.forEach { (state, stateModelResource) ->
-        func(ModelDiscoveryContext(this, state, stateModelResource))
-    }
-}
-
-/**
- * Information about a single [BlockState] and all the [UnbakedModel]s it could render as.
- */
-class ModelDiscoveryContext(
-    loader: ModelLoader,
-    val state: BlockState,
-    val modelId: ModelIdentifier
-) {
-    val models = loader.unwrapVariants(loader.getOrLoadModel(modelId) to modelId)
-        .filter { it.second != loader.getOrLoadModel(ModelLoader.MISSING) }
-
-    fun ModelLoader.unwrapVariants(modelAndLoc: Pair<UnbakedModel, Identifier>): List<Pair<UnbakedModel, Identifier>> = when(val model = modelAndLoc.first) {
-        is WeightedUnbakedModel -> (model.variants as List<ModelVariant>).flatMap {
-            variant -> unwrapVariants(getOrLoadModel(variant.location) to variant.location)
-        }
-        is JsonUnbakedModel -> listOf(modelAndLoc)
-        else -> emptyList()
-    }
-}
-
-interface ModelDiscovery {
-    fun discover(loader: ModelLoader, atlas: Consumer<Identifier>): Map<BlockState, BlockRenderKey>
-}
-
-abstract class ModelDiscoveryBase : ModelDiscovery, HasLogger {
-    override fun discover(loader: ModelLoader, atlas: Consumer<Identifier>): Map<BlockState, BlockRenderKey> {
-        val keys = mutableMapOf<BlockState, BlockRenderKey>()
-        var errors = 0
-
-        loader.iterateModels { ctx ->
-            try {
-                val result = processModel(ctx, atlas)
-                result?.let { keys[ctx.state] = it }
-            } catch (e: Exception) {
-                errors++
-            }
-        }
-        log("${keys.size} BlockStates discovered, $errors errors")
-        return keys
-    }
-
-    abstract fun processModel(ctx: ModelDiscoveryContext, atlas: Consumer<Identifier>): BlockRenderKey?
-}
-
-
@@ -0,0 +1,61 @@
+package mods.betterfoliage.resource.discovery
+
+import mods.betterfoliage.model.WeightedModelWrapper
+import mods.betterfoliage.model.WrappedBakedModel
+import mods.betterfoliage.model.WrappedMeshModel
+import mods.betterfoliage.util.HasLogger
+import net.fabricmc.fabric.api.renderer.v1.material.BlendMode
+import net.minecraft.client.render.model.BakedModel
+import net.minecraft.client.render.model.BasicBakedModel
+import net.minecraft.client.render.model.json.WeightedUnbakedModel
+import net.minecraft.util.Identifier
+import org.apache.logging.log4j.Level.INFO
+import org.apache.logging.log4j.Level.WARN
+
+class WeightedUnbakedKey(
+    val replacements: Map<Identifier, ModelBakingKey>
+) : ModelBakingKey {
+
+    override fun bake(ctx: ModelBakingContext): BakedModel? {
+        val unbaked = ctx.getUnbaked()
+        if (unbaked !is WeightedUnbakedModel) return super.bake(ctx)
+
+        // bake all variants, replace as needed
+        val bakedModels = unbaked.variants.mapNotNull {
+            val variantCtx = ctx.copy(location = it.location, transform = it)
+            val replacement = replacements[it.location]
+            val baked = replacement?.let { replacement ->
+                ctx.logger.log(INFO, "Baking replacement for variant [${variantCtx.getUnbaked()::class.java.simpleName}] ${variantCtx.location} -> $replacement")
+                replacement.bake(variantCtx)
+            } ?: variantCtx.getBaked()
+            when(baked) {
+                is WrappedBakedModel -> it to baked
+                // just in case we replaced some variants in the list, but not others
+                // this should not realistically happen, this is just a best-effort fallback
+                is BasicBakedModel -> it to WrappedMeshModel.converter(
+                    state = null, unshade = false, noDiffuse = true, blendModeOverride = BlendMode.CUTOUT_MIPPED
+                ).convert(baked)!!
+                else -> null
+            }
+        }
+
+        // something fishy going on, possibly unknown model type
+        // let it through unchanged
+        if (bakedModels.isEmpty()) return super.bake(ctx)
+
+        if (bakedModels.size < unbaked.variants.size) {
+            detailLogger.log(
+                WARN,
+                "Dropped ${unbaked.variants.size - bakedModels.size} variants from model ${ctx.location}"
+            )
+        }
+        val weightedSpecials = bakedModels.map { (variant, model) ->
+            WeightedModelWrapper.WeightedModel(model, variant.weight)
+        }
+        return WeightedModelWrapper(weightedSpecials, weightedSpecials[0].model)
+    }
+
+    override fun toString() = "[WeightedUnbakedKey, ${replacements.size} replacements]"
+
+    companion object : HasLogger()
+}
@@ -13,7 +13,7 @@ data class CenteredSprite(val sprite: Identifier, val atlas: Atlas = Atlas.BLOCK
    fun register(pack: GeneratedBlockTexturePack) = pack.register(this, this::draw)

    fun draw(resourceManager: ResourceManager): ByteArray {
-        val baseTexture = resourceManager.loadSprite(atlas.wrap(sprite))
+        val baseTexture = resourceManager.loadSprite(atlas.file(sprite))

        val frameWidth = baseTexture.width
        val frameHeight = baseTexture.width * aspectHeight / aspectWidth
@@ -10,6 +10,8 @@ import net.minecraft.resource.metadata.ResourceMetadataReader
 import net.minecraft.text.LiteralText
 import net.minecraft.util.Identifier
 import net.minecraft.util.profiler.Profiler
+import org.apache.logging.log4j.Level
+import org.apache.logging.log4j.Level.INFO
 import org.apache.logging.log4j.Logger
 import java.io.IOException
 import java.lang.IllegalStateException
@@ -29,7 +31,9 @@ import java.util.function.Supplier
 * @param[packDesc] Description of pack
 * @param[logger] Logger to log to when generating resources
 */
-class GeneratedBlockTexturePack(val reloadId: Identifier, val nameSpace: String, val packName: String, val packDesc: String, override val logger: Logger) : HasLogger, ResourcePack {
+class GeneratedBlockTexturePack(
+    val reloadId: Identifier, val nameSpace: String, val packName: String, val packDesc: String
+) : HasLogger(), ResourcePack {

    override fun getName() = reloadId.toString()
    override fun getNamespaces(type: ResourceType) = setOf(nameSpace)
@@ -51,8 +55,8 @@ class GeneratedBlockTexturePack(val reloadId: Identifier, val nameSpace: String,
        val resource = func(manager!!)

        identifiers[key] = id
-        resources[Atlas.BLOCKS.wrap(id)] = resource
-        log("generated resource $key -> $id")
+        resources[Atlas.BLOCKS.file(id)] = resource
+        detailLogger.log(INFO, "generated resource $key -> $id")
        return id
    }

@@ -17,7 +17,7 @@ data class GeneratedGrassSprite(val sprite: Identifier, val isSnowed: Boolean, v
    fun register(pack: GeneratedBlockTexturePack) = pack.register(this, this::draw)

    fun draw(resourceManager: ResourceManager): ByteArray {
-        val baseTexture = resourceManager.loadSprite(atlas.wrap(sprite))
+        val baseTexture = resourceManager.loadSprite(atlas.file(sprite))

        val result = BufferedImage(baseTexture.width, baseTexture.height, BufferedImage.TYPE_4BYTE_ABGR)
        val graphics = result.createGraphics()
@@ -20,7 +20,7 @@ data class GeneratedLeafSprite(val sprite: Identifier, val leafType: String, val
    fun register(pack: GeneratedBlockTexturePack) = pack.register(this, this::draw)

    fun draw(resourceManager: ResourceManager): ByteArray {
-        val baseTexture = resourceManager.loadSprite(atlas.wrap(sprite))
+        val baseTexture = resourceManager.loadSprite(atlas.file(sprite))

        val size = baseTexture.width
        val frames = baseTexture.height / size
@@ -67,7 +67,7 @@ data class GeneratedLeafSprite(val sprite: Identifier, val leafType: String, val
     * @param[maxSize] Preferred mask size.
     */
    fun getLeafMask(type: String, maxSize: Int) = getMultisizeTexture(maxSize) { size ->
-        Atlas.BLOCKS.wrap(Identifier(BetterFoliage.MOD_ID, "blocks/leafmask_${size}_${type}"))
+        Atlas.BLOCKS.file(Identifier(BetterFoliage.MOD_ID, "blocks/leafmask_${size}_${type}"))
    }

    /**