rhdl_dsig_scalable_core/

slice_merkle.rs

//! A slice's share of the batch Merkle tree: one subtree per batch.
//!
//! # The problem
//!
//! DSig pays for one expensive Ed25519 signature per *batch* of 128 one-time
//! keys, not per key. To make that one signature cover 128 keys, the keys'
//! public commitments (`pk_hash`, one 32-byte hash per key) become the leaves
//! of a Merkle tree: a binary tree where every parent is the hash of its two
//! children. Ed25519 signs only the root. A signature then proves that its key
//! belongs to the batch by carrying the 7 *siblings* on the path from its leaf
//! to the root (the *inclusion proof*); a verifier rehashes up the path and
//! compares with the signed root.
//!
//! The 128 keys of a batch are not made in one place. Each
//! [slice](crate::slice) makes its own share, so each slice can build the
//! part of the tree that sits above its own keys, without talking to any
//! other slice. Only the last few levels need everyone's results; they are
//! built by the shared [cert unit](crate::cert_unit).
//!
//! ```text
//!   one batch, 8-slice build (128 keys, 7 levels)
//!
//!   level 7                          batch root ──▶ Ed25519 (cert unit)
//!   levels 5..7          ┌──── Merkle top: 4 + 2 + 1 hashes (cert unit) ────┐
//!   level 4             r0     r1     r2     r3     r4     r5     r6     r7   subtree roots
//!   levels 1..4        [s0]   [s1]   [s2]   [s3]   [s4]   [s5]   [s6]   [s7]  15 parents each
//!   level 0 (leaves)   0..15 16..31 32..47  ...                        112..127  keys
//!                      └─ this module, one copy per slice ─┘
//! ```
//!
//! # What this module does
//!
//! A batch of 128 keys has a 7-level Merkle tree. Slice `s` owns keys
//! `K·s .. K·s + K` with `K =` [`KEYS_PER_SLICE`] (16, 32 or 64 for 8, 4 or 2
//! slices), which form exactly one subtree of height [`SUBTREE_LEVELS`]. For
//! the 8-slice build:
//!
//! ```text
//!   level 4                         root ──▶ shared Merkle top (3 more levels)
//!   level 3               ┌─────────┴─────────┐
//!   level 2          ┌────┴────┐         ┌────┴────┐
//!   level 1        ┌─┴─┐     ┌─┴─┐     ┌─┴─┐     ┌─┴─┐
//!   level 0 (leaf) j=0 1     2   3  …                 15     pk_hash of each key
//! ```
//!
//! With 4 slices the subtree is one level taller and the certificate unit's
//! top one level shorter; the inclusion proof is always 7 siblings.
//!
//! [`SubtreeBuilder`] receives each finished key's `pk_hash` from the
//! [`slice_factory`](crate::slice_factory), stores it, and as soon as both
//! children of a node exist it queues that parent. Queued parents go into the
//! factory's Merkle slot one at a time (one per 68-cycle period, enough for
//! `K − 1` parents per `K` keys). A finished top-level node leaves as the
//! slice's **subtree root**.
//!
//! # Why there is no BLAKE3 core here
//!
//! A parent is one BLAKE3 compression of `left ‖ right`. The factory's
//! [timetable](crate::slice_schedule) already reserves one free "Merkle slot"
//! on its BLAKE3 lane every 68-cycle key period. A subtree of `K` keys needs
//! only `K − 1` parents per `K` keys, so that one slot is enough, and the
//! slice saves a whole hash core. This module only decides *what* to hash
//! and stores the results.
//!
//! # When a parent is ready
//!
//! Keys finish in order `j = 0, 1, 2, …`. When an odd leaf `j` arrives, its
//! even sibling `j − 1` is already stored, so parent `j / 2` at level 1 can be
//! queued at once. The same rule repeats one level up: a parent result with an
//! odd index means its even sibling finished earlier (jobs are served in
//! order), so the grandparent is queued. When the level-[`SUBTREE_LEVELS`]
//! result arrives, it is the subtree root.
//!
//! ```text
//!   leaf j=0 ─ store
//!   leaf j=1 ─ store, queue (level 1, index 0)
//!   leaf j=2 ─ store
//!   leaf j=3 ─ store, queue (level 1, index 1)
//!   result (1,0) ─ store
//!   result (1,1) ─ store, queue (level 2, index 0)   … and so on up to the root
//! ```
//!
//! For each queued job the builder reads the two children from its node RAM
//! (one cycle each), then offers `left ‖ right` to the factory until the
//! factory takes it. The 16-bit request tag carries `(ring, level, index)` so
//! the result can be stored where it belongs when it comes back.
//!
//! # Storage
//!
//! Every node is kept, because a signature's inclusion proof needs one sibling
//! per level. The signer reads siblings through the separate proof port. Nodes
//! live in a ring of [`NODE_RING_BATCHES`] batches, the same number of batches
//! the key cache holds. The slice's admission limit keeps one whole batch of
//! slack, so a batch's nodes are never overwritten while one of its keys can
//! still sign (see [`slice`](crate::slice)).
//!
//! Node address = `ring · 2K + NODE_BASE[level] + index`, with
//! `NODE_BASE[level] = 2K − 2K / 2^level` (for 16 keys: 0, 16, 24, 28, 30).
//! All `2K − 1` nodes of a batch fit in its `2K` slots, and ring entries ×
//! slots is always 1,024, a 10-bit address.
//!
//! ```text
//!   one ring entry, 16 keys (32 slots of 32 bytes)
//!   slot:  0 ............ 15 | 16 ...... 23 | 24 .. 27 | 28 29 | 30 | 31
//!          level 0 (leaves)  | level 1      | level 2  | lvl 3 | root | spare
//! ```
//!
//! The node RAM exists twice ([`SubtreeBuilder`]'s `job_nodes` and
//! `proof_nodes`) with identical writes. A synchronous RAM has one read port;
//! the builder and the signer read at unrelated times, so each gets its own
//! copy instead of arbitrating for one port.

