tokenspeed-pd 0.1.0

// Copyright (c) 2026 LightSeek Foundation
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

//! The cache transfer contract (`pd/cache_protocol.py`'s
//! `CacheTransferContract`): one rank's cache memory plan, its groups'
//! semantics and its fields' TP partitioning, as the PD wire carries it.
//!
//! The Python cache recipes build the contract and serialize it once at
//! startup; this side only parses it, checks it the way the Python
//! constructors do, and answers the questions transfer planning asks.

use std::collections::{BTreeMap, HashMap, HashSet};

use serde_json::{Map, Value};

use super::LayoutError;

/// The largest contract the wire carries.
pub const MAX_CONTRACT_WIRE_BYTES: usize = 256 << 10;

/// A field's element width by dtype name (`plan._CACHE_DTYPE_BYTES`).
fn dtype_bytes(dtype: &str) -> Option<i64> {
    Some(match dtype {
        "uint8" | "int8" | "float8_e4m3fn" | "float8_e5m2" | "float8_e8m0fnu" => 1,
        "bfloat16" | "float16" => 2,
        "int32" | "float32" => 4,
        "int64" | "float64" => 8,
        _ => return None,
    })
}

/// How many CacheBlocks of one group a parent holds, and how many it has.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GroupLayout {
    /// The group.
    pub group_id: String,
    /// CacheBlocks per LCM parent.
    pub cache_blocks_per_lcm_block: i64,
    /// Local pages, the null page included.
    pub page_count: i64,
}

/// One plane of the arena.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PlaneLayout {
    /// The plane.
    pub plane_id: String,
    /// Its bytes per LCM parent.
    pub bytes_per_lcm_block: i64,
    /// Where it starts in the arena.
    pub arena_offset_bytes: i64,
}

/// One field's geometry.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FieldLayout {
    /// Its group.
    pub group_id: String,
    /// The field (`layer.<id>.<plane>`).
    pub field_id: String,
    /// Its plane.
    pub plane_id: String,
    /// The rank-local shape of one page's payload.
    pub shape: Vec<i64>,
    /// The dtype the arena views it as.
    pub dtype: String,
    /// Its offset within a page of the plane.
    pub field_offset_bytes: i64,
    /// Bytes between consecutive pages.
    pub page_stride_bytes: i64,
}

impl FieldLayout {
    /// Bytes per element.
    pub fn element_size(&self) -> i64 {
        dtype_bytes(&self.dtype).expect("checked when parsed")
    }

    /// Bytes of one page's payload.
    pub fn payload_bytes(&self) -> i64 {
        self.shape.iter().product::<i64>() * self.element_size()
    }
}

/// One cache group's scheduler-facing semantics (`CacheGroupSpec`).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GroupSpec {
    /// The group.
    pub group_id: String,
    /// `full_history`, `sliding_window` or `state`.
    pub attn_kind: String,
    /// Rows per page; row-geometry groups only.
    pub rows_per_page: Option<i64>,
    /// Raw tokens per row; row-geometry groups only.
    pub entry_stride_tokens: Option<i64>,
    /// The retained window of a sliding-window group.
    pub sliding_window_tokens: Option<i64>,
    /// A sliding-window group regenerated by bounded replay.
    pub replayable: bool,
    /// `full_suffix` or `latest_snapshot`.
    pub transfer_policy: Option<String>,
    /// Cyclic owners of the group's virtual blocks; 1 is replicated.
    pub shard_count: i64,
}

impl GroupSpec {
    /// The raw-token span of one block-table slot (`block_granularity`);
    /// a state group has none.
    pub fn block_granularity(&self) -> Result<i64, LayoutError> {
        match (self.rows_per_page, self.entry_stride_tokens) {
            (Some(rows), Some(stride)) => Ok(rows * stride),
            _ => Err(LayoutError::new(format!(
                "group {:?}: block_granularity is undefined for a state group",
                self.group_id
            ))),
        }
    }

    /// Whether the group transfers its retained history suffix (else its
    /// latest snapshot).
    pub fn full_suffix(&self) -> bool {
        self.transfer_policy.as_deref() == Some("full_suffix")
    }

