package unit;

import java.util.*;
import java.util.stream.*;

/**
 * Standalone unit test for trino-0001:
 * PushDownDereferenceThroughJoin — List.contains() in stream filter → O(N²).
 *
 * Simulates the two hot loops in PushDownDereferenceThroughJoin.apply():
 *   Loop 1: forEach over dereferenceAssignments, calling outputSymbols.contains()
 *   Loop 2: stream.filter(s -> outputSymbols.contains(s))
 *
 * Compile: javac -d . PushDownDerefJoinAlgorithm.java
 * Run:     java unit.PushDownDerefJoinAlgorithm
 */
public class PushDownDerefJoinAlgorithm {

    // ── Node ─────────────────────────────────────────────────────────────────

    static class Symbol {
        final String name;
        Symbol(String name) { this.name = name; }

        @Override public boolean equals(Object o) {
            return o instanceof Symbol && name.equals(((Symbol) o).name);
        }
        @Override public int hashCode() { return name.hashCode(); }
        @Override public String toString() { return name; }
    }

    // ── Defective: uses List.contains() O(N) per look-up ────────────────────

    static class DefectivePushDown {
        List<Symbol> filterSymbols(
                List<Symbol> referred,
                List<Symbol> nodeOutputSymbols) {       // List → O(N) .contains
            // Mirrors: referredSymbols.stream().filter(s -> nodeOutput.contains(s))
            return referred.stream()
                    .filter(s -> nodeOutputSymbols.contains(s))   // CWE-407 site
                    .collect(Collectors.toList());
        }

        Map<Symbol, Symbol> classifyDerefs(
                List<Symbol> derefs,
                List<Symbol> leftOutput,                // List → O(N) .contains
                List<Symbol> rightOutput) {             // List → O(N) .contains
            Map<Symbol, Symbol> result = new LinkedHashMap<>();
            for (Symbol d : derefs) {
                if (leftOutput.contains(d)) {            // CWE-407 site
                    result.put(d, new Symbol("left:" + d.name));
                } else if (rightOutput.contains(d)) {    // CWE-407 site
                    result.put(d, new Symbol("right:" + d.name));
                }
            }
            return result;
        }
    }

    // ── Fixed: converts to Set once, then O(1) per look-up ──────────────────

    static class FixedPushDown {
        List<Symbol> filterSymbols(
                List<Symbol> referred,
                List<Symbol> nodeOutputSymbols) {
            Set<Symbol> outputSet = new HashSet<>(nodeOutputSymbols);  // O(N) once
            return referred.stream()
                    .filter(outputSet::contains)                        // O(1) each
                    .collect(Collectors.toList());
        }

        Map<Symbol, Symbol> classifyDerefs(
                List<Symbol> derefs,
                List<Symbol> leftOutput,
                List<Symbol> rightOutput) {
            Set<Symbol> leftSet  = new HashSet<>(leftOutput);   // O(N) once
            Set<Symbol> rightSet = new HashSet<>(rightOutput);  // O(N) once
            Map<Symbol, Symbol> result = new LinkedHashMap<>();
            for (Symbol d : derefs) {
                if (leftSet.contains(d)) {            // O(1)
                    result.put(d, new Symbol("left:" + d.name));
                } else if (rightSet.contains(d)) {    // O(1)
                    result.put(d, new Symbol("right:" + d.name));
                }
            }
            return result;
        }
    }

    // ── Result container ─────────────────────────────────────────────────────

    static class Result {
        final String label;
        final List<Symbol> filteredLeft;
        final List<Symbol> filteredRight;
        final Map<Symbol, Symbol> classified;
        final long ns;

        Result(String label, List<Symbol> fl, List<Symbol> fr, Map<Symbol, Symbol> cl, long ns) {
            this.label = label; this.filteredLeft = fl; this.filteredRight = fr;
            this.classified = cl; this.ns = ns;
        }
    }

    // ── Test driver ──────────────────────────────────────────────────────────

    static List<Symbol> symbols(int n, String prefix) {
        List<Symbol> out = new ArrayList<>(n);
        for (int i = 0; i < n; i++) out.add(new Symbol(prefix + i));
        return out;
    }

