typebit 0.1.8

//! Per-torrent piece/block state tracking: pieces of `piece_length` bytes
//! (last possibly shorter), each split into 16 KiB blocks. Tracks verified,
//! partially-received, and in-flight block state.

use crate::bitfield::Bitfield;
use crate::consts::BLOCK_LEN;
use crate::error::Result;
use alloc::collections::BTreeMap;
use alloc::vec::Vec;

/// Exact received-byte accounting for one assembling piece.
///
/// Standard block bitmaps cannot express partial or misaligned blocks: an
/// 8 KiB block would mark a 16 KiB slot "received", leaving the missing
/// half unrequestable (permanently stuck) or — worse — letting a zero-filled
/// hole reach the hasher, which then fails verification forever. This
/// tracker records the exact set of non-overlapping byte ranges actually
/// copied into the piece buffer, so a piece is "data-complete" iff every
/// byte is present. Real-world clients that send tail-remainder or
/// odd-sized blocks are absorbed safely instead of poisoning the piece.
#[derive(Debug, Clone, Default)]
pub struct PieceBytes {
    /// Sorted, non-overlapping received ranges `[start, end)`.
    ranges: Vec<(u32, u32)>,
    /// Total distinct bytes received.
    total: u64,
}

impl PieceBytes {
    /// Create an empty tracker.
    pub fn new() -> Self {
        Self::default()
    }

    /// Add the byte window `[start, end)`; returns the number of bytes that
    /// were previously missing (0 = fully duplicate). The window is merged
    /// into the sorted range list.
    pub fn add(&mut self, start: u32, end: u32) -> u64 {
        if end <= start {
            return 0;
        }
        let mut new_bytes = (end - start) as u64;
        // Subtract bytes already covered by existing ranges (ranges are
        // sorted and non-overlapping, so this is exact).
        for &(rs, re) in &self.ranges {
            if re <= start {
                continue;
            }
            if rs >= end {
                break;
            }
            let ov_s = core::cmp::max(start, rs);
            let ov_e = core::cmp::min(end, re);
            new_bytes = new_bytes.saturating_sub((ov_e - ov_s) as u64);
        }
        if new_bytes == 0 {
            return 0;
        }
        self.ranges.push((start, end));
        self.ranges.sort_unstable();
        let mut merged: Vec<(u32, u32)> = Vec::with_capacity(self.ranges.len());
        for (s, e) in self.ranges.drain(..) {
            if let Some(last) = merged.last_mut() {
                if s <= last.1 {
                    last.1 = core::cmp::max(last.1, e);
                    continue;
                }
            }
            merged.push((s, e));
        }
        self.ranges = merged;
        self.total = self.total.saturating_add(new_bytes);
        new_bytes
    }

    /// Distinct bytes present.
    pub fn total(&self) -> u64 {
        self.total
    }

    /// Whether the window `[start, end)` is fully covered by received bytes.
    pub fn covers(&self, start: u32, end: u32) -> bool {
        for &(rs, re) in &self.ranges {
            if rs <= start && re >= end {
                return true;
            }
            if rs > start {
                break;
            }
        }
        false
    }

    /// The sub-windows of `[start, end)` that are **not** yet covered.
    /// Used so an arriving block only ever writes previously-missing bytes —
    /// a hostile peer must not be able to rewrite bytes another peer already
    /// supplied (that would corrupt a good piece and frame an innocent
    /// supplier).
    pub fn new_ranges(&self, start: u32, end: u32) -> Vec<(u32, u32)> {
        let mut out: Vec<(u32, u32)> = Vec::new();
        let mut cur = start;
        for &(rs, re) in &self.ranges {
            if re <= cur {
                continue;
            }
            if rs >= end {
                break;
            }
            if rs > cur {
                out.push((cur, core::cmp::min(rs, end)));
            }
            cur = core::cmp::max(cur, core::cmp::min(re, end));
            if cur >= end {
                break;
            }
        }
        if cur < end {
            out.push((cur, end));
        }
        out
    }
}

/// State of a single piece.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PieceState {
    /// Nothing received.
    Missing,
    /// Some blocks received, not yet verified.
    Partial,
    /// Fully received and hash-verified.
    Have,
}

