clay-codes 0.3.0

//! Single-node repair for Clay codes
//!
//! This module implements the repair algorithm from the FAST'18 paper and
//! extends it to work around missing helpers. Write (x0, y0) for the lost
//! node and P(x) for the β planes whose digit y0 is x.
//!
//! - With every helper available, the d helpers send P(x0).
//! - When same-section helpers M are missing, the k other-section helpers
//!   also send P(x) for each x in M.
//! - When only r < k other-section helpers are available, s = k − r
//!   same-section helpers B make up the shortfall. B and the r other-section
//!   helpers send the wide set P(x0), P(B) and P(M).
//!
//! The decode takes P(B) first, then P(x0) and P(M), then the pairwise
//! transforms within the lost section.

use std::ops::Range;

use crate::coords::LayerGeometry;
use crate::decode::{decode_uncoupled_layer, DecodeParams, NodeRows, RsCodec};
use crate::error::ClayError;
use crate::transforms::{compute_c_into, compute_cstar_into, compute_u_into, pft_into, prt_into};

/// Parameters needed for repair (alias to DecodeParams)
pub type RepairParams = DecodeParams;

/// A node's part in one planned repair
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum RepairRole {
    /// The node under repair
    Lost,
    /// A same-section helper that sends the wide plane set
    Wide,
    /// A same-section helper that sends the narrow set P(x0), or a virtual
    /// same-section node
    Narrow,
    /// A same-section node outside the plan
    Missing,
    /// An other-section helper, or a virtual other-section node
    Helper,
    /// An other-section node outside the plan
    Aloof,
}

/// The lost node's y-section and the arithmetic of its plane sets
struct Section {
    geometry: LayerGeometry,
    q: usize,
    beta: usize,
    /// Place value of digit y0
    place: usize,
    x0: usize,
    y0: usize,
    /// Internal index of the lost node
    lost: usize,
}

impl Section {
    fn new(params: &RepairParams, lost_node: usize) -> Self {
        let q = params.q;
        let lost = to_internal(params, lost_node);
        let geometry = LayerGeometry::new(q, params.t, params.sub_chunk_no);
        let place = geometry.q_powers[params.t - 1 - lost / q];
        Section {
            geometry,
            q,
            beta: params.sub_chunk_no / q,
            place,
            x0: lost % q,
            y0: lost / q,
            lost,
        }
    }

    /// Internal index of the section's node at digit x
    fn node(&self, x: usize) -> usize {
        self.y0 * self.q + x
    }

    /// Whether an internal node sits in the section
    fn holds(&self, node: usize) -> bool {
        node / self.q == self.y0
    }

    /// The plane of P(x) at `rank`, counting the planes of P(x) in order
    fn plane(&self, x: usize, rank: usize) -> usize {
        (rank / self.place) * self.q * self.place + x * self.place + rank % self.place
    }

    /// Digit y0 of plane z
    fn digit(&self, z: usize) -> usize {
        z / self.place % self.q
    }

    /// Rank of plane z within P(digit(z)), the inverse of `plane`
    fn rank(&self, z: usize) -> usize {
        z / (self.q * self.place) * self.place + z % self.place
    }
}

/// The plane sets a plan decodes, in order: P(B), P(x0), then P(M)
struct PlaneSets {
    /// Digits in decode order: B, then x0, then M
    digits: Vec<usize>,
    /// How many of the digits belong to B
    wide_count: usize,
    /// Slot of each digit in `digits`
    slot: Vec<usize>,
    /// Index of each digit among `digits` in ascending order, the order a
    /// wide sender sends its planes in
    sent: Vec<usize>,
}

impl PlaneSets {
    fn new(section: &Section, roles: &[RepairRole]) -> Self {
        let mut digits: Vec<usize> = Vec::with_capacity(section.q);
        let mut missing: Vec<usize> = Vec::new();
        for x in 0..section.q {
            match roles[section.node(x)] {
                RepairRole::Wide => digits.push(x),
                RepairRole::Missing => missing.push(x),
                RepairRole::Lost | RepairRole::Narrow | RepairRole::Helper | RepairRole::Aloof => {}
            }
        }
        let wide_count = digits.len();
        digits.push(section.x0);
        digits.extend_from_slice(&missing);

        let mut slot = vec![0; section.q];
        for (index, &x) in digits.iter().enumerate() {
            slot[x] = index;
        }
        let mut ascending = digits.clone();
        ascending.sort_unstable();
        let mut sent = vec![0; section.q];
        for (index, &x) in ascending.iter().enumerate() {
            sent[x] = index;
        }
        PlaneSets {
            digits,
            wide_count,
            slot,
            sent,
        }
    }

