tokenizer: index added tokens by first char and length

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Fixed
 
+- **Added-token matching no longer scans every added token at every position** (`skainet-frontend-text`, `AddedTokenMatcher` behind `SpecialTokenSplitter` and the inline matchers). Gemma's GGUF declares ~6,400 added tokens; a 70-character chat turn cost ~450k region compares per `encode`. Candidates are now indexed by first character and length (one map lookup per distinct length). Gemma-3 270M fine-tune vocab, JVM: 2.0 ms → 36 µs per chat-turn encode, ids unchanged; on a 32-bit ARM Android device the same encode took 36–166 ms.
 - **`CoroutineSchedule.forRange` could call its body with an empty, inverted range** (`skainet-backend-cpu`): the chunk length is rounded up, so for some `n` and task counts the last chunk started past `n` (63 elements on 12 tasks: chunks of 6, the twelfth at 66..63). The engine's own bodies loop over the range and did nothing for it; a body that passes the range length to a kernel saw a negative size. The task count now follows the chunk length, and every range is non-empty and covers `[0, n)` as the `Schedule` contract states.
 
 ### Changed

skainet-frontend/skainet-frontend-text/src/commonMain/kotlin/sk/ainet/frontend/tokenizer/TextPipeline.kt

@@ -37,9 +37,16 @@ public interface RunDecoder {
  */
 internal class AddedTokenMatcher(tokens: Collection<AddedToken>) {
     private val byContent: Map<String, AddedToken> = tokens.filter { it.content.isNotEmpty() }.associateBy { it.content }
-    private val byLengthDesc: List<String> = byContent.keys.sortedByDescending { it.length }
+    /**
+     * The distinct lengths of the added tokens per first character, longest first. Only a token starting with
+     * `text[from]` can match at `from`, and at most one per length, so a position costs one map lookup per
+     * length. Gemma's GGUF declares ~6,400 added tokens; comparing each of them at every position cost ~450k
+     * region compares for a 70-character chat turn.
+     */
+    private val lengthsByFirstChar: Map<Char, IntArray> =
+        byContent.keys.groupBy { it[0] }.mapValues { (_, c) -> c.map { it.length }.distinct().sortedDescending().toIntArray() }
 
-    val isEmpty: Boolean get() = byLengthDesc.isEmpty()
+    val isEmpty: Boolean get() = byContent.isEmpty()
     val ids: Set<Int> = byContent.values.map { it.id }.toSet()
     val specialIds: Set<Int> = byContent.values.filter { it.special }.map { it.id }.toSet()
     val idToContent: Map<Int, String> = byContent.values.associate { it.id to it.content }
@@ -64,12 +71,12 @@ internal class AddedTokenMatcher(tokens: Collection<AddedToken>) {
     }
 
     private fun matchAt(text: String, from: Int): AddedToken? {
-        for (content in byLengthDesc) {
-            if (!text.regionMatches(from, content, 0, content.length)) continue
-            val tok = byContent.getValue(content)
+        for (length in lengthsByFirstChar[text[from]] ?: return null) {
+            if (from + length > text.length) continue
+            val tok = byContent[text.substring(from, from + length)] ?: continue
             if (tok.singleWord) {
                 val before = if (from > 0) text[from - 1] else ' '
-                val afterIdx = from + content.length
+                val afterIdx = from + length
                 val after = if (afterIdx < text.length) text[afterIdx] else ' '
                 if (before.isLetterOrDigit() || after.isLetterOrDigit()) continue
             }

skainet-frontend/skainet-frontend-text/src/commonTest/kotlin/sk/ainet/frontend/tokenizer/SpecialTokenSplitterTest.kt

@@ -83,6 +83,21 @@ class SpecialTokenSplitterTest {
         assertEquals("a<|im_start|>b", splitter.decode(ids))
     }
 
+    @Test
+    fun `a rejected single_word match falls back to the next shorter token with the same first char`() {
+        val splitter = SpecialTokenSplitter(
+            CharBase,
+            listOf(
+                AddedToken(content = "word", id = 5, special = false, singleWord = true),
+                AddedToken(content = "wo", id = 6, special = false),
+                AddedToken(content = "wordy-long", id = 9, special = false),   // longer than what is left of the text
+            ),
+        )
+        // "word" inside "swordx" is not a whole word, so the shorter "wo" matches instead
+        assertContentEquals(intArrayOf('s'.code, 6, 'r'.code, 'd'.code, 'x'.code), splitter.encode("swordx"))
+        assertContentEquals(intArrayOf('a'.code, ' '.code, 5), splitter.encode("a word"))
+    }
+
     @Test
     fun `consecutive specials emit no spurious base segments`() {
         val specials = mapOf("<a>" to 1, "<b>" to 2)