/// Tracks pieces and their blocks for one torrent.
#[derive(Debug, Clone)]
pub struct PieceTracker {
    piece_count: u32,
    piece_length: u32,
    /// Verified-complete pieces.
    have: Bitfield,
    /// Pieces with at least one block requested or received (in-flight).
    in_flight: Bitfield,
    /// Received blocks per partial piece (bitmap).
    partial: BTreeMap<u32, Bitfield>,
    /// Blocks requested from peers but not yet received (in-flight).
    requested: BTreeMap<u32, Bitfield>,
    /// Total blocks received (for stats).
    blocks_received_total: u64,
}

impl PieceTracker {
    /// Create with `piece_count` pieces of `piece_length` bytes.
    pub fn new(piece_count: u32, piece_length: u32) -> Self {
        PieceTracker {
            piece_count,
            piece_length,
            have: Bitfield::new(piece_count),
            in_flight: Bitfield::new(piece_count),
            partial: BTreeMap::new(),
            requested: BTreeMap::new(),
            blocks_received_total: 0,
        }
    }

    /// Number of pieces.
    pub fn piece_count(&self) -> u32 {
        self.piece_count
    }

    /// Piece length (bytes).
    pub fn piece_length(&self) -> u32 {
        self.piece_length
    }

    /// Number of blocks in a given piece (last piece may be shorter).
    pub fn block_count(&self, piece: u32) -> u16 {
        if piece >= self.piece_count {
            return 0;
        }
        let len = if piece + 1 == self.piece_count && self.piece_length > 0 {
            // caller must know the true last-piece size; we use piece_length
            // as an upper bound — engines override via `set_last_piece_len`.
            self.piece_length
        } else {
            self.piece_length
        };
        (len as u64).div_ceil(BLOCK_LEN as u64) as u16
    }

    /// Whether a piece is fully verified.
    pub fn is_have(&self, piece: u32) -> bool {
        self.have.get(piece)
    }

    /// Whether a piece is in-flight (blocks requested).
    pub fn is_in_flight(&self, piece: u32) -> bool {
        self.in_flight.get(piece)
    }

    /// Set in-flight flag.
    pub fn set_in_flight(&mut self, piece: u32, on: bool) {
        if on {
            self.in_flight.set(piece);
        } else {
            self.in_flight.clear(piece);
        }
    }

    /// Per-piece state.
    pub fn state(&self, piece: u32) -> PieceState {
        if self.have.get(piece) {
            PieceState::Have
        } else if self.partial.contains_key(&piece) {
            PieceState::Partial
        } else {
            PieceState::Missing
        }
    }

    /// Number of received blocks in a piece.
    pub fn blocks_received(&self, piece: u32) -> u16 {
        self.partial
            .get(&piece)
            .map(|b| b.count() as u16)
            .unwrap_or(0)
    }

    /// Whether a block has been received.
    pub fn block_received(&self, piece: u32, block: u16) -> bool {
        self.partial
            .get(&piece)
            .map(|b| b.get(block as u32))
            .unwrap_or(false)
    }

    /// Whether a block is currently requested (in-flight).
    pub fn block_requested(&self, piece: u32, block: u16) -> bool {
        self.requested
            .get(&piece)
            .map(|b| b.get(block as u32))
            .unwrap_or(false)
    }

    /// Mark a block as requested (in-flight).
    pub fn mark_block_requested(&mut self, piece: u32, block: u16, total_blocks: u16) {
        if piece >= self.piece_count || block >= total_blocks {
            return;
        }
        let entry = self
            .requested
            .entry(piece)
            .or_insert_with(|| Bitfield::new(total_blocks as u32));
        entry.set(block as u32);
        self.in_flight.set(piece);
    }

    /// Clear a block's in-flight flag (block arrived, was cancelled, or the
    /// request failed).
    pub fn clear_block_requested(&mut self, piece: u32, block: u16, total_blocks: u16) {
        if piece >= self.piece_count || block >= total_blocks {
            return;
        }
        if let Some(entry) = self.requested.get_mut(&piece) {
            entry.clear(block as u32);
            if entry.count() == 0 {
                self.requested.remove(&piece);
                self.in_flight.clear(piece);
            }
        }
    }

    /// Clear all in-flight flags for a piece (e.g. peer disconnected).
    pub fn clear_piece_requests(&mut self, piece: u32) {
        self.requested.remove(&piece);
        self.in_flight.clear(piece);
    }

    /// Mark a block as received. Returns true if it was newly received.
    pub fn mark_block_received(&mut self, piece: u32, block: u16, total_blocks: u16) -> bool {
        if piece >= self.piece_count || block >= total_blocks {
            return false;
        }
        let entry = self
            .partial
            .entry(piece)
            .or_insert_with(|| Bitfield::new(total_blocks as u32));
        if entry.get(block as u32) {
            return false;
        }
        entry.set(block as u32);
        self.blocks_received_total += 1;
        true
    }

    /// Whether all blocks of a piece have arrived (not yet hash-verified).
    pub fn piece_data_complete(&self, piece: u32, total_blocks: u16) -> bool {
        self.partial
            .get(&piece)
            .map(|b| b.count() == total_blocks as u32)
            .unwrap_or(false)
    }

    /// Mark a piece as verified-complete; clears partial bookkeeping.
    pub fn mark_piece_have(&mut self, piece: u32) {
        self.have.set(piece);
        self.in_flight.clear(piece);
        self.partial.remove(&piece);
        self.requested.remove(&piece);
    }

    /// Mark a piece as missing (e.g. after a hash failure reset).
    pub fn reset_piece(&mut self, piece: u32) {
        self.in_flight.clear(piece);
        self.partial.remove(&piece);
        self.requested.remove(&piece);
    }

    /// Count of verified pieces.
    pub fn have_count(&self) -> u32 {
        self.have.count()
    }

    /// Verified piece bitfield (network bytes).
    pub fn have_bitfield(&self) -> &Bitfield {
        &self.have
    }

    /// Total bytes downloaded (blocks received × block len).
    pub fn downloaded_bytes(&self) -> u64 {
        self.blocks_received_total * BLOCK_LEN as u64
    }

    /// Number of pieces that still need blocks (missing or partial).
    pub fn outstanding_pieces(&self) -> u32 {
        self.piece_count - self.have.count()
    }

    /// Restore `have` from network bytes.
    pub fn set_have_from_bytes(&mut self, bytes: &[u8], count: u32) -> Result<()> {
        let mut b = Bitfield::new(0);
        b.from_bytes(bytes, count)?;
        self.have = b;
        Ok(())
    }

    /// Iterate pieces that are neither have nor in-flight (candidates).
    pub fn candidates(&self) -> CandidateIter<'_> {
        CandidateIter {
            inner: self,
            next: 0,
        }
    }

    /// Serialize `have` + partial state for persistence.
    pub fn snapshot(&self) -> (Vec<u8>, Vec<(u32, Vec<u8>)>) {
        let have = self.have.to_bytes();
        let partial = self
            .partial
            .iter()
            .map(|(p, b)| (*p, b.to_bytes()))
            .collect();
        (have, partial)
    }

    /// Restore from a snapshot.
    pub fn restore(&mut self, have: &[u8], partial: &[(u32, Vec<u8>)]) -> Result<()> {
        self.have.from_bytes(have, self.piece_count)?;
        self.partial.clear();
        for (p, bytes) in partial {
            if *p >= self.piece_count {
                continue;
            }
            let total = self.block_count(*p);
            let mut b = Bitfield::new(0);
            b.from_bytes(bytes, total as u32)?;
            self.partial.insert(*p, b);
        }
        self.blocks_received_total =
            self.partial.values().map(|b| b.count() as u64).sum::<u64>() * BLOCK_LEN as u64;
        Ok(())
    }
}

/// Iterator over pieces that need downloading.
pub struct CandidateIter<'a> {
    inner: &'a PieceTracker,
    next: u32,
}

impl<'a> Iterator for CandidateIter<'a> {
    type Item = u32;
    fn next(&mut self) -> Option<u32> {
        while self.next < self.inner.piece_count {
            let i = self.next;
            self.next += 1;
            if !self.inner.have.get(i) && !self.inner.in_flight.get(i) {
                return Some(i);
            }
        }
        None
    }
}

/// Compute the number of 16 KiB blocks in a piece of `piece_len` bytes.
pub fn block_count_for(piece_len: u32) -> u16 {
    (piece_len as u64).div_ceil(BLOCK_LEN as u64) as u16
}

/// Given a piece length and its data size, the block index containing
/// `offset` and its length.
pub fn block_at(piece_len: u32, offset: u32, block_len: u32) -> (u16, u32) {
    let idx = offset / block_len;
    let start = idx * block_len;
    let len = core::cmp::min(block_len, piece_len - start);
    (idx as u16, len)
}

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

    #[test]
    fn block_accounting() {
        let mut t = PieceTracker::new(4, 32 * 1024);
        assert_eq!(t.block_count(0), 2);
        assert!(t.mark_block_received(0, 0, 2)); // first time → new
        assert!(!t.mark_block_received(0, 0, 2)); // duplicate
        assert!(t.mark_block_received(0, 1, 2));
        assert!(t.piece_data_complete(0, 2));
        assert!(!t.piece_data_complete(1, 2));
        t.mark_piece_have(0);
        assert!(t.is_have(0));
        assert_eq!(t.state(1), PieceState::Missing);
        assert_eq!(t.have_count(), 1);
        assert_eq!(t.outstanding_pieces(), 3);
    }

    #[test]
    fn snapshot_roundtrip() {
        let mut t = PieceTracker::new(8, 16 * 1024);
        t.mark_block_received(2, 0, 1);
        t.mark_piece_have(5);
        let (have, partial) = t.snapshot();
        let mut t2 = PieceTracker::new(8, 16 * 1024);
        t2.restore(&have, &partial).unwrap();
        assert_eq!(t.have_bitfield(), t2.have_bitfield());
        assert_eq!(t.blocks_received(2), t2.blocks_received(2));
    }

    #[test]
    fn candidate_iteration() {
        let mut t = PieceTracker::new(5, 16 * 1024);
        t.mark_piece_have(0);
        t.set_in_flight(1, true);
        let cands: Vec<u32> = t.candidates().collect();
        assert_eq!(cands, vec![2, 3, 4]);
    }

    #[test]
    fn piece_bytes_counts_only_new_bytes() {
        let mut p = PieceBytes::new();
        // First block: full 16 KiB.
        assert_eq!(p.add(0, 16 * 1024), 16 * 1024);
        assert_eq!(p.total(), 16 * 1024);
        // Exact duplicate → nothing new.
        assert_eq!(p.add(0, 16 * 1024), 0);
        assert_eq!(p.total(), 16 * 1024);
        // Adjoining second block → 16 KiB new, ranges merge into one.
        assert_eq!(p.add(16 * 1024, 32 * 1024), 16 * 1024);
        assert_eq!(p.total(), 32 * 1024);
        assert!(p.covers(0, 32 * 1024));
        assert!(!p.covers(32 * 1024, 48 * 1024));
    }

    #[test]
    fn piece_bytes_handles_partial_and_misaligned_blocks() {
        let mut p = PieceBytes::new();
        // An 8 KiB half-block (short block / tail remainder style).
        assert_eq!(p.add(0, 8 * 1024), 8 * 1024);
        assert_eq!(p.total(), 8 * 1024);
        assert!(!p.covers(0, 16 * 1024), "half block is not a full block");
        // The missing half arrives later → counted, never double-counted.
        assert_eq!(p.add(8 * 1024, 16 * 1024), 8 * 1024);
        assert_eq!(p.total(), 16 * 1024);
        assert!(p.covers(0, 16 * 1024));
        // Misaligned tail block overlapping already-received bytes.
        assert_eq!(p.add(12 * 1024, 20 * 1024), 4 * 1024);
        assert_eq!(p.total(), 20 * 1024);
        assert!(p.covers(12 * 1024, 20 * 1024));
        // Whole-piece coverage after a final fragment.
        assert_eq!(p.add(20 * 1024, 32 * 1024), 12 * 1024);
        assert_eq!(p.total(), 32 * 1024);
        assert!(p.covers(0, 32 * 1024));
    }
}