<commit_before>use std::collections::HashMap;
use std::{mem, time};
use std::sync::Arc;
use torrent::{Bitfield, Info, Peer};
use control::cio;

mod rarest;
mod sequential;

#[cfg(test)]
mod tests;

#[derive(Clone, Debug)]
pub struct Picker {
    /// Number of blocks per piece
    scale: u32,
    last_piece_scale: u32,
    last_piece: u32,
    /// Number of detected seeders
    seeders: u16,
    /// Set of pieces which have blocks waiting. These should be prioritized.
    downloading: HashMap<u32, Vec<Downloading>>,
    /// Bitfield of unpicked pieces, not in progress or
    /// completed yet. A set bit is picked, unset is unpicked.
    unpicked: Bitfield,
    /// The current picker in use
    picker: PickerKind,
    /// Piece priorities
    priorities: Vec<u8>,
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct Block {
    pub index: u32,
    pub offset: u32,
}

/// Pickers act solely as piece picking algorithm.
/// They will select the optimal next piece for a peer,
/// and can be told when a piece is complete(or invalid).
#[derive(Clone, Debug)]
enum PickerKind {
    Rarest(rarest::Picker),
    Sequential(sequential::Picker),
}

/// A downloading block and the peers it has been
/// requested from.
#[derive(Clone, Debug)]
struct Downloading {
    offset: u32,
    completed: bool,
    requested: Vec<Request>,
}

/// A request to a peer and the time it was initiated.
#[derive(Clone, Debug)]
struct Request {
    peer: usize,
    requested_at: time::Instant,
}

const MAX_DUP_REQS: usize = 3;
const MAX_DL_Q: usize = 100;
const MAX_PC_SIZE: usize = 50;
const REQ_TIMEOUT: u64 = 10;

impl Picker {
    /// Creates a new picker, which will select over
    /// the given pieces. The algorithm used for selection
    /// will vary based on the current swarm state, but
    /// will default to rarest first.
    pub fn new(info: &Arc<Info>, pieces: &Bitfield, priorities: &[u8]) -> Picker {
        let scale = info.piece_len / 16_384;
        let picker = rarest::Picker::new(pieces);
        let last_piece = info.pieces().saturating_sub(1);
        let lpl = info.piece_len(last_piece);
        let last_piece_scale = if lpl % 16_384 == 0 {
            lpl / 16_384
        } else {
            lpl / 16_384 + 1
        };
        let downloading = if pieces.complete() {
            HashMap::with_capacity(0)
        } else {
            HashMap::with_capacity(info.pieces() as usize)
        };
        let mut picker = Picker {
            picker: PickerKind::Rarest(picker),
            scale,
            last_piece,
            last_piece_scale,
            seeders: 0,
            unpicked: pieces.clone(),
            downloading,
            priorities: vec![3; info.pieces() as usize],
        };
        picker.set_priorities(priorities, info);
        picker
    }

    /// Returns true if the current picker algorithm is sequential
    pub fn is_sequential(&self) -> bool {
        match self.picker {
            PickerKind::Sequential(_) => true,
            _ => false,
        }
    }

    pub fn tick(&mut self) {
        let mut expired = 0;
        for chunks in self.downloading.values_mut() {
            for chunk in chunks.iter_mut() {
                if !chunk.completed {
                    chunk.requested.retain(|req| {
                        let unexp = req.requested_at.elapsed().as_secs() < REQ_TIMEOUT;
                        if !unexp {
                            expired += 1;
                        }
                        unexp
                    })
                }
            }
        }
        if expired != 0 {
            debug!("Expired {} chunks!", expired);
        }
    }

    /// Attempts to select a block for a peer.
    pub fn pick<T: cio::CIO>(&mut self, peer: &mut Peer<T>) -> Option<Block> {
        if let Some(b) = self.pick_expired(peer) {
            return Some(b);
        }

        let piece = match self.picker {
            PickerKind::Sequential(ref mut p) => p.pick(peer),
            PickerKind::Rarest(ref mut p) => p.pick(peer),
        };
        let res = piece
            .and_then(|p| self.pick_piece(p, peer.id()))
            .or_else(|| self.pick_dl(peer));
        if res.is_none() && peer.pieces().complete() && !self.unpicked.complete() {
            debug_assert!(
                false,
                "Couldnt pick {:?} from a seeder for some reason!",
                piece
            );
        }
        res
    }

    /// Attempts to pick an expired block
    fn pick_expired<T: cio::CIO>(&mut self, _: &Peer<T>) -> Option<Block> {
        // TODO: Use some form of heuristic here to say "we expect to have
        // downloaded some pieces by X, hit the picker with a tick which checks
        // that, flags shit as invalid, and then does a double request
        None
    }

    /// Picks a block from a given piece for a peer
    fn pick_piece(&mut self, piece: u32, id: usize) -> Option<Block> {
        self.downloading.entry(piece).or_insert_with(|| vec![]);
        let dl = self.downloading.get_mut(&piece).unwrap();
        let offset = dl.len() as u32 * 16_384;
        dl.push(Downloading {
            offset,
            completed: false,
            requested: vec![Request::new(id)],
        });

        if dl.len() == self.scale as usize
            || (piece == self.last_piece && dl.len() == self.last_piece_scale as usize)
        {
            match self.picker {
                PickerKind::Sequential(ref mut p) => p.completed(piece),
                PickerKind::Rarest(ref mut p) => p.completed(piece),
            }
            self.unpicked.set_bit(u64::from(piece));
        }
        Some(Block {
            index: piece,
            offset,
        })
    }

    /// Attempts to pick the highest priority piece in the dl q
    fn pick_dl<T: cio::CIO>(&mut self, peer: &Peer<T>) -> Option<Block> {
        for (idx, dl) in self.downloading.iter_mut().take(MAX_DL_Q) {
            if peer.pieces().has_bit(u64::from(*idx)) {
                let r = dl.iter_mut()
                    .find(|r| {
                        !r.completed && r.requested.len() < MAX_DUP_REQS
                            && r.requested.iter().all(|req| req.peer != peer.id())
                    })
                    .map(|r| {
                        r.requested.push(Request::new(peer.id()));
                        Block::new(*idx, r.offset)
                    });
                if r.is_some() {
                    return r;
                }
            }
        }
        None
    }

    /// Marks a block as completed. Returns a result indicating if the block
    /// was actually requested, the success value containing a bool indicating
    /// if the block is complete, and a vector of peers from which the block
    /// was requested(for cancellation).
    pub fn completed(&mut self, b: Block) -> Result<(bool, Vec<usize>), ()> {
        // Find the block in our downloading blocks, mark as true,
        // and extract the current peer list for return.
        let res = self.downloading
            .get_mut(&b.index)
            .and_then(|dl| {
                dl.iter_mut()
                    .find(|r| r.offset == b.offset)
                    .map(|r| r.complete())
            })
            .map(|r| r.into_iter().map(|e| e.peer).collect());

        // If we've requested every single block for this piece and they're all complete, remove it
        // and report completion
        let scale = self.scale;
        let lp = self.last_piece;
        let lps = self.last_piece_scale;
        let complete = self.downloading
            .get_mut(&b.index)
            .map(|r| {
                (r.len() as u32 == scale || (b.index == lp && r.len() as u32 == lps))
                    && r.iter().all(|d| d.completed)
            })
            .unwrap_or(false);

        if complete {
            self.downloading.remove(&b.index);
        }

        res.map(|r| (complete, r)).ok_or(())
    }

    pub fn have_block(&mut self, b: Block) -> bool {
        self.downloading
            .get_mut(&b.index)
            .and_then(|dl| {
                dl.iter()
                    .find(|r| r.offset == b.offset)
                    .map(|r| r.completed)
            })
            .unwrap_or(false)
    }

    /// Invalidates a piece
    pub fn invalidate_piece(&mut self, idx: u32) {
        match self.picker {
            PickerKind::Sequential(ref mut p) => p.incomplete(idx),
            PickerKind::Rarest(ref mut p) => p.incomplete(idx),
        }
        self.unpicked.unset_bit(u64::from(idx));
        self.downloading.remove(&idx);
    }

    pub fn piece_available(&mut self, idx: u32) {
        if let PickerKind::Rarest(ref mut p) = self.picker {
            p.piece_available(idx);
        }
    }

    pub fn add_peer<T: cio::CIO>(&mut self, peer: &Peer<T>) {
        if peer.pieces().complete() {
            self.seeders += 1;
        } else {
            if let PickerKind::Rarest(ref mut p) = self.picker {
                p.add_peer(peer);
            }
        }
    }

    pub fn remove_peer<T: cio::CIO>(&mut self, peer: &Peer<T>) {
        // Have to consider situation where a peer became a seeder but joined as leecher.
        if peer.pieces().complete() && self.seeders > 0 {
            self.seeders -= 1;
        } else {
            if let PickerKind::Rarest(ref mut p) = self.picker {
                p.remove_peer(peer);
            }
        }

        for piece in self.downloading.values_mut() {
            for block in piece {
                block.requested.retain(|req| req.peer != peer.id())
            }
        }
    }

    /// Alters the picker to sequential/non sequential. If changing
    /// from sequential to non sequential, peer state will need to be loaded
    /// after this.
    pub fn change_picker(&mut self, sequential: bool, pieces: &Bitfield) {
        self.unpicked = pieces.clone();
        self.picker = if sequential {
            PickerKind::Sequential(sequential::Picker::new(&self.unpicked))
        } else {
            PickerKind::Rarest(rarest::Picker::new(&self.unpicked))
        };
        self.downloading.clear();
    }

    pub fn set_priorities(&mut self, pri: &[u8], info: &Arc<Info>) {
        self.unapply_priorities();
        self.priorities = generate_piece_pri(pri, info);
        self.apply_priorities();
    }

    pub fn apply_priorities(&mut self) {
        if self.is_sequential() {
            let mut pieces = [vec![], vec![], vec![], vec![], vec![], vec![]];
            for (piece, pri) in self.priorities.iter().enumerate() {
                pieces[*pri as usize].push(piece as u32);
            }
            self.picker = PickerKind::Sequential(sequential::Picker::new_pri(pieces));
        } else {
            for (piece, pri) in self.priorities.iter().enumerate() {
                if let PickerKind::Rarest(ref mut p) = self.picker {
                    for _ in 0..*pri {
                        p.piece_unavailable(piece as u32);
                    }
                }

                if *pri == 0 && !self.unpicked.has_bit(piece as u64) {
                    match self.picker {
                        PickerKind::Rarest(ref mut p) => p.completed(piece as u32),
                        _ => unreachable!(),
                    }
                }
            }
        }
    }

    pub fn unapply_priorities(&mut self) {
        if !self.is_sequential() {
            for (piece, pri) in self.priorities.iter().enumerate() {
                if let PickerKind::Rarest(ref mut p) = self.picker {
                    for _ in 0..*pri {
                        p.piece_available(piece as u32);
                    }
                }

                if *pri == 0 && !self.unpicked.has_bit(piece as u64) {
                    match self.picker {
                        PickerKind::Rarest(ref mut p) => p.incomplete(piece as u32),
                        _ => unreachable!(),
                    }
                }
            }
        }
    }
}

fn generate_piece_pri(pri: &[u8], info: &Arc<Info>) -> Vec<u8> {
    // Map piece -> priority
    let mut priorities = Vec::with_capacity(info.pieces() as usize);
    // If a piece is completely in a file, just assign that pri.
    // Otherwise mark it as the higher pri piece
    for p in 0..info.pieces() {
        let max = Info::piece_disk_locs(&info, p)
            .map(|loc| pri[loc.file])
            .max()
            .expect("Piece must have locations!");
        priorities.push(max);
    }
    priorities
}

#[cfg(test)]
impl Picker {
    pub fn new_rarest(info: &Info, pieces: &Bitfield) -> Picker {
        Picker::new(
            &Arc::new(info.clone()),
            pieces,
            &vec![3u8; info.files.len()],
        )
    }

    pub fn new_sequential(info: &Info, pieces: &Bitfield) -> Picker {
        let mut p = Picker::new(
            &Arc::new(info.clone()),
            pieces,
            &vec![3u8; info.files.len()],
        );
        p.change_picker(true, pieces);
        p
    }
}

impl Block {
    pub fn new(index: u32, offset: u32) -> Block {
        Block { index, offset }
    }
}

impl Request {
    fn new(peer: usize) -> Request {
        Request {
            peer,
            requested_at: time::Instant::now(),
        }
    }
}

impl Downloading {
    fn complete(&mut self) -> Vec<Request> {
        self.completed = true;
        mem::replace(&mut self.requested, Vec::with_capacity(0))
    }
}
<commit_msg>Prevent download queue from filling unnecessarily<commit_after>use std::collections::HashMap;
use std::{mem, time};
use std::sync::Arc;
use torrent::{Bitfield, Info, Peer};
use control::cio;

mod rarest;
mod sequential;

#[cfg(test)]
mod tests;

#[derive(Clone, Debug)]
pub struct Picker {
    /// Number of blocks per piece
    scale: u32,
    last_piece_scale: u32,
    last_piece: u32,
    /// Number of detected seeders
    seeders: u16,
    /// Set of pieces which have blocks waiting. These should be prioritized.
    downloading: HashMap<u32, Vec<Downloading>>,
    /// Bitfield of unpicked pieces, not in progress or
    /// completed yet. A set bit is picked, unset is unpicked.
    unpicked: Bitfield,
    /// The current picker in use
    picker: PickerKind,
    /// Piece priorities
    priorities: Vec<u8>,
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct Block {
    pub index: u32,
    pub offset: u32,
}

/// Pickers act solely as piece picking algorithm.
/// They will select the optimal next piece for a peer,
/// and can be told when a piece is complete(or invalid).
#[derive(Clone, Debug)]
enum PickerKind {
    Rarest(rarest::Picker),
    Sequential(sequential::Picker),
}

/// A downloading block and the peers it has been
/// requested from.
#[derive(Clone, Debug)]
struct Downloading {
    offset: u32,
    completed: bool,
    requested: Vec<Request>,
}

/// A request to a peer and the time it was initiated.
#[derive(Clone, Debug)]
struct Request {
    peer: usize,
    requested_at: time::Instant,
}

const MAX_DUP_REQS: usize = 3;
const MAX_DL_Q: usize = 100;
const MAX_PC_SIZE: usize = 50;
const REQ_TIMEOUT: u64 = 10;

impl Picker {
    /// Creates a new picker, which will select over
    /// the given pieces. The algorithm used for selection
    /// will vary based on the current swarm state, but
    /// will default to rarest first.
    pub fn new(info: &Arc<Info>, pieces: &Bitfield, priorities: &[u8]) -> Picker {
        let scale = info.piece_len / 16_384;
        let picker = rarest::Picker::new(pieces);
        let last_piece = info.pieces().saturating_sub(1);
        let lpl = info.piece_len(last_piece);
        let last_piece_scale = if lpl % 16_384 == 0 {
            lpl / 16_384
        } else {
            lpl / 16_384 + 1
        };
        let downloading = if pieces.complete() {
            HashMap::with_capacity(0)
        } else {
            HashMap::with_capacity(info.pieces() as usize)
        };
        let mut picker = Picker {
            picker: PickerKind::Rarest(picker),
            scale,
            last_piece,
            last_piece_scale,
            seeders: 0,
            unpicked: pieces.clone(),
            downloading,
            priorities: vec![3; info.pieces() as usize],
        };
        picker.set_priorities(priorities, info);
        picker
    }

    /// Returns true if the current picker algorithm is sequential
    pub fn is_sequential(&self) -> bool {
        match self.picker {
            PickerKind::Sequential(_) => true,
            _ => false,
        }
    }

    pub fn tick(&mut self) {
        let mut expired = 0;
        for (piece, chunks) in &mut self.downloading {
            let cl = chunks.len();
            for chunk in chunks.iter_mut() {
                if !chunk.completed {
                    chunk.requested.retain(|req| {
                        let unexp = req.requested_at.elapsed().as_secs() < REQ_TIMEOUT;
                        if !unexp {
                            expired += 1;
                        }
                        unexp
                    });

                    // Rare situation where we pick from a slow peer, choose every other the chunk
                    // in the piece then get stuck: need to unpick here and try again
                    if chunk.requested.is_empty() && cl == self.scale as usize
                        || (*piece == self.last_piece && cl == self.last_piece_scale as usize)
                    {
                        match self.picker {
                            PickerKind::Sequential(ref mut p) => p.incomplete(*piece),
                            PickerKind::Rarest(ref mut p) => p.incomplete(*piece),
                        }
                    }
                }
            }
        }
        if expired != 0 {
            debug!("Expired {} chunks!", expired);
        }
    }

    /// Attempts to select a block for a peer.
    pub fn pick<T: cio::CIO>(&mut self, peer: &mut Peer<T>) -> Option<Block> {
        if let Some(b) = self.pick_expired(peer) {
            return Some(b);
        }

        let piece = match self.picker {
            PickerKind::Sequential(ref mut p) => p.pick(peer),
            PickerKind::Rarest(ref mut p) => p.pick(peer),
        };
        let res = piece
            .and_then(|p| self.pick_piece(p, peer.id()))
            .or_else(|| self.pick_dl(peer));
        if res.is_none() && peer.pieces().complete() && !self.unpicked.complete() {
            debug_assert!(
                false,
                "Couldnt pick {:?} from a seeder for some reason!",
                piece
            );
        }
        res
    }

    /// Attempts to pick an expired block
    fn pick_expired<T: cio::CIO>(&mut self, _: &Peer<T>) -> Option<Block> {
        // TODO: Use some form of heuristic here to say "we expect to have
        // downloaded some pieces by X, hit the picker with a tick which checks
        // that, flags shit as invalid, and then does a double request
        None
    }

    /// Picks a block from a given piece for a peer
    fn pick_piece(&mut self, piece: u32, id: usize) -> Option<Block> {
        self.downloading.entry(piece).or_insert_with(|| vec![]);
        let dl = self.downloading.get_mut(&piece).unwrap();
        let offset = dl.len() as u32 * 16_384;
        dl.push(Downloading {
            offset,
            completed: false,
            requested: vec![Request::new(id)],
        });

        if dl.len() == self.scale as usize
            || (piece == self.last_piece && dl.len() == self.last_piece_scale as usize)
        {
            match self.picker {
                PickerKind::Sequential(ref mut p) => p.completed(piece),
                PickerKind::Rarest(ref mut p) => p.completed(piece),
            }
            self.unpicked.set_bit(u64::from(piece));
        }
        Some(Block {
            index: piece,
            offset,
        })
    }

    /// Attempts to pick the highest priority piece in the dl q
    fn pick_dl<T: cio::CIO>(&mut self, peer: &Peer<T>) -> Option<Block> {
        for (idx, dl) in self.downloading.iter_mut().take(MAX_DL_Q) {
            if peer.pieces().has_bit(u64::from(*idx)) {
                let r = dl.iter_mut()
                    .find(|r| {
                        !r.completed && r.requested.len() < MAX_DUP_REQS
                            && r.requested.iter().all(|req| req.peer != peer.id())
                    })
                    .map(|r| {
                        r.requested.push(Request::new(peer.id()));
                        Block::new(*idx, r.offset)
                    });
                if r.is_some() {
                    return r;
                }
            }
        }
        None
    }

    /// Marks a block as completed. Returns a result indicating if the block
    /// was actually requested, the success value containing a bool indicating
    /// if the block is complete, and a vector of peers from which the block
    /// was requested(for cancellation).
    pub fn completed(&mut self, b: Block) -> Result<(bool, Vec<usize>), ()> {
        // Find the block in our downloading blocks, mark as true,
        // and extract the current peer list for return.
        let res = self.downloading
            .get_mut(&b.index)
            .and_then(|dl| {
                dl.iter_mut()
                    .find(|r| r.offset == b.offset)
                    .map(|r| r.complete())
            })
            .map(|r| r.into_iter().map(|e| e.peer).collect());

        // If we've requested every single block for this piece and they're all complete, remove it
        // and report completion
        let scale = self.scale;
        let lp = self.last_piece;
        let lps = self.last_piece_scale;
        let complete = self.downloading
            .get_mut(&b.index)
            .map(|r| {
                (r.len() as u32 == scale || (b.index == lp && r.len() as u32 == lps))
                    && r.iter().all(|d| d.completed)
            })
            .unwrap_or(false);

        if complete {
            self.downloading.remove(&b.index);
        }

        res.map(|r| (complete, r)).ok_or(())
    }

    pub fn have_block(&mut self, b: Block) -> bool {
        self.downloading
            .get_mut(&b.index)
            .and_then(|dl| {
                dl.iter()
                    .find(|r| r.offset == b.offset)
                    .map(|r| r.completed)
            })
            .unwrap_or(false)
    }

    /// Invalidates a piece
    pub fn invalidate_piece(&mut self, idx: u32) {
        match self.picker {
            PickerKind::Sequential(ref mut p) => p.incomplete(idx),
            PickerKind::Rarest(ref mut p) => p.incomplete(idx),
        }
        self.unpicked.unset_bit(u64::from(idx));
        self.downloading.remove(&idx);
    }

    pub fn piece_available(&mut self, idx: u32) {
        if let PickerKind::Rarest(ref mut p) = self.picker {
            p.piece_available(idx);
        }
    }

    pub fn add_peer<T: cio::CIO>(&mut self, peer: &Peer<T>) {
        if peer.pieces().complete() {
            self.seeders += 1;
        } else {
            if let PickerKind::Rarest(ref mut p) = self.picker {
                p.add_peer(peer);
            }
        }
    }

    pub fn remove_peer<T: cio::CIO>(&mut self, peer: &Peer<T>) {
        // Have to consider situation where a peer became a seeder but joined as leecher.
        if peer.pieces().complete() && self.seeders > 0 {
            self.seeders -= 1;
        } else {
            if let PickerKind::Rarest(ref mut p) = self.picker {
                p.remove_peer(peer);
            }
        }

        for piece in self.downloading.values_mut() {
            for block in piece {
                block.requested.retain(|req| req.peer != peer.id())
            }
        }
    }

    /// Alters the picker to sequential/non sequential. If changing
    /// from sequential to non sequential, peer state will need to be loaded
    /// after this.
    pub fn change_picker(&mut self, sequential: bool, pieces: &Bitfield) {
        self.unpicked = pieces.clone();
        self.picker = if sequential {
            PickerKind::Sequential(sequential::Picker::new(&self.unpicked))
        } else {
            PickerKind::Rarest(rarest::Picker::new(&self.unpicked))
        };
        self.downloading.clear();
    }

    pub fn set_priorities(&mut self, pri: &[u8], info: &Arc<Info>) {
        self.unapply_priorities();
        self.priorities = generate_piece_pri(pri, info);
        self.apply_priorities();
    }

    pub fn apply_priorities(&mut self) {
        if self.is_sequential() {
            let mut pieces = [vec![], vec![], vec![], vec![], vec![], vec![]];
            for (piece, pri) in self.priorities.iter().enumerate() {
                pieces[*pri as usize].push(piece as u32);
            }
            self.picker = PickerKind::Sequential(sequential::Picker::new_pri(pieces));
        } else {
            for (piece, pri) in self.priorities.iter().enumerate() {
                if let PickerKind::Rarest(ref mut p) = self.picker {
                    for _ in 0..*pri {
                        p.piece_unavailable(piece as u32);
                    }
                }

                if *pri == 0 && !self.unpicked.has_bit(piece as u64) {
                    match self.picker {
                        PickerKind::Rarest(ref mut p) => p.completed(piece as u32),
                        _ => unreachable!(),
                    }
                }
            }
        }
    }

    pub fn unapply_priorities(&mut self) {
        if !self.is_sequential() {
            for (piece, pri) in self.priorities.iter().enumerate() {
                if let PickerKind::Rarest(ref mut p) = self.picker {
                    for _ in 0..*pri {
                        p.piece_available(piece as u32);
                    }
                }

                if *pri == 0 && !self.unpicked.has_bit(piece as u64) {
                    match self.picker {
                        PickerKind::Rarest(ref mut p) => p.incomplete(piece as u32),
                        _ => unreachable!(),
                    }
                }
            }
        }
    }
}

fn generate_piece_pri(pri: &[u8], info: &Arc<Info>) -> Vec<u8> {
    // Map piece -> priority
    let mut priorities = Vec::with_capacity(info.pieces() as usize);
    // If a piece is completely in a file, just assign that pri.
    // Otherwise mark it as the higher pri piece
    for p in 0..info.pieces() {
        let max = Info::piece_disk_locs(&info, p)
            .map(|loc| pri[loc.file])
            .max()
            .expect("Piece must have locations!");
        priorities.push(max);
    }
    priorities
}

#[cfg(test)]
impl Picker {
    pub fn new_rarest(info: &Info, pieces: &Bitfield) -> Picker {
        Picker::new(
            &Arc::new(info.clone()),
            pieces,
            &vec![3u8; info.files.len()],
        )
    }

    pub fn new_sequential(info: &Info, pieces: &Bitfield) -> Picker {
        let mut p = Picker::new(
            &Arc::new(info.clone()),
            pieces,
            &vec![3u8; info.files.len()],
        );
        p.change_picker(true, pieces);
        p
    }
}

impl Block {
    pub fn new(index: u32, offset: u32) -> Block {
        Block { index, offset }
    }
}

impl Request {
    fn new(peer: usize) -> Request {
        Request {
            peer,
            requested_at: time::Instant::now(),
        }
    }
}

impl Downloading {
    fn complete(&mut self) -> Vec<Request> {
        self.completed = true;
        mem::replace(&mut self.requested, Vec::with_capacity(0))
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::{mem, time};
use std::sync::Arc;
use torrent::{Info, Peer, Bitfield};
use control::cio;

mod rarest;
mod sequential;

#[cfg(test)]
mod tests;

#[derive(Clone, Debug)]
pub struct Picker {
    /// Number of blocks per piece
    scale: u32,
    last_piece_scale: u32,
    last_piece: u32,
    /// Number of detected seeders
    seeders: u16,
    /// Set of pieces which have blocks waiting. These should be prioritized.
    downloading: HashMap<u32, Vec<Downloading>>,
    /// Bitfield of unpicked pieces, not in progress or
    /// completed yet. A set bit is picked, unset is unpicked.
    unpicked: Bitfield,
    /// The current picker in use
    picker: PickerKind,
    /// Piece priorities
    priorities: Vec<u8>,
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct Block {
    pub index: u32,
    pub offset: u32,
}

/// Pickers act solely as piece picking algorithm.
/// They will select the optimal next piece for a peer,
/// and can be told when a piece is complete(or invalid).
#[derive(Clone, Debug)]
enum PickerKind {
    Rarest(rarest::Picker),
    Sequential(sequential::Picker),
}

/// A downloading block and the peers it has been
/// requested from.
#[derive(Clone, Debug)]
struct Downloading {
    offset: u32,
    completed: bool,
    requested: Vec<Request>,
}

/// A request to a peer and the time it was initiated.
#[derive(Clone, Debug)]
struct Request {
    peer: usize,
    requested_at: time::Instant,
}

const MAX_DUP_REQS: usize = 5;

impl Picker {
    /// Creates a new picker, which will select over
    /// the given pieces. The algorithm used for selection
    /// will vary based on the current swarm state, but
    /// will default to rarest first.
    pub fn new(info: &Arc<Info>, pieces: &Bitfield, priorities: &[u8]) -> Picker {
        let scale = info.piece_len / 16_384;
        let picker = rarest::Picker::new(pieces);
        let last_piece = info.pieces().saturating_sub(1);
        let lpl = info.piece_len(last_piece);
        let last_piece_scale = if lpl % 16_384 == 0 {
            lpl / 16_384
        } else {
            lpl / 16_384 + 1
        };
        let mut picker = Picker {
            picker: PickerKind::Rarest(picker),
            scale,
            last_piece,
            last_piece_scale,
            seeders: 0,
            unpicked: pieces.clone(),
            downloading: HashMap::new(),
            priorities: vec![3; info.pieces() as usize],
        };
        picker.set_priorities(priorities, info);
        picker
    }

    /// Returns true if the current picker algorithm is sequential
    pub fn is_sequential(&self) -> bool {
        match self.picker {
            PickerKind::Sequential(_) => true,
            _ => false,
        }
    }

    /// Attempts to select a block for a peer.
    pub fn pick<T: cio::CIO>(&mut self, peer: &Peer<T>) -> Option<Block> {
        if let Some(b) = self.pick_expired(peer) {
            return Some(b);
        }

        let piece = match self.picker {
            PickerKind::Sequential(ref mut p) => p.pick(peer),
            PickerKind::Rarest(ref mut p) => p.pick(peer),
        };
        piece.and_then(|p| self.pick_piece(p, peer.id())).or_else(
            || {
                self.pick_downloading(peer)
            },
        )
    }

    /// Attempts to pick an expired block
    fn pick_expired<T: cio::CIO>(&mut self, peer: &Peer<T>) -> Option<Block> {
        // TODO: Use some form of heuristic here to say "we expect to have
        // downloaded some pieces by X, hit the picker with a tick which checks
        // that, flags shit as invalid, and then does a double request
        None
    }

    /// Picks a block from a given piece for a peer
    fn pick_piece(&mut self, piece: u32, id: usize) -> Option<Block> {
        self.downloading.entry(piece).or_insert_with(|| vec![]);
        let dl = self.downloading.get_mut(&piece).unwrap();
        let offset = dl.len() as u32 * 16_384;
        dl.push(Downloading {
            offset,
            completed: false,
            requested: vec![Request::new(id)],
        });

        if dl.len() == self.scale as usize ||
            (piece == self.last_piece && dl.len() == self.last_piece_scale as usize)
        {
            match self.picker {
                PickerKind::Sequential(ref mut p) => p.completed(piece),
                PickerKind::Rarest(ref mut p) => p.completed(piece),
            }
            self.unpicked.set_bit(u64::from(piece));
        }
        Some(Block {
            index: piece,
            offset,
        })
    }

    /// Attempts to pick an already requested block
    fn pick_downloading<T: cio::CIO>(&mut self, peer: &Peer<T>) -> Option<Block> {
        for (idx, dl) in &mut self.downloading {
            if peer.pieces().has_bit(u64::from(*idx)) {
                let r = dl.iter_mut()
                    .find(|r| {
                        !r.completed && r.requested.len() < MAX_DUP_REQS &&
                            r.requested.iter().all(|req| req.peer != peer.id())
                    })
                    .map(|r| {
                        r.requested.push(Request::new(peer.id()));
                        Block::new(*idx, r.offset)
                    });
                if r.is_some() {
                    return r;
                }
            }
        }
        None
    }

    /// Marks a block as completed. Returns a result indicating if the block
    /// was actually requested, the success value containing a bool indicating
    /// if the block is complete, and a vector of peers from which the block
    /// was requested(for cancellation).
    pub fn completed(&mut self, b: Block) -> Result<(bool, Vec<usize>), ()> {
        // Find the block in our downloading blocks, mark as true,
        // and extract the current peer list for return.
        let res = self.downloading
            .get_mut(&b.index)
            .and_then(|dl| {
                dl.iter_mut().find(|r| r.offset == b.offset).map(
                    |r| r.complete(),
                )
            })
            .map(|r| r.into_iter().map(|e| e.peer).collect());

        // If we've requested every single block for this piece and they're all complete, remove it
        // and report completion
        let scale = self.scale;
        let lp = self.last_piece;
        let lps = self.last_piece_scale;
        let complete = self.downloading
            .get_mut(&b.index)
            .map(|r| {
                (r.len() as u32 == scale || (b.index == lp && r.len() as u32 == lps)) &&
                    r.iter().all(|d| d.completed)
            })
            .unwrap_or(false);

        if complete {
            self.downloading.remove(&b.index);
        }

        res.map(|r| (complete, r)).ok_or(())
    }

    pub fn have_block(&mut self, b: Block) -> bool {
        self.downloading
            .get_mut(&b.index)
            .and_then(|dl| {
                dl.iter().find(|r| r.offset == b.offset).map(
                    |r| r.completed,
                )
            })
            .unwrap_or(false)
    }

    /// Invalidates a piece
    pub fn invalidate_piece(&mut self, idx: u32) {
        match self.picker {
            PickerKind::Sequential(ref mut p) => p.incomplete(idx),
            PickerKind::Rarest(ref mut p) => p.incomplete(idx),
        }
        self.unpicked.unset_bit(u64::from(idx));
    }

    pub fn piece_available(&mut self, idx: u32) {
        if let PickerKind::Rarest(ref mut p) = self.picker {
            p.piece_available(idx);
        }
    }

    pub fn add_peer<T: cio::CIO>(&mut self, peer: &Peer<T>) {
        if peer.pieces().complete() {
            self.seeders += 1;
        }
        if let PickerKind::Rarest(ref mut p) = self.picker {
            p.add_peer(peer);
        }
    }

    pub fn remove_peer<T: cio::CIO>(&mut self, peer: &Peer<T>) {
        if let PickerKind::Rarest(ref mut p) = self.picker {
            p.remove_peer(peer);
        }

        for piece in self.downloading.values_mut() {
            for block in piece {
                block.requested.retain(|req| req.peer != peer.id())
            }
        }
    }

    /// Alters the picker to sequential/non sequential. If changing
    /// from sequential to non sequential, peer state will need to be loaded
    /// after this.
    pub fn change_picker(&mut self, sequential: bool) {
        self.picker = if sequential {
            PickerKind::Sequential(sequential::Picker::new(&self.unpicked))
        } else {
            PickerKind::Rarest(rarest::Picker::new(&self.unpicked))
        };
        self.apply_priorities();
        for i in self.downloading.keys() {
            if let PickerKind::Rarest(ref mut p) = self.picker {
                p.inc_pri(*i);
            }
        }
    }

    pub fn set_priorities(&mut self, pri: &[u8], info: &Arc<Info>) {
        self.unapply_priorities();
        self.priorities = generate_piece_pri(pri, info);
        self.apply_priorities();
    }

    pub fn apply_priorities(&mut self) {
        if self.is_sequential() {
            let mut pieces = [vec![], vec![], vec![], vec![], vec![], vec![]];
            for (piece, pri) in self.priorities.iter().enumerate() {
                pieces[*pri as usize].push(piece as u32);
            }
            self.picker = PickerKind::Sequential(sequential::Picker::new_pri(pieces));
        } else {
            for (piece, pri) in self.priorities.iter().enumerate() {
                for _ in 0..*pri {
                    if let PickerKind::Rarest(ref mut p) = self.picker {
                        p.piece_unavailable(piece as u32);
                    }
                }

                if *pri == 0 {
                    match self.picker {
                        PickerKind::Rarest(ref mut p) => p.completed(piece as u32),
                        _ => unreachable!(),
                    }
                }
            }
        }
    }

    pub fn unapply_priorities(&mut self) {
        if !self.is_sequential() {
            for (piece, pri) in self.priorities.iter().enumerate() {
                for _ in 0..*pri {
                    if let PickerKind::Rarest(ref mut p) = self.picker {
                        p.piece_available(piece as u32);
                    }
                }

                if *pri == 0 {
                    match self.picker {
                        PickerKind::Rarest(ref mut p) => p.incomplete(piece as u32),
                        _ => unreachable!(),
                    }
                }
            }
        }
    }
}

fn generate_piece_pri(pri: &[u8], info: &Arc<Info>) -> Vec<u8> {
    // Map piece -> priority
    let mut priorities = Vec::with_capacity(info.pieces() as usize);
    // If a piece is completely in a file, just assign that pri.
    // Otherwise mark it as the higher pri piece
    for p in 0..info.pieces() {
        let max = Info::piece_disk_locs(&info, p)
            .map(|loc| pri[loc.file])
            .max()
            .expect("Piece must have locations!");
        priorities.push(max);
    }
    priorities
}

#[cfg(test)]
impl Picker {
    pub fn new_rarest(info: &Info, pieces: &Bitfield) -> Picker {
        Picker::new(
            &Arc::new(info.clone()),
            pieces,
            &vec![3u8; info.files.len()],
        )
    }

    pub fn new_sequential(info: &Info, pieces: &Bitfield) -> Picker {
        let mut p = Picker::new(
            &Arc::new(info.clone()),
            pieces,
            &vec![3u8; info.files.len()],
        );
        p.change_picker(true);
        p
    }
}
impl Block {
    pub fn new(index: u32, offset: u32) -> Block {
        Block { index, offset }
    }
}

impl Request {
    fn new(peer: usize) -> Request {
        Request {
            peer,
            requested_at: time::Instant::now(),
        }
    }
}

impl Downloading {
    fn complete(&mut self) -> Vec<Request> {
        self.completed = true;
        mem::replace(&mut self.requested, Vec::with_capacity(0))
    }
}
<commit_msg>Prioritize downloading pieces more<commit_after>use std::collections::HashMap;
use std::{mem, time};
use std::sync::Arc;
use torrent::{Info, Peer, Bitfield};
use control::cio;

mod rarest;
mod sequential;

#[cfg(test)]
mod tests;

#[derive(Clone, Debug)]
pub struct Picker {
    /// Number of blocks per piece
    scale: u32,
    last_piece_scale: u32,
    last_piece: u32,
    /// Number of detected seeders
    seeders: u16,
    /// Set of pieces which have blocks waiting. These should be prioritized.
    downloading: HashMap<u32, Vec<Downloading>>,
    /// Bitfield of unpicked pieces, not in progress or
    /// completed yet. A set bit is picked, unset is unpicked.
    unpicked: Bitfield,
    /// The current picker in use
    picker: PickerKind,
    /// Piece priorities
    priorities: Vec<u8>,
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct Block {
    pub index: u32,
    pub offset: u32,
}

/// Pickers act solely as piece picking algorithm.
/// They will select the optimal next piece for a peer,
/// and can be told when a piece is complete(or invalid).
#[derive(Clone, Debug)]
enum PickerKind {
    Rarest(rarest::Picker),
    Sequential(sequential::Picker),
}

/// A downloading block and the peers it has been
/// requested from.
#[derive(Clone, Debug)]
struct Downloading {
    offset: u32,
    completed: bool,
    requested: Vec<Request>,
}

/// A request to a peer and the time it was initiated.
#[derive(Clone, Debug)]
struct Request {
    peer: usize,
    requested_at: time::Instant,
}

const MAX_DUP_REQS: usize = 5;
const MAX_DL_Q: usize = 15;

impl Picker {
    /// Creates a new picker, which will select over
    /// the given pieces. The algorithm used for selection
    /// will vary based on the current swarm state, but
    /// will default to rarest first.
    pub fn new(info: &Arc<Info>, pieces: &Bitfield, priorities: &[u8]) -> Picker {
        let scale = info.piece_len / 16_384;
        let picker = rarest::Picker::new(pieces);
        let last_piece = info.pieces().saturating_sub(1);
        let lpl = info.piece_len(last_piece);
        let last_piece_scale = if lpl % 16_384 == 0 {
            lpl / 16_384
        } else {
            lpl / 16_384 + 1
        };
        let mut picker = Picker {
            picker: PickerKind::Rarest(picker),
            scale,
            last_piece,
            last_piece_scale,
            seeders: 0,
            unpicked: pieces.clone(),
            downloading: HashMap::new(),
            priorities: vec![3; info.pieces() as usize],
        };
        picker.set_priorities(priorities, info);
        picker
    }

    /// Returns true if the current picker algorithm is sequential
    pub fn is_sequential(&self) -> bool {
        match self.picker {
            PickerKind::Sequential(_) => true,
            _ => false,
        }
    }

    /// Attempts to select a block for a peer.
    pub fn pick<T: cio::CIO>(&mut self, peer: &Peer<T>) -> Option<Block> {
        if let Some(b) = self.pick_expired(peer) {
            return Some(b);
        }
        if self.downloading.len() > MAX_DL_Q {
            if let Some(b) =  self.pick_downloading(peer) {
                return Some(b)
            }
        }

        let piece = match self.picker {
            PickerKind::Sequential(ref mut p) => p.pick(peer),
            PickerKind::Rarest(ref mut p) => p.pick(peer),
        };
        piece.and_then(|p| self.pick_piece(p, peer.id())).or_else(
            || {
                self.pick_downloading(peer)
            },
        )
    }

    /// Attempts to pick an expired block
    fn pick_expired<T: cio::CIO>(&mut self, peer: &Peer<T>) -> Option<Block> {
        // TODO: Use some form of heuristic here to say "we expect to have
        // downloaded some pieces by X, hit the picker with a tick which checks
        // that, flags shit as invalid, and then does a double request
        None
    }

    /// Picks a block from a given piece for a peer
    fn pick_piece(&mut self, piece: u32, id: usize) -> Option<Block> {
        self.downloading.entry(piece).or_insert_with(|| vec![]);
        let dl = self.downloading.get_mut(&piece).unwrap();
        let offset = dl.len() as u32 * 16_384;
        dl.push(Downloading {
            offset,
            completed: false,
            requested: vec![Request::new(id)],
        });

        if dl.len() == self.scale as usize ||
            (piece == self.last_piece && dl.len() == self.last_piece_scale as usize)
        {
            match self.picker {
                PickerKind::Sequential(ref mut p) => p.completed(piece),
                PickerKind::Rarest(ref mut p) => p.completed(piece),
            }
            self.unpicked.set_bit(u64::from(piece));
        }
        Some(Block {
            index: piece,
            offset,
        })
    }

    /// Attempts to pick an already requested block
    fn pick_downloading<T: cio::CIO>(&mut self, peer: &Peer<T>) -> Option<Block> {
        for (idx, dl) in &mut self.downloading {
            if peer.pieces().has_bit(u64::from(*idx)) {
                let r = dl.iter_mut()
                    .find(|r| {
                        !r.completed && r.requested.len() < MAX_DUP_REQS &&
                            r.requested.iter().all(|req| req.peer != peer.id())
                    })
                    .map(|r| {
                        r.requested.push(Request::new(peer.id()));
                        Block::new(*idx, r.offset)
                    });
                if r.is_some() {
                    return r;
                }
            }
        }
        None
    }

    /// Marks a block as completed. Returns a result indicating if the block
    /// was actually requested, the success value containing a bool indicating
    /// if the block is complete, and a vector of peers from which the block
    /// was requested(for cancellation).
    pub fn completed(&mut self, b: Block) -> Result<(bool, Vec<usize>), ()> {
        // Find the block in our downloading blocks, mark as true,
        // and extract the current peer list for return.
        let res = self.downloading
            .get_mut(&b.index)
            .and_then(|dl| {
                dl.iter_mut().find(|r| r.offset == b.offset).map(
                    |r| r.complete(),
                )
            })
            .map(|r| r.into_iter().map(|e| e.peer).collect());

        // If we've requested every single block for this piece and they're all complete, remove it
        // and report completion
        let scale = self.scale;
        let lp = self.last_piece;
        let lps = self.last_piece_scale;
        let complete = self.downloading
            .get_mut(&b.index)
            .map(|r| {
                (r.len() as u32 == scale || (b.index == lp && r.len() as u32 == lps)) &&
                    r.iter().all(|d| d.completed)
            })
            .unwrap_or(false);

        if complete {
            self.downloading.remove(&b.index);
        }

        res.map(|r| (complete, r)).ok_or(())
    }

    pub fn have_block(&mut self, b: Block) -> bool {
        self.downloading
            .get_mut(&b.index)
            .and_then(|dl| {
                dl.iter().find(|r| r.offset == b.offset).map(
                    |r| r.completed,
                )
            })
            .unwrap_or(false)
    }

    /// Invalidates a piece
    pub fn invalidate_piece(&mut self, idx: u32) {
        match self.picker {
            PickerKind::Sequential(ref mut p) => p.incomplete(idx),
            PickerKind::Rarest(ref mut p) => p.incomplete(idx),
        }
        self.unpicked.unset_bit(u64::from(idx));
    }

    pub fn piece_available(&mut self, idx: u32) {
        if let PickerKind::Rarest(ref mut p) = self.picker {
            p.piece_available(idx);
        }
    }

    pub fn add_peer<T: cio::CIO>(&mut self, peer: &Peer<T>) {
        if peer.pieces().complete() {
            self.seeders += 1;
        }
        if let PickerKind::Rarest(ref mut p) = self.picker {
            p.add_peer(peer);
        }
    }

    pub fn remove_peer<T: cio::CIO>(&mut self, peer: &Peer<T>) {
        if let PickerKind::Rarest(ref mut p) = self.picker {
            p.remove_peer(peer);
        }

        for piece in self.downloading.values_mut() {
            for block in piece {
                block.requested.retain(|req| req.peer != peer.id())
            }
        }
    }

    /// Alters the picker to sequential/non sequential. If changing
    /// from sequential to non sequential, peer state will need to be loaded
    /// after this.
    pub fn change_picker(&mut self, sequential: bool) {
        self.picker = if sequential {
            PickerKind::Sequential(sequential::Picker::new(&self.unpicked))
        } else {
            PickerKind::Rarest(rarest::Picker::new(&self.unpicked))
        };
        self.apply_priorities();
        for i in self.downloading.keys() {
            if let PickerKind::Rarest(ref mut p) = self.picker {
                p.inc_pri(*i);
            }
        }
    }

    pub fn set_priorities(&mut self, pri: &[u8], info: &Arc<Info>) {
        self.unapply_priorities();
        self.priorities = generate_piece_pri(pri, info);
        self.apply_priorities();
    }

    pub fn apply_priorities(&mut self) {
        if self.is_sequential() {
            let mut pieces = [vec![], vec![], vec![], vec![], vec![], vec![]];
            for (piece, pri) in self.priorities.iter().enumerate() {
                pieces[*pri as usize].push(piece as u32);
            }
            self.picker = PickerKind::Sequential(sequential::Picker::new_pri(pieces));
        } else {
            for (piece, pri) in self.priorities.iter().enumerate() {
                for _ in 0..*pri {
                    if let PickerKind::Rarest(ref mut p) = self.picker {
                        p.piece_unavailable(piece as u32);
                    }
                }

                if *pri == 0 {
                    match self.picker {
                        PickerKind::Rarest(ref mut p) => p.completed(piece as u32),
                        _ => unreachable!(),
                    }
                }
            }
        }
    }

    pub fn unapply_priorities(&mut self) {
        if !self.is_sequential() {
            for (piece, pri) in self.priorities.iter().enumerate() {
                for _ in 0..*pri {
                    if let PickerKind::Rarest(ref mut p) = self.picker {
                        p.piece_available(piece as u32);
                    }
                }

                if *pri == 0 {
                    match self.picker {
                        PickerKind::Rarest(ref mut p) => p.incomplete(piece as u32),
                        _ => unreachable!(),
                    }
                }
            }
        }
    }
}

fn generate_piece_pri(pri: &[u8], info: &Arc<Info>) -> Vec<u8> {
    // Map piece -> priority
    let mut priorities = Vec::with_capacity(info.pieces() as usize);
    // If a piece is completely in a file, just assign that pri.
    // Otherwise mark it as the higher pri piece
    for p in 0..info.pieces() {
        let max = Info::piece_disk_locs(&info, p)
            .map(|loc| pri[loc.file])
            .max()
            .expect("Piece must have locations!");
        priorities.push(max);
    }
    priorities
}

#[cfg(test)]
impl Picker {
    pub fn new_rarest(info: &Info, pieces: &Bitfield) -> Picker {
        Picker::new(
            &Arc::new(info.clone()),
            pieces,
            &vec![3u8; info.files.len()],
        )
    }

    pub fn new_sequential(info: &Info, pieces: &Bitfield) -> Picker {
        let mut p = Picker::new(
            &Arc::new(info.clone()),
            pieces,
            &vec![3u8; info.files.len()],
        );
        p.change_picker(true);
        p
    }
}
impl Block {
    pub fn new(index: u32, offset: u32) -> Block {
        Block { index, offset }
    }
}

impl Request {
    fn new(peer: usize) -> Request {
        Request {
            peer,
            requested_at: time::Instant::now(),
        }
    }
}

impl Downloading {
    fn complete(&mut self) -> Vec<Request> {
        self.completed = true;
        mem::replace(&mut self.requested, Vec::with_capacity(0))
    }
}
<|endoftext|>


<commit_before>// Make me compile! Scroll down for hints :)

pub fn main() {
    let vec0 = Vec::new();

    let vec1 = fill_vec(vec0);

    println!("{} has length {} content `{:?}`", "vec1", vec1.len(), vec1);

    vec1.push(88);

    println!("{} has length {} content `{:?}`", "vec1", vec1.len(), vec1);

}

fn fill_vec(vec: Vec<i32>) -> Vec<i32> {
    let mut vec = vec;

    vec.push(22);
    vec.push(44);
    vec.push(66);

    vec
}

// So you've got the "cannot borrow immutable local variable `vec1` as mutable" error on line 8,
// right? The fix for this is going to be adding one keyword, and the addition is NOT on line 8
// where the error is.
<commit_msg>Fix typo in hint for move_semantics1.rs.<commit_after>// Make me compile! Scroll down for hints :)

pub fn main() {
    let vec0 = Vec::new();

    let vec1 = fill_vec(vec0);

    println!("{} has length {} content `{:?}`", "vec1", vec1.len(), vec1);

    vec1.push(88);

    println!("{} has length {} content `{:?}`", "vec1", vec1.len(), vec1);

}

fn fill_vec(vec: Vec<i32>) -> Vec<i32> {
    let mut vec = vec;

    vec.push(22);
    vec.push(44);
    vec.push(66);

    vec
}

// So you've got the "cannot borrow immutable local variable `vec1` as mutable" error on line 10,
// right? The fix for this is going to be adding one keyword, and the addition is NOT on line 10
// where the error is.
<|endoftext|>


<commit_before>use std::mem;
use std::ops;
use std::ptr;
use std::slice;

use nodrop::NoDrop;

use num_traits::{Unsigned, Zero, One};

pub struct OverflowError<T1, T2>(pub T1, pub T2);

pub enum InsertError<T1, T2> {
    Overflow(T1, T2),
    OutOfBounds(T1, T2)
}

pub unsafe trait Array {
    type Item;
    type Index: Index;
    fn as_ptr(&self) -> *const Self::Item;
    fn as_mut_ptr(&mut self) -> *mut Self::Item;
    fn capacity(&self) -> usize;
}

pub trait Index : Unsigned + ops::AddAssign + ops::SubAssign + Copy {
    fn as_usize(self) -> usize;
    fn from_usize(value: usize) -> Self;
}

macro_rules! impl_index {
    ($index_type:ty) => (
        impl Index for $index_type {
            fn as_usize(self) -> usize {
                self as usize
            }

            fn from_usize(value: usize) -> Self {
                value as $index_type
            }
        }
    )
}

impl_index!(u8);
impl_index!(u16);
impl_index!(u32);
impl_index!(usize);

macro_rules! impl_array {
    ($index_type:ty, $len:expr ) => (
        unsafe impl<T> Array for [T; $len] {
            type Item = T;
            type Index = $index_type;
            #[inline(always)]
            fn as_ptr(&self) -> *const T {
                self as *const _ as *const _
            }

            #[inline(always)]
            fn as_mut_ptr(&mut self) -> *mut T {
                self as *mut _ as *mut _
            }

            #[inline(always)]
            fn capacity(&self) -> usize {
                $len
            }
        }
    )
}

macro_rules! impl_array_many {
    ($index_type:ty, $($len:expr),+) => (
        $(
            impl_array!($index_type, $len);
        )*
    );
}

impl_array_many!(u8, 1, 2, 7, 31, 256);

pub struct LockstepArray<A1: Array, A2: Array> {
    len: A1::Index,
    array1: NoDrop<A1>,
    array2: NoDrop<A2>
}

impl<A1, A2> LockstepArray<A1, A2> where A1: Array, A2: Array {
    pub fn new() -> LockstepArray<A1, A2> {
        LockstepArray {
            len: Zero::zero(),
            array1: NoDrop::new(unsafe {mem::uninitialized()}),
            array2: NoDrop::new(unsafe {mem::uninitialized()})
        }
    }

    pub fn from_slices(slice1: &[A1::Item], slice2: &[A2::Item]) -> LockstepArray<A1, A2>{
        let mut lockstep = LockstepArray::<A1, A2>::new();
        assert!(slice1.len() <= lockstep.array1.capacity());
        assert!(slice2.len() <= lockstep.array2.capacity());
        assert_eq!(slice1.len(), slice2.len());
        lockstep
    }

    pub fn push(&mut self, item1: A1::Item, item2: A2::Item)
                -> Result<(), OverflowError<A1::Item, A2::Item>> {
        if self.len.as_usize() == self.array1.capacity() {
            return Err(OverflowError(item1, item2));
        }
        unsafe {
            let p1 = self.array1.as_mut_ptr();
            let p2 = self.array2.as_mut_ptr();
        }
        self.len += One::one();
        Ok(())
    }

    pub fn pop(&mut self) -> Option<(A1::Item, A2::Item)> {
        if self.len == Zero::zero() {
            None
        } else {
            unsafe {
                // Calculate item 1's pointer
                let src1 = self.array1.as_mut_ptr()
                    .offset(self.len.as_usize() as isize);
                // Move item 1 out of the pointer
                let item1 = ptr::read(src1);
                // Calculate item 1's pointer
                let src2 = self.array2.as_mut_ptr()
                    .offset(self.len.as_usize() as isize);
                // Move item 1 out of the pointer
                let item2 = ptr::read(src2);
                // Update length
                self.len -= One::one();
                Some((item1, item2))
            }
        }
    }

    pub fn insert(&mut self, index: usize, item1: A1::Item, item2: A2::Item)
                  -> Result<(), InsertError<A1::Item, A2::Item>> {
        if index >= self.len.as_usize() {
            Err(InsertError::OutOfBounds(item1, item2))
        } else if self.len.as_usize() == self.array1.capacity() {
            Err(InsertError::Overflow(item1, item2))
        } else {
            unsafe {
                // TODO: Tricky math, check indices and such
                let dest1 = self.array1.as_mut_ptr().offset(index as isize);
                ptr::copy(dest1, dest1.offset(1), self.len.as_usize() - index);
                ptr::write(dest1, item1);
                let dest2 = self.array2.as_mut_ptr().offset(index as isize);
                ptr::copy(dest2, dest2.offset(1), self.len.as_usize() - index);
                ptr::write(dest2, item2);
            }
            self.len += One::one();
            Ok(())
        }
    }

    pub fn remove(&mut self, index: usize) -> Option<(A1::Item, A2::Item)> {
        if index >= self.len.as_usize() {
            None
        } else {
            let item1;
            let item2;
            unsafe {
                // Calculate item 1's pointer
                let src1 = self.array1.as_mut_ptr().offset(index as isize);
                // Move item 1 out of the pointer
                item1 = ptr::read(src1);
                // Shift down array 1 down by 1
                // Is self.len - index correct?
                ptr::copy(src1.offset(1), src1, self.len.as_usize() - index);
                // Calculate item 1's pointer
                let src2 = self.array2.as_mut_ptr().offset(index as isize);
                // Move item 1 out of the pointer
                item2 = ptr::read(src2);
                // Shift down array 1 down by 1
                // Is self.len - index correct?
                ptr::copy(src2, src2.offset(1), self.len.as_usize() - index);
                // Update length
                self.len -= One::one();
            }
            Some((item1, item2))
        }
    }

    pub fn is_empty(&self) -> bool {
        return self.len != Zero::zero()
    }

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

    pub fn array1(&self) -> &[A1::Item] {
        unsafe {
            let ptr = self.array1.as_ptr();
            let len = self.len.as_usize();
            slice::from_raw_parts(ptr, len)
        }
    }

    pub fn array2(&self) -> &[A2::Item] {
        unsafe {
            let ptr = self.array2.as_ptr();
            let len = self.len.as_usize();
            slice::from_raw_parts(ptr, len)
        }
    }

    pub fn array1_mut(&mut self) -> &mut [A1::Item] {
        unsafe {
            let ptr = self.array1.as_mut_ptr();
            let len = self.len.as_usize();
            slice::from_raw_parts_mut(ptr, len)
        }
    }

    pub fn array2_mut(&mut self) -> &mut [A2::Item] {
        unsafe {
            let ptr = self.array2.as_mut_ptr();
            let len = self.len.as_usize();
            slice::from_raw_parts_mut(ptr, len)
        }
    }

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

impl<A1, A2> Default for LockstepArray<A1, A2> where A1: Array, A2: Array {
    fn default() -> Self {
        Self::new()
    }
}
<commit_msg>Implement missing write into LockstepArray<commit_after>use std::mem;
use std::ops;
use std::ptr;
use std::slice;

use nodrop::NoDrop;

use num_traits::{Unsigned, Zero, One};

pub struct OverflowError<T1, T2>(pub T1, pub T2);

pub enum InsertError<T1, T2> {
    Overflow(T1, T2),
    OutOfBounds(T1, T2)
}

pub unsafe trait Array {
    type Item;
    type Index: Index;
    fn as_ptr(&self) -> *const Self::Item;
    fn as_mut_ptr(&mut self) -> *mut Self::Item;
    fn capacity(&self) -> usize;
}

pub trait Index : Unsigned + ops::AddAssign + ops::SubAssign + Copy {
    fn as_usize(self) -> usize;
    fn from_usize(value: usize) -> Self;
}

macro_rules! impl_index {
    ($index_type:ty) => (
        impl Index for $index_type {
            fn as_usize(self) -> usize {
                self as usize
            }

            fn from_usize(value: usize) -> Self {
                value as $index_type
            }
        }
    )
}

impl_index!(u8);
impl_index!(u16);
impl_index!(u32);
impl_index!(usize);

macro_rules! impl_array {
    ($index_type:ty, $len:expr ) => (
        unsafe impl<T> Array for [T; $len] {
            type Item = T;
            type Index = $index_type;
            #[inline(always)]
            fn as_ptr(&self) -> *const T {
                self as *const _ as *const _
            }

            #[inline(always)]
            fn as_mut_ptr(&mut self) -> *mut T {
                self as *mut _ as *mut _
            }

            #[inline(always)]
            fn capacity(&self) -> usize {
                $len
            }
        }
    )
}

macro_rules! impl_array_many {
    ($index_type:ty, $($len:expr),+) => (
        $(
            impl_array!($index_type, $len);
        )*
    );
}

impl_array_many!(u8, 1, 2, 7, 31, 256);

pub struct LockstepArray<A1: Array, A2: Array> {
    len: A1::Index,
    array1: NoDrop<A1>,
    array2: NoDrop<A2>
}

impl<A1, A2> LockstepArray<A1, A2> where A1: Array, A2: Array {
    pub fn new() -> LockstepArray<A1, A2> {
        LockstepArray {
            len: Zero::zero(),
            array1: NoDrop::new(unsafe {mem::uninitialized()}),
            array2: NoDrop::new(unsafe {mem::uninitialized()})
        }
    }

    pub fn from_slices(slice1: &[A1::Item], slice2: &[A2::Item]) -> LockstepArray<A1, A2>{
        let mut lockstep = LockstepArray::<A1, A2>::new();
        assert!(slice1.len() <= lockstep.array1.capacity());
        assert!(slice2.len() <= lockstep.array2.capacity());
        assert_eq!(slice1.len(), slice2.len());
        // TODO: Implement from_slices, will require unsafe code
        lockstep
    }

    pub fn push(&mut self, item1: A1::Item, item2: A2::Item)
                -> Result<(), OverflowError<A1::Item, A2::Item>> {
        if self.len.as_usize() == self.array1.capacity() {
            return Err(OverflowError(item1, item2));
        }
        unsafe {
            let p1 = self.array1.as_mut_ptr();
            ptr::write(p1.offset(self.len.as_usize() as isize), item1);
            let p2 = self.array2.as_mut_ptr();
            ptr::write(p2.offset(self.len.as_usize() as isize), item2);
        }
        self.len += One::one();
        Ok(())
    }

    pub fn pop(&mut self) -> Option<(A1::Item, A2::Item)> {
        if self.len == Zero::zero() {
            None
        } else {
            unsafe {
                // Calculate item 1's pointer
                let src1 = self.array1.as_mut_ptr()
                    .offset(self.len.as_usize() as isize);
                // Move item 1 out of the pointer
                let item1 = ptr::read(src1);
                // Calculate item 1's pointer
                let src2 = self.array2.as_mut_ptr()
                    .offset(self.len.as_usize() as isize);
                // Move item 1 out of the pointer
                let item2 = ptr::read(src2);
                // Update length
                self.len -= One::one();
                Some((item1, item2))
            }
        }
    }

    pub fn insert(&mut self, index: usize, item1: A1::Item, item2: A2::Item)
                  -> Result<(), InsertError<A1::Item, A2::Item>> {
        if index >= self.len.as_usize() {
            Err(InsertError::OutOfBounds(item1, item2))
        } else if self.len.as_usize() == self.array1.capacity() {
            Err(InsertError::Overflow(item1, item2))
        } else {
            unsafe {
                // TODO: Tricky math, check indices and such
                let dest1 = self.array1.as_mut_ptr().offset(index as isize);
                ptr::copy(dest1, dest1.offset(1), self.len.as_usize() - index);
                ptr::write(dest1, item1);
                let dest2 = self.array2.as_mut_ptr().offset(index as isize);
                ptr::copy(dest2, dest2.offset(1), self.len.as_usize() - index);
                ptr::write(dest2, item2);
            }
            self.len += One::one();
            Ok(())
        }
    }

    pub fn remove(&mut self, index: usize) -> Option<(A1::Item, A2::Item)> {
        if index >= self.len.as_usize() {
            None
        } else {
            let item1;
            let item2;
            unsafe {
                // Calculate item 1's pointer
                let src1 = self.array1.as_mut_ptr().offset(index as isize);
                // Move item 1 out of the pointer
                item1 = ptr::read(src1);
                // Shift down array 1 down by 1
                // Is self.len - index correct?
                ptr::copy(src1.offset(1), src1, self.len.as_usize() - index);
                // Calculate item 1's pointer
                let src2 = self.array2.as_mut_ptr().offset(index as isize);
                // Move item 1 out of the pointer
                item2 = ptr::read(src2);
                // Shift down array 1 down by 1
                // Is self.len - index correct?
                ptr::copy(src2, src2.offset(1), self.len.as_usize() - index);
                // Update length
                self.len -= One::one();
            }
            Some((item1, item2))
        }
    }

    pub fn is_empty(&self) -> bool {
        return self.len != Zero::zero()
    }

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

    pub fn array1(&self) -> &[A1::Item] {
        unsafe {
            let ptr = self.array1.as_ptr();
            let len = self.len.as_usize();
            slice::from_raw_parts(ptr, len)
        }
    }

    pub fn array2(&self) -> &[A2::Item] {
        unsafe {
            let ptr = self.array2.as_ptr();
            let len = self.len.as_usize();
            slice::from_raw_parts(ptr, len)
        }
    }

    pub fn array1_mut(&mut self) -> &mut [A1::Item] {
        unsafe {
            let ptr = self.array1.as_mut_ptr();
            let len = self.len.as_usize();
            slice::from_raw_parts_mut(ptr, len)
        }
    }

    pub fn array2_mut(&mut self) -> &mut [A2::Item] {
        unsafe {
            let ptr = self.array2.as_mut_ptr();
            let len = self.len.as_usize();
            slice::from_raw_parts_mut(ptr, len)
        }
    }

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

impl<A1, A2> Default for LockstepArray<A1, A2> where A1: Array, A2: Array {
    fn default() -> Self {
        Self::new()
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::io::Read;

use crate::common::LineIter;

type EdgeMap = HashMap<String, Vec<String>>;

fn read_edges(input: &mut dyn Read) -> EdgeMap {
    let mut reader = LineIter::new(input);
    let mut edges = EdgeMap::new();

    while let Some(line) = reader.next() {
        let (from, to) = line.split_once('-').unwrap();

        edges
            .entry(from.to_owned())
            .or_default()
            .push(to.to_owned());

        edges
            .entry(to.to_owned())
            .or_default()
            .push(from.to_owned());
    }

    edges
}

fn is_small(cave: &str) -> bool {
    cave.chars().all(|c| c.is_ascii_lowercase())
}

fn dfs_routes<'a>(edges: &'a EdgeMap, route: &'_ mut Vec<&'a str>, pos: &'a str) -> usize {
    if is_small(pos) && route.contains(&pos) {
        return 0;
    }

    if pos == "end" {
        return 1;
    }

    route.push(pos);

    let routes = edges[pos]
        .iter()
        .map(|new_pos| dfs_routes(edges, route, new_pos))
        .sum();

    route.pop();

    routes
}

pub fn part1(input: &mut dyn Read) -> String {
    let edges = read_edges(input);
    let mut route = Vec::new();

    dfs_routes(&edges, &mut route, "start").to_string()
}

pub fn part2(_input: &mut dyn Read) -> String {
    todo!()
}

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

    use crate::test_implementation;

    const SAMPLE1: &[u8] = include_bytes!("samples/12.1.txt");
    const SAMPLE2: &[u8] = include_bytes!("samples/12.2.txt");
    const SAMPLE3: &[u8] = include_bytes!("samples/12.3.txt");

    #[test]
    fn sample_part1() {
        test_implementation(part1, SAMPLE1, 10);
        test_implementation(part1, SAMPLE2, 19);
        test_implementation(part1, SAMPLE3, 226);
    }
}
<commit_msg>Implement day 12 part 2<commit_after>use std::collections::HashMap;
use std::io::Read;

use crate::common::LineIter;

type EdgeMap = HashMap<String, Vec<String>>;

fn read_edges(input: &mut dyn Read) -> EdgeMap {
    let mut reader = LineIter::new(input);
    let mut edges = EdgeMap::new();

    while let Some(line) = reader.next() {
        let (from, to) = line.split_once('-').unwrap();

        edges
            .entry(from.to_owned())
            .or_default()
            .push(to.to_owned());

        edges
            .entry(to.to_owned())
            .or_default()
            .push(from.to_owned());
    }

    edges
}

fn is_small(cave: &str) -> bool {
    cave.chars().all(|c| c.is_ascii_lowercase())
}

fn dfs_routes<'a>(
    edges: &'a EdgeMap,
    route: &'_ mut Vec<&'a str>,
    pos: &'a str,
    mut small_twice: bool,
) -> usize {
    match pos {
        "end" => return 1,
        "start" if !route.is_empty() => return 0,
        pos if is_small(pos) && route.contains(&pos) => {
            if small_twice {
                small_twice = false;
            } else {
                return 0;
            }
        }
        _ => (),
    }

    route.push(pos);

    let routes = edges[pos]
        .iter()
        .map(|new_pos| dfs_routes(edges, route, new_pos, small_twice))
        .sum();

    route.pop();

    routes
}

fn parts_common(input: &mut dyn Read, small_twice: bool) -> String {
    let edges = read_edges(input);
    let mut route = Vec::new();

    dfs_routes(&edges, &mut route, "start", small_twice).to_string()
}

pub fn part1(input: &mut dyn Read) -> String {
    parts_common(input, false)
}

pub fn part2(input: &mut dyn Read) -> String {
    parts_common(input, true)
}

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

    use crate::test_implementation;

    const SAMPLE1: &[u8] = include_bytes!("samples/12.1.txt");
    const SAMPLE2: &[u8] = include_bytes!("samples/12.2.txt");
    const SAMPLE3: &[u8] = include_bytes!("samples/12.3.txt");

    #[test]
    fn sample_part1() {
        test_implementation(part1, SAMPLE1, 10);
        test_implementation(part1, SAMPLE2, 19);
        test_implementation(part1, SAMPLE3, 226);
    }

    #[test]
    fn sample_part2() {
        test_implementation(part2, SAMPLE1, 36);
        test_implementation(part2, SAMPLE2, 103);
        test_implementation(part2, SAMPLE3, 3509);
    }
}
<|endoftext|>