check_late_key_partition.py

#!/usr/bin/env python3
"""Finite checker for the guarded late-key accessibility partition function.

The enumerators in this file construct blocks, head intervals, and their
payload intersections from coordinates.  Closed forms occur only on the
right-hand side of assertions checked against those enumerations.
"""

from __future__ import annotations

import itertools
import math
from dataclasses import dataclass
from typing import Iterable


@dataclass(frozen=True)
class Instance:
    """A guarded late-key instance with positions numbered 0 through T."""

    R: int
    d: int
    sigma: int

    @property
    def h(self) -> int:
        # A head can read p,...,p+h over sigma selector and d query ticks.
        return self.sigma + self.d - 1

    @property
    def T(self) -> int:
        return self.R * (self.h + self.d)

    @property
    def positions(self) -> range:
        return range(self.T + 1)

    @property
    def blocks(self) -> tuple[frozenset[int], ...]:
        result = []
        for selector in range(self.R):
            start = selector * (self.h + self.d) + self.h + 1
            result.append(frozenset(range(start, start + self.d)))
        return tuple(result)

    @property
    def payload(self) -> frozenset[int]:
        return frozenset().union(*self.blocks)

    def interval(self, position: int) -> frozenset[int]:
        assert 0 <= position <= self.T
        return frozenset(range(position, position + self.h + 1))

    def accessible(self, head_positions: Iterable[int]) -> frozenset[int]:
        cells: set[int] = set()
        for position in head_positions:
            cells.update(self.interval(position) & self.payload)
        return frozenset(cells)

    def one_head_weight(self, position: int) -> int:
        return 1 << len(self.accessible((position,)))

    def z_exact(self, k: int) -> int:
        """Enumerate the labeled k-head partition function."""

        total = 0
        for head_positions in itertools.product(self.positions, repeat=k):
            total += 1 << len(self.accessible(head_positions))
        return total


def ceil_log2(value: int) -> int:
    assert value >= 1
    return (value - 1).bit_length()


def z1_enumerated(instance: Instance) -> int:
    return sum(instance.one_head_weight(p) for p in instance.positions)


def z1_claimed(instance: Instance) -> int:
    # This is deliberately separate from the coordinate enumerator.
    return 1 + instance.R * (
        (instance.sigma + 3) * (1 << instance.d) - 4
    )


def verify_one_head_geometry(instance: Instance) -> None:
    """Check separation and the full overlap-size histogram by coordinates."""

    histogram = {overlap: 0 for overlap in range(instance.d + 1)}
    for position in instance.positions:
        interval = instance.interval(position)
        met_blocks = [block for block in instance.blocks if interval & block]
        assert len(met_blocks) <= 1, (
            instance,
            position,
            tuple(sorted(interval)),
            met_blocks,
        )
        overlap = len(interval & instance.payload)
        assert 0 <= overlap <= instance.d
        histogram[overlap] += 1

    # The enumerated coordinates must account for every legal position.
    assert sum(histogram.values()) == instance.T + 1

    # These assertions check the claimed profile, but the histogram above was
    # obtained without using it.
    assert histogram[0] == 1, (instance, histogram)
    assert histogram[instance.d] == instance.R * (instance.sigma + 1), (
        instance,
        histogram,
    )
    for overlap in range(1, instance.d):
        assert histogram[overlap] == 2 * instance.R, (
            instance,
            histogram,
        )

    enumerated = z1_enumerated(instance)
    claimed = z1_claimed(instance)
    assert enumerated == claimed, (instance, enumerated, claimed, histogram)


def verify_zk(instance: Instance, max_k: int = 3) -> list[tuple[int, int, int]]:
    """Compare exact multihead Z_k to the product upper bound Z_1^k."""

    z1 = z1_enumerated(instance)
    rows = []
    for k in range(max_k + 1):
        exact = instance.z_exact(k)
        product_bound = z1**k
        assert exact <= product_bound, (
            instance,
            k,
            exact,
            product_bound,
        )
        if k <= 1:
            assert exact == product_bound, (
                instance,
                k,
                exact,
                product_bound,
            )
        rows.append((k, exact, product_bound))
    return rows


def project_bits(value: int, coordinates: Iterable[int]) -> tuple[int, ...]:
    """Project an Rd-bit integer onto zero-based logical payload coordinates."""

    return tuple((value >> coordinate) & 1 for coordinate in coordinates)


