SDR-Rover/protocol/fec_interleaver.py

"""
Packet interleaving for consecutive packet-erasure FEC blocks.

The interleaver is deliberately independent from GF(256).  It accepts fully
encoded outer symbols from ``packet_erasure_fec`` and only changes their wire
order.  A group contains up to ``depth`` consecutive blocks.  Symbols are
emitted by symbol index across the blocks:

    block0[0], block1[0], ..., block0[1], block1[1], ...

If a block is shorter than the others, its missing symbol index is skipped.
No packet is added, removed, or modified.
"""

from __future__ import annotations

from collections.abc import Sequence

from protocol.packet_erasure_fec import (
    OuterSymbol,
    decode_outer_symbol,
)


class InterleaverError(ValueError):
    """An interleaving group or outer-symbol sequence is inconsistent."""


def _decode_block(
    wire_block: Sequence[bytes],
) -> tuple[int, tuple[bytes, ...], tuple[OuterSymbol, ...]]:
    if not wire_block:
        raise InterleaverError("FEC block is empty")

    packets = tuple(wire_block)
    symbols = tuple(decode_outer_symbol(packet) for packet in packets)
    first = symbols[0]
    expected_count = first.source_count + first.parity_count
    if len(symbols) != expected_count:
        raise InterleaverError(
            f"block {first.block_id} has {len(symbols)} symbols, "
            f"expected {expected_count}"
        )
    metadata = (
        first.block_id,
        first.source_count,
        first.parity_count,
        first.symbol_size,
    )
    for expected_index, symbol in enumerate(symbols):
        current = (
            symbol.block_id,
            symbol.source_count,
            symbol.parity_count,
            symbol.symbol_size,
        )
        if current != metadata:
            raise InterleaverError("outer symbols have inconsistent metadata")
        if symbol.symbol_index != expected_index:
            raise InterleaverError(
                f"block {first.block_id} is not in symbol-index order"
            )
    return first.block_id, packets, symbols


def interleave_group(
    wire_blocks: Sequence[Sequence[bytes]],
) -> tuple[bytes, ...]:
    """
    Interleave one non-empty group of complete consecutive FEC blocks.

    Source symbols remain before parity symbols inside every individual block
    because each block is traversed by increasing ``symbol_index``.
    """

    if not wire_blocks:
        raise InterleaverError("interleaving group is empty")

    decoded = tuple(_decode_block(block) for block in wire_blocks)
    block_ids = tuple(item[0] for item in decoded)
    if len(set(block_ids)) != len(block_ids):
        raise InterleaverError("interleaving group repeats a block_id")
    if any(
        right != left + 1
        for left, right in zip(block_ids, block_ids[1:])
    ):
        raise InterleaverError("FEC blocks are not consecutive")

    packets_by_block = tuple(item[1] for item in decoded)
    maximum_symbols = max(len(block) for block in packets_by_block)
    return tuple(
        block[symbol_index]
        for symbol_index in range(maximum_symbols)
        for block in packets_by_block
        if symbol_index < len(block)
    )


def interleave_blocks(
    wire_blocks: Sequence[Sequence[bytes]],
    depth: int,
) -> tuple[bytes, ...]:
    """
    Interleave a stream in groups of ``depth`` blocks.

    The final group may contain fewer than ``depth`` blocks and is emitted
    without synthetic padding.
    """

    if not isinstance(depth, int) or isinstance(depth, bool) or depth < 1:
        raise ValueError("interleaving depth must be a positive integer")
    blocks = tuple(tuple(block) for block in wire_blocks)
    if not blocks:
        return ()

    block_ids = tuple(_decode_block(block)[0] for block in blocks)
    if any(
        right != left + 1
        for left, right in zip(block_ids, block_ids[1:])
    ):
        raise InterleaverError("FEC block stream is not consecutive")

    result = []
    for start in range(0, len(blocks), depth):
        group = blocks[start:start + depth]
        result.extend(interleave_group(group))
    return tuple(result)


def deinterleave_symbols(
    wire_symbols: Sequence[bytes],
) -> tuple[tuple[bytes, ...], ...]:
    """
    Group received symbols by ``block_id`` and restore symbol-index order.

    Lost symbols are simply absent, so the returned blocks may be incomplete.
    Blocks are returned in first-seen order.  A conflicting duplicate symbol
    is rejected; a byte-identical duplicate is retained only once.
    """

    block_order = []
    grouped: dict[int, dict[int, bytes]] = {}
    metadata: dict[int, tuple[int, int, int]] = {}
    for wire_packet in wire_symbols:
        symbol = decode_outer_symbol(wire_packet)
        if symbol.block_id not in grouped:
            block_order.append(symbol.block_id)
            grouped[symbol.block_id] = {}
            metadata[symbol.block_id] = (
                symbol.source_count,
                symbol.parity_count,
                symbol.symbol_size,
            )
        current_metadata = (
            symbol.source_count,
            symbol.parity_count,
            symbol.symbol_size,
        )
        if current_metadata != metadata[symbol.block_id]:
            raise InterleaverError(
                f"block {symbol.block_id} metadata changed"
            )
        existing = grouped[symbol.block_id].get(symbol.symbol_index)
        if existing is not None and existing != wire_packet:
            raise InterleaverError(
                f"conflicting duplicate block={symbol.block_id}, "
                f"symbol={symbol.symbol_index}"
            )
        grouped[symbol.block_id][symbol.symbol_index] = wire_packet

    return tuple(
        tuple(
            symbols[index]
            for index in sorted(symbols)
        )
        for block_id in block_order
        for symbols in (grouped[block_id],)
    )