package unit;

import java.util.*;

/**
 * Unit test for solc-0002: Assembly calculateMaxStackHeight RJUMP target CWE-407.
 *
 * Defect: calculateMaxStackHeight() (EOF CFG builder) resolves each RJUMP/CRJUMP
 * target via std::find(items.begin(), items.end(), item.tag()) — an O(N) scan
 * over all assembly items per jump. With J jumps over N items, total cost is
 * O(J × N).
 *
 * File:   libevmasm/Assembly.cpp:1077
 * Symbol: calculateMaxStackHeight — `std::find(items.begin(), items.end(), item.tag())`
 *
 * Fix:    Build an unordered_map<u256, size_t> (tagIndex) once before the
 *         worklist loop, mapping each Tag item's data value to its index.
 *         Each RJUMP resolution becomes O(1). Total cost O(J + N).
 *
 * Modeled here in Java:
 *   - AssemblyItem ≡ Item record (type + label)
 *   - tag lookup ≡ linear scan of items[] vs HashMap<Integer,Integer>
 *   - comparisons counted at each scan step
 *
 * Expected at J=N=100:
 *   defective ≈ J × N     = 10000
 *   fixed     ≈ J + N     = 200   (index build + lookups)
 *   ratio     > 10×
 */
public class SolcAssemblyRjumpTest {

    static final int TYPE_TAG          = 0;
    static final int TYPE_RJUMP        = 1;
    static final int TYPE_COND_RJUMP   = 2;
    static final int TYPE_OTHER        = 3;

    // ── Defective: std::find linear scan per jump ─────────────────────────────

    static class DefectiveStackCalc {
        final int[] itemTypes;   // TYPE_* per item
        final int[] itemLabels;  // label for Tag and Jump items (−1 if not applicable)
        long comparisons = 0;

        DefectiveStackCalc(int[] itemTypes, int[] itemLabels) {
            this.itemTypes  = itemTypes;
            this.itemLabels = itemLabels;
        }

        /** Returns index of the Tag item whose label equals targetLabel, −1 if not found. */
        int findTag(int targetLabel) {
            for (int i = 0; i < itemTypes.length; i++) {
                comparisons++;
                if (itemTypes[i] == TYPE_TAG && itemLabels[i] == targetLabel)
                    return i;
            }
            return -1;
        }

        /** Process all jumps, resolving targets via linear scan. */
        void resolveAllJumps() {
            for (int idx = 0; idx < itemTypes.length; idx++) {
                int type = itemTypes[idx];
                if (type == TYPE_RJUMP || type == TYPE_COND_RJUMP) {
                    int target = itemLabels[idx];
                    int pos = findTag(target);
                    assert pos >= 0 : "Tag not found for label " + target;
                }
            }
        }
    }

    // ── Fixed: pre-built HashMap index, O(1) lookup ───────────────────────────

    static class FixedStackCalc {
        final int[] itemTypes;
        final int[] itemLabels;
        final Map<Integer, Integer> tagIndex;
        long comparisons = 0;

        FixedStackCalc(int[] itemTypes, int[] itemLabels) {
            this.itemTypes  = itemTypes;
            this.itemLabels = itemLabels;
            // Build index once: O(N)
            tagIndex = new HashMap<>();
            for (int i = 0; i < itemTypes.length; i++) {
                if (itemTypes[i] == TYPE_TAG) {
                    comparisons++;     // count the index-build work
                    tagIndex.put(itemLabels[i], i);
                }
            }
        }

        void resolveAllJumps() {
            for (int idx = 0; idx < itemTypes.length; idx++) {
                int type = itemTypes[idx];
                if (type == TYPE_RJUMP || type == TYPE_COND_RJUMP) {
                    int target = itemLabels[idx];
                    comparisons++;    // O(1) hash lookup
                    Integer pos = tagIndex.get(target);
                    assert pos != null : "Tag not found for label " + target;
                }
            }
        }
    }

    // ── Graph builders ────────────────────────────────────────────────────────