use crate::slice_factory::{CACHE_SLOTS, KeyDone, SlotRequest, SlotResult};
use crate::{KEY_BITS, KEY_ONE, KEY_ZERO, KEYS_PER_SLICE, KeyIndex, NODE_LEVELS, NODE_SLOT_BITS, SUBTREE_LEVELS};
use rhdl::prelude::*;
use rhdl_fpga::core::{
    dff::DFF,
    ram::synchronous::{In as RamIn, SyncBRAM, Write as RamWrite},
};

/// Batches whose subtree nodes are kept (ring size): one per cache batch.
///
/// 512 cache slots ÷ 16 keys = 32 entries in the 8-slice build. Matching the
/// key cache means a batch's proof nodes live exactly as long as its keys.
pub const NODE_RING_BATCHES: usize = CACHE_SLOTS / KEYS_PER_SLICE;
/// Node slots per ring entry: `2·KEYS_PER_SLICE` (`2K − 1` nodes, one spare).
pub const NODES_PER_ENTRY: usize = 2 * KEYS_PER_SLICE;
/// Mask turning a batch number into its ring entry.
///
/// The ring size is a power of two, so `batch mod NODE_RING_BATCHES` is just
/// the low bits of the batch number (RHDL kernels avoid `%`).
pub const RING_MASK: b5 = b5((NODE_RING_BATCHES - 1) as u128);
/// Level of the subtree root (`log2 KEYS_PER_SLICE`).
pub const ROOT_LEVEL: b3 = b3(SUBTREE_LEVELS as u128);

/// First node slot of each level inside a ring entry:
/// `2K − 2K / 2^level`. For 16 keys: 0, 16, 24, 28, 30.
///
/// [`NODE_LEVELS`] entries (seven in every batch-engine build) so that any
/// level up to [`KEY_BITS`] has a defined base; levels above
/// [`ROOT_LEVEL`] are never used.
pub const NODE_BASE: [NodeSlot; NODE_LEVELS] = node_bases();
/// A node slot within one ring entry ([`NODE_SLOT_BITS`] wide).
pub type NodeSlot = Bits<NODE_SLOT_BITS>;

/// Compute [`NODE_BASE`] at compile time: level `l` starts after all the
/// nodes of levels `0..l` (`2K − 2K >> l`).
const fn node_bases() -> [NodeSlot; NODE_LEVELS] {
    let mut bases = [bits(0); NODE_LEVELS];
    let mut level = 0;
    while level < NODE_LEVELS {
        bases[level] = bits((NODES_PER_ENTRY - (NODES_PER_ENTRY >> level)) as u128);
        level += 1;
    }
    bases
}

// Ring entries × node slots always fill exactly the 10-bit node address.
const _: () = assert!(NODE_RING_BATCHES * NODES_PER_ENTRY == 1 << 10);
// Levels are 3-bit fields, and the index sits at bits 8.. of the 16-bit tag.
const _: () = assert!(NODE_LEVELS <= 8 && KEY_BITS <= 8);
// The ring entry number must fit in the 5-bit `ring` fields.
const _: () = assert!(NODE_RING_BATCHES <= 32);
/// Pending-parent queue depth.
///
/// Parents are produced in short bursts (a leaf can complete a whole chain
/// of ancestors) but drained at one per 68-cycle period. Eight entries absorb
/// the bursts; if the queue ever fills, `overflow` latches as a design error.
pub const JOB_QUEUE: usize = 8;
/// `JOB_QUEUE - 1`, named because RHDL loop bounds must be a literal or a const.
const JOB_SHIFTS: usize = JOB_QUEUE - 1;