    /// Digits of B
    fn wide(&self) -> &[usize] {
        &self.digits[..self.wide_count]
    }

    /// Digits of M
    fn missing(&self) -> &[usize] {
        &self.digits[self.wide_count + 1..]
    }
}

/// The nodes the RS decode solves for, and its passes in score order
struct Erasures {
    /// Whether the decode solves for each node's U
    mask: Vec<bool>,
    count: usize,
    /// Plane ranks of each pass, by intersection score
    passes: Vec<Vec<usize>>,
}

impl Erasures {
    fn new(section: &Section, roles: &[RepairRole]) -> Self {
        let mask: Vec<bool> = roles
            .iter()
            .map(|role| {
                matches!(
                    role,
                    RepairRole::Lost | RepairRole::Narrow | RepairRole::Missing | RepairRole::Aloof
                )
            })
            .collect();
        let count = mask.iter().filter(|&&is_erased| is_erased).count();

        // A helper whose partner is aloof needs the partner's U from a plane
        // that scores one lower, so the passes run in score order
        let aloof: Vec<usize> = (0..roles.len())
            .filter(|&node| roles[node] == RepairRole::Aloof)
            .collect();
        let mut passes = vec![Vec::new(); aloof.len() + 1];
        for rank in 0..section.beta {
            let digits = section.geometry.plane_digits(section.plane(section.x0, rank));
            let score = aloof
                .iter()
                .filter(|&&node| digits[node / section.q] == node % section.q)
                .count();
            passes[score].push(rank);
        }
        Erasures { mask, count, passes }
    }
}

/// A plan from plan_repair, checked and laid out for the decode
struct Layout {
    section: Section,
    roles: Vec<RepairRole>,
    sets: PlaneSets,
    erasures: Erasures,
}

impl Layout {
    /// Check a plan from plan_repair and lay out its planes
    fn from_plan(
        params: &RepairParams,
        lost_node: usize,
        plan: &[(usize, Vec<usize>)],
    ) -> Result<Self, ClayError> {
        let section = Section::new(params, lost_node);
        let narrow = planes_with_digits(params, section.y0, &[section.x0]);
        let (roles, other_count) = assign_roles(params, lost_node, plan, &narrow)?;
        let sets = PlaneSets::new(&section, &roles);

        let wide_count = sets.wide_count;
        let virtual_count = (params.k..params.k + params.nu)
            .filter(|&node| section.holds(node))
            .count();
        let is_fit = if virtual_count == 0 {
            other_count + wide_count == params.k
        } else {
            wide_count == 0 && sets.missing().is_empty() && other_count == params.k + virtual_count
        };
        if !is_fit {
            return Err(ClayError::InvalidParameters(format!(
                "Plan does not fit node {}: {} other-section and {} wide same-section helpers",
                lost_node, other_count, wide_count
            )));
        }

        let wide = planes_with_digits(params, section.y0, &sets.digits);
        for (node, planes) in plan {
            let role = roles[to_internal(params, *node)];
            if matches!(role, RepairRole::Wide | RepairRole::Helper) && *planes != wide {
                return Err(ClayError::InvalidParameters(format!(
                    "Plan helper {} sends the wrong sub-chunks for node {}",
                    node, lost_node
                )));
            }
        }

        let erasures = Erasures::new(&section, &roles);
        Ok(Layout {
            section,
            roles,
            sets,
            erasures,
        })
    }

    /// Where plane z sits in the U rows, in sub-chunks
    fn row(&self, z: usize) -> usize {
        self.sets.slot[self.section.digit(z)] * self.section.beta + self.section.rank(z)
    }

    /// Where plane z sits in a wide sender's sub-chunks
    fn payload(&self, z: usize) -> usize {
        let (q, place) = (self.section.q, self.section.place);
        let width = self.sets.digits.len();
        z / (q * place) * width * place + self.sets.sent[self.section.digit(z)] * place + z % place
    }
}

