java-topology/defects/tidb/unit/TidbTest.java

import java.util.*;

/**
 * Java simulation of TiDB CWE-407 defects.
 *
 * tidb-0001: mergeInAndNotEQLists – removeValues []int slice + slices.Contains O(P²)
 *   pkg/planner/core/rule/rule_predicate_simplification.go
 *
 * tidb-0002: ListPartitionGroup.intersect – findGroupIdx slices.Contains O(G²)
 *   pkg/table/tables/partition.go
 */
public class TidbTest {

    // ---------------------------------------------------------------
    // tidb-0001: predicate removeValues dedup
    // ---------------------------------------------------------------

    /** Unpatched: accumulate remove indices in a list, then filter with list.contains – O(P²) */
    static List<Integer> mergeFilterUnpatched(List<Integer> predicates) {
        List<Integer> removeValues = new ArrayList<>();
        for (int i = 0; i < predicates.size(); i++) {
            for (int j = i + 1; j < predicates.size(); j++) {
                // Simulate: if ith is NE predicate and jth is IN predicate
                if (predicates.get(i) < 0 && predicates.get(j) >= 0) {
                    removeValues.add(i);          // O(1) append
                }
            }
        }
        List<Integer> result = new ArrayList<>();
        for (int i = 0; i < predicates.size(); i++) {
            if (!removeValues.contains(i)) {      // O(R) linear scan — the defect
                result.add(predicates.get(i));
            }
        }
        return result;
    }

    /** Patched: use HashSet for O(1) lookup */
    static List<Integer> mergeFilterPatched(List<Integer> predicates) {
        Set<Integer> removeSet = new HashSet<>();
        for (int i = 0; i < predicates.size(); i++) {
            for (int j = i + 1; j < predicates.size(); j++) {
                if (predicates.get(i) < 0 && predicates.get(j) >= 0) {
                    removeSet.add(i);
                }
            }
        }
        List<Integer> result = new ArrayList<>();
        for (int i = 0; i < predicates.size(); i++) {
            if (!removeSet.contains(i)) {         // O(1) hash lookup — the fix
                result.add(predicates.get(i));
            }
        }
        return result;
    }

    // ---------------------------------------------------------------
    // tidb-0002: ListPartitionGroup.intersect
    // ---------------------------------------------------------------

    /** Unpatched: for each gidx in other, call slices.Contains(pg.GroupIdxs) – O(G²) */
    static List<Integer> intersectUnpatched(List<Integer> pgIdxs, List<Integer> otherIdxs) {
        List<Integer> result = new ArrayList<>();
        for (int gidx : otherIdxs) {
            if (pgIdxs.contains(gidx)) {          // O(G) linear scan — the defect
                result.add(gidx);
            }
        }
        return result;
    }

    /** Patched: build HashSet from pg.GroupIdxs first, then O(1) per lookup */
    static List<Integer> intersectPatched(List<Integer> pgIdxs, List<Integer> otherIdxs) {
        Set<Integer> existing = new HashSet<>(pgIdxs);
        List<Integer> result = new ArrayList<>();
        for (int gidx : otherIdxs) {
            if (existing.contains(gidx)) {        // O(1) — the fix
                result.add(gidx);
            }
        }
        return result;
    }

    // ---------------------------------------------------------------
    // Correctness assertions
    // ---------------------------------------------------------------

    static void assertEquals(Object a, Object b, String msg) {
        if (!a.equals(b)) throw new AssertionError(msg + ": expected " + a + " got " + b);
        System.out.println("PASS  " + msg);
    }

    // ---------------------------------------------------------------
    // Benchmark helpers
    // ---------------------------------------------------------------

    static long benchMergeUnpatched(int p) {
        List<Integer> predicates = new ArrayList<>();
        for (int i = 0; i < p; i++) {
            predicates.add(i % 3 == 0 ? -(i + 1) : i + 1);
        }
        long t0 = System.nanoTime();
        mergeFilterUnpatched(predicates);
        return System.nanoTime() - t0;
    }

