package unit;

import java.util.ArrayList;
import java.util.HashSet;

/**
 * SpiderMonkeySimpleSetTest
 *
 * Models the CWE-407 defect in UnrollLoops.cpp SimpleSet<T,N,AP>:
 *
 *   Defective: mozilla::Vector<T,N,AP> backed set.
 *              contains() iterates all N elements linearly.
 *              add() calls contains() before appending.
 *              Called inside triple-nested loop (copies × blocks × values).
 *
 *   Fixed:     HashSet<T> backed set.
 *              contains() / add() are O(1) amortised.
 *
 * Models ValueSet (inline cap = 64) and its use in the unrolling phase.
 * "Value" is modelled as an Integer ID. Comparison counts are instrumented
 * explicitly — not wall-clock timing.
 *
 * Run: javac -d . SpiderMonkeySimpleSetTest.java && java -ea unit.SpiderMonkeySimpleSetTest
 */
public class SpiderMonkeySimpleSetTest {

    // -----------------------------------------------------------------------
    // Defective SimpleSet: ArrayList<Integer> with linear contains()
    // -----------------------------------------------------------------------

    static class DefectiveSimpleSet {
        private final ArrayList<Integer> vec = new ArrayList<>();
        long comparisons = 0;

        boolean contains(int t) {
            for (int existing : vec) {
                comparisons++;
                if (existing == t) return true;
            }
            return false;
        }

        boolean add(int t) {
            if (contains(t)) return true;
            vec.add(t);
            return true;
        }

        int size() { return vec.size(); }
    }

    // -----------------------------------------------------------------------
    // Fixed SimpleSet: HashSet<Integer> with O(1) contains()
    // -----------------------------------------------------------------------

    static class FixedSimpleSet {
        private final HashSet<Integer> set = new HashSet<>();
        long lookups = 0;

        boolean contains(int t) {
            lookups++;
            return set.contains(t);
        }

        boolean add(int t) {
            lookups++;
            set.add(t);
            return true;
        }

        int size() { return set.size(); }
    }

    // -----------------------------------------------------------------------
    // Simulate the unrolling inner loop:
    //   unrollFactor copies × numBlocks blocks × numSuccessors successors
    //   exitTargetBlocks.contains(succ) called per successor per block per copy
    //   exitingValues.contains(value) called per value per block per copy
    //
    // Parameters:
    //   unrollFactor   - number of loop body copies (typically 2-4)
    //   numBlocks      - basic blocks in the loop body
    //   numSuccessors  - successors per block's last instruction
    //   numExitTargets - size of exitTargetBlocks set (B, up to 8)
    //   numValues      - size of exitingValues set (V, up to 64)
    // -----------------------------------------------------------------------

    static long simulateUnrollDefective(int unrollFactor, int numBlocks,
                                         int numSuccessors, int numExitTargets,
                                         int numValues) {
        DefectiveSimpleSet exitTargetBlocks = new DefectiveSimpleSet();
        DefectiveSimpleSet exitingValues    = new DefectiveSimpleSet();

        // Populate sets
        for (int i = 0; i < numExitTargets; i++) exitTargetBlocks.add(i);
        for (int i = 0; i < numValues; i++)       exitingValues.add(i);

        long ops = 0;

        for (int cix = 0; cix < unrollFactor; cix++) {
            for (int bix = 0; bix < numBlocks; bix++) {
                // Successor scan — check if successor is an exit target
                for (int s = 0; s < numSuccessors; s++) {
                    int succ = s % (numExitTargets + 2); // most are not exits
                    exitTargetBlocks.contains(succ);
                }
                // Value scan — check if each value is an exiting value
                for (int v = 0; v < numValues; v++) {
                    exitingValues.contains(v);
                }
            }
        }

        ops = exitTargetBlocks.comparisons + exitingValues.comparisons;
        return ops;
    }

    static long simulateUnrollFixed(int unrollFactor, int numBlocks,
                                     int numSuccessors, int numExitTargets,
                                     int numValues) {
        FixedSimpleSet exitTargetBlocks = new FixedSimpleSet();
        FixedSimpleSet exitingValues    = new FixedSimpleSet();

        for (int i = 0; i < numExitTargets; i++) exitTargetBlocks.add(i);
        for (int i = 0; i < numValues; i++)       exitingValues.add(i);

        long ops = 0;

        for (int cix = 0; cix < unrollFactor; cix++) {
            for (int bix = 0; bix < numBlocks; bix++) {
                for (int s = 0; s < numSuccessors; s++) {
                    int succ = s % (numExitTargets + 2);
                    exitTargetBlocks.contains(succ);
                }
                for (int v = 0; v < numValues; v++) {
                    exitingValues.contains(v);
                }
            }
        }

        ops = exitTargetBlocks.lookups + exitingValues.lookups;
        return ops;
    }

    // -----------------------------------------------------------------------
    // Test 1 — correctness: defective and fixed sets agree on membership
    // -----------------------------------------------------------------------

    static void test1_correctness() {
        DefectiveSimpleSet slow = new DefectiveSimpleSet();
        FixedSimpleSet     fast = new FixedSimpleSet();

        int[] values = {10, 20, 30, 10, 40, 20};
        for (int v : values) {
            slow.add(v);
            fast.add(v);
        }

        // Both should have 4 distinct values
        assert slow.size() == 4 : "defective size expected 4, got " + slow.size();
        assert fast.size() == 4 : "fixed size expected 4, got " + fast.size();

        // Both should contain all inserted values
        for (int v : new int[]{10, 20, 30, 40}) {
            assert slow.contains(v) : "defective missing " + v;
            assert fast.contains(v) : "fixed missing " + v;
        }

        // Both should not contain absent values
        assert !slow.contains(99) : "defective should not contain 99";
        assert !fast.contains(99) : "fixed should not contain 99";

        System.out.printf("test1: size=%d (both agree), membership checks PASS%n",
                slow.size());
    }