/// Plan the sub-chunks that repair a lost node from the available nodes
///
/// `available` is in the caller's order of preference. The plan takes every
/// available same-section node and the first k available other-section nodes.
/// With r < k other-section nodes available, the first k − r same-section
/// nodes send the wide plane set. Repeats and the lost node itself are
/// skipped. A lost node whose section holds v virtual nodes needs every real
/// same-section node and the first k + v other-section nodes.
pub fn plan_repair(
    params: &RepairParams,
    lost_node: usize,
    available: &[usize],
) -> Result<Vec<(usize, Vec<usize>)>, ClayError> {
    check_lost(params, lost_node)?;
    let q = params.q;
    let lost = to_internal(params, lost_node);
    let (x0, y0) = (lost % q, lost / q);

    let mut is_available = vec![false; params.n];
    let mut order: Vec<usize> = Vec::with_capacity(available.len());
    let mut other_count = 0;
    for &node in available {
        if node >= params.n {
            return Err(ClayError::InvalidParameters(format!(
                "Available node {} out of range [0, {})",
                node, params.n
            )));
        }
        if node == lost_node || is_available[node] {
            continue;
        }
        is_available[node] = true;
        order.push(node);
        other_count += (to_internal(params, node) / q != y0) as usize;
    }
    if order.len() < params.k {
        return Err(ClayError::InsufficientHelpers {
            needed: params.k,
            provided: order.len(),
        });
    }

    let mut digits = vec![x0];
    let mut first_missing = None;
    let mut virtual_count = 0;
    for x in 0..q {
        match to_external(params, y0 * q + x) {
            None => virtual_count += 1,
            Some(node) if node != lost_node && !is_available[node] => {
                digits.push(x);
                first_missing.get_or_insert(node);
            }
            Some(_) => {}
        }
    }
    let quota = params.k + virtual_count;
    if virtual_count > 0 {
        if let Some(missing_helper) = first_missing {
            return Err(ClayError::MissingYSectionHelper {
                lost_node,
                missing_helper,
            });
        }
        if other_count < quota {
            return Err(ClayError::InsufficientHelpers {
                needed: params.k + q - 1,
                provided: order.len(),
            });
        }
    }

    let mut shortfall = quota - other_count.min(quota);
    let mut is_wide = vec![false; params.n];
    for &node in &order {
        let internal = to_internal(params, node);
        if shortfall > 0 && internal / q == y0 {
            is_wide[node] = true;
            digits.push(internal % q);
            shortfall -= 1;
        }
    }
    digits.sort_unstable();
    let narrow = planes_with_digits(params, y0, &[x0]);
    let wide = planes_with_digits(params, y0, &digits);

    let mut plan = Vec::with_capacity(params.k + q - 1);
    let mut others_taken = 0;
    for &node in &order {
        if to_internal(params, node) / q == y0 {
            let planes = if is_wide[node] { &wide } else { &narrow };
            plan.push((node, planes.clone()));
        } else if others_taken < quota {
            others_taken += 1;
            plan.push((node, wide.clone()));
        }
    }
    Ok(plan)
}

