Choosing HE Parameters for Spatial Circuits

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Sequence


class DepthExceeded(RuntimeError):
    """A circuit consumed more levels than its parameters provide."""


# Minimum ring dimension per total coefficient-modulus size, at 128-bit security.
# Values follow the homomorphic-encryption security standard's tables.
SECURITY_TABLE: dict[int, int] = {
    109: 4096, 218: 8192, 438: 16384, 881: 32768,
}


@dataclass(frozen=True)
class HEParams:
    ring_dimension: int
    modulus_bits: int
    scale_bits: int
    depth: int

    @property
    def slots(self) -> int:
        return self.ring_dimension // 2

    def multiply_cost_ms(self) -> float:
        """Empirical scaling: cost grows roughly as N log N per multiplication."""
        n = self.ring_dimension
        return 3.1 * (n / 4096) * math.log2(n) / math.log2(4096)

    def per_slot_cost_us(self) -> float:
        return self.multiply_cost_ms() * 1000.0 / self.slots


def derive_params(depth: int, *, scale_bits: int = 40, security_bits: int = 128) -> HEParams:
    """Derive parameters FROM the circuit depth, in the only sound order.

    The chain needs one prime per level plus a larger first and last prime; the
    total modulus then selects the smallest ring dimension the security tables
    permit. Picking the ring dimension first and hoping is how deployments end up
    unable to state their security level.
    """
    if depth < 0:
        raise ValueError("depth cannot be negative")
    if security_bits != 128:
        raise ValueError("only the 128-bit table is encoded here")
    total_bits = scale_bits * depth + 2 * (scale_bits + 20)
    for limit in sorted(SECURITY_TABLE):
        if total_bits <= limit:
            return HEParams(SECURITY_TABLE[limit], total_bits, scale_bits, depth)
    raise ValueError(
        f"a depth-{depth} circuit needs {total_bits} modulus bits, beyond the "
        "tabulated range — restructure the circuit rather than growing the ring"
    )


@dataclass
class DepthCounter:
    """Wrap the evaluator so depth is measured in code, not tracked in a document."""

    budget: int
    used: int = 0

    def multiply(self) -> None:
        self.used += 1
        if self.used > self.budget:
            raise DepthExceeded(
                f"circuit used {self.used} levels, parameters provide {self.budget} — "
                "decryption would return a plausible but meaningless value"
            )

    def add(self) -> None:
        """Additions consume no level. Present so call sites read symmetrically."""

    def remaining(self) -> int:
        return self.budget - self.used


def squared_distance_circuit(counter: DepthCounter, n_points: int) -> int:
    """Depth-1 circuit: (x-a)^2 + (y-b)^2 against a PLAINTEXT centre."""
    counter.add()          # x - a, free
    counter.multiply()     # (x - a)^2
    counter.add()          # + (y - b)^2
    return n_points


def threshold_circuit(counter: DepthCounter, degree: int) -> None:
    """Polynomial threshold of the given degree, evaluated by Horner's rule."""
    for _ in range(int(math.ceil(math.log2(max(2, degree))))):
        counter.multiply()

def _validate() -> None:
    # 1. Parameters must grow with depth, and the ring must follow the modulus.
    shallow = derive_params(1)
    deep = derive_params(5)
    assert deep.modulus_bits > shallow.modulus_bits
    assert deep.ring_dimension >= shallow.ring_dimension

    # 2. A bigger ring is slower per multiplication and cheaper per slot.
    assert deep.multiply_cost_ms() >= shallow.multiply_cost_ms()
    assert deep.per_slot_cost_us() <= shallow.per_slot_cost_us()

    # 3. A depth-1 circuit fits depth-1 parameters exactly.
    counter = DepthCounter(budget=shallow.depth)
    squared_distance_circuit(counter, n_points=1_000)
    assert counter.remaining() == 0

    # 4. Adding a threshold step to depth-1 parameters must RAISE, not silently
    #    return a wrong answer — the failure this whole page exists to prevent.
    counter = DepthCounter(budget=shallow.depth)
    squared_distance_circuit(counter, n_points=1_000)
    try:
        threshold_circuit(counter, degree=4)
    except DepthExceeded as exc:
        assert "plausible but meaningless" in str(exc)
    else:
        raise AssertionError("exceeding the level budget must be detected")

    # 5. Sized correctly, the same circuit fits.
    params = derive_params(3)
    counter = DepthCounter(budget=params.depth)
    squared_distance_circuit(counter, n_points=1_000)
    threshold_circuit(counter, degree=4)
    assert counter.remaining() >= 0

    # 6. Additions never consume a level, however many there are.
    counter = DepthCounter(budget=0)
    for _ in range(10_000):
        counter.add()
    assert counter.remaining() == 0

    # 7. An unreachable depth must fail loudly rather than silently downgrade security.
    try:
        derive_params(40)
    except ValueError as exc:
        assert "restructure the circuit" in str(exc)
    else:
        raise AssertionError("an out-of-range depth must raise")

    print("HE parameter derivation: all assertions passed")


_validate()