    // -----------------------------------------------------------------------
    // Test 2 — ratio >= 5x at V=64, unrollFactor=4, numBlocks=20
    // -----------------------------------------------------------------------

    static void test2_ratioAtMaxValueSet() {
        int unrollFactor   = 4;
        int numBlocks      = 20;
        int numSuccessors  = 3;
        int numExitTargets = 8;   // BlockSet inline cap
        int numValues      = 64;  // ValueSet inline cap (MaxValuesForPeel)

        long defectOps = simulateUnrollDefective(unrollFactor, numBlocks,
                numSuccessors, numExitTargets, numValues);
        long fixedOps  = simulateUnrollFixed(unrollFactor, numBlocks,
                numSuccessors, numExitTargets, numValues);

        double ratio = (double) defectOps / Math.max(1, fixedOps);
        System.out.printf("test2: unroll=%d blocks=%d V=%d  defect=%d fixed=%d ratio=%.1fx%n",
                unrollFactor, numBlocks, numValues, defectOps, fixedOps, ratio);

        assert ratio >= 5.0 : "expected ratio >= 5x, got " + ratio;
    }

    // -----------------------------------------------------------------------
    // Test 3 — worst-case simulation: V=64, 4 copies, 50 blocks
    // -----------------------------------------------------------------------

    static void test3_worstCase() {
        int unrollFactor   = 4;
        int numBlocks      = 50;
        int numSuccessors  = 4;
        int numExitTargets = 8;
        int numValues      = 64;

        long defectOps = simulateUnrollDefective(unrollFactor, numBlocks,
                numSuccessors, numExitTargets, numValues);
        long fixedOps  = simulateUnrollFixed(unrollFactor, numBlocks,
                numSuccessors, numExitTargets, numValues);

        double ratio = (double) defectOps / Math.max(1, fixedOps);
        System.out.printf("test3: worst-case unroll=%d blocks=%d V=%d  defect=%d fixed=%d ratio=%.1fx%n",
                unrollFactor, numBlocks, numValues, defectOps, fixedOps, ratio);

        assert ratio >= 10.0 : "expected ratio >= 10x at worst case, got " + ratio;
    }

    // -----------------------------------------------------------------------
    // Test 4 — scaling: doubling V roughly doubles defect ops, fixed stays constant
    // -----------------------------------------------------------------------

    static void test4_linearScalingDefect() {
        int unrollFactor   = 4;
        int numBlocks      = 20;
        int numSuccessors  = 2;
        int numExitTargets = 4;

        long d32  = simulateUnrollDefective(unrollFactor, numBlocks,
                numSuccessors, numExitTargets, 32);
        long d64  = simulateUnrollDefective(unrollFactor, numBlocks,
                numSuccessors, numExitTargets, 64);
        long f32  = simulateUnrollFixed(unrollFactor, numBlocks,
                numSuccessors, numExitTargets, 32);
        long f64  = simulateUnrollFixed(unrollFactor, numBlocks,
                numSuccessors, numExitTargets, 64);

        double defectGrowth = (double) d64 / Math.max(1, d32);
        double fixedGrowth  = (double) f64 / Math.max(1, f32);

        System.out.printf("test4: V=32→64  defect %d→%d (%.2fx)  fixed %d→%d (%.2fx)%n",
                d32, d64, defectGrowth, f32, f64, fixedGrowth);

        // Defect ops should grow faster than fixed when V doubles
        assert defectGrowth > fixedGrowth
                : "defect should grow faster than fixed when V doubles";
        // Fixed should grow roughly linearly in V (more lookups, each O(1))
        assert fixedGrowth <= 2.5
                : "fixed growth when V doubles should be near 2x, got " + fixedGrowth;
    }

    // -----------------------------------------------------------------------
    // Test 5 — add() dedup correctness under repeated inserts
    // -----------------------------------------------------------------------

    static void test5_addDedup() {
        DefectiveSimpleSet slow = new DefectiveSimpleSet();
        FixedSimpleSet     fast = new FixedSimpleSet();

        // Insert values 0..9 three times each
        for (int round = 0; round < 3; round++) {
            for (int v = 0; v < 10; v++) {
                slow.add(v);
                fast.add(v);
            }
        }

        // Both sets should contain exactly 10 distinct values
        assert slow.size() == 10
                : "defective: expected 10 unique values, got " + slow.size();
        assert fast.size() == 10
                : "fixed: expected 10 unique values, got " + fast.size();

        System.out.printf("test5: add() dedup correct, size=%d (both)%n", slow.size());
    }

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

    public static void main(String[] args) {
        System.out.println("=== SpiderMonkeySimpleSetTest ===");
        System.out.println("Modelling CWE-407 sm-0003: SimpleSet::contains() linear scan vs HashSet O(1)");
        System.out.println("Location: js/src/jit/UnrollLoops.cpp SimpleSet<T,N,AP>");
        System.out.println();

        test1_correctness();
        System.out.println("  PASS test1_correctness");

        test2_ratioAtMaxValueSet();
        System.out.println("  PASS test2_ratioAtMaxValueSet");

        test3_worstCase();
        System.out.println("  PASS test3_worstCase");

        test4_linearScalingDefect();
        System.out.println("  PASS test4_linearScalingDefect");

        test5_addDedup();
        System.out.println("  PASS test5_addDedup");

        System.out.println();
        System.out.println("5/5 PASS");
    }
}