/// Repair a lost chunk from the sub-chunks a plan from `plan_repair` asked for
///
/// `helpers` holds one slot per node. Each planned helper's slot holds its
/// sub-chunks in the plan's order, and the other slots are ignored.
pub fn repair(
    params: &RepairParams,
    rs: &RsCodec,
    lost_node: usize,
    plan: &[(usize, Vec<usize>)],
    helpers: &[Option<&[u8]>],
    chunk_size: usize,
) -> Result<Vec<u8>, ClayError> {
    check_lost(params, lost_node)?;
    if helpers.len() != params.n {
        return Err(ClayError::InvalidParameters(format!(
            "Expected {} helper slots (n), got {}",
            params.n,
            helpers.len()
        )));
    }
    if chunk_size == 0 || !chunk_size.is_multiple_of(params.sub_chunk_no) {
        return Err(ClayError::InvalidChunkSize {
            expected: params.sub_chunk_no,
            actual: chunk_size,
        });
    }
    let sub_chunk_size = chunk_size / params.sub_chunk_no;
    let layout = Layout::from_plan(params, lost_node, plan)?;

    // Virtual nodes send zeros, and a node that sends nothing keeps an empty slice
    let set_count = layout.sets.digits.len();
    let row_len = set_count * layout.section.beta * sub_chunk_size;
    let zero_row = vec![0u8; row_len];
    let mut sources: Vec<&[u8]> = vec![&[]; params.q * params.t];
    for source in sources.iter_mut().take(params.k + params.nu).skip(params.k) {
        *source = &zero_row;
    }
    for (node, planes) in plan {
        let expected = planes.len() * sub_chunk_size;
        let data = helpers[*node].unwrap_or(&[]);
        if data.len() != expected {
            return Err(ClayError::InsufficientHelperData {
                helper: *node,
                expected,
                actual: data.len(),
            });
        }
        sources[to_internal(params, *node)] = data;
    }

    let mut u_buf = NodeRows::new(params.q * params.t, row_len);
    let mut recovered = vec![0u8; chunk_size];
    let wide_count = layout.sets.wide_count;
    if wide_count > 0 {
        couple_wide(&layout, &sources, &mut u_buf, sub_chunk_size);
        decode_sets(params, rs, &layout, &sources, 0..wide_count, &mut u_buf, sub_chunk_size)?;
        recover_wide(&layout, &sources, &mut u_buf, &mut recovered, sub_chunk_size);
    }

    let section_sets = wide_count..set_count;
    decode_sets(params, rs, &layout, &sources, section_sets, &mut u_buf, sub_chunk_size)?;
    recover_section(&layout, &sources, &u_buf, &mut recovered, sub_chunk_size);

    Ok(recovered)
}

/// Each node's role under a plan, and the count of real other-section helpers
fn assign_roles(
    params: &RepairParams,
    lost_node: usize,
    plan: &[(usize, Vec<usize>)],
    narrow: &[usize],
) -> Result<(Vec<RepairRole>, usize), ClayError> {
    let q = params.q;
    let y0 = to_internal(params, lost_node) / q;
    let mut roles: Vec<RepairRole> = Vec::with_capacity(q * params.t);
    for node in 0..q * params.t {
        let is_virtual = node >= params.k && node < params.k + params.nu;
        roles.push(match (node / q == y0, is_virtual) {
            (true, true) => RepairRole::Narrow,
            (true, false) => RepairRole::Missing,
            (false, true) => RepairRole::Helper,
            (false, false) => RepairRole::Aloof,
        });
    }
    roles[to_internal(params, lost_node)] = RepairRole::Lost;

    // Missing marks every same-section node and Aloof every other-section
    // node that no plan entry has claimed yet
    let mut other_count = 0;
    for (node, planes) in plan {
        if *node >= params.n || *node == lost_node {
            return Err(ClayError::InvalidParameters(format!(
                "Plan helper {} is out of range or the lost node",
                node
            )));
        }
        let internal = to_internal(params, *node);
        roles[internal] = match roles[internal] {
            RepairRole::Missing if planes.as_slice() == narrow => RepairRole::Narrow,
            RepairRole::Missing => RepairRole::Wide,
            RepairRole::Aloof => {
                other_count += 1;
                RepairRole::Helper
            }
            RepairRole::Lost | RepairRole::Wide | RepairRole::Narrow | RepairRole::Helper => {
                return Err(ClayError::InvalidParameters(format!(
                    "Plan lists helper {} twice",
                    node
                )))
            }
        };
    }
    Ok((roles, other_count))
}