    fn validate(&self) -> Result<(), LayoutError> {
        let id = &self.group_id;
        if self.shard_count <= 0 {
            return Err(LayoutError::new("shard_count must be a positive integer"));
        }
        if !matches!(
            self.attn_kind.as_str(),
            "full_history" | "sliding_window" | "state"
        ) {
            return Err(LayoutError::new(format!(
                "group {id:?}: unknown cache-group kind {:?}",
                self.attn_kind
            )));
        }
        if self.attn_kind == "sliding_window" {
            if self.sliding_window_tokens.is_none_or(|tokens| tokens <= 0) {
                return Err(LayoutError::new(format!(
                    "group {id:?}: a sliding-window group needs a positive sliding_window_tokens"
                )));
            }
        } else if self.sliding_window_tokens.is_some() {
            return Err(LayoutError::new(format!(
                "group {id:?}: only a sliding-window group carries sliding_window_tokens"
            )));
        }
        if self.replayable && self.attn_kind != "sliding_window" {
            return Err(LayoutError::new(format!(
                "group {id:?}: only a sliding-window group can be regenerated by bounded replay"
            )));
        }
        let has_rows = self.rows_per_page.is_some() || self.entry_stride_tokens.is_some();
        if self.attn_kind == "state" {
            if has_rows {
                return Err(LayoutError::new(format!(
                    "group {id:?}: a state group declares no row geometry"
                )));
            }
            return Ok(());
        }
        match (self.rows_per_page, self.entry_stride_tokens) {
            (Some(rows), Some(stride)) if rows > 0 && stride > 0 => Ok(()),
            (Some(_), Some(_)) => Err(LayoutError::new(format!(
                "group {id:?}: rows_per_page and entry_stride_tokens must be > 0"
            ))),
            _ => Err(LayoutError::new(format!(
                "group {id:?}: a {} group declares row geometry",
                self.attn_kind
            ))),
        }
    }
}

/// How one field's global extent splits over TP ranks
/// (`CacheFieldPartition`).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Partition {
    /// The partitioned axis of the local shape.
    pub axis: i64,
    /// Its global extent.
    pub global_extent: i64,
    /// Sub-ranges partitioned separately (a fused QKV), or empty.
    pub global_parts: Vec<i64>,
}

impl Partition {
    fn validate(&self, field_id: &str, local_shape: &[i64]) -> Result<(), LayoutError> {
        let invalid = || {
            LayoutError::new(format!(
                "cache field {field_id:?} has invalid partition geometry"
            ))
        };
        let axis = usize::try_from(self.axis).map_err(|_| invalid())?;
        let local_extent = *local_shape.get(axis).ok_or_else(invalid)?;
        if self.global_extent < local_extent || local_extent <= 0 {
            return Err(invalid());
        }
        if self.global_extent % local_extent != 0 {
            return Err(LayoutError::new(format!(
                "cache field {field_id:?} global partition extent is not divisible by its \
                 local extent"
            )));
        }
        let distinct_shards = self.global_extent / local_extent;
        if !self.global_parts.is_empty()
            && (self.global_parts.len() < 2
                || self.global_parts.len() > 16
                || self
                    .global_parts
                    .iter()
                    .any(|&part| part <= 0 || part % distinct_shards != 0)
                || self.global_parts.iter().sum::<i64>() != self.global_extent)
        {
            return Err(LayoutError::new(format!(
                "cache field {field_id:?} has invalid global partition parts"
            )));
        }
        Ok(())
    }
}

/// One rank's static cache byte geometry (`CacheMemoryPlan`).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Plan {
    /// Tokens per prefix-matching unit.
    pub prefix_granularity: i64,
    /// Bytes of one LCM parent.
    pub lcm_block_bytes: i64,
    /// Parents, the null one excluded.
    pub num_lcm_blocks: i64,
    /// The groups.
    pub groups: Vec<GroupLayout>,
    /// The planes.
    pub planes: Vec<PlaneLayout>,
    /// The fields.
    pub fields: Vec<FieldLayout>,
    /// The retained planes' bytes per parent, when narrowed to a layer
    /// window.
    pub resident_block_bytes: Option<i64>,
}

