package unit;

import java.util.*;

/**
 * BevyTest — bevy-0001
 *
 * Proves CWE-407 in Bevy's slab_allocator.rs:
 *   bevy-0001: free_empty_slabs() — Vec::iter().position() inside nested layout scan
 *              O(E × L × S) vs O(E) with reverse slab_id→layout HashMap
 *
 * Run: javac -d . BevyTest.java && java -ea unit.BevyTest
 */
public class BevyTest {

    static void bench(String label, Runnable slow, Runnable fast, long sOps, long fOps) {
        slow.run(); fast.run(); // warmup
        long t0 = System.nanoTime(); slow.run(); long sMs = (System.nanoTime() - t0) / 1_000_000;
        long t1 = System.nanoTime(); fast.run(); long fMs = (System.nanoTime() - t1) / 1_000_000;
        double r = fOps > 0 ? (double) sOps / fOps : 0;
        System.out.printf("  %-52s  slow:%4dms (%,d ops)  fast:%4dms (%,d ops)  speedup:%.0fx%n",
                label, sMs, sOps, fMs, fOps, r);
    }

    // ── bevy-0001: slab_allocator free_empty_slabs ───────────────────────────

    /**
     * SLOW: simulates free_empty_slabs() — for each freed slab, scan ALL layout
     * buckets (Vec<SlabId>) with linear search to find which bucket contains it.
     * O(E × L × S): E freed slabs, L layouts, S slabs per layout.
     */
    static long freeSlabsSlow(int layoutCount, int slabsPerLayout, int freeCount) {
        // slab_layouts: Map<layout, List<slabId>>
        // slab IDs are integers 0..layoutCount*slabsPerLayout
        List<List<Integer>> layouts = new ArrayList<>(layoutCount);
        List<Integer> allSlabs = new ArrayList<>();
        int slabId = 0;
        for (int l = 0; l < layoutCount; l++) {
            List<Integer> bucket = new ArrayList<>(slabsPerLayout);
            for (int s = 0; s < slabsPerLayout; s++) {
                bucket.add(slabId);
                allSlabs.add(slabId);
                slabId++;
            }
            layouts.add(bucket);
        }

        // Simulate freeing the last `freeCount` slabs
        long ops = 0;
        for (int f = allSlabs.size() - 1; f >= allSlabs.size() - freeCount; f--) {
            int emptySlab = allSlabs.get(f);
            // Scan all layout buckets — O(L * S)
            for (List<Integer> bucket : layouts) {
                for (int i = 0; i < bucket.size(); i++) {
                    ops++;
                    if (bucket.get(i).equals(emptySlab)) {
                        bucket.remove(i);
                        break;
                    }
                }
            }
        }
        return ops;
    }

    /**
     * FAST: simulates fixed free_empty_slabs() — maintain reverse map slabId→layoutIndex
     * for O(1) layout lookup. Only scans the one bucket that actually contains the slab.
     * O(E): one map lookup + one swap-remove per freed slab.
     */
    static long freeSlabsFast(int layoutCount, int slabsPerLayout, int freeCount) {
        List<List<Integer>> layouts = new ArrayList<>(layoutCount);
        Map<Integer, Integer> slabToLayout = new HashMap<>(); // reverse map: slabId → layoutIndex
        List<Integer> allSlabs = new ArrayList<>();
        int slabId = 0;
        for (int l = 0; l < layoutCount; l++) {
            List<Integer> bucket = new ArrayList<>(slabsPerLayout);
            for (int s = 0; s < slabsPerLayout; s++) {
                bucket.add(slabId);
                slabToLayout.put(slabId, l); // maintain reverse map on insertion
                allSlabs.add(slabId);
                slabId++;
            }
            layouts.add(bucket);
        }

        long ops = 0;
        for (int f = allSlabs.size() - 1; f >= allSlabs.size() - freeCount; f--) {
            int emptySlab = allSlabs.get(f);
            ops++;
            Integer layoutIdx = slabToLayout.remove(emptySlab); // O(1) reverse lookup
            if (layoutIdx != null) {
                List<Integer> bucket = layouts.get(layoutIdx);
                // swap-remove: O(1)
                for (int i = 0; i < bucket.size(); i++) {
                    ops++;
                    if (bucket.get(i).equals(emptySlab)) {
                        bucket.set(i, bucket.get(bucket.size() - 1));
                        bucket.remove(bucket.size() - 1);
                        break;
                    }
                }
            }
        }
        return ops;
    }

    public static void main(String[] args) {
        System.out.println("=== UNIT bevy-0001: Bevy CWE-407 — slab_allocator free_empty_slabs ===");

        // Small: 50 layouts, 20 slabs each, free 50 slabs
        final int L1 = 50, S1 = 20, E1 = 50;
        long slow1 = freeSlabsSlow(L1, S1, E1);
        long fast1 = freeSlabsFast(L1, S1, E1);
        bench(String.format("bevy-0001 free_empty_slabs L=%d S=%d E=%d", L1, S1, E1),
                () -> freeSlabsSlow(L1, S1, E1),
                () -> freeSlabsFast(L1, S1, E1),
                slow1, fast1);

        // Medium: 100 layouts, 50 slabs each, free 200 slabs
        final int L2 = 100, S2 = 50, E2 = 200;
        long slow2 = freeSlabsSlow(L2, S2, E2);
        long fast2 = freeSlabsFast(L2, S2, E2);
        bench(String.format("bevy-0001 free_empty_slabs L=%d S=%d E=%d", L2, S2, E2),
                () -> freeSlabsSlow(L2, S2, E2),
                () -> freeSlabsFast(L2, S2, E2),
                slow2, fast2);

        // Large: 200 layouts, 100 slabs each, free 500 slabs
        final int L3 = 200, S3 = 100, E3 = 500;
        long slow3 = freeSlabsSlow(L3, S3, E3);
        long fast3 = freeSlabsFast(L3, S3, E3);
        bench(String.format("bevy-0001 free_empty_slabs L=%d S=%d E=%d", L3, S3, E3),
                () -> freeSlabsSlow(L3, S3, E3),
                () -> freeSlabsFast(L3, S3, E3),
                slow3, fast3);

        System.out.println();

        // Assertions: slow must generate significantly more operations than fast
        int pass = 0;
        assert slow1 > fast1 * 5 :
                "bevy-0001 (small) expected slow ops > 5× fast ops, got slow=" + slow1 + " fast=" + fast1;
        pass++;
        assert slow2 > fast2 * 5 :
                "bevy-0001 (medium) expected slow ops > 5× fast ops, got slow=" + slow2 + " fast=" + fast2;
        pass++;
        assert slow3 > fast3 * 5 :
                "bevy-0001 (large) expected slow ops > 5× fast ops, got slow=" + slow3 + " fast=" + fast3;
        pass++;

        System.out.printf("%d/3 PASS%n", pass);
    }
}