def verify_sharpness_layout(instance: Instance, k: int) -> None:
    """Enumerate the fixed-k-block layout attaining H(C)+L=Rd.

    The first k blocks are exposed through fixed heads.  The remaining block
    bits form the control-state advice.  We enumerate all payloads and check
    that (advice, accessed bits) is injective and has the claimed uniform
    cardinalities.
    """

    assert 0 <= k <= instance.R
    selected_blocks = tuple(range(k))
    head_positions = tuple(
        min(instance.blocks[block_index]) for block_index in selected_blocks
    )
    physical_access = instance.accessible(head_positions)
    expected_access = frozenset().union(
        *(instance.blocks[index] for index in selected_blocks)
    )
    assert physical_access == expected_access, (
        instance,
        k,
        physical_access,
        expected_access,
    )

    # Logical payload coordinates are block-major: 0,...,Rd-1.
    selected_logical = tuple(
        block * instance.d + offset
        for block in selected_blocks
        for offset in range(instance.d)
    )
    selected_set = set(selected_logical)
    stored_logical = tuple(
        coordinate
        for coordinate in range(instance.R * instance.d)
        if coordinate not in selected_set
    )

    seen_pairs: set[tuple[tuple[int, ...], tuple[int, ...]]] = set()
    seen_states: set[tuple[int, ...]] = set()
    for payload_value in range(1 << (instance.R * instance.d)):
        state_advice = project_bits(payload_value, stored_logical)
        accessed_bits = project_bits(payload_value, selected_logical)
        pair = (state_advice, accessed_bits)
        assert pair not in seen_pairs, (instance, k, payload_value, pair)
        seen_pairs.add(pair)
        seen_states.add(state_advice)

    assert len(seen_pairs) == 1 << (instance.R * instance.d)
    assert len(seen_states) == 1 << (instance.d * (instance.R - k))
    assert len(selected_logical) == instance.d * k

    state_entropy = math.log2(len(seen_states))
    access_length = len(selected_logical)
    assert state_entropy + access_length == instance.R * instance.d


def verify_strict_example() -> None:
    instance = Instance(R=2, d=3, sigma=1)
    exact_z1 = z1_enumerated(instance)
    assert instance.T == 12
    assert exact_z1 == 57

    payload_count = 1 << (instance.R * instance.d)
    required_m = (payload_count + exact_z1 - 1) // exact_z1
    assert required_m == 2

    coarse_capacity_at_m1 = (instance.T + 1) * (1 << instance.d)
    assert coarse_capacity_at_m1 == 104
    assert coarse_capacity_at_m1 >= payload_count

    print(
        "strict example: "
        f"2^(Rd)={payload_count}, Z1={exact_z1}, "
        f"ceil(2^(Rd)/Z1)={required_m}; "
        f"coarse M=1 capacity={coarse_capacity_at_m1}"
    )


def main() -> None:
    geometry_cases = 0
    zk_cases = 0
    sharpness_cases = 0
    sample_rows: list[tuple[int, int, int]] | None = None

    # Exercise shortest binary selectors and two longer selector windows.
    for R in range(2, 5):
        shortest_selector = ceil_log2(R)
        for d in range(1, 5):
            for sigma in range(shortest_selector, shortest_selector + 3):
                instance = Instance(R=R, d=d, sigma=sigma)
                verify_one_head_geometry(instance)
                geometry_cases += 1

                # k <= 3 remains a small exhaustive enumeration here.
                rows = verify_zk(instance, max_k=3)
                zk_cases += len(rows)
                if instance == Instance(R=2, d=3, sigma=1):
                    sample_rows = rows

                # Keep payload enumeration modest while covering k=0, partial,
                # and all-block layouts.
                if instance.R * instance.d <= 12:
                    for k in sorted({0, min(1, instance.R), instance.R}):
                        verify_sharpness_layout(instance, k)
                        sharpness_cases += 1

    verify_strict_example()
    assert sample_rows is not None
    print("exact Z_k versus Z1^k for R=2, d=3, sigma=1:")
    for k, exact, product_bound in sample_rows:
        relation = "=" if exact == product_bound else "<"
        print(f"  k={k}: {exact} {relation} {product_bound}")

    print(
        "all checks passed: "
        f"{geometry_cases} geometry/Z1 cases, "
        f"{zk_cases} exact Zk comparisons, "
        f"{sharpness_cases} equality-layout checks"
    )


if __name__ == "__main__":
    main()