    static long benchMergePatched(int p) {
        List<Integer> predicates = new ArrayList<>();
        for (int i = 0; i < p; i++) {
            predicates.add(i % 3 == 0 ? -(i + 1) : i + 1);
        }
        long t0 = System.nanoTime();
        mergeFilterPatched(predicates);
        return System.nanoTime() - t0;
    }

    static long benchIntersectUnpatched(int g) {
        List<Integer> pg = new ArrayList<>();
        List<Integer> other = new ArrayList<>();
        for (int i = 0; i < g; i++) { pg.add(i); other.add(g - 1 - i); }
        long t0 = System.nanoTime();
        intersectUnpatched(pg, other);
        return System.nanoTime() - t0;
    }

    static long benchIntersectPatched(int g) {
        List<Integer> pg = new ArrayList<>();
        List<Integer> other = new ArrayList<>();
        for (int i = 0; i < g; i++) { pg.add(i); other.add(g - 1 - i); }
        long t0 = System.nanoTime();
        intersectPatched(pg, other);
        return System.nanoTime() - t0;
    }

    // ---------------------------------------------------------------
    // Main
    // ---------------------------------------------------------------

    public static void main(String[] args) {
        System.out.println("=== tidb-0001: mergeInAndNotEQLists removeValues ===");

        // Correctness
        List<Integer> preds = Arrays.asList(-1, 2, -3, 4, -5, 6);
        List<Integer> r1 = mergeFilterUnpatched(preds);
        List<Integer> r2 = mergeFilterPatched(preds);
        assertEquals(r1, r2, "tidb-0001 correctness (unpatched==patched output)");

        // Warmup
        for (int i = 0; i < 3; i++) { benchMergeUnpatched(200); benchMergePatched(200); }

        // Benchmark P=500 predicates
        int P = 500;
        long u1 = 0, p1 = 0;
        int rounds = 5;
        for (int i = 0; i < rounds; i++) { u1 += benchMergeUnpatched(P); p1 += benchMergePatched(P); }
        u1 /= rounds; p1 /= rounds;
        double ratio1 = (double) u1 / Math.max(p1, 1);
        System.out.printf("  P=%d  unpatched=%,d ns  patched=%,d ns  ratio=%.1fx%n", P, u1, p1, ratio1);
        if (ratio1 < 2.0) System.out.println("  WARN: ratio below 2x (small N may not show O(N²) effect)");
        System.out.println("PASS  tidb-0001 benchmark");

        System.out.println();
        System.out.println("=== tidb-0002: ListPartitionGroup.intersect ===");

        // Correctness
        List<Integer> pg = Arrays.asList(0, 1, 2, 3, 4);
        List<Integer> other = Arrays.asList(2, 3, 5, 6);
        List<Integer> r3 = intersectUnpatched(pg, other);
        List<Integer> r4 = intersectPatched(pg, other);
        assertEquals(r3, r4, "tidb-0002 correctness (unpatched==patched output)");

        // Warmup
        for (int i = 0; i < 3; i++) { benchIntersectUnpatched(200); benchIntersectPatched(200); }

        // Benchmark G=1000 group indices
        int G = 1000;
        long u2 = 0, p2 = 0;
        for (int i = 0; i < rounds; i++) { u2 += benchIntersectUnpatched(G); p2 += benchIntersectPatched(G); }
        u2 /= rounds; p2 /= rounds;
        double ratio2 = (double) u2 / Math.max(p2, 1);
        System.out.printf("  G=%d  unpatched=%,d ns  patched=%,d ns  ratio=%.1fx%n", G, u2, p2, ratio2);
        if (ratio2 < 2.0) System.out.println("  WARN: ratio below 2x");
        System.out.println("PASS  tidb-0002 benchmark");

        System.out.println();
        System.out.println("ALL PASS");
    }
}