/// A rank's cache transfer contract.
#[derive(Debug, Clone)]
pub struct Contract {
    /// The memory plan.
    pub plan: Plan,
    /// The groups' semantics, in the scheduler's order.
    pub group_specs: Vec<GroupSpec>,
    /// The partitioned fields, by field.
    partitions: HashMap<String, Partition>,
    /// The fields' landing addresses, on an endpoint's contract; absent in
    /// the rank-independent bootstrap layout.
    pub field_addresses: Option<BTreeMap<String, u64>>,
    groups_by_id: HashMap<String, usize>,
    specs_by_id: HashMap<String, usize>,
    fields_by_id: HashMap<String, usize>,
    planes_by_id: HashMap<String, usize>,
    /// Each group's fields, sorted by field id (`fields_for_group`).
    group_fields: HashMap<String, Vec<usize>>,
}

impl PartialEq for Contract {
    fn eq(&self, other: &Self) -> bool {
        self.plan == other.plan
            && self.group_specs == other.group_specs
            && self.partitions == other.partitions
            && self.field_addresses == other.field_addresses
    }
}

impl Eq for Contract {}

/// The keys a dataclass built from `**mapping` accepts: anything else, or a
/// missing required key, is refused, as the Python constructor refuses it.
fn keyword_fields<'a>(
    value: &'a Value,
    name: &str,
    required: &[&str],
    optional: &[&str],
) -> Result<&'a Map<String, Value>, LayoutError> {
    let object = value
        .as_object()
        .ok_or_else(|| LayoutError::new(format!("{name} must be a JSON object")))?;
    if let Some(unknown) = object
        .keys()
        .find(|key| !required.contains(&key.as_str()) && !optional.contains(&key.as_str()))
    {
        return Err(LayoutError::new(format!(
            "{name} has an unknown field {unknown:?}"
        )));
    }
    if let Some(missing) = required.iter().find(|key| !object.contains_key(**key)) {
        return Err(LayoutError::new(format!("{name} misses {missing:?}")));
    }
    Ok(object)
}

fn key<'a>(
    object: &'a Map<String, Value>,
    name: &str,
    key: &str,
) -> Result<&'a Value, LayoutError> {
    object
        .get(key)
        .ok_or_else(|| LayoutError::new(format!("{name} misses {key:?}")))
}

fn int(value: &Value, what: &str) -> Result<i64, LayoutError> {
    value
        .as_i64()
        .ok_or_else(|| LayoutError::new(format!("{what} must be an integer")))
}

fn opt_int(object: &Map<String, Value>, field: &str) -> Result<Option<i64>, LayoutError> {
    match object.get(field) {
        None | Some(Value::Null) => Ok(None),
        Some(value) => int(value, field).map(Some),
    }
}

fn string(value: &Value, what: &str) -> Result<String, LayoutError> {
    value
        .as_str()
        .map(str::to_owned)
        .ok_or_else(|| LayoutError::new(format!("{what} must be a string")))
}

fn list<'a>(value: &'a Value, what: &str) -> Result<&'a Vec<Value>, LayoutError> {
    value
        .as_array()
        .ok_or_else(|| LayoutError::new(format!("{what} must be a list")))
}

/// A wire payload as a JSON object (`_load_wire_json`).
pub(crate) fn load_wire_json(
    raw: &[u8],
    name: &str,
    maximum: usize,
) -> Result<Map<String, Value>, LayoutError> {
    if raw.is_empty() || raw.len() > maximum {
        return Err(LayoutError::new(format!(
            "{name} payload must be 1..{maximum} bytes"
        )));
    }
    match serde_json::from_slice::<Value>(raw) {
        Ok(Value::Object(object)) => Ok(object),
        Ok(_) => Err(LayoutError::new(format!(
            "{name} payload must be a JSON object"
        ))),
        Err(_) => Err(LayoutError::new(format!("invalid {name} JSON"))),
    }
}

