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

package unit;
import java.util.*;

/**
 * TiDB CWE-407 unit tests — standalone, no JUnit.
 *
 *   tidb-0001  getEnforcedMergeJoin / getNewJoinKeys offsets slices.Contains
 *              pkg/planner/core/operator/physicalop/physical_merge_join.go:173,513,527
 *
 *   tidb-0002  mergeInAndNotEQLists removeValues slices.Contains
 *              pkg/planner/core/rule/rule_predicate_simplification.go:267
 */
public class TiDBTest {

    // -----------------------------------------------------------------------
    // tidb-0001: getNewJoinKeysByOffsets — slices.Contains vs map lookup
    //
    // Models: given N join keys and K offsets (already-placed keys), iterate
    // all N keys and skip those whose position is in offsets.
    // Slow: slices.Contains — O(K) per key, O(N*K) total.
    // Fast: map/set — O(1) per key, O(N) total.
    // -----------------------------------------------------------------------

    /** Slow: scan offsets slice for every key position. Returns op count. */
    static long mergeJoinSlowOps(int numJoinKeys, int numOffsets) {
        List<Integer> offsets = new ArrayList<>();
        Random rng = new Random(7);
        Set<Integer> added = new HashSet<>();
        while (offsets.size() < numOffsets) {
            int o = rng.nextInt(numJoinKeys);
            if (added.add(o)) offsets.add(o);
        }

        long ops = 0;
        // getNewJoinKeysByOffsets inner loop: for each pos, scan offsets
        for (int pos = 0; pos < numJoinKeys; pos++) {
            for (int i = 0; i < offsets.size(); i++) {
                ops++;
                if (offsets.get(i) == pos) break;
            }
        }
        return ops;
    }

    /** Fast: build map once, then O(1) lookup per key. Returns op count. */
    static long mergeJoinFastOps(int numJoinKeys, int numOffsets) {
        List<Integer> offsets = new ArrayList<>();
        Random rng = new Random(7);
        Set<Integer> added = new HashSet<>();
        while (offsets.size() < numOffsets) {
            int o = rng.nextInt(numJoinKeys);
            if (added.add(o)) offsets.add(o);
        }

        long ops = 0;
        // Build map: O(K)
        Set<Integer> offsetSet = new HashSet<>(offsets);
        ops += offsets.size(); // building cost

        // getNewJoinKeysByOffsets inner loop: O(1) per pos
        for (int pos = 0; pos < numJoinKeys; pos++) {
            ops++; // one map lookup
            offsetSet.contains(pos);
        }
        return ops;
    }

    // -----------------------------------------------------------------------
    // tidb-0002: mergeInAndNotEQLists — removeValues slice vs map
    //
    // Models the filter pass: given N predicates and R indices to remove,
    // iterate all N and check membership in the remove-set.
    // Slow: slices.Contains — O(R) per predicate, O(N*R) total.
    // Fast: map — O(1) per predicate, O(N) total.
    // -----------------------------------------------------------------------

    /** Slow: scan removeValues list for every predicate. Returns op count. */
    static long predicateSimplifySlowOps(int numPredicates, int numRemoved) {
        // Simulate a double-loop that marks some indices for removal
        List<Integer> removeValues = new ArrayList<>();
        // mark the first numRemoved predicates for removal (worst case: all removed come late)
        for (int i = numPredicates - numRemoved; i < numPredicates; i++) {
            removeValues.add(i);
        }

        long ops = 0;
        // Filter pass: for each predicate, scan removeValues
        for (int i = 0; i < numPredicates; i++) {
            for (int j = 0; j < removeValues.size(); j++) {
                ops++;
                if (removeValues.get(j) == i) break;
            }
        }
        return ops;
    }

    /** Fast: use a HashSet for O(1) removal check. Returns op count. */
    static long predicateSimplifyFastOps(int numPredicates, int numRemoved) {
        Set<Integer> removeSet = new HashSet<>();
        for (int i = numPredicates - numRemoved; i < numPredicates; i++) {
            removeSet.add(i);
        }

        long ops = 0;
        // Filter pass: O(1) per predicate
        for (int i = 0; i < numPredicates; i++) {
            ops++; // one set lookup
            removeSet.contains(i);
        }
        return ops;
    }

    // -----------------------------------------------------------------------
    // Test runner
    // -----------------------------------------------------------------------

    public static void main(String[] args) {
        int passed = 0;
        int failed = 0;

        // tidb-0001: join key counts
        int[][] joinSizes = {{20, 10}, {50, 25}, {100, 50}, {200, 100}};
        for (int[] sz : joinSizes) {
            int n = sz[0], k = sz[1];
            long slow = mergeJoinSlowOps(n, k);
            long fast = mergeJoinFastOps(n, k);
            // Slow should be significantly more: O(N*K) vs O(N+K)
            // At N=20,K=10: slow~100, fast~30 → ratio ~3x minimum
            boolean pass = slow > fast * 2;
            System.out.printf("tidb-0001 N=%-4d K=%-4d  slow=%6d  fast=%4d  ratio=%5.1fx  %s%n",
                n, k, slow, fast, (double) slow / fast, pass ? "PASS" : "FAIL");
            if (pass) passed++; else failed++;
        }

        // tidb-0002: predicate + removed counts
        int[][] predSizes = {{50, 25}, {100, 50}, {200, 100}, {500, 250}};
        for (int[] sz : predSizes) {
            int n = sz[0], r = sz[1];
            long slow = predicateSimplifySlowOps(n, r);
            long fast = predicateSimplifyFastOps(n, r);
            // Slow = O(N*R), fast = O(N). At N=50,R=25: slow~625, fast~50 → 12.5x
            boolean pass = slow > fast * 5;
            System.out.printf("tidb-0002 N=%-4d R=%-4d  slow=%6d  fast=%4d  ratio=%5.1fx  %s%n",
                n, r, slow, fast, (double) slow / fast, pass ? "PASS" : "FAIL");
            if (pass) passed++; else failed++;
        }

        System.out.printf("%nTotal: %d/%d PASS%n", passed, passed + failed);
        if (failed > 0) System.exit(1);
    }
}