java-topology/defects/elasticsearch/unit/MMRDiversificationContains.java

package unit;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * CWE-407 unit test: elasticsearch-001
 * MMRResultDiversification.java:63 — selectedDocRanks.contains() (ArrayList) inside
 * double loop — O(n²) membership test.
 *
 * Slow path: List.contains() — O(selected) per inner-loop iteration.
 * Fast path: HashSet.contains() — O(1) per inner-loop iteration.
 *
 * Compile: javac -d . MMRDiversificationContains.java
 * Run:     java -ea unit.MMRDiversificationContains
 */
public class MMRDiversificationContains {

    /**
     * Simulates the defective MMR selection loop.
     * Returns total number of list element comparisons performed.
     */
    static long slowMMRSelect(int[] docRanks, int topK) {
        List<Integer> selectedDocRanks = new ArrayList<>();
        long comparisons = 0;

        // seed with first doc
        selectedDocRanks.add(docRanks[0]);

        for (int x = 0; x < topK && selectedDocRanks.size() < topK && selectedDocRanks.size() < docRanks.length; x++) {
            int bestRank = -1;

            for (int docRank : docRanks) {
                // O(selectedDocRanks.size()) scan — the defect
                comparisons += selectedDocRanks.size();
                boolean alreadySelected = selectedDocRanks.contains(docRank);
                if (alreadySelected) {
                    continue;
                }
                bestRank = docRank; // simplified: just pick the last unselected
            }

            if (bestRank >= 0) {
                selectedDocRanks.add(bestRank);
            }
        }

        return comparisons;
    }

    /**
     * Simulates the fixed MMR selection loop using HashSet for O(1) membership.
     * Returns total number of hash lookups performed.
     */
    static long fastMMRSelect(int[] docRanks, int topK) {
        List<Integer> selectedDocRanks = new ArrayList<>();
        Set<Integer> selectedSet = new HashSet<>();
        long operations = 0;

        // seed with first doc
        selectedDocRanks.add(docRanks[0]);
        selectedSet.add(docRanks[0]);

        for (int x = 0; x < topK && selectedDocRanks.size() < topK && selectedDocRanks.size() < docRanks.length; x++) {
            int bestRank = -1;

            for (int docRank : docRanks) {
                operations++; // O(1) HashSet lookup
                if (selectedSet.contains(docRank)) {
                    continue;
                }
                bestRank = docRank;
            }

            if (bestRank >= 0) {
                selectedDocRanks.add(bestRank);
                selectedSet.add(bestRank);
            }
        }

        return operations;
    }

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

        // Test 1: small case — verify correctness
        {
            total++;
            int[] docs = {1, 2, 3, 4, 5};
            long slowCost = slowMMRSelect(docs, 3);
            long fastCost = fastMMRSelect(docs, 3);
            // slow must do more comparisons than fast for non-trivial input
            assert slowCost > 0 : "slow path should do comparisons";
            assert fastCost > 0 : "fast path should do operations";
            System.out.printf("Test 1 (small): slow=%d slow-ops, fast=%d hash-ops%n", slowCost, fastCost);
            passed++;
        }

        // Test 2: medium case — verify quadratic vs linear growth
        {
            total++;
            int n = 200;
            int[] docs = new int[n];
            for (int i = 0; i < n; i++) docs[i] = i;
            int topK = 50;

            long slowCost = slowMMRSelect(docs, topK);
            long fastCost = fastMMRSelect(docs, topK);

            // slow should be significantly more expensive than fast
            assert slowCost > fastCost * 5 :
                String.format("Expected slow >> fast, got slow=%d fast=%d", slowCost, fastCost);
            System.out.printf("Test 2 (n=%d, topK=%d): slow=%d, fast=%d, ratio=%.1fx%n",
                n, topK, slowCost, fastCost, (double) slowCost / fastCost);
            passed++;
        }

        // Test 3: large case — measure speedup at realistic window
        {
            total++;
            int n = 1000;
            int[] docs = new int[n];
            for (int i = 0; i < n; i++) docs[i] = i;
            int topK = 100;

            long slowCost = slowMMRSelect(docs, topK);
            long fastCost = fastMMRSelect(docs, topK);

            double ratio = (double) slowCost / fastCost;
            assert ratio > 20.0 :
                String.format("Expected >20x speedup at n=%d topK=%d, got %.1fx", n, topK, ratio);
            System.out.printf("Test 3 (n=%d, topK=%d): slow=%d, fast=%d, speedup=%.1fx%n",
                n, topK, slowCost, fastCost, ratio);
            passed++;
        }

        // Test 4: verify that both paths produce consistent selection behavior
        {
            total++;
            int[] docs = {10, 20, 30, 40, 50, 60};
            // Both paths should complete without error with full topK
            long s = slowMMRSelect(docs, 4);
            long f = fastMMRSelect(docs, 4);
            assert s > 0 && f > 0;
            System.out.printf("Test 4 (correctness): slow=%d fast=%d%n", s, f);
            passed++;
        }

        System.out.printf("%n%d/%d PASS%n", passed, total);
    }
}