/// U of the B helpers on P(B), which their own sub-chunks pin down
fn couple_wide(layout: &Layout, sources: &[&[u8]], u_buf: &mut NodeRows, sub_chunk_size: usize) {
    let (section, sc) = (&layout.section, sub_chunk_size);
    let wide = layout.sets.wide();
    for (slot, &x) in wide.iter().enumerate() {
        let node = section.node(x);
        for rank in 0..section.beta {
            let z = section.plane(x, rank);
            let offset = (slot * section.beta + rank) * sc;
            u_buf.row_mut(node)[offset..offset + sc]
                .copy_from_slice(&sources[node][layout.payload(z) * sc..][..sc]);

            // (partner, z) couples with (x, z_partner), and both sent their C
            for (partner_slot, &partner_x) in wide.iter().enumerate().skip(slot + 1) {
                let partner_node = section.node(partner_x);
                let z_partner = section.plane(partner_x, rank);
                let partner_offset = (partner_slot * section.beta + rank) * sc;
                let (u_partner, u_node) = u_buf.pair_mut(partner_node, node);
                prt_into(
                    &sources[partner_node][layout.payload(z) * sc..][..sc],
                    &sources[node][layout.payload(z_partner) * sc..][..sc],
                    &mut u_partner[offset..offset + sc],
                    &mut u_node[partner_offset..partner_offset + sc],
                );
            }
        }
    }
}

/// Decode the plane sets in `slots`, one pass per intersection score
fn decode_sets(
    params: &RepairParams,
    rs: &RsCodec,
    layout: &Layout,
    sources: &[&[u8]],
    slots: Range<usize>,
    u_buf: &mut NodeRows,
    sub_chunk_size: usize,
) -> Result<(), ClayError> {
    let (section, erasures) = (&layout.section, &layout.erasures);
    // Each entry is a plane and its position in the U rows
    let mut planes: Vec<(usize, usize)> = Vec::with_capacity(slots.len() * section.beta);
    for ranks in &erasures.passes {
        if ranks.is_empty() {
            continue;
        }
        planes.clear();
        for slot in slots.clone() {
            for &rank in ranks {
                let z = section.plane(layout.sets.digits[slot], rank);
                couple_plane(layout, sources, z, u_buf, sub_chunk_size);
                planes.push((z, slot * section.beta + rank));
            }
        }

        // RS over GF(2^8) is byte-column independent and every plane shares
        // the erasure set, so each run of adjacent positions is one RS call
        let mut start = 0;
        while start < planes.len() {
            let mut end = start;
            while end + 1 < planes.len() && planes[end + 1].1 == planes[end].1 + 1 {
                end += 1;
            }
            decode_uncoupled_layer(
                params,
                rs,
                &erasures.mask,
                erasures.count,
                planes[start].0,
                planes[start].1 * sub_chunk_size,
                (end - start + 1) * sub_chunk_size,
                u_buf,
            )?;
            start = end + 1;
        }
    }
    Ok(())
}

/// U of every other-section helper on plane z
fn couple_plane(
    layout: &Layout,
    sources: &[&[u8]],
    z: usize,
    u_buf: &mut NodeRows,
    sub_chunk_size: usize,
) {
    let (section, sc) = (&layout.section, sub_chunk_size);
    let digits = section.geometry.plane_digits(z);
    let offset = layout.row(z) * sc;
    let c_offset = layout.payload(z) * sc;
    for (y, &z_y) in digits.iter().enumerate() {
        if y == section.y0 {
            continue;
        }
        for x in 0..section.q {
            let node = y * section.q + x;
            if layout.roles[node] == RepairRole::Aloof {
                continue;
            }
            let c = &sources[node][c_offset..c_offset + sc];
            if z_y == x {
                u_buf.row_mut(node)[offset..offset + sc].copy_from_slice(c);
                continue;
            }

            // A pair of helpers is coupled once, from the side with the larger x
            let partner = y * section.q + z_y;
            let is_aloof_partner = layout.roles[partner] == RepairRole::Aloof;
            if !is_aloof_partner && z_y > x {
                continue;
            }
            let z_partner = section.geometry.companion_layer(z, x, y, z_y);
            let partner_offset = layout.row(z_partner) * sc;
            if is_aloof_partner {
                // An earlier pass decoded the aloof partner's U
                let (u_node, u_partner) = u_buf.pair_mut(node, partner);
                compute_u_into(
                    c,
                    &u_partner[partner_offset..partner_offset + sc],
                    &mut u_node[offset..offset + sc],
                );
            } else {
                let c_partner = &sources[partner][layout.payload(z_partner) * sc..][..sc];
                let (u_node, u_partner) = u_buf.pair_mut(node, partner);
                prt_into(
                    c,
                    c_partner,
                    &mut u_node[offset..offset + sc],
                    &mut u_partner[partner_offset..partner_offset + sc],
                );
            }
        }
    }
}