    /**
     * Build an items array with N tags interleaved with J jumps.
     * Layout: [TAG_0, OTHER, OTHER, RJUMP→TAG_0, TAG_1, OTHER, RJUMP→TAG_1, …]
     * Ensures every jump target exists as a Tag item.
     */
    static int[][] buildItems(int N, int J) {
        // Each tag gets an index slot; jumps are interspersed.
        // Simple layout: slots 0..N-1 are Tags, slots N..N+J-1 are RJUMPs
        // targeting tag (slot % N).
        int total = N + J;
        int[] types  = new int[total];
        int[] labels = new int[total];

        for (int i = 0; i < N; i++) {
            types[i]  = TYPE_TAG;
            labels[i] = i;       // label == index for simplicity
        }
        for (int j = 0; j < J; j++) {
            types[N + j]  = TYPE_RJUMP;
            labels[N + j] = j % N;   // jump to tag j%N
        }
        return new int[][]{ types, labels };
    }

    // ── Simulation helpers ────────────────────────────────────────────────────

    public static long simulateDefective(int N, int J) {
        int[][] items = buildItems(N, J);
        DefectiveStackCalc calc = new DefectiveStackCalc(items[0], items[1]);
        calc.resolveAllJumps();
        return calc.comparisons;
    }

    public static long simulateFixed(int N, int J) {
        int[][] items = buildItems(N, J);
        FixedStackCalc calc = new FixedStackCalc(items[0], items[1]);
        calc.resolveAllJumps();
        return calc.comparisons;
    }

    // ── Tests ─────────────────────────────────────────────────────────────────

    static void testCorrectnessMatch() {
        // Both must resolve same targets without assertion failure
        int N = 10, J = 10;
        int[][] items = buildItems(N, J);
        DefectiveStackCalc def = new DefectiveStackCalc(items[0], items[1]);
        FixedStackCalc     fix = new FixedStackCalc(items[0], items[1]);

        // just verify they don't throw
        def.resolveAllJumps();
        fix.resolveAllJumps();
        System.out.println("PASS testCorrectnessMatch");
    }

    static void testDefectiveGrowsQuadratically() {
        long prev = -1;
        for (int S : new int[]{10, 20, 40}) {
            long c = simulateDefective(S, S);
            if (prev > 0) {
                double ratio = (double) c / prev;
                assert ratio > 3.0
                    : "defective should grow >3x when N=J doubled; got " + ratio + " at N=J=" + S;
            }
            prev = c;
        }
        System.out.println("PASS testDefectiveGrowsQuadratically");
    }

    static void testFixedGrowsLinearly() {
        long prev = -1;
        for (int S : new int[]{10, 20, 40}) {
            long c = simulateFixed(S, S);
            if (prev > 0) {
                double ratio = (double) c / prev;
                assert ratio < 2.5
                    : "fixed should grow ~2x when N=J doubled; got " + ratio + " at N=J=" + S;
            }
            prev = c;
        }
        System.out.println("PASS testFixedGrowsLinearly");
    }

    static void testRatioAtScale() {
        int N = 100, J = 100;
        long defComp = simulateDefective(N, J);
        long fixComp = simulateFixed(N, J);
        double ratio  = (double) defComp / fixComp;

        // defective: each of J=100 jumps scans up to N+J=200 items before finding tag.
        // Tags are at positions 0..99, jumps at 100..199; worst-case each jump scans
        // past all 200 items. In our layout tags are first so average scan is ~N/2.
        // Minimum expected: J * 1 = 100 comparisons. We assert > 10x fixed.
        assert ratio > 10.0
            : "ratio should be >10x at N=J=100; got " + ratio +
              " (defective=" + defComp + ", fixed=" + fixComp + ")";

        System.out.printf(
            "PASS testRatioAtScale  (defective=%d, fixed=%d, ratio=%.1fx)%n",
            defComp, fixComp, ratio);
    }

    public static void main(String[] args) {
        testCorrectnessMatch();
        testDefectiveGrowsQuadratically();
        testFixedGrowsLinearly();
        testRatioAtScale();
        System.out.println("All solc-0002 tests passed.");
    }
}