    static Result runDefective(int N, int D) {
        List<Symbol> leftOutput  = symbols(N, "lout");
        List<Symbol> rightOutput = symbols(N, "rout");
        // derefs: first half from left, second from right
        List<Symbol> derefs = new ArrayList<>();
        for (int i = 0; i < D / 2; i++) derefs.add(leftOutput.get(i % N));
        for (int i = 0; i < D / 2; i++) derefs.add(rightOutput.get(i % N));

        List<Symbol> referred = new ArrayList<>();
        referred.addAll(leftOutput.subList(0, Math.min(N / 2, N)));
        referred.addAll(rightOutput.subList(0, Math.min(N / 2, N)));

        DefectivePushDown alg = new DefectivePushDown();
        long t0 = System.nanoTime();
        Map<Symbol, Symbol> classified = alg.classifyDerefs(derefs, leftOutput, rightOutput);
        List<Symbol> fl = alg.filterSymbols(referred, leftOutput);
        List<Symbol> fr = alg.filterSymbols(referred, rightOutput);
        long ns = System.nanoTime() - t0;
        return new Result("DEFECTIVE", fl, fr, classified, ns);
    }

    static Result runFixed(int N, int D) {
        List<Symbol> leftOutput  = symbols(N, "lout");
        List<Symbol> rightOutput = symbols(N, "rout");
        List<Symbol> derefs = new ArrayList<>();
        for (int i = 0; i < D / 2; i++) derefs.add(leftOutput.get(i % N));
        for (int i = 0; i < D / 2; i++) derefs.add(rightOutput.get(i % N));

        List<Symbol> referred = new ArrayList<>();
        referred.addAll(leftOutput.subList(0, Math.min(N / 2, N)));
        referred.addAll(rightOutput.subList(0, Math.min(N / 2, N)));

        FixedPushDown alg = new FixedPushDown();
        long t0 = System.nanoTime();
        Map<Symbol, Symbol> classified = alg.classifyDerefs(derefs, leftOutput, rightOutput);
        List<Symbol> fl = alg.filterSymbols(referred, leftOutput);
        List<Symbol> fr = alg.filterSymbols(referred, rightOutput);
        long ns = System.nanoTime() - t0;
        return new Result("FIXED", fl, fr, classified, ns);
    }

    // ── Assertions ───────────────────────────────────────────────────────────

    static void assertEquals(Object expected, Object actual, String msg) {
        if (!expected.equals(actual)) {
            throw new AssertionError(msg + ": expected=" + expected + " actual=" + actual);
        }
    }

    static void assertTrue(boolean cond, String msg) {
        if (!cond) throw new AssertionError(msg);
    }

    public static void main(String[] args) {
        System.out.println("=== trino-0001: PushDownDerefJoinAlgorithm ===");

        // Small correctness test
        {
            int N = 20, D = 10;
            Result d = runDefective(N, D);
            Result f = runFixed(N, D);

            // Both should classify the same derefs
            assertEquals(d.classified.size(), f.classified.size(), "classified size");
            assertEquals(d.filteredLeft.size(), f.filteredLeft.size(), "filteredLeft size");
            assertEquals(d.filteredRight.size(), f.filteredRight.size(), "filteredRight size");
            System.out.println("PASS  correctness (N=20, D=10)");
        }

        // Benchmark at N=500, D=200
        {
            int N = 500, D = 200;
            // warm up
            for (int i = 0; i < 5; i++) { runDefective(N, D); runFixed(N, D); }

            long defNs = 0, fixNs = 0;
            int reps = 20;
            for (int i = 0; i < reps; i++) {
                defNs += runDefective(N, D).ns;
                fixNs += runFixed(N, D).ns;
            }
            defNs /= reps; fixNs /= reps;
            double ratio = (double) defNs / fixNs;

            System.out.printf("BENCH N=500 D=200 reps=%d%n", reps);
            System.out.printf("  defective avg: %,d ns%n", defNs);
            System.out.printf("  fixed     avg: %,d ns%n", fixNs);
            System.out.printf("  speedup:       %.1fx%n", ratio);

            assertTrue(ratio >= 2.0, "Expected fixed to be at least 2x faster, got " + ratio + "x");
            System.out.println("PASS  speedup >= 2x");
        }

        System.out.println("=== ALL PASS ===");
    }
}