coordinates.py

"""Coordinate index types for the IR envelope.

These types bridge Policy IR nodes to their Evidence IR (blob) origins via
byte offsets, node indices, and structural hashes.  The coordinate index
is populated after both lanes produce their outputs, during envelope assembly.

Types:

- ``NodeCoordinate``:  per-node blob location (indices + offsets).
- ``CoordinateIndex``: full node-to-blob bridge + op-table index.
- ``SectionHashes``:   per-section SHA-256 for the three blob regions.
- ``HashBoundaries``:  structural hashes at three granularities.

The digest module (``pawl.structure.inventory.digests.api``) computes
section hashes from raw blobs; this module defines the contract-level
representations that the envelope consumes.  Subgraph hashes come from
normalized ``Policy`` metadata (``structural_hash_entry`` written by
``pawl.forward.layer2.normalize.canonicalize_operation``).
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Mapping, Sequence

from pawl.structure.core.blob_geometry import (
    NODE_BYTE_SIZE,
    node_index_to_absolute_offset,
)


@dataclass(frozen=True)
class NodeCoordinate:
    """Map a single Policy IR node to its Evidence IR (blob) origin.

    Tuples (not scalars) because the compiler shares nodes across
    operations — a single ``NodeTest`` may map to multiple blob offsets.
    """

    node_indices: tuple[int, ...] = ()
    byte_offsets: tuple[int, ...] = ()
    literal_pool_offset: int | None = None
    absence_reason: str | None = None

    def as_dict(self) -> dict[str, object]:
        out: dict[str, object] = {}
        if self.node_indices:
            out["node_indices"] = list(self.node_indices)
        if self.byte_offsets:
            out["byte_offsets"] = list(self.byte_offsets)
        if self.literal_pool_offset is not None:
            out["literal_pool_offset"] = self.literal_pool_offset
        if self.absence_reason is not None:
            out["absence_reason"] = self.absence_reason
        return out


@dataclass(frozen=True)
class CoordinateIndex:
    """Bridge between Policy IR nodes and Evidence IR blob locations.

    ``nodes`` maps each node_id to its blob coordinate.  Every Policy IR
    node must either have a populated coordinate or an explicit
    ``absence_reason``.

    ``op_table`` maps operation names to their op-table entry index in
    the compiled blob.
    """

    nodes: Mapping[str, NodeCoordinate] = field(default_factory=dict)
    op_table: Mapping[str, int] = field(default_factory=dict)

    def as_dict(self) -> dict[str, object]:
        return {
            "nodes": {k: v.as_dict() for k, v in self.nodes.items()},
            "op_table": dict(self.op_table),
        }


@dataclass(frozen=True)
class SectionHashes:
    """SHA-256 hashes for the three structural sections of a compiled blob.

    Field names match the IR envelope contract.  A structurally identical
    type exists in ``pawl.structure.inventory.digests.api`` for the digest
    pipeline; this is the contract-level representation.
    """

    op_table: str
    node_stream: str
    literal_pool: str

    def as_dict(self) -> dict[str, object]:
        return {
            "op_table": self.op_table,
            "node_stream": self.node_stream,
            "literal_pool": self.literal_pool,
        }


@dataclass(frozen=True)
class HashBoundaries:
    """Structural hashes at three granularities.

    - ``profile_hash``: SHA-256 of the full compiled blob.
      Null when no blob is available (Source-Evaluate path).
    - ``section_hashes``: per-section SHA-256 (op_table, node_stream,
      literal_pool).  Null when no blob is available.
    - ``subgraph_hashes``: per-operation SHA-256 of the reachable
      subgraph.  Computed from normalized ``Policy`` metadata
      (``structural_hash_entry`` written by ``canonicalize_operation``).
    """

    profile_hash: str | None = None
    section_hashes: SectionHashes | None = None
    subgraph_hashes: Mapping[str, str] = field(default_factory=dict)

    def as_dict(self) -> dict[str, object]:
        return {
            "profile_hash": self.profile_hash,
            "section_hashes": (
                self.section_hashes.as_dict()
                if self.section_hashes is not None
                else None
            ),
            "subgraph_hashes": dict(self.subgraph_hashes),
        }


# ---------------------------------------------------------------------------
# Builder helpers
# ---------------------------------------------------------------------------


def coordinate_from_provenance(
    origin_offsets: Sequence[int],
    nodes_start_offset: int,
    node_byte_size: int = NODE_BYTE_SIZE,
    *,
    literal_pool_offset: int | None = None,
) -> NodeCoordinate:
    """Build a ``NodeCoordinate`` from provenance origin offsets.

    ``origin_offsets`` are node-stream indices (not byte offsets).
    ``nodes_start_offset`` is the absolute byte offset where the node
    stream begins in the blob.
    """
    byte_offsets = tuple(
        node_index_to_absolute_offset(
            nodes_start_offset,
            idx,
            node_byte_size=node_byte_size,
        )
        for idx in origin_offsets
    )
    return NodeCoordinate(
        node_indices=tuple(origin_offsets),
        byte_offsets=byte_offsets,
        literal_pool_offset=literal_pool_offset,
    )


def build_coordinate_index(
    node_origins: Mapping[str, Sequence[int]],
    op_table_indices: Mapping[str, int],
    nodes_start_offset: int,
    node_byte_size: int = NODE_BYTE_SIZE,
    *,
    literal_pool_offsets: Mapping[str, int] | None = None,
    absence_reasons: Mapping[str, str] | None = None,
) -> CoordinateIndex:
    """Assemble a ``CoordinateIndex`` from pre-extracted provenance data.

    Parameters
    ----------
    node_origins:
        ``{node_id: origin_offsets}`` — node-stream indices from
        ``NodeProvenance.origin_offsets`` for each Policy IR node.
    op_table_indices:
        ``{operation_name: index}`` — op-table entry index per operation.
    nodes_start_offset:
        Absolute byte offset where the blob's node stream begins
        (from ``SectionOffsets.nodes_start``).
    literal_pool_offsets:
        Optional ``{node_id: offset}`` for nodes that reference a
        literal pool entry.
    absence_reasons:
        Optional ``{node_id: reason}`` for nodes without blob
        coordinates (e.g. ``source_evaluate_no_blob``).
    """
    lpo = literal_pool_offsets or {}
    reasons = absence_reasons or {}
    coords: dict[str, NodeCoordinate] = {}

    for node_id, offsets in node_origins.items():
        if offsets:
            coords[node_id] = coordinate_from_provenance(
                offsets,
                nodes_start_offset,
                node_byte_size,
                literal_pool_offset=lpo.get(node_id),
            )
        else:
            coords[node_id] = NodeCoordinate(
                absence_reason=reasons.get(node_id, "no_provenance"),
            )

    # Nodes mentioned in absence_reasons but not in node_origins.
    for node_id, reason in reasons.items():
        if node_id not in coords:
            coords[node_id] = NodeCoordinate(absence_reason=reason)

    return CoordinateIndex(nodes=coords, op_table=op_table_indices)


__all__ = [
    "NODE_BYTE_SIZE",
    "NodeCoordinate",
    "CoordinateIndex",
    "SectionHashes",
    "HashBoundaries",
    "coordinate_from_provenance",
    "build_coordinate_index",
]