feat: Add generic component matcher

3 files changed, 296 insertions, 223 deletions

diff --git a/src/texturePacks/java/moe/nea/firmament/features/texturepack/CustomModelOverrideParser.kt b/src/texturePacks/java/moe/nea/firmament/features/texturepack/CustomModelOverrideParser.kt
index fe7951c..491a344 100644
--- a/src/texturePacks/java/moe/nea/firmament/features/texturepack/CustomModelOverrideParser.kt
+++ b/src/texturePacks/java/moe/nea/firmament/features/texturepack/CustomModelOverrideParser.kt
@@ -17,6 +17,7 @@ import moe.nea.firmament.features.texturepack.predicates.AndPredicate
 import moe.nea.firmament.features.texturepack.predicates.CastPredicate
 import moe.nea.firmament.features.texturepack.predicates.DisplayNamePredicate
 import moe.nea.firmament.features.texturepack.predicates.ExtraAttributesPredicate
+import moe.nea.firmament.features.texturepack.predicates.GenericComponentPredicate
 import moe.nea.firmament.features.texturepack.predicates.ItemPredicate
 import moe.nea.firmament.features.texturepack.predicates.LorePredicate
 import moe.nea.firmament.features.texturepack.predicates.NotPredicate
@@ -63,6 +64,7 @@ object CustomModelOverrideParser {
 		registerPredicateParser("item", ItemPredicate.Parser)
 		registerPredicateParser("extra_attributes", ExtraAttributesPredicate.Parser)
 		registerPredicateParser("pet", PetPredicate.Parser)
+		registerPredicateParser("component", GenericComponentPredicate.Parser)
 	}
 
 	private val neverPredicate = listOf(
diff --git a/src/texturePacks/java/moe/nea/firmament/features/texturepack/predicates/ExtraAttributesPredicate.kt b/src/texturePacks/java/moe/nea/firmament/features/texturepack/predicates/ExtraAttributesPredicate.kt
index 3c8023d..6fea1c3 100644
--- a/src/texturePacks/java/moe/nea/firmament/features/texturepack/predicates/ExtraAttributesPredicate.kt
+++ b/src/texturePacks/java/moe/nea/firmament/features/texturepack/predicates/ExtraAttributesPredicate.kt
@@ -22,194 +22,200 @@ import net.minecraft.nbt.NbtString
 import moe.nea.firmament.util.extraAttributes
 
 fun interface NbtMatcher {
-    fun matches(nbt: NbtElement): Boolean
+	fun matches(nbt: NbtElement): Boolean
 
-    object Parser {
-        fun parse(jsonElement: JsonElement): NbtMatcher? {
-            if (jsonElement is JsonPrimitive) {
-                if (jsonElement.isString) {
-                    val string = jsonElement.asString
-                    return MatchStringExact(string)
-                }
-                if (jsonElement.isNumber) {
-                    return MatchNumberExact(jsonElement.asLong) //TODO: parse generic number
-                }
-            }
-            if (jsonElement is JsonObject) {
-                var encounteredParser: NbtMatcher? = null
-                for (entry in ExclusiveParserType.entries) {
-                    val data = jsonElement[entry.key] ?: continue
-                    if (encounteredParser != null) {
-                        // TODO: warn
-                        return null
-                    }
-                    encounteredParser = entry.parse(data) ?: return null
-                }
-                return encounteredParser
-            }
-            return null
-        }
+	object Parser {
+		fun parse(jsonElement: JsonElement): NbtMatcher? {
+			if (jsonElement is JsonPrimitive) {
+				if (jsonElement.isString) {
+					val string = jsonElement.asString
+					return MatchStringExact(string)
+				}
+				if (jsonElement.isNumber) {
+					return MatchNumberExact(jsonElement.asLong) // TODO: parse generic number
+				}
+			}
+			if (jsonElement is JsonObject) {
+				var encounteredParser: NbtMatcher? = null
+				for (entry in ExclusiveParserType.entries) {
+					val data = jsonElement[entry.key] ?: continue
+					if (encounteredParser != null) {
+						// TODO: warn
+						return null
+					}
+					encounteredParser = entry.parse(data) ?: return null
+				}
+				return encounteredParser
+			}
+			return null
+		}
 
-        enum class ExclusiveParserType(val key: String) {
-            STRING("string") {
-                override fun parse(element: JsonElement): NbtMatcher? {
-                    return MatchString(StringMatcher.parse(element))
-                }
-            },
-            INT("int") {
-                override fun parse(element: JsonElement): NbtMatcher? {
-                    return parseGenericNumber(element,
-                                              { it.asInt },
-                                              { (it as? NbtInt)?.intValue() },
-                                              { a, b ->
-                                                  if (a == b) Comparison.EQUAL
-                                                  else if (a < b) Comparison.LESS_THAN
-                                                  else Comparison.GREATER
-                                              })
-                }
-            },
-            FLOAT("float") {
-                override fun parse(element: JsonElement): NbtMatcher? {
-                    return parseGenericNumber(element,
-                                              { it.asFloat },
-                                              { (it as? NbtFloat)?.floatValue() },
-                                              { a, b ->
-                                                  if (a == b) Comparison.EQUAL
-                                                  else if (a < b) Comparison.LESS_THAN
-                                                  else Comparison.GREATER
-                                              })
-                }
-            },
-            DOUBLE("double") {
-                override fun parse(element: JsonElement): NbtMatcher? {
-                    return parseGenericNumber(element,
-                                              { it.asDouble },
-                                              { (it as? NbtDouble)?.doubleValue() },
-                                              { a, b ->
-                                                  if (a == b) Comparison.EQUAL
-                                                  else if (a < b) Comparison.LESS_THAN
-                                                  else Comparison.GREATER
-                                              })
-                }
-            },
-            LONG("long") {
-                override fun parse(element: JsonElement): NbtMatcher? {
-                    return parseGenericNumber(element,
-                                              { it.asLong },
-                                              { (it as? NbtLong)?.longValue() },
-                                              { a, b ->
-                                                  if (a == b) Comparison.EQUAL
-                                                  else if (a < b) Comparison.LESS_THAN
-                                                  else Comparison.GREATER
-                                              })
-                }
-            },
-            SHORT("short") {
-                override fun parse(element: JsonElement): NbtMatcher? {
-                    return parseGenericNumber(element,
-                                              { it.asShort },
-                                              { (it as? NbtShort)?.shortValue() },
-                                              { a, b ->
-                                                  if (a == b) Comparison.EQUAL
-                                                  else if (a < b) Comparison.LESS_THAN
-                                                  else Comparison.GREATER
-                                              })
-                }
-            },
-            BYTE("byte") {
-                override fun parse(element: JsonElement): NbtMatcher? {
-                    return parseGenericNumber(element,
-                                              { it.asByte },
-                                              { (it as? NbtByte)?.byteValue() },
-                                              { a, b ->
-                                                  if (a == b) Comparison.EQUAL
-                                                  else if (a < b) Comparison.LESS_THAN
-                                                  else Comparison.GREATER
-                                              })
-                }
-            },
-            ;
+		enum class ExclusiveParserType(val key: String) {
+			STRING("string") {
+				override fun parse(element: JsonElement): NbtMatcher? {
+					return MatchString(StringMatcher.parse(element))
+				}
+			},
+			INT("int") {
+				override fun parse(element: JsonElement): NbtMatcher? {
+					return parseGenericNumber(
+						element,
+						{ it.asInt },
+						{ (it as? NbtInt)?.intValue() },
+						{ a, b ->
+							if (a == b) Comparison.EQUAL
+							else if (a < b) Comparison.LESS_THAN
+							else Comparison.GREATER
+						})
+				}
+			},
+			FLOAT("float") {
+				override fun parse(element: JsonElement): NbtMatcher? {
+					return parseGenericNumber(
+						element,
+						{ it.asFloat },
+						{ (it as? NbtFloat)?.floatValue() },
+						{ a, b ->
+							if (a == b) Comparison.EQUAL
+							else if (a < b) Comparison.LESS_THAN
+							else Comparison.GREATER
+						})
+				}
+			},
+			DOUBLE("double") {
+				override fun parse(element: JsonElement): NbtMatcher? {
+					return parseGenericNumber(
+						element,
+						{ it.asDouble },
+						{ (it as? NbtDouble)?.doubleValue() },
+						{ a, b ->
+							if (a == b) Comparison.EQUAL
+							else if (a < b) Comparison.LESS_THAN
+							else Comparison.GREATER
+						})
+				}
+			},
+			LONG("long") {
+				override fun parse(element: JsonElement): NbtMatcher? {
+					return parseGenericNumber(
+						element,
+						{ it.asLong },
+						{ (it as? NbtLong)?.longValue() },
+						{ a, b ->
+							if (a == b) Comparison.EQUAL
+							else if (a < b) Comparison.LESS_THAN
+							else Comparison.GREATER
+						})
+				}
+			},
+			SHORT("short") {
+				override fun parse(element: JsonElement): NbtMatcher? {
+					return parseGenericNumber(
+						element,
+						{ it.asShort },
+						{ (it as? NbtShort)?.shortValue() },
+						{ a, b ->
+							if (a == b) Comparison.EQUAL
+							else if (a < b) Comparison.LESS_THAN
+							else Comparison.GREATER
+						})
+				}
+			},
+			BYTE("byte") {
+				override fun parse(element: JsonElement): NbtMatcher? {
+					return parseGenericNumber(
+						element,
+						{ it.asByte },
+						{ (it as? NbtByte)?.byteValue() },
+						{ a, b ->
+							if (a == b) Comparison.EQUAL
+							else if (a < b) Comparison.LESS_THAN
+							else Comparison.GREATER
+						})
+				}
+			},
+			;
 
-            abstract fun parse(element: JsonElement): NbtMatcher?
-        }
+			abstract fun parse(element: JsonElement): NbtMatcher?
+		}
 
-        enum class Comparison {
-            LESS_THAN, EQUAL, GREATER
-        }
+		enum class Comparison {
+			LESS_THAN, EQUAL, GREATER
+		}
 
-        inline fun <T : Any> parseGenericNumber(
-            jsonElement: JsonElement,
-            primitiveExtractor: (JsonPrimitive) -> T?,
-            crossinline nbtExtractor: (NbtElement) -> T?,
-            crossinline compare: (T, T) -> Comparison
-        ): NbtMatcher? {
-            if (jsonElement is JsonPrimitive) {
-                val expected = primitiveExtractor(jsonElement) ?: return null
-                return NbtMatcher {
-                    val actual = nbtExtractor(it) ?: return@NbtMatcher false
-                    compare(actual, expected) == Comparison.EQUAL
-                }
-            }
-            if (jsonElement is JsonObject) {
-                val minElement = jsonElement.getAsJsonPrimitive("min")
-                val min = if (minElement != null) primitiveExtractor(minElement) ?: return null else null
-                val minExclusive = jsonElement.get("minExclusive")?.asBoolean ?: false
-                val maxElement = jsonElement.getAsJsonPrimitive("max")
-                val max = if (maxElement != null) primitiveExtractor(maxElement) ?: return null else null
-                val maxExclusive = jsonElement.get("maxExclusive")?.asBoolean ?: true
-                if (min == null && max == null) return null
-                return NbtMatcher {
-                    val actual = nbtExtractor(it) ?: return@NbtMatcher false
-                    if (max != null) {
-                        val comp = compare(actual, max)
-                        if (comp == Comparison.GREATER) return@NbtMatcher false
-                        if (comp == Comparison.EQUAL && maxExclusive) return@NbtMatcher false
-                    }
-                    if (min != null) {
-                        val comp = compare(actual, min)
-                        if (comp == Comparison.LESS_THAN) return@NbtMatcher false
-                        if (comp == Comparison.EQUAL && minExclusive) return@NbtMatcher false
-                    }
-                    return@NbtMatcher true
-                }
-            }
-            return null
+		inline fun <T : Any> parseGenericNumber(
+			jsonElement: JsonElement,
+			primitiveExtractor: (JsonPrimitive) -> T?,
+			crossinline nbtExtractor: (NbtElement) -> T?,
+			crossinline compare: (T, T) -> Comparison
+		): NbtMatcher? {
+			if (jsonElement is JsonPrimitive) {
+				val expected = primitiveExtractor(jsonElement) ?: return null
+				return NbtMatcher {
+					val actual = nbtExtractor(it) ?: return@NbtMatcher false
+					compare(actual, expected) == Comparison.EQUAL
+				}
+			}
+			if (jsonElement is JsonObject) {
+				val minElement = jsonElement.getAsJsonPrimitive("min")
+				val min = if (minElement != null) primitiveExtractor(minElement) ?: return null else null
+				val minExclusive = jsonElement.get("minExclusive")?.asBoolean ?: false
+				val maxElement = jsonElement.getAsJsonPrimitive("max")
+				val max = if (maxElement != null) primitiveExtractor(maxElement) ?: return null else null
+				val maxExclusive = jsonElement.get("maxExclusive")?.asBoolean ?: true
+				if (min == null && max == null) return null
+				return NbtMatcher {
+					val actual = nbtExtractor(it) ?: return@NbtMatcher false
+					if (max != null) {
+						val comp = compare(actual, max)
+						if (comp == Comparison.GREATER) return@NbtMatcher false
+						if (comp == Comparison.EQUAL && maxExclusive) return@NbtMatcher false
+					}
+					if (min != null) {
+						val comp = compare(actual, min)
+						if (comp == Comparison.LESS_THAN) return@NbtMatcher false
+						if (comp == Comparison.EQUAL && minExclusive) return@NbtMatcher false
+					}
+					return@NbtMatcher true
+				}
+			}
+			return null
 
-        }
-    }
+		}
+	}
 
-    class MatchNumberExact(val number: Long) : NbtMatcher {
-        override fun matches(nbt: NbtElement): Boolean {
-            return when (nbt) {
-                is NbtByte -> nbt.byteValue().toLong() == number
-                is NbtInt -> nbt.intValue().toLong() == number
-                is NbtShort -> nbt.shortValue().toLong() == number
-                is NbtLong -> nbt.longValue().toLong() == number
-                else -> false
-            }
-        }
+	class MatchNumberExact(val number: Long) : NbtMatcher {
+		override fun matches(nbt: NbtElement): Boolean {
+			return when (nbt) {
+				is NbtByte -> nbt.byteValue().toLong() == number
+				is NbtInt -> nbt.intValue().toLong() == number
+				is NbtShort -> nbt.shortValue().toLong() == number
+				is NbtLong -> nbt.longValue().toLong() == number
+				else -> false
+			}
+		}
 
-    }
+	}
 
-    class MatchStringExact(val string: String) : NbtMatcher {
-        override fun matches(nbt: NbtElement): Boolean {
-            return nbt is NbtString && nbt.asString() == string
-        }
+	class MatchStringExact(val string: String) : NbtMatcher {
+		override fun matches(nbt: NbtElement): Boolean {
+			return nbt.asString() == string
+		}
 
-        override fun toString(): String {
-            return "MatchNbtStringExactly($string)"
-        }
-    }
+		override fun toString(): String {
+			return "MatchNbtStringExactly($string)"
+		}
+	}
 
-    class MatchString(val string: StringMatcher) : NbtMatcher {
-        override fun matches(nbt: NbtElement): Boolean {
-            return nbt is NbtString && string.matches(nbt.asString())
-        }
+	class MatchString(val string: StringMatcher) : NbtMatcher {
+		override fun matches(nbt: NbtElement): Boolean {
+			return nbt.asString().let(string::matches)
+		}
 
-        override fun toString(): String {
-            return "MatchNbtString($string)"
-        }
-    }
+		override fun toString(): String {
+			return "MatchNbtString($string)"
+		}
+	}
 }
 
 data class ExtraAttributesPredicate(
@@ -217,55 +223,63 @@ data class ExtraAttributesPredicate(
 	val matcher: NbtMatcher,
 ) : FirmamentModelPredicate {
 
-    object Parser : FirmamentModelPredicateParser {
-        override fun parse(jsonElement: JsonElement): FirmamentModelPredicate? {
-            if (jsonElement !is JsonObject) return null
-            val path = jsonElement.get("path") ?: return null
-            val pathSegments = if (path is JsonArray) {
-                path.map { (it as JsonPrimitive).asString }
-            } else if (path is JsonPrimitive && path.isString) {
-                path.asString.split(".")
-            } else return null
-            val matcher = NbtMatcher.Parser.parse(jsonElement.get("match") ?: jsonElement)
-                ?: return null
-            return ExtraAttributesPredicate(NbtPrism(pathSegments), matcher)
-        }
-    }
+	object Parser : FirmamentModelPredicateParser {
+		override fun parse(jsonElement: JsonElement): FirmamentModelPredicate? {
+			if (jsonElement !is JsonObject) return null
+			val path = jsonElement.get("path") ?: return null
+			val prism = NbtPrism.fromElement(path) ?: return null
+			val matcher = NbtMatcher.Parser.parse(jsonElement.get("match") ?: jsonElement)
+				?: return null
+			return ExtraAttributesPredicate(prism, matcher)
+		}
+	}
 
-    override fun test(stack: ItemStack): Boolean {
-        return path.access(stack.extraAttributes)
-            .any { matcher.matches(it) }
-    }
+	override fun test(stack: ItemStack): Boolean {
+		return path.access(stack.extraAttributes)
+			.any { matcher.matches(it) }
+	}
 }
 
 class NbtPrism(val path: List<String>) {
-    override fun toString(): String {
-        return "Prism($path)"
-    }
-    fun access(root: NbtElement): Collection<NbtElement> {
-        var rootSet = mutableListOf(root)
-        var switch = mutableListOf<NbtElement>()
-        for (pathSegment in path) {
-            if (pathSegment == ".") continue
-            for (element in rootSet) {
-                if (element is NbtList) {
-                    if (pathSegment == "*")
-                        switch.addAll(element)
-                    val index = pathSegment.toIntOrNull() ?: continue
-                    if (index !in element.indices) continue
-                    switch.add(element[index])
-                }
-                if (element is NbtCompound) {
-                    if (pathSegment == "*")
-                        element.keys.mapTo(switch) { element.get(it)!! }
-                    switch.add(element.get(pathSegment) ?: continue)
-                }
-            }
-            val temp = switch
-            switch = rootSet
-            rootSet = temp
-            switch.clear()
-        }
-        return rootSet
-    }
+	companion object {
+		fun fromElement(path: JsonElement): NbtPrism? {
+			if (path is JsonArray) {
+				return NbtPrism(path.map { (it as JsonPrimitive).asString })
+			} else if (path is JsonPrimitive && path.isString) {
+				return NbtPrism(path.asString.split("."))
+			}
+			return null
+		}
+	}
+
+	override fun toString(): String {
+		return "Prism($path)"
+	}
+
+	fun access(root: NbtElement): Collection<NbtElement> {
+		var rootSet = mutableListOf(root)
+		var switch = mutableListOf<NbtElement>()
+		for (pathSegment in path) {
+			if (pathSegment == ".") continue
+			for (element in rootSet) {
+				if (element is NbtList) {
+					if (pathSegment == "*")
+						switch.addAll(element)
+					val index = pathSegment.toIntOrNull() ?: continue
+					if (index !in element.indices) continue
+					switch.add(element[index])
+				}
+				if (element is NbtCompound) {
+					if (pathSegment == "*")
+						element.keys.mapTo(switch) { element.get(it)!! }
+					switch.add(element.get(pathSegment) ?: continue)
+				}
+			}
+			val temp = switch
+			switch = rootSet
+			rootSet = temp
+			switch.clear()
+		}
+		return rootSet
+	}
 }
diff --git a/src/texturePacks/java/moe/nea/firmament/features/texturepack/predicates/GenericComponentPredicate.kt b/src/texturePacks/java/moe/nea/firmament/features/texturepack/predicates/GenericComponentPredicate.kt
new file mode 100644
index 0000000..a99c724
--- /dev/null
+++ b/src/texturePacks/java/moe/nea/firmament/features/texturepack/predicates/GenericComponentPredicate.kt
@@ -0,0 +1,57 @@
+package moe.nea.firmament.features.texturepack.predicates
+
+import com.google.gson.JsonElement
+import com.google.gson.JsonObject
+import com.mojang.serialization.Codec
+import kotlin.jvm.optionals.getOrNull
+import net.minecraft.component.ComponentType
+import net.minecraft.component.type.NbtComponent
+import net.minecraft.entity.LivingEntity
+import net.minecraft.item.ItemStack
+import net.minecraft.nbt.NbtOps
+import net.minecraft.registry.RegistryKey
+import net.minecraft.registry.RegistryKeys
+import net.minecraft.util.Identifier
+import moe.nea.firmament.features.texturepack.FirmamentModelPredicate
+import moe.nea.firmament.features.texturepack.FirmamentModelPredicateParser
+import moe.nea.firmament.util.MC
+
+data class GenericComponentPredicate<T>(
+	val componentType: ComponentType<T>,
+	val codec: Codec<T>,
+	val path: NbtPrism,
+	val matcher: NbtMatcher,
+) : FirmamentModelPredicate {
+	constructor(componentType: ComponentType<T>, path: NbtPrism, matcher: NbtMatcher)
+		: this(componentType, componentType.codecOrThrow, path, matcher)
+
+	override fun test(stack: ItemStack, holder: LivingEntity?): Boolean {
+		val component = stack.get(componentType) ?: return false
+		// TODO: cache this
+		val nbt =
+			if (component is NbtComponent) component.nbt
+			else codec.encodeStart(NbtOps.INSTANCE, component)
+				.resultOrPartial().getOrNull() ?: return false
+		return path.access(nbt).any { matcher.matches(it) }
+	}
+
+	object Parser : FirmamentModelPredicateParser {
+		override fun parse(jsonElement: JsonElement): GenericComponentPredicate<*>? {
+			if (jsonElement !is JsonObject) return null
+			val path = jsonElement.get("path") ?: return null
+			val prism = NbtPrism.fromElement(path) ?: return null
+			val matcher = NbtMatcher.Parser.parse(jsonElement.get("match") ?: jsonElement)
+				?: return null
+			val component = MC.currentOrDefaultRegistries
+				.getOrThrow(RegistryKeys.DATA_COMPONENT_TYPE)
+				.getOrThrow(
+					RegistryKey.of(
+						RegistryKeys.DATA_COMPONENT_TYPE,
+						Identifier.of(jsonElement.get("component").asString)
+					)
+				).value()
+			return GenericComponentPredicate(component, prism, matcher)
+		}
+	}
+
+}