PSI Cardinality-Only Joins at Scale

from __future__ import annotations

import hashlib
import math
import random
import secrets
from dataclasses import dataclass
from typing import Iterable, Sequence

# A prime-order group stand-in. Production code uses an elliptic curve with a
# hash-to-curve; the structure — commutative blinding — is identical.
P = (1 << 521) - 1
Q = (P - 1) // 2


def _hash_to_group(token: str) -> int:
    digest = hashlib.blake2b(token.encode(), digest_size=64).digest()
    return pow(4, int.from_bytes(digest, "big") % Q, P)


def blind(values: Iterable[int], key: int) -> list[int]:
    """Raise each group element to the party's secret exponent."""
    return [pow(v, key, P) for v in values]


def bucket_of(value: int, bucket_bits: int) -> int:
    """Partition on the BLINDED value, so the bucket reveals nothing about the token."""
    return int.from_bytes(
        hashlib.blake2b(str(value).encode(), digest_size=8).digest(), "big"
    ) % (1 << bucket_bits)


@dataclass(frozen=True)
class CardinalityResult:
    exact: int
    released: float
    epsilon: float
    buckets: int


def psi_cardinality(
    a_tokens: Sequence[str],
    b_tokens: Sequence[str],
    *,
    bucket_bits: int = 8,
    epsilon: float = 1.0,
    min_set_size: int = 500,
    seed: int = 0,
) -> CardinalityResult:
    """Run a bucketed DH-style PSI and release only a noised cardinality.

    Both parties blind with their own secret exponent; because exponentiation
    commutes, a token present in both sets produces the same double-blinded value on
    both sides. Bucketing on the double-blinded value turns a quadratic comparison
    into 2^bucket_bits small ones without revealing anything extra.
    """
    if len(a_tokens) < min_set_size or len(b_tokens) < min_set_size:
        raise ValueError(
            f"set sizes {len(a_tokens)}/{len(b_tokens)} below the floor {min_set_size} — "
            "a cardinality answer over a tiny set is close to membership disclosure"
        )
    key_a, key_b = secrets.randbelow(Q) + 1, secrets.randbelow(Q) + 1

    a_once = blind((_hash_to_group(t) for t in a_tokens), key_a)
    b_once = blind((_hash_to_group(t) for t in b_tokens), key_b)
    a_twice = blind(a_once, key_b)          # B blinds A's set
    b_twice = blind(b_once, key_a)          # A blinds B's set

    buckets_a: dict[int, set[int]] = {}
    for v in a_twice:
        buckets_a.setdefault(bucket_of(v, bucket_bits), set()).add(v)
    exact = 0
    for v in b_twice:
        if v in buckets_a.get(bucket_of(v, bucket_bits), ()):
            exact += 1

    rng = random.Random(seed)
    u = rng.random() - 0.5
    noise = -(1.0 / epsilon) * math.copysign(1.0, u) * math.log1p(-2 * abs(u))
    return CardinalityResult(exact, max(0.0, exact + noise), epsilon, 1 << bucket_bits)


def comparison_cost(n: int, bucket_bits: int) -> int:
    """Comparisons performed: quadratic without bucketing, near-linear with it."""
    buckets = 1 << bucket_bits
    per_bucket = max(1, n // buckets)
    return buckets * per_bucket * per_bucket

def _validate() -> None:
    rng = random.Random(7)
    shared = [f"cell-{i}" for i in range(1_200)]
    a = shared + [f"a-only-{i}" for i in range(2_000)]
    b = shared + [f"b-only-{i}" for i in range(3_000)]

    result = psi_cardinality(a, b, bucket_bits=8, epsilon=2.0, seed=1)

    # 1. The protocol must find exactly the true intersection size.
    assert result.exact == 1_200, result.exact

    # 2. The released value is noised, so it is not the exact count.
    assert result.released != float(result.exact)
    assert abs(result.released - result.exact) < 20

    # 3. Disjoint sets must produce a cardinality of zero before noise.
    disjoint = psi_cardinality(
        [f"x-{i}" for i in range(800)], [f"y-{i}" for i in range(800)],
        bucket_bits=8, epsilon=2.0, min_set_size=500, seed=2,
    )
    assert disjoint.exact == 0

    # 4. Bucketing must not change the answer — only the cost.
    coarse = psi_cardinality(a, b, bucket_bits=2, epsilon=2.0, seed=1)
    assert coarse.exact == result.exact

    # 5. Bucketing must actually reduce the comparison count substantially.
    assert comparison_cost(100_000, 8) < comparison_cost(100_000, 0) / 100

    # 6. Small sets must be refused, not answered.
    try:
        psi_cardinality(["a"] * 10, ["a"] * 10, min_set_size=500)
    except ValueError as exc:
        assert "below the floor" in str(exc)
    else:
        raise AssertionError("tiny sets must be refused")

    print("PSI cardinality: all assertions passed")


_validate()