java-topology/defects/llamacpp-0001/test/LlamacppGrammarStacksDedupTest.java

import java.util.*;

/**
 * Unit test for llamacpp-0001: llama_grammar_advance_stack / llama_grammar_accept_token
 * new_stacks / stacks_new dedup via std::find on vector<vector<ptr>> — O(S^2) per token.
 *
 * In llama.cpp's grammar-constrained sampling (used for JSON schema output, regex grammars,
 * structured generation) each sampled token triggers:
 *   1. llama_grammar_advance_stack — appends terminal-stacks to new_stacks, deduplicating
 *      with std::find: O(|new_stacks|) per candidate.  Called S times (once per stack in
 *      grammar.stacks), so total dedup cost = O(S^2) per token.
 *   2. llama_grammar_accept_token — surviving_stack dedup with std::find on stacks_new:
 *      another O(S^2) term.
 *
 * For a complex JSON grammar with S=150 grammar stacks and G=1000 output tokens the defect
 * costs ~22.5M list-scan steps per generation; each scan compares stack vectors.
 *
 * Fix: maintain a companion HashSet<List<Integer>> alongside new_stacks / stacks_new.
 * HashSet.add() gives O(1) amortized dedup, reducing total cost to O(G * S).
 *
 * This test models the dedup pattern in Java using ArrayList vs HashSet:
 *   - Defect: List.contains() inside an append-loop  → O(S^2)
 *   - Fix:    HashSet.add() tracks seen stacks        → O(S) amortized
 */
public class LlamacppGrammarStacksDedupTest {

    // A "stack" is modelled as a list of integers (pointer addresses in C++ become int IDs here).

    // --- Defect: linear scan dedup (std::find equivalent) ---
    static List<List<Integer>> dedupLinear(List<List<Integer>> candidates) {
        List<List<Integer>> result = new ArrayList<>();
        for (List<Integer> stack : candidates) {
            if (!result.contains(stack)) {   // O(|result|) — the defect
                result.add(stack);
            }
        }
        return result;
    }

    // --- Fix: hash set dedup (std::set / std::unordered_set insert equivalent) ---
    static List<List<Integer>> dedupHash(List<List<Integer>> candidates) {
        Set<List<Integer>> seen = new HashSet<>();
        List<List<Integer>> result = new ArrayList<>();
        for (List<Integer> stack : candidates) {
            if (seen.add(stack)) {           // O(1) amortized — the fix
                result.add(stack);
            }
        }
        return result;
    }

    /**
     * Build S candidate stacks with depth elements each.
     * Half are unique, half are duplicates of existing stacks (simulating grammar branches
     * that converge to the same continuation after rule expansion).
     */
    static List<List<Integer>> buildCandidates(int S, int depth) {
        List<List<Integer>> candidates = new ArrayList<>(S * 2);
        // S unique stacks
        for (int i = 0; i < S; i++) {
            List<Integer> stack = new ArrayList<>(depth);
            for (int d = 0; d < depth; d++) {
                stack.add(i * 100 + d);
            }
            candidates.add(stack);
        }
        // S duplicate stacks (mirrors of unique stacks — forces the dedup to scan all existing)
        for (int i = 0; i < S; i++) {
            candidates.add(new ArrayList<>(candidates.get(i)));
        }
        return candidates;
    }

    public static void main(String[] args) {
        // --- Correctness check ---
        int smallS = 12;
        List<List<Integer>> smallCandidates = buildCandidates(smallS, 3);
        List<List<Integer>> linearResult = dedupLinear(smallCandidates);
        List<List<Integer>> hashResult   = dedupHash(smallCandidates);
        Set<List<Integer>> linearSet = new HashSet<>(linearResult);
        Set<List<Integer>> hashSet   = new HashSet<>(hashResult);
        assert linearSet.equals(hashSet) : "Correctness failed: results differ at S=" + smallS;
        assert linearResult.size() == smallS : "Expected " + smallS + " unique stacks, got " + linearResult.size();
        System.out.println("Correctness OK: S=" + smallS + " unique=" + linearResult.size());

        // --- Performance benchmark ---
        // S=300 simulates a complex JSON grammar (object with many optional fields / deep nesting).
        // reps=800 simulates 800 accepted tokens during a structured-output generation.
        // At S=300: defect costs S^2=90,000 list comparisons per token (each comparing S-deep lists).
        int S = 300;
        int depth = 5;
        int reps = 800;
        List<List<Integer>> candidates = buildCandidates(S, depth);

        // Correctness at full scale
        List<List<Integer>> lR = dedupLinear(candidates);
        List<List<Integer>> hR = dedupHash(candidates);
        assert new HashSet<>(lR).equals(new HashSet<>(hR)) : "Correctness failed at S=" + S;

        // Warmup
        for (int w = 0; w < 5; w++) {
            dedupLinear(candidates);
            dedupHash(candidates);
        }

        long t0 = System.nanoTime();
        for (int r = 0; r < reps; r++) dedupLinear(candidates);
        long linearNs = System.nanoTime() - t0;

        long t1 = System.nanoTime();
        for (int r = 0; r < reps; r++) dedupHash(candidates);
        long hashNs = System.nanoTime() - t1;

        double ratio = (double) linearNs / hashNs;
        System.out.printf("Grammar stack dedup: S=%d (candidates=%d) depth=%d reps=%d%n",
                S, candidates.size(), depth, reps);
        System.out.printf("  Defect (List.contains O(S^2)): %,d ns%n", linearNs);
        System.out.printf("  Fix    (HashSet.add   O(S)):    %,d ns%n", hashNs);
        System.out.printf("  speedup: %.1fx%n", ratio);

        assert ratio > 3.0 : "Expected >3x speedup, got " + ratio;
        System.out.println("PASS");
    }
}