/// One queued parent: ring entry, level (`1..=SUBTREE_LEVELS`) and index
/// within the level.
///
/// Together these name a node uniquely, so the job also serves as the
/// request tag ([`job_tag_kernel`]) that comes back with the hash result.
#[derive(Clone, Copy, Debug, PartialEq, Digital, Default)]
pub struct Job {
    /// Ring entry of the batch.
    pub ring: b5,
    /// Tree level of the parent.
    pub level: b3,
    /// Index of the parent within its level.
    pub index: KeyIndex,
}

/// Inputs from the factory and the signer.
#[derive(Clone, Copy, Debug, PartialEq, Digital, Default)]
pub struct SubtreeInput {
    /// Finished keys (leaves).
    pub key_done: KeyDone,
    /// The factory took the offered Merkle request this cycle.
    pub merkle_taken: bool,
    /// Results of the factory's slot requests (non-message ones are ours).
    pub slot_result: SlotResult,
    /// Signer proof read address (`ring << 5 | node`), data next cycle.
    pub proof_address: b10,
}

/// A finished subtree root.
#[derive(Clone, Copy, Debug, PartialEq, Digital, Default)]
pub struct SubtreeRoot {
    /// Root present this cycle.
    pub valid: bool,
    /// Batch number (relative to the run's first batch).
    pub batch: b32,
    /// The level-4 node.
    pub root: [b8; 32],
}

/// Outputs to the factory, the shared Merkle top and the signer.
#[derive(Clone, Copy, Debug, PartialEq, Digital, Default)]
pub struct SubtreeOutput {
    /// Merkle request offered to the factory.
    pub merkle: SlotRequest,
    /// Finished subtree root.
    pub root: SubtreeRoot,
    /// Node read through the proof port.
    pub proof_data: [b8; 32],
    /// The job queue overflowed (a design error; should never assert).
    pub overflow: bool,
}

/// Resettable control.
#[derive(Clone, Copy, Debug, PartialEq, Digital, Default)]
pub struct SubtreeControl {
    /// Pending parents, oldest at index 0.
    pub jobs: [Job; JOB_QUEUE],
    /// Number of pending parents.
    pub count: b4,
    /// Fetch phase of the head job: 0 idle, 1 left read, 2 right read, 3 offered.
    pub phase: b2,
    /// Batch number of each ring entry.
    ///
    /// Recorded when leaf `j = 0` of a batch arrives, so the root can be
    /// labeled with its full 32-bit batch number later.
    pub batches: [b32; NODE_RING_BATCHES],
    /// Sticky overflow.
    pub overflow: bool,
}

/// Resetless datapath: the two children of the job being offered.
#[derive(Clone, Copy, Debug, PartialEq, Digital)]
pub struct SubtreeData {
    /// Left child.
    pub left: [b8; 32],
    /// Right child.
    pub right: [b8; 32],
}

/// The subtree builder.
///
/// Two node RAMs (one per reader) plus a small control register and the two
/// child registers of the job being offered. Timing: a leaf or parent
/// result is written the cycle it arrives; a queued job takes four cycles to
/// fetch its children and then waits for the factory's Merkle slot.
#[derive(Clone, Debug, Synchronous, SynchronousDQ)]
pub struct SubtreeBuilder {
    /// Copy read by the parent jobs.
    job_nodes: SyncBRAM<[b8; 32], 10>,
    /// Copy read by the signer.
    proof_nodes: SyncBRAM<[b8; 32], 10>,
    /// Queue, fetch phase, ring labels and overflow flag (reset to empty).
    control: DFF<SubtreeControl>,
    /// Children of the job being offered (no reset needed).
    data: DFF<SubtreeData>,
}

impl Default for SubtreeBuilder {
    /// Empty queue, idle phase; node RAMs start with unknown contents (every
    /// node is written before anyone reads it).
    fn default() -> Self {
        Self {
            job_nodes: SyncBRAM::default(),
            proof_nodes: SyncBRAM::default(),
            control: DFF::new(SubtreeControl::default()),
            data: DFF::resetless(),
        }
    }
}

/// Input: leaves, parent results and proof reads. Output: Merkle requests,
/// subtree roots and proof data.
impl SynchronousIO for SubtreeBuilder {
    type I = SubtreeInput;
    type O = SubtreeOutput;
    type Kernel = subtree_kernel;
}

