java-topology/defects/duckdb/unit/DuckDbCorrelatedColumnsAlgorithm.java

package unit;

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

/**
 * Unit test for CWE-407 defect in DuckDB CorrelatedColumns dedup.
 *
 * Defect: CorrelatedColumns is backed by vector<CorrelatedColumnInfo>. AddCorrelatedColumn()
 * and ExtractCorrelatedColumns() each do std::find (O(n)) before inserting. Called from:
 *   - MergeCorrelatedColumns: for loop over other → AddCorrelatedColumn → O(n²)
 *   - ExtractCorrelatedColumns: recursive traversal × O(n) per ref
 *   - HasCorrelatedExpressions::VisitReplace: inner loop × O(n)
 *
 * Fix: add column_binding_set_t (unordered_set with ColumnBindingHashFunction) to CorrelatedColumns
 * as a shadow set. CorrelatedColumns::contains() becomes O(1).
 *
 * This test models MergeCorrelatedColumns: merging S sets of C columns each into an accumulator.
 * slow(): O(n) contains per insert → O(C * accumulated_size) total
 * fast(): O(1) set contains → O(C * S) total
 *
 * Asserts slow ops > fast ops * 10x at N=300.
 */
public class DuckDbCorrelatedColumnsAlgorithm {

    // -----------------------------------------------------------------------
    // Node simulates CorrelatedColumnInfo (equality by binding integer id).
    // -----------------------------------------------------------------------
    static class Node {
        final int binding;
        Node(int b) { this.binding = b; }

        @Override
        public boolean equals(Object o) {
            return o instanceof Node && ((Node) o).binding == this.binding;
        }

        @Override
        public int hashCode() { return Integer.hashCode(binding); }
    }

    // -----------------------------------------------------------------------
    // slow(): simulates vector-backed CorrelatedColumns with O(n) contains.
    // Merges `numSets` sets of `sizePerSet` columns (with overlap to trigger dedup).
    // Returns total comparison ops.
    // -----------------------------------------------------------------------
    static Result slow(int sizePerSet, int numSets) {
        List<Node> accumulator = new ArrayList<>();
        long ops = 0;

        for (int s = 0; s < numSets; s++) {
            for (int c = 0; c < sizePerSet; c++) {
                int bindingId = c; // overlap: same columns each set → all deduped
                // O(n) membership test (std::find)
                boolean found = false;
                for (Node n : accumulator) {
                    ops++;
                    if (n.binding == bindingId) {
                        found = true;
                        break;
                    }
                }
                if (!found) {
                    accumulator.add(new Node(bindingId));
                }
            }
        }
        return new Result(ops);
    }

    // -----------------------------------------------------------------------
    // fast(): simulates CorrelatedColumns with shadow set contains() → O(1).
    // Returns total lookup ops.
    // -----------------------------------------------------------------------
    static Result fast(int sizePerSet, int numSets) {
        List<Node> accumulator = new ArrayList<>();
        Set<Integer> shadowSet = new HashSet<>(); // column_binding_set_t
        long ops = 0;

        for (int s = 0; s < numSets; s++) {
            for (int c = 0; c < sizePerSet; c++) {
                int bindingId = c;
                ops++; // O(1) hash lookup
                if (!shadowSet.contains(bindingId)) {
                    shadowSet.add(bindingId);
                    accumulator.add(new Node(bindingId));
                }
            }
        }
        return new Result(ops);
    }

    // -----------------------------------------------------------------------
    static class Result {
        final long ops;
        Result(long ops) { this.ops = ops; }
    }

    // -----------------------------------------------------------------------
    public static void main(String[] args) {
        int N = 300;   // columns per set
        int S = 50;    // number of subquery levels (MergeCorrelatedColumns calls)
        int NX = 10;

        Result s = slow(N, S);
        Result f = fast(N, S);

        System.out.printf("slow ops=%d  fast ops=%d  ratio=%.1fx%n",
                s.ops, f.ops, (double) s.ops / f.ops);

        if (s.ops <= f.ops * NX) {
            System.out.printf("FAIL: expected slow(%d) > fast(%d) * %d%n", s.ops, f.ops, NX);
            System.exit(1);
        }

        System.out.printf("1/1 PASS  (slow=%d >> fast=%d, N=%d S=%d)%n",
                s.ops, f.ops, N, S);
    }
}