fn parse_group(value: &Value) -> Result<GroupLayout, LayoutError> {
    let name = "cache group layout";
    let object = keyword_fields(
        value,
        name,
        &["group_id", "cache_blocks_per_lcm_block", "page_count"],
        &[],
    )?;
    Ok(GroupLayout {
        group_id: string(key(object, name, "group_id")?, "group_id")?,
        cache_blocks_per_lcm_block: int(
            key(object, name, "cache_blocks_per_lcm_block")?,
            "cache_blocks_per_lcm_block",
        )?,
        page_count: int(key(object, name, "page_count")?, "page_count")?,
    })
}

fn parse_plane(value: &Value) -> Result<PlaneLayout, LayoutError> {
    let name = "cache plane layout";
    let object = keyword_fields(
        value,
        name,
        &["plane_id", "bytes_per_lcm_block", "arena_offset_bytes"],
        &[],
    )?;
    Ok(PlaneLayout {
        plane_id: string(key(object, name, "plane_id")?, "plane_id")?,
        bytes_per_lcm_block: int(
            key(object, name, "bytes_per_lcm_block")?,
            "bytes_per_lcm_block",
        )?,
        arena_offset_bytes: int(
            key(object, name, "arena_offset_bytes")?,
            "arena_offset_bytes",
        )?,
    })
}

fn parse_field(value: &Value) -> Result<FieldLayout, LayoutError> {
    let name = "cache field layout";
    let object = keyword_fields(
        value,
        name,
        &[
            "group_id",
            "field_id",
            "plane_id",
            "shape",
            "dtype",
            "field_offset_bytes",
            "page_stride_bytes",
        ],
        &[],
    )?;
    let field = FieldLayout {
        group_id: string(key(object, name, "group_id")?, "group_id")?,
        field_id: string(key(object, name, "field_id")?, "field_id")?,
        plane_id: string(key(object, name, "plane_id")?, "plane_id")?,
        shape: list(key(object, name, "shape")?, "shape")?
            .iter()
            .map(|dimension| int(dimension, "shape"))
            .collect::<Result<_, _>>()?,
        dtype: string(key(object, name, "dtype")?, "dtype")?,
        field_offset_bytes: int(
            key(object, name, "field_offset_bytes")?,
            "field_offset_bytes",
        )?,
        page_stride_bytes: int(key(object, name, "page_stride_bytes")?, "page_stride_bytes")?,
    };
    if dtype_bytes(&field.dtype).is_none() {
        return Err(LayoutError::new(format!(
            "unsupported cache dtype {:?}",
            field.dtype
        )));
    }
    Ok(field)
}

fn parse_spec(value: &Value) -> Result<GroupSpec, LayoutError> {
    let name = "cache group spec";
    let object = keyword_fields(
        value,
        name,
        &["group_id", "attn_kind", "replayable"],
        &[
            "rows_per_page",
            "entry_stride_tokens",
            "sliding_window_tokens",
            "transfer_policy",
            "shard_count",
        ],
    )?;
    let spec = GroupSpec {
        group_id: string(key(object, name, "group_id")?, "group_id")?,
        attn_kind: string(key(object, name, "attn_kind")?, "attn_kind")?,
        rows_per_page: opt_int(object, "rows_per_page")?,
        entry_stride_tokens: opt_int(object, "entry_stride_tokens")?,
        sliding_window_tokens: opt_int(object, "sliding_window_tokens")?,
        replayable: key(object, name, "replayable")?
            .as_bool()
            .ok_or_else(|| LayoutError::new("replayable must be a bool"))?,
        transfer_policy: match object.get("transfer_policy") {
            None | Some(Value::Null) => None,
            Some(value) => Some(string(value, "transfer_policy")?),
        },
        shard_count: match object.get("shard_count") {
            None => 1,
            Some(value) => int(value, "shard_count")?,
        },
    };
    spec.validate()?;
    Ok(spec)
}

fn parse_partition(value: &Value) -> Result<Partition, LayoutError> {
    let name = "cache field partition";
    let object = value
        .as_object()
        .ok_or_else(|| LayoutError::new(format!("{name} must be a JSON object")))?;
    Ok(Partition {
        axis: int(key(object, name, "axis")?, "axis")?,
        global_extent: int(key(object, name, "global_extent")?, "global_extent")?,
        global_parts: list(key(object, name, "global_parts")?, "global_parts")?
            .iter()
            .map(|part| int(part, "global_parts"))
            .collect::<Result<_, _>>()?,
    })
}

impl Contract {
    /// Parse a contract off the wire (`CacheTransferContract.from_wire_bytes`).
    pub fn from_wire(raw: &[u8]) -> Result<Self, LayoutError> {
        let name = "cache transfer contract";
        let payload = load_wire_json(raw, name, MAX_CONTRACT_WIRE_BYTES)?;
        let invalid = |error: LayoutError| {
            LayoutError::new(format!("invalid cache transfer contract: {error}"))
        };
        let parse = || -> Result<Self, LayoutError> {
            let plan = key(&payload, name, "plan")?
                .as_object()
                .ok_or_else(|| LayoutError::new("plan must be a JSON object"))?;
            let plan = Plan {
                prefix_granularity: int(
                    key(plan, "plan", "prefix_granularity")?,
                    "prefix_granularity",
                )?,
                lcm_block_bytes: int(key(plan, "plan", "lcm_block_bytes")?, "lcm_block_bytes")?,
                num_lcm_blocks: int(key(plan, "plan", "num_lcm_blocks")?, "num_lcm_blocks")?,
                groups: list(key(plan, "plan", "groups")?, "groups")?
                    .iter()
                    .map(parse_group)
                    .collect::<Result<_, _>>()?,
                planes: list(key(plan, "plan", "planes")?, "planes")?
                    .iter()
                    .map(parse_plane)
                    .collect::<Result<_, _>>()?,
                fields: list(key(plan, "plan", "fields")?, "fields")?
                    .iter()
                    .map(parse_field)
                    .collect::<Result<_, _>>()?,
                resident_block_bytes: opt_int(plan, "resident_block_bytes")?,
            };
            let field_addresses = match payload.get("field_addresses") {
                None | Some(Value::Null) => None,
                Some(Value::Object(addresses)) => Some(
                    addresses
                        .iter()
                        .map(|(field, address)| {
                            let address =
                                address.as_u64().filter(|&address| address > 0).ok_or_else(
                                    || LayoutError::new("field addresses must be nonzero uint64"),
                                )?;
                            Ok((field.clone(), address))
                        })
                        .collect::<Result<_, LayoutError>>()?,
                ),
                Some(_) => return Err(LayoutError::new("field addresses must be a mapping")),
            };
            let group_specs = list(key(&payload, name, "group_specs")?, "group_specs")?
                .iter()
                .map(parse_spec)
                .collect::<Result<Vec<_>, _>>()?;
            let schema = key(&payload, name, "transfer_schema")?
                .as_object()
                .ok_or_else(|| LayoutError::new("transfer_schema must be a JSON object"))?;
            let mut partitions = Vec::new();
            for entry in list(key(schema, "transfer_schema", "fields")?, "fields")? {
                let entry = entry
                    .as_object()
                    .ok_or_else(|| LayoutError::new("a transfer field must be a JSON object"))?;
                let field_id = string(key(entry, "transfer field", "field_id")?, "field_id")?;
                if field_id.is_empty() {
                    return Err(LayoutError::new(
                        "cache transfer field_id must be non-empty",
                    ));
                }
                partitions.push((
                    field_id,
                    parse_partition(key(entry, "transfer field", "partition")?)?,
                ));
            }
            Self::new(plan, group_specs, partitions, field_addresses)
        };
        parse().map_err(invalid)
    }

    /// A contract from its parts, checked as the Python constructor checks
    /// it.
    pub fn new(
        plan: Plan,
        group_specs: Vec<GroupSpec>,
        partitions: Vec<(String, Partition)>,
        field_addresses: Option<BTreeMap<String, u64>>,
    ) -> Result<Self, LayoutError> {
        let plan_groups: HashSet<&str> = plan.groups.iter().map(|g| g.group_id.as_str()).collect();
        let spec_groups: HashSet<&str> = group_specs.iter().map(|s| s.group_id.as_str()).collect();
        if plan_groups != spec_groups {
            return Err(LayoutError::new(
                "cache plan and scheduler group IDs disagree",
            ));
        }
        let mut seen = HashSet::new();
        if partitions
            .iter()
            .any(|(field, _)| !seen.insert(field.clone()))
        {
            return Err(LayoutError::new(
                "cache transfer schema contains a duplicate field",
            ));
        }
        // Later entries win, as in a Python dict.
        let fields_by_id: HashMap<String, usize> = plan
            .fields
            .iter()
            .enumerate()
            .map(|(index, field)| (field.field_id.clone(), index))
            .collect();
        for (field_id, partition) in &partitions {
            let Some(&index) = fields_by_id.get(field_id) else {
                return Err(LayoutError::new(format!(
                    "cache transfer schema references unknown fields [{field_id:?}]"
                )));
            };
            partition.validate(field_id, &plan.fields[index].shape)?;
        }
        if let Some(addresses) = &field_addresses {
            let planned: HashSet<&str> = plan.fields.iter().map(|f| f.field_id.as_str()).collect();
            let addressed: HashSet<&str> = addresses.keys().map(String::as_str).collect();
            if planned != addressed {
                return Err(LayoutError::new(
                    "field addresses must cover the complete plan",
                ));
            }
        }
        let groups_by_id = plan
            .groups
            .iter()
            .enumerate()
            .map(|(index, group)| (group.group_id.clone(), index))
            .collect();
        let specs_by_id = group_specs
            .iter()
            .enumerate()
            .map(|(index, spec)| (spec.group_id.clone(), index))
            .collect();
        let planes_by_id = plan
            .planes
            .iter()
            .enumerate()
            .map(|(index, plane)| (plane.plane_id.clone(), index))
            .collect();
        let mut group_fields: HashMap<String, Vec<usize>> = HashMap::new();
        for (index, field) in plan.fields.iter().enumerate() {
            group_fields
                .entry(field.group_id.clone())
                .or_default()
                .push(index);
        }
        for indexes in group_fields.values_mut() {
            indexes.sort_by(|&a, &b| plan.fields[a].field_id.cmp(&plan.fields[b].field_id));
        }
        Ok(Self {
            plan,
            group_specs,
            partitions: partitions.into_iter().collect(),
            field_addresses,
            groups_by_id,
            specs_by_id,
            fields_by_id,
            planes_by_id,
            group_fields,
        })
    }

    /// The same contract without its landing addresses: the logical layout
    /// every rank of a node registers with the bootstrap server.
    pub fn without_addresses(&self) -> Self {
        Self {
            field_addresses: None,
            ..self.clone()
        }
    }

    /// The group's semantics.
    pub fn group_spec(&self, group_id: &str) -> Result<&GroupSpec, LayoutError> {
        self.specs_by_id
            .get(group_id)
            .map(|&index| &self.group_specs[index])
            .ok_or_else(|| LayoutError::new(format!("unknown cache group {group_id:?}")))
    }

    /// The group's geometry.
    pub fn group(&self, group_id: &str) -> Result<&GroupLayout, LayoutError> {
        self.groups_by_id
            .get(group_id)
            .map(|&index| &self.plan.groups[index])
            .ok_or_else(|| LayoutError::new(format!("unknown cache group {group_id:?}")))
    }

    /// A field's geometry.
    pub fn field(&self, field_id: &str) -> Result<&FieldLayout, LayoutError> {
        self.fields_by_id
            .get(field_id)
            .map(|&index| &self.plan.fields[index])
            .ok_or_else(|| LayoutError::new(format!("unknown cache field {field_id:?}")))
    }

    /// The group's fields, sorted by field id (`fields_for_group`).
    pub fn fields_for_group(&self, group_id: &str) -> Vec<&FieldLayout> {
        self.group_fields
            .get(group_id)
            .map(|indexes| {
                indexes
                    .iter()
                    .map(|&index| &self.plan.fields[index])
                    .collect()
            })
            .unwrap_or_default()
    }

    /// The group's cyclic owner count; one means replicated.
    pub fn shard_count(&self, group_id: &str) -> Result<i64, LayoutError> {
        Ok(self.group_spec(group_id)?.shard_count)
    }

    /// The exclusive bound of the scheduler block ids a manifest of the
    /// group may carry (`virtual_block_count`).
    pub fn virtual_block_count(&self, group_id: &str) -> Result<i64, LayoutError> {
        let pages = self.group(group_id)?.page_count;
        Ok(1 + (pages - 1) * self.shard_count(group_id)?)
    }

    /// A field's TP partition, if it has one.
    pub fn partition_for(&self, field_id: &str) -> Option<&Partition> {
        self.partitions.get(field_id)
    }

    /// One field page's byte offset in the arena (`field_page_byte_offset`).
    pub fn field_page_byte_offset(&self, field_id: &str, page_id: i64) -> Result<i64, LayoutError> {
        let field = self.field(field_id)?;
        let group = self.group(&field.group_id)?;
        if page_id < 0 || page_id >= group.page_count {
            return Err(LayoutError::new(format!(
                "page_id {page_id} outside [0, {}) for group {:?}",
                group.page_count, group.group_id
            )));
        }
        let plane = self
            .planes_by_id
            .get(&field.plane_id)
            .map(|&index| &self.plan.planes[index])
            .ok_or_else(|| LayoutError::new(format!("unknown cache plane {:?}", field.plane_id)))?;
        Ok(
            plane.arena_offset_bytes + plane.bytes_per_lcm_block - field.page_stride_bytes
                + page_id * field.page_stride_bytes
                + field.field_offset_bytes,
        )
    }

    /// Where a field's page 0 lands for an arena at `base`: its advertised
    /// address on an endpoint's contract, else `base` plus its offset.
    pub fn field_address(&self, base: u64, field_id: &str) -> Result<i128, LayoutError> {
        if let Some(addresses) = &self.field_addresses {
            return addresses
                .get(field_id)
                .map(|&address| i128::from(address))
                .ok_or_else(|| LayoutError::new(format!("unknown cache field {field_id:?}")));
        }
        Ok(i128::from(base) + i128::from(self.field_page_byte_offset(field_id, 0)?))
    }
}

/// The field's global shape: its local shape with the partitioned axis at
/// its global extent.
fn global_shape(field: &FieldLayout, partition: Option<&Partition>) -> Vec<i64> {
    let mut shape = field.shape.clone();
    if let Some(partition) = partition {
        if let Some(extent) = usize::try_from(partition.axis)
            .ok()
            .and_then(|axis| shape.get_mut(axis))
        {
            *extent = partition.global_extent;
        }
    }
    shape
}

/// Check that a peer's contract describes the same cache semantically
/// (`validate_cache_peer_layout`): the same prefix granularity, groups in
/// the same order with the same semantics, and the same fields per group
/// with the same dtype, partitioning and global shape. Shard counts may
/// differ: each side shards its groups on its own.
pub fn validate_peer(layout: &Contract, peer: &Contract) -> Result<(), LayoutError> {
    if layout.plan.prefix_granularity != peer.plan.prefix_granularity {
        return Err(LayoutError::new(
            "Cache P/D contract mismatch: prefix_granularity",
        ));
    }
    let local_ids: Vec<&str> = layout
        .group_specs
        .iter()
        .map(|s| s.group_id.as_str())
        .collect();
    let peer_ids: Vec<&str> = peer
        .group_specs
        .iter()
        .map(|s| s.group_id.as_str())
        .collect();
    if local_ids != peer_ids {
        return Err(LayoutError::new("Cache P/D contract mismatch: group order"));
    }
    for (local, remote) in layout.group_specs.iter().zip(&peer.group_specs) {
        if local.attn_kind != remote.attn_kind
            || local.rows_per_page != remote.rows_per_page
            || local.entry_stride_tokens != remote.entry_stride_tokens
            || local.sliding_window_tokens != remote.sliding_window_tokens
            || local.replayable != remote.replayable
            || local.transfer_policy != remote.transfer_policy
        {
            return Err(LayoutError::new(format!(
                "Cache P/D contract mismatch: group {:?} semantics",
                local.group_id
            )));
        }
        let local_fields = layout.fields_for_group(&local.group_id);
        let peer_fields = peer.fields_for_group(&remote.group_id);
        let local_field_ids: Vec<&str> = local_fields.iter().map(|f| f.field_id.as_str()).collect();
        let peer_field_ids: Vec<&str> = peer_fields.iter().map(|f| f.field_id.as_str()).collect();
        if local_field_ids != peer_field_ids {
            return Err(LayoutError::new(format!(
                "Cache P/D contract mismatch: group {:?} transfer field order",
                local.group_id
            )));
        }
        for (local_field, peer_field) in local_fields.iter().zip(&peer_fields) {
            let local_partition = layout.partition_for(&local_field.field_id);
            let peer_partition = peer.partition_for(&peer_field.field_id);
            if local_field.dtype != peer_field.dtype
                || local_partition != peer_partition
                || global_shape(local_field, local_partition)
                    != global_shape(peer_field, peer_partition)
            {
                return Err(LayoutError::new(format!(
                    "Cache P/D contract mismatch: field {:?} semantics",
                    local_field.field_id
                )));
            }
        }
    }
    Ok(())
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;

    pub(crate) fn wire() -> String {
        serde_json::json!({
            "plan": {
                "prefix_granularity": 2,
                "lcm_block_bytes": 32,
                "num_lcm_blocks": 15,
                "groups": [
                    {"group_id": "history", "cache_blocks_per_lcm_block": 1, "page_count": 16},
                    {"group_id": "state", "cache_blocks_per_lcm_block": 1, "page_count": 16}
                ],
                "planes": [
                    {"plane_id": "p0", "bytes_per_lcm_block": 32, "arena_offset_bytes": 0},
                    {"plane_id": "p1", "bytes_per_lcm_block": 32, "arena_offset_bytes": 16}
                ],
                "fields": [
                    {"group_id": "history", "field_id": "layer.0.kv", "plane_id": "p0",
                     "shape": [8], "dtype": "bfloat16", "field_offset_bytes": 0,
                     "page_stride_bytes": 32},
                    {"group_id": "state", "field_id": "layer.1.state", "plane_id": "p1",
                     "shape": [8], "dtype": "bfloat16", "field_offset_bytes": 0,
                     "page_stride_bytes": 32}
                ],
                "resident_block_bytes": null
            },
            "group_specs": [
                {"group_id": "history", "attn_kind": "full_history", "rows_per_page": 2,
                 "entry_stride_tokens": 1, "sliding_window_tokens": null, "replayable": false,
                 "transfer_policy": "full_suffix", "shard_count": 1},
                {"group_id": "state", "attn_kind": "state", "rows_per_page": null,
                 "entry_stride_tokens": null, "sliding_window_tokens": null, "replayable": false,
                 "transfer_policy": "latest_snapshot", "shard_count": 1}
            ],
            "transfer_schema": {"fields": []},
            "field_addresses": null
        })
        .to_string()
    }

    #[test]
    fn a_contract_parses_and_answers_its_geometry() {
        let contract = Contract::from_wire(wire().as_bytes()).unwrap();
        assert_eq!(contract.virtual_block_count("history").unwrap(), 16);
        assert_eq!(
            contract.field_address(0x1000, "layer.1.state").unwrap(),
            0x1000 + 16
        );
        assert_eq!(contract.field("layer.0.kv").unwrap().payload_bytes(), 16);
        assert!(contract.group_spec("history").unwrap().full_suffix());
        assert!(contract
            .group_spec("state")
            .unwrap()
            .block_granularity()
            .is_err());
    }

    #[test]
    fn malformed_contracts_are_refused() {
        let mutate = |edit: fn(&mut Value)| {
            let mut value: Value = serde_json::from_str(&wire()).unwrap();
            edit(&mut value);
            Contract::from_wire(value.to_string().as_bytes())
        };
        assert!(mutate(|v| v["plan"]["fields"][0]["dtype"] = "float7".into()).is_err());
        assert!(mutate(|v| v["plan"]["groups"][0]["extra"] = 1.into()).is_err());
        assert!(mutate(|v| v["group_specs"][1]["rows_per_page"] = 2.into()).is_err());
        assert!(mutate(|v| v["field_addresses"] = serde_json::json!({"layer.0.kv": 1})).is_err());
        assert!(Contract::from_wire(b"").is_err());
        assert!(Contract::from_wire(b"[]").is_err());
    }
}