/// The lost sub-chunks on P(B), then U of the B helpers on P(x0) and P(M)
fn recover_wide(
    layout: &Layout,
    sources: &[&[u8]],
    u_buf: &mut NodeRows,
    recovered: &mut [u8],
    sub_chunk_size: usize,
) {
    let (section, sc) = (&layout.section, sub_chunk_size);
    for &x in layout.sets.wide() {
        let node = section.node(x);
        let source = sources[node];
        for rank in 0..section.beta {
            // (lost, z) couples with (x, z0), whose C x sent
            let z = section.plane(x, rank);
            let z0 = section.plane(section.x0, rank);
            let c_z0 = &source[layout.payload(z0) * sc..][..sc];
            compute_c_into(
                &u_buf.row(section.lost)[layout.row(z) * sc..][..sc],
                c_z0,
                &mut recovered[z * sc..(z + 1) * sc],
            );

            // U = C + γC* is the same map as C = U + γC*
            let offset = layout.row(z0) * sc;
            compute_c_into(
                c_z0,
                &recovered[z * sc..(z + 1) * sc],
                &mut u_buf.row_mut(node)[offset..offset + sc],
            );

            // (x, z_missing) couples with (missing, z), whose U the P(B) decode gave
            for &missing in layout.sets.missing() {
                let z_missing = section.plane(missing, rank);
                let offset = layout.row(z_missing) * sc;
                let (u_node, u_missing) = u_buf.pair_mut(node, section.node(missing));
                compute_u_into(
                    &source[layout.payload(z_missing) * sc..][..sc],
                    &u_missing[layout.row(z) * sc..][..sc],
                    &mut u_node[offset..offset + sc],
                );
            }
        }
    }
}

/// The lost sub-chunks on P(x0), on P(M) and on each narrow helper's planes
fn recover_section(
    layout: &Layout,
    sources: &[&[u8]],
    u_buf: &NodeRows,
    recovered: &mut [u8],
    sub_chunk_size: usize,
) {
    let (section, sc) = (&layout.section, sub_chunk_size);
    let lost_row = u_buf.row(section.lost);
    let mut scratch = vec![0u8; sc];
    for rank in 0..section.beta {
        let z0 = section.plane(section.x0, rank);
        let offset = layout.row(z0) * sc;
        recovered[z0 * sc..(z0 + 1) * sc].copy_from_slice(&lost_row[offset..offset + sc]);

        for x in 0..section.q {
            let node = section.node(x);
            let z = section.plane(x, rank);
            match layout.roles[node] {
                // A narrow helper's C and U at z0 pin down the lost C at z
                RepairRole::Narrow => compute_cstar_into(
                    &sources[node][rank * sc..(rank + 1) * sc],
                    &u_buf.row(node)[offset..offset + sc],
                    &mut recovered[z * sc..(z + 1) * sc],
                ),
                // (lost, z) couples with (missing, z0), and the decodes gave both U
                RepairRole::Missing => pft_into(
                    &lost_row[layout.row(z) * sc..][..sc],
                    &u_buf.row(node)[offset..offset + sc],
                    &mut recovered[z * sc..(z + 1) * sc],
                    &mut scratch,
                ),
                RepairRole::Lost | RepairRole::Wide | RepairRole::Helper | RepairRole::Aloof => {}
            }
        }
    }
}

/// Planes whose digit y is in `digits`, in order
fn planes_with_digits(params: &RepairParams, y: usize, digits: &[usize]) -> Vec<usize> {
    let place = params.sub_chunk_no / params.q.pow(y as u32 + 1);
    let mut is_wanted = vec![false; params.q];
    for &x in digits {
        is_wanted[x] = true;
    }
    let mut planes = Vec::with_capacity(digits.len() * params.sub_chunk_no / params.q);
    for z in 0..params.sub_chunk_no {
        if is_wanted[z / place % params.q] {
            planes.push(z);
        }
    }
    planes
}

/// Refuse a lost node index past n
fn check_lost(params: &RepairParams, lost_node: usize) -> Result<(), ClayError> {
    if lost_node >= params.n {
        return Err(ClayError::InvalidParameters(format!(
            "Invalid lost node index: {} >= {}",
            lost_node, params.n
        )));
    }
    Ok(())
}