/// The ring entry holding batch `batch`'s nodes.
///
/// `batch mod NODE_RING_BATCHES`, done by keeping the low bits.
#[kernel]
pub fn ring_of_kernel(batch: b32) -> b5 {
    batch.resize::<5>() & RING_MASK
}

/// Node address of `(level, index)` in ring entry `ring`:
/// `ring · NODES_PER_ENTRY + NODE_BASE[level] + index`.
///
/// Used by the builder (writes and child reads) and by the signer (sibling
/// reads, through [`sibling_address_kernel`](crate::slice_signer::sibling_address_kernel)).
/// The multiply is a shift because `NODES_PER_ENTRY = 2^(SUBTREE_LEVELS + 1)`.
#[allow(clippy::needless_range_loop)]
#[kernel]
pub fn node_address_kernel(ring: b5, level: b3, index: KeyIndex) -> b10 {
    // Look up NODE_BASE[level]; a loop over constants because RHDL kernels
    // need constant array indices.
    let mut base: NodeSlot = bits(0);
    for l in 0..NODE_LEVELS {
        if level == b3(l as u128) {
            base = NODE_BASE[l];
        }
    }
    (ring.resize::<10>() << ((SUBTREE_LEVELS + 1) as u128))
        | (base + index.resize::<NODE_SLOT_BITS>()).resize::<10>()
}

/// Pack a job into the 16-bit slot tag and back.
///
/// Layout: bits 0..5 ring, 5..8 level, 8.. index ([`KEY_BITS`] wide). The factory returns the
/// tag unchanged with the result; [`subtree_kernel`] unpacks it with shifts.
#[kernel]
pub fn job_tag_kernel(job: Job) -> b16 {
    job.ring.resize::<16>() | (job.level.resize::<16>() << 5) | (job.index.resize::<16>() << 8)
}

/// Words of `left ‖ right`, little-endian.
///
/// The 64-byte block of one Merkle parent hash, as the sixteen 32-bit words
/// the BLAKE3 lane takes.
#[allow(clippy::needless_range_loop)]
#[kernel]
pub fn pair_block_kernel(left: [b8; 32], right: [b8; 32]) -> [b32; 16] {
    let mut bytes = [b8(0); 64];
    for b in 0..32 {
        bytes[b] = left[b];
        bytes[32 + b] = right[b];
    }
    crate::slice_factory::words_from_bytes_kernel(bytes)
}

/// The builder's wiring.
///
/// Each cycle it does three things:
///
/// 1. **Write** at most one node: a new leaf from `key_done`, else a parent
///    result from the factory. Maybe queue the next parent, or emit the root.
/// 2. **Fetch/offer** the head job: phase 0 → read left child, 1 → read right
///    child, 2 → hold both, 3 → offer to the factory until `merkle_taken`.
/// 3. **Update the queue**: pop the offered job when taken, push a new one.
///
/// A leaf and a parent result never arrive in the same cycle: both are
/// results of the factory's single BLAKE3 lane, which returns at most one
/// result per cycle. The `if`/`else if` below still gives the leaf priority.
#[allow(clippy::needless_range_loop)]
#[kernel]
pub fn subtree_kernel(
    clock_reset: ClockReset,
    input: SubtreeInput,
    q: SubtreeBuilderQ,
) -> (SubtreeOutput, SubtreeBuilderD) {
    let resetting = clock_reset.reset.any();
    // Working copies of the registers; written back at the end.
    let mut control = q.control;
    let mut data = q.data;

    // ------------------------------------------------ writes (leaf / parent)
    // The leaf (if any) and the ring entry of its batch.
    let leaf = input.key_done;
    let leaf_ring = ring_of_kernel(leaf.batch);
    // A non-message slot result is one of our parent hashes.
    let parent = input.slot_result.valid && !input.slot_result.message;
    // Unpack the tag (see `job_tag_kernel`) to find where the parent goes.
    let p_ring: b5 = input.slot_result.tag.resize();
    let p_level: b3 = (input.slot_result.tag >> 5).resize();
    let p_index: KeyIndex = (input.slot_result.tag >> 8).resize();
    // The parent node is the first 32 bytes of the 64-byte hash output.
    let p_bytes = crate::slice_factory::bytes_from_words_kernel(input.slot_result.words);
    let mut p_node = [b8(0); 32];
    for b in 0..32 {
        p_node[b] = p_bytes[b];
    }
    // The single node write of this cycle (same for both RAM copies).
    let mut write_address = b10(0);
    let mut write_value = leaf.pk_hash;
    let mut write_enable = false;
    // A new parent to queue (from the leaf or the parent result).
    let mut push = false;
    let mut pushed = Job::default();
    // The subtree root, if this cycle's parent is the top node.
    let mut root = SubtreeRoot::default();
    if leaf.valid {
        write_address = node_address_kernel(leaf_ring, b3(0), leaf.j);
        write_enable = true;
        // First leaf of a batch: remember which batch now owns this entry.
        if leaf.j == KEY_ZERO {
            for e in 0..NODE_RING_BATCHES {
                if leaf_ring == b5(e as u128) {
                    control.batches[e] = leaf.batch;
                }
            }
        }
        // Odd leaf: its even sibling is stored, so the level-1 parent is ready.
        if (leaf.j & KEY_ONE) == KEY_ONE {
            push = true;
            pushed = Job {
                ring: leaf_ring,
                level: b3(1),
                index: leaf.j >> 1,
            };
        }
    } else if parent {
        write_address = node_address_kernel(p_ring, p_level, p_index);
        write_value = p_node;
        write_enable = true;
        if p_level == ROOT_LEVEL {
            // Top of the subtree: label it with the entry's batch and send it.
            let mut batch = b32(0);
            for e in 0..NODE_RING_BATCHES {
                if p_ring == b5(e as u128) {
                    batch = q.control.batches[e];
                }
            }
            root = SubtreeRoot {
                valid: true,
                batch,
                root: p_node,
            };
        } else if (p_index & KEY_ONE) == KEY_ONE {
            // Odd parent: its sibling is done, so queue the next level up.
            push = true;
            pushed = Job {
                ring: p_ring,
                level: p_level + b3(1),
                index: p_index >> 1,
            };
        }
    }

    // ------------------------------------------------ head job fetch/offer
    // The oldest queued parent; children are at level − 1, indices 2i, 2i + 1.
    let head = q.control.jobs[0];
    let mut read_address = b10(0);
    // The head job was taken by the factory this cycle.
    let mut pop = false;
    if q.control.phase == b2(0) && q.control.count != b4(0) {
        // Read the left child.
        read_address = node_address_kernel(head.ring, head.level - b3(1), head.index << 1);
        control.phase = b2(1);
    } else if q.control.phase == b2(1) {
        // Left data arrives; read the right child.
        data.left = q.job_nodes;
        read_address =
            node_address_kernel(head.ring, head.level - b3(1), (head.index << 1) | KEY_ONE);
        control.phase = b2(2);
    } else if q.control.phase == b2(2) {
        // Right data arrives; the request is complete from next cycle on.
        data.right = q.job_nodes;
        control.phase = b2(3);
    } else if q.control.phase == b2(3) && input.merkle_taken {
        control.phase = b2(0);
        pop = true;
    }
    // The request held on the factory's Merkle slot while in phase 3.
    let offered = SlotRequest {
        valid: q.control.phase == b2(3) && !resetting,
        block: pair_block_kernel(q.data.left, q.data.right),
        tag: job_tag_kernel(head),
        tail: false,
        cv: [b32(0); 8],
    };

    // ------------------------------------------------ queue update
    // Pop first (shift everything down), then push at the new tail, so a pop
    // and a push in the same cycle work.
    let mut jobs = q.control.jobs;
    let mut count = q.control.count;
    if pop {
        for k in 0..JOB_SHIFTS {
            jobs[k] = jobs[k + 1];
        }
        count = count - b4(1);
    }
    if push {
        for k in 0..JOB_QUEUE {
            if count == b4(k as u128) {
                jobs[k] = pushed;
            }
        }
        // Full queue: drop the job and latch the fault (should never happen).
        if count == b4(JOB_QUEUE as u128) {
            control.overflow = true;
        } else {
            count = count + b4(1);
        }
    }
    control.jobs = jobs;
    control.count = count;

    // Both RAM copies get the same write; only their read addresses differ.
    let d = SubtreeBuilderD {
        job_nodes: RamIn::<[b8; 32], 10> {
            read_addr: read_address,
            write: RamWrite::<[b8; 32], 10> {
                addr: write_address,
                value: write_value,
                enable: write_enable,
            },
        },
        proof_nodes: RamIn::<[b8; 32], 10> {
            read_addr: input.proof_address,
            write: RamWrite::<[b8; 32], 10> {
                addr: write_address,
                value: write_value,
                enable: write_enable,
            },
        },
        control,
        data,
    };
    (
        SubtreeOutput {
            merkle: offered,
            root,
            proof_data: q.proof_nodes,
            overflow: q.control.overflow,
        },
        d,
    )
}