rsurl 0.1.18

//! Piece bookkeeping: a [`Bitfield`] and a rarest-first [`Picker`].

use std::collections::HashSet;

/// A BitTorrent bitfield (BEP 3): bit 0 is the most-significant bit of byte 0.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Bitfield {
    bits: Vec<u8>,
    len: usize,
}

impl Bitfield {
    pub fn new(len: usize) -> Self {
        Bitfield {
            bits: vec![0u8; len.div_ceil(8)],
            len,
        }
    }

    /// Build from stored bytes (e.g. resume state); extra bytes/bits beyond
    /// `len` are ignored. Use [`from_wire`](Self::from_wire) for a peer's
    /// `bitfield` message, which must be validated instead.
    pub fn from_bytes(bytes: &[u8], len: usize) -> Self {
        let mut bf = Bitfield::new(len);
        let n = bytes.len().min(bf.bits.len());
        bf.bits[..n].copy_from_slice(&bytes[..n]);
        // Clear spare bits so `as_bytes` stays canonical.
        if !len.is_multiple_of(8) {
            if let Some(last) = bf.bits.last_mut() {
                *last &= 0xffu8 << (8 - len % 8);
            }
        }
        bf
    }

    /// Parse a peer's `bitfield` message payload (BEP 3): exactly
    /// `ceil(len / 8)` bytes with every spare trailing bit clear. Returns
    /// `None` for a malformed bitfield, on which the peer must be dropped.
    pub fn from_wire(bytes: &[u8], len: usize) -> Option<Self> {
        if bytes.len() != len.div_ceil(8) {
            return None;
        }
        if !len.is_multiple_of(8) {
            let spare = 0xffu8 >> (len % 8);
            if bytes.last().is_some_and(|&b| b & spare != 0) {
                return None;
            }
        }
        Some(Bitfield {
            bits: bytes.to_vec(),
            len,
        })
    }

    /// Set every bit that is set in `other` (same length).
    pub fn merge(&mut self, other: &Bitfield) {
        for (a, b) in self.bits.iter_mut().zip(&other.bits) {
            *a |= *b;
        }
    }

    pub fn len(&self) -> usize {
        self.len
    }

    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    pub fn has(&self, i: usize) -> bool {
        if i >= self.len {
            return false;
        }
        self.bits[i / 8] & (0x80 >> (i % 8)) != 0
    }

    pub fn set(&mut self, i: usize) {
        if i < self.len {
            self.bits[i / 8] |= 0x80 >> (i % 8);
        }
    }

    pub fn unset(&mut self, i: usize) {
        if i < self.len {
            self.bits[i / 8] &= !(0x80 >> (i % 8));
        }
    }

    /// Spare bits past `len` are always clear (every constructor and setter
    /// keeps them so), so a popcount over the bytes is exact.
    pub fn count(&self) -> usize {
        self.bits.iter().map(|b| b.count_ones() as usize).sum()
    }

    pub fn is_complete(&self) -> bool {
        self.count() == self.len
    }

    pub fn as_bytes(&self) -> &[u8] {
        &self.bits
    }
}

/// Rarest-first piece picker. The engine feeds it each peer's advertised pieces
/// (so it knows availability), and asks it which piece to fetch next.
pub struct Picker {
    availability: Vec<u32>,
}

impl Picker {
    pub fn new(num_pieces: usize) -> Self {
        Picker {
            availability: vec![0; num_pieces],
        }
    }

    pub fn add_have(&mut self, index: usize) {
        if let Some(c) = self.availability.get_mut(index) {
            *c += 1;
        }
    }

    pub fn add_bitfield(&mut self, bf: &Bitfield) {
        for i in 0..bf.len() {
            if bf.has(i) {
                self.add_have(i);
            }
        }
    }

    /// A peer left; decrement availability for the pieces it had.
    pub fn remove_bitfield(&mut self, bf: &Bitfield) {
        for i in 0..bf.len() {
            if bf.has(i) {
                if let Some(c) = self.availability.get_mut(i) {
                    *c = c.saturating_sub(1);
                }
            }
        }
    }

    /// Pick the rarest piece we lack (`!ours`) that `peer_has`, excluding pieces
    /// already being fetched (`exclude`). Ties break toward the lowest index.
    pub fn pick(
        &self,
        ours: &Bitfield,
        peer_has: &Bitfield,
        exclude: &HashSet<usize>,
    ) -> Option<usize> {
        let mut best: Option<(u32, usize)> = None;
        for i in 0..self.availability.len() {
            if ours.has(i) || !peer_has.has(i) || exclude.contains(&i) {
                continue;
            }
            let avail = self.availability[i];
            if best.is_none_or(|(b, _)| avail < b) {
                best = Some((avail, i));
            }
        }
        best.map(|(_, i)| i)
    }
}

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

    #[test]
    fn bitfield_msb_first() {
        let mut bf = Bitfield::new(10);
        assert!(!bf.has(0));
        bf.set(0);
        bf.set(9);
        assert!(bf.has(0));
        assert!(bf.has(9));
        assert!(!bf.has(5));
        assert_eq!(bf.count(), 2);
        // piece 0 => high bit of byte 0; piece 9 => bit 1 of byte 1.
        assert_eq!(bf.as_bytes(), &[0x80, 0x40]);
        assert!(!bf.is_complete());
    }

    #[test]
    fn from_bytes_round_trips() {
        let bf = Bitfield::from_bytes(&[0b1010_0000], 3);
        assert!(bf.has(0));
        assert!(!bf.has(1));
        assert!(bf.has(2));
    }

    #[test]
    fn from_wire_validates_length_and_spare_bits() {
        // 10 pieces => exactly 2 bytes, low 6 bits of byte 1 are spare.
        assert!(Bitfield::from_wire(&[0xff, 0xc0], 10).is_some());
        assert!(Bitfield::from_wire(&[0xff], 10).is_none(), "too short");
        assert!(
            Bitfield::from_wire(&[0xff, 0xc0, 0], 10).is_none(),
            "too long"
        );
        assert!(
            Bitfield::from_wire(&[0xff, 0xc1], 10).is_none(),
            "spare bit set"
        );
        // A multiple of 8 has no spare bits.
        assert!(Bitfield::from_wire(&[0xff], 8).is_some());
        assert!(Bitfield::from_wire(&[], 0).is_some());
    }

    #[test]
    fn merge_and_unset() {
        let mut a = Bitfield::new(10);
        a.set(1);
        let mut b = Bitfield::new(10);
        b.set(9);
        a.merge(&b);
        assert!(a.has(1) && a.has(9));
        a.unset(1);
        assert!(!a.has(1) && a.has(9));
    }

    #[test]
    fn from_bytes_clears_spare_bits() {
        let bf = Bitfield::from_bytes(&[0xff, 0xff], 10);
        assert_eq!(bf.as_bytes(), &[0xff, 0xc0]);
        assert_eq!(bf.count(), 10);
    }

    #[test]
    fn picker_prefers_rarest() {
        let mut p = Picker::new(4);
        let ours = Bitfield::new(4);
        // availability: piece0=3, piece1=1, piece2=2, piece3=0(peer lacks)
        let mut a = Bitfield::new(4);
        a.set(0);
        a.set(1);
        a.set(2);
        p.add_bitfield(&a); // all three +1
        p.add_have(0);
        p.add_have(0);
        p.add_have(2);
        let mut peer = Bitfield::new(4);
        peer.set(0);
        peer.set(1);
        peer.set(2);
        let exclude = HashSet::new();
        // rarest among {0:3,1:1,2:2} is piece 1
        assert_eq!(p.pick(&ours, &peer, &exclude), Some(1));
        // exclude piece 1 → next rarest is piece 2
        let exclude: HashSet<usize> = [1].into_iter().collect();
        assert_eq!(p.pick(&ours, &peer, &exclude), Some(2));
    }

    #[test]
    fn picker_skips_owned_and_unavailable() {
        let p = Picker::new(3);
        let mut ours = Bitfield::new(3);
        ours.set(0);
        let mut peer = Bitfield::new(3);
        peer.set(0); // we already have it
        peer.set(1); // peer has it, we lack it
                     // piece 2 nobody has
        assert_eq!(p.pick(&ours, &peer, &HashSet::new()), Some(1));
    }
}