/// Internal index of a node: the virtual nodes sit after the data nodes
fn to_internal(params: &RepairParams, node: usize) -> usize {
    if node < params.k {
        node
    } else {
        node + params.nu
    }
}

/// Node index of an internal node, None for a virtual node
fn to_external(params: &RepairParams, internal: usize) -> Option<usize> {
    if internal < params.k {
        Some(internal)
    } else if internal < params.k + params.nu {
        None
    } else {
        Some(internal - params.nu)
    }
}

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

    fn test_params() -> RepairParams {
        RepairParams {
            k: 4,
            m: 2,
            n: 6,
            q: 2,
            t: 3,
            nu: 0,
            sub_chunk_no: 8,
            original_count: 4,
        }
    }

    // a digit set takes beta planes per digit, each carrying a wanted digit
    #[test]
    fn plane_sets() {
        let params = test_params();

        for y in 0..params.t {
            let place = params.sub_chunk_no / params.q.pow(y as u32 + 1);
            for x in 0..params.q {
                let planes = planes_with_digits(&params, y, &[x]);
                assert_eq!(planes.len(), params.sub_chunk_no / params.q);
                assert!(planes.iter().all(|z| z / place % params.q == x));
            }
            let every: Vec<usize> = (0..params.sub_chunk_no).collect();
            assert_eq!(planes_with_digits(&params, y, &[0, 1]), every);
        }
    }

    // with every node available the plan takes d helpers at beta each
    #[test]
    fn all_available() {
        let params = test_params();

        let plan = plan_repair(&params, 0, &[1, 2, 3, 4, 5]).expect("plan");

        assert_eq!(plan.len(), 5);
        assert!(plan.iter().all(|(_, planes)| planes.len() == 4));
        assert!(plan.iter().any(|(node, _)| *node == 1));
    }

    // fewer than k available nodes give no plan
    #[test]
    fn below_k() {
        let params = test_params();

        let result = plan_repair(&params, 0, &[1, 2, 3]);

        assert_eq!(
            result,
            Err(ClayError::InsufficientHelpers {
                needed: 4,
                provided: 3
            })
        );
    }

    // the layout puts P(x0) first when nothing is missing
    #[test]
    fn all_available_layout() {
        let params = test_params();
        let plan = plan_repair(&params, 0, &[1, 2, 3, 4, 5]).expect("plan");

        let layout = Layout::from_plan(&params, 0, &plan).expect("layout");

        assert_eq!(layout.sets.digits, vec![0]);
        assert_eq!(layout.erasures.count, params.m);
        for (rank, &z) in plan[0].1.iter().enumerate() {
            assert_eq!(layout.section.plane(0, rank), z);
            assert_eq!(layout.section.rank(z), rank);
            assert_eq!(layout.row(z), rank);
            assert_eq!(layout.payload(z), rank);
        }
    }

    // a wide sender's sub-chunks run in plane order, and the U rows in decode order
    #[test]
    fn plane_offsets() {
        let params = RepairParams {
            k: 3,
            m: 3,
            n: 6,
            q: 3,
            t: 2,
            nu: 0,
            sub_chunk_no: 9,
            original_count: 3,
        };
        // node 5 is down, so node 2 makes up for it, and P(2) decodes before P(0)
        let plan = plan_repair(&params, 0, &[2, 1, 3, 4]).expect("plan");

        let layout = Layout::from_plan(&params, 0, &plan).expect("layout");

        assert_eq!(layout.sets.digits, vec![2, 0]);
        let wide = planes_with_digits(&params, 0, &[0, 2]);
        assert_eq!(wide, vec![0, 1, 2, 6, 7, 8]);
        let rows: Vec<usize> = wide.iter().map(|&z| layout.row(z)).collect();
        assert_eq!(rows, vec![3, 4, 5, 0, 1, 2]);
        let payloads: Vec<usize> = wide.iter().map(|&z| layout.payload(z)).collect();
        assert_eq!(payloads, vec![0, 1, 2, 3, 4, 5]);
        for &z in &wide {
            assert_eq!(layout.section.plane(layout.section.digit(z), layout.section.rank(z)), z);
        }
    }
}