On-Device Budget Management for LDP Clients

from __future__ import annotations

import hashlib
import hmac
from dataclasses import dataclass, field
from typing import Callable, Mapping


class BudgetExhausted(RuntimeError):
    """Raised when the client has already reported its cap of distinct values."""


@dataclass
class ClientLedger:
    """The on-device budget: memoised outputs plus a distinct-value cap.

    The ledger is authoritative — there is no server-side equivalent, because under
    local differential privacy the loss happens before anything is transmitted.
    """

    epsilon_per_value: float
    max_distinct_values: int
    epoch_id: str
    device_key: bytes
    _memo: dict[int, int] = field(default_factory=dict)

    def spent(self) -> float:
        """Epoch privacy loss so far: one epsilon per DISTINCT value reported."""
        return self.epsilon_per_value * len(self._memo)

    def remaining(self) -> int:
        return self.max_distinct_values - len(self._memo)

    def rotate(self, new_epoch_id: str) -> None:
        """Start a new epoch. Memoised outputs must NOT survive the boundary.

        Carrying memos across epochs would keep the guarantee intact but would make
        the released value linkable across an unbounded horizon; dropping them costs
        budget in the new epoch, which is the intended trade.
        """
        if new_epoch_id == self.epoch_id:
            return
        self.epoch_id = new_epoch_id
        self._memo.clear()

    def report(self, cell: int, randomise: Callable[[int], int]) -> int:
        """Return the value to transmit for `cell`, memoising the randomisation.

        Repeat reports of an already-seen cell return the SAME randomised output and
        cost nothing: the adversary sees no new randomness, so no new information.
        A previously unseen cell costs one unit of the cap.
        """
        if cell in self._memo:
            return self._memo[cell]
        if self.remaining() <= 0:
            raise BudgetExhausted(
                f"epoch {self.epoch_id}: cap of {self.max_distinct_values} distinct "
                "cells reached; the client must stop reporting until rotation"
            )
        value = randomise(cell)
        self._memo[cell] = value
        return value

    def memo_digest(self) -> str:
        """A keyed digest of the memo table, for local integrity checks only.

        Never transmit this: the digest is derived from the set of visited cells and
        would be a linkable identifier across reports.
        """
        payload = ",".join(f"{k}:{v}" for k, v in sorted(self._memo.items()))
        return hmac.new(self.device_key, payload.encode(), hashlib.blake2b).hexdigest()

def _validate() -> None:
    calls = {"n": 0}

    def randomise(cell: int) -> int:
        calls["n"] += 1
        return (cell * 7919 + calls["n"]) % 64   # a stand-in for the real mechanism

    ledger = ClientLedger(
        epsilon_per_value=1.0, max_distinct_values=3,
        epoch_id="2026-08-11", device_key=b"test-key",
    )

    # 1. Repeat reports of the same cell must be byte-identical and free.
    first = ledger.report(12, randomise)
    again = ledger.report(12, randomise)
    assert first == again
    assert calls["n"] == 1, "a memoised value must not be re-randomised"
    assert ledger.spent() == 1.0

    # 2. Distinct cells consume the cap.
    ledger.report(30, randomise)
    ledger.report(44, randomise)
    assert ledger.remaining() == 0
    assert ledger.spent() == 3.0

    # 3. Reporting a new cell past the cap must raise, not silently randomise.
    try:
        ledger.report(51, randomise)
    except BudgetExhausted:
        pass
    else:
        raise AssertionError("the cap must be enforced, not advisory")

    # 4. A memoised cell is still reportable after the cap — it costs nothing.
    assert ledger.report(12, randomise) == first

    # 5. Rotation clears memos and restores the cap.
    ledger.rotate("2026-08-12")
    assert ledger.remaining() == 3
    fresh = ledger.report(12, randomise)
    assert fresh != first or calls["n"] == 4, "a new epoch must re-randomise"

    # 6. Rotating to the same epoch id must be a no-op, not a reset.
    ledger.rotate("2026-08-12")
    assert ledger.remaining() == 2

    print("on-device LDP ledger: all assertions passed")


_validate()