Add extra attribute item predicate

4 files changed, 285 insertions, 3 deletions

diff --git a/src/main/kotlin/moe/nea/firmament/features/texturepack/CustomModelOverrideParser.kt b/src/main/kotlin/moe/nea/firmament/features/texturepack/CustomModelOverrideParser.kt
index 2f5667b..458f8ff 100644
--- a/src/main/kotlin/moe/nea/firmament/features/texturepack/CustomModelOverrideParser.kt
+++ b/src/main/kotlin/moe/nea/firmament/features/texturepack/CustomModelOverrideParser.kt
@@ -44,6 +44,7 @@ object CustomModelOverrideParser {
         registerPredicateParser("any", OrPredicate.Parser)
         registerPredicateParser("not", NotPredicate.Parser)
         registerPredicateParser("item", ItemPredicate.Parser)
+        registerPredicateParser("extra_attributes", ExtraAttributesPredicate.Parser)
     }
 
     private val neverPredicate = listOf(
diff --git a/src/main/kotlin/moe/nea/firmament/features/texturepack/ExtraAttributesPredicate.kt b/src/main/kotlin/moe/nea/firmament/features/texturepack/ExtraAttributesPredicate.kt
new file mode 100644
index 0000000..4a726cc
--- /dev/null
+++ b/src/main/kotlin/moe/nea/firmament/features/texturepack/ExtraAttributesPredicate.kt
@@ -0,0 +1,273 @@
+/*
+ * SPDX-FileCopyrightText: 2024 Linnea Gräf <nea@nea.moe>
+ *
+ * SPDX-License-Identifier: GPL-3.0-or-later
+ */
+
+package moe.nea.firmament.features.texturepack
+
+import com.google.gson.JsonArray
+import com.google.gson.JsonElement
+import com.google.gson.JsonObject
+import com.google.gson.JsonPrimitive
+import net.minecraft.item.ItemStack
+import net.minecraft.nbt.NbtByte
+import net.minecraft.nbt.NbtCompound
+import net.minecraft.nbt.NbtDouble
+import net.minecraft.nbt.NbtElement
+import net.minecraft.nbt.NbtFloat
+import net.minecraft.nbt.NbtInt
+import net.minecraft.nbt.NbtList
+import net.minecraft.nbt.NbtLong
+import net.minecraft.nbt.NbtShort
+import net.minecraft.nbt.NbtString
+import moe.nea.firmament.util.extraAttributes
+
+fun interface NbtMatcher {
+    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
+        }
+
+        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?
+        }
+
+        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
+
+        }
+    }
+
+    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
+        }
+
+        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())
+        }
+
+        override fun toString(): String {
+            return "MatchNbtString($string)"
+        }
+    }
+}
+
+data class ExtraAttributesPredicate(
+    val path: NbtPrism,
+    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)
+        }
+    }
+
+    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
+    }
+}
diff --git a/src/main/kotlin/moe/nea/firmament/features/texturepack/OrPredicate.kt b/src/main/kotlin/moe/nea/firmament/features/texturepack/OrPredicate.kt
index c171367..cab918e 100644
--- a/src/main/kotlin/moe/nea/firmament/features/texturepack/OrPredicate.kt
+++ b/src/main/kotlin/moe/nea/firmament/features/texturepack/OrPredicate.kt
@@ -24,7 +24,7 @@ class OrPredicate(val children: Array<FirmamentModelPredicate>) : FirmamentModel
                         CustomModelOverrideParser.parsePredicates(it as JsonObject)
                     }
                     .toTypedArray()
-            return AndPredicate(children)
+            return OrPredicate(children)
         }
 
     }
diff --git a/src/main/kotlin/moe/nea/firmament/features/texturepack/StringMatcher.kt b/src/main/kotlin/moe/nea/firmament/features/texturepack/StringMatcher.kt
index 1c2675f..3fc3ab9 100644
--- a/src/main/kotlin/moe/nea/firmament/features/texturepack/StringMatcher.kt
+++ b/src/main/kotlin/moe/nea/firmament/features/texturepack/StringMatcher.kt
@@ -46,13 +46,21 @@ interface StringMatcher {
         override fun matches(string: String): Boolean {
             return expected == (if (stripColorCodes) string.removeColorCodes() else string)
         }
+
+        override fun toString(): String {
+            return "Equals($expected, stripColorCodes = $stripColorCodes)"
+        }
     }
 
-    class Pattern(patternWithColorCodes: String, val stripColorCodes: Boolean) : StringMatcher {
+    class Pattern(val patternWithColorCodes: String, val stripColorCodes: Boolean) : StringMatcher {
         private val regex: Predicate<String> = patternWithColorCodes.toPattern().asMatchPredicate()
         override fun matches(string: String): Boolean {
             return regex.test(if (stripColorCodes) string.removeColorCodes() else string)
         }
+
+        override fun toString(): String {
+            return "Pattern($patternWithColorCodes, stripColorCodes = $stripColorCodes)"
+        }
     }
 
     object Serializer : KSerializer<StringMatcher> {
@@ -73,7 +81,7 @@ interface StringMatcher {
     }
 
     companion object {
-        fun serialize(stringMatcher: StringMatcher):JsonElement {
+        fun serialize(stringMatcher: StringMatcher): JsonElement {
             TODO("Cannot serialize string matchers rn")
         }