<commit_before>// Lumol, an extensible molecular simulation engine
// Copyright (C) 2015-2016 G. Fraux — BSD license

//! Metropolis Monte-Carlo propagator implementation
use rand::{self, SeedableRng};

use consts::K_BOLTZMANN;
use sys::{System, EnergyCache};
use sim::{Propagator, TemperatureStrategy};

use super::MCMove;

/// Metropolis Monte-Carlo propagator
pub struct MonteCarlo {
    /// Boltzmann factor: beta = 1/(kB * T)
    beta: f64,
    /// List of possible Monte-Carlo moves
    moves: Vec<(Box<MCMove>, MoveCounter)>,
    /// Cummulative frequencies of the Monte-Carlo moves
    frequencies: Vec<f64>,
    /// Specifies the number of moves after which an update of a move's
    /// amplitude is performed.
    update_frequency: u64,
    /// Random number generator for the simulation. All random state will be
    /// taken from this.
    rng: Box<rand::Rng>,
    /// Cache for faster energy computation
    cache: EnergyCache,
    /// Flag checking if the moves frequencies has been converted to
    /// cumulative frequencies or not yet.
    initialized: bool,
}

impl MonteCarlo {
    /// Create a new Monte-Carlo propagator at temperature `T`.
    pub fn new(temperature: f64) -> MonteCarlo {
        let mut rng = Box::new(rand::XorShiftRng::new_unseeded());
        rng.reseed([2015u32, 42u32, 3u32, 12u32]);
        return MonteCarlo::from_rng(temperature, rng);
    }

    /// Create a Monte-Carlo propagator at temperature `T`, using the `rng`
    /// random number generator.
    pub fn from_rng(temperature: f64, rng: Box<rand::Rng>) -> MonteCarlo {
        assert!(temperature >= 0.0, "Monte-Carlo temperature must be positive");
        MonteCarlo {
            beta: 1.0 / (K_BOLTZMANN * temperature),
            moves: Vec::new(),
            frequencies: Vec::new(),
            update_frequency: 0,
            rng: rng,
            cache: EnergyCache::new(),
            initialized: false,
        }
    }

    /// Add the `mcmove` Monte-Carlo move to this propagator, with frequency
    /// `frequency`. All calls to this function should happen before any
    /// simulation run.
    ///
    /// # Panics
    ///
    /// If called after a simulation run.
    pub fn add(&mut self, mcmove: Box<MCMove>, frequency: f64) {
        if self.initialized {
            fatal_error!(
                "Monte-Carlo simulation has already been initialized, \
                we can not add new moves."
            );
        }
        self.moves.push((mcmove, MoveCounter::new(None)));
        self.frequencies.push(frequency);
    }

    /// Add the `mcmove` Monte-Carlo move to the propagator.
    /// `frequency` describes how frequent a move is called, `target_acceptance`
    /// is the desired acceptance ratio of the move.
    ///
    /// # Panics
    ///
    /// If called after a simulation run.
    /// If `target_acceptance` is either negative or larger than one.
    pub fn add_move_with_acceptance(&mut self, mcmove: Box<MCMove>, frequency: f64, target_acceptance: f64) {
        if self.initialized {
            fatal_error!(
                "Monte-Carlo simulation has already been initialized, \
                we can not add new moves."
            );
        }
        if target_acceptance >= 1.0 || target_acceptance < 0.0 {
            fatal_error!(
                "The target acceptance ratio of the move has to be a positive \
                value smaller equal than 1."
            );
        }
        self.moves.push((mcmove, MoveCounter::new(Some(target_acceptance))));
        self.frequencies.push(frequency);
    }

    /// Set the number of times a move has to be called before its amplitude
    /// is updated. This value is applied to all moves.
    pub fn set_amplitude_update_frequency(&mut self, frequency: u64) {
        self.update_frequency = frequency;
    }

    /// Get the temperature of the simulation
    pub fn temperature(&self) -> f64 {
        1.0 / (self.beta * K_BOLTZMANN)
    }

    /// Set the temperature of the simulation
    pub fn set_temperature(&mut self, temperature: f64) {
        self.beta = 1.0 / (temperature * K_BOLTZMANN);
    }

    fn normalize_frequencies(&mut self) {
        assert!(self.frequencies.len() == self.moves.len());
        if self.frequencies.is_empty() {
            warn!(
                "No move in the Monte-Carlo simulation, \
                did you forget to specify them?"
            );
            return;
        }

        if self.initialized {
            error!("This Monte-Carlo simulation has already been initialized.");
            return;
        }

        self.initialized = true;
        // Normalize the frequencies
        let sum = self.frequencies.iter().fold(0.0, |sum, &f| sum + f);
        for frequency in &mut self.frequencies {
            *frequency /= sum;
        }
        // Make the frequencies vector contain cumulative frequencies
        for i in 1..self.frequencies.len() {
            self.frequencies[i] += self.frequencies[i - 1];
        }
        let last = self.frequencies.len() - 1;
        self.frequencies[last] = 1.0;
    }
}

impl Propagator for MonteCarlo {
    fn temperature_strategy(&self) -> TemperatureStrategy {
        TemperatureStrategy::External(self.temperature())
    }

    fn setup(&mut self, system: &System) {
        self.normalize_frequencies();
        self.cache.init(system);
    }

    fn propagate(&mut self, system: &mut System) {
        let mcmove = {
            let probability = self.rng.next_f64();
            // Get the index of the first move with frequency >= probability.
            let (i, _) = self.frequencies.iter()
                                         .enumerate()
                                         .find(|&(_, f)| probability <= *f)
                                         .expect("Could not find a move in MonteCarlo moves list");
            &mut self.moves[i]
        };
        trace!("Selected move is '{}'", mcmove.0.describe());

        if !mcmove.0.prepare(system, &mut self.rng) {
            trace!("    --> Can not perform the move");
            return;
        }

        // attempt the move: increase counter
        mcmove.1.ncalled += 1;
        mcmove.1.nattempted += 1;

        // compute cost
        let cost = mcmove.0.cost(system, self.beta, &mut self.cache);
        trace!("    --> Move cost is {}", cost);

        // apply metropolis criterion
        let accepted = cost <= 0.0 || self.rng.next_f64() < f64::exp(-cost);

        if accepted {
            trace!("    --> Move was accepted");
            mcmove.0.apply(system);
            self.cache.update(system);
            mcmove.1.naccepted += 1
        } else {
            trace!("    --> Move was rejected");
            mcmove.0.restore(system);
        }

        // Do the adjustments for the selected move as needed
        if mcmove.1.nattempted == self.update_frequency {
            mcmove.0.update_amplitude(mcmove.1.compute_scaling_factor());
            // Set `nattempted` and `naccepted` to zero.
            // This way, only the running acceptance since the last update is considered.
            mcmove.1.naccepted = 0;
            mcmove.1.nattempted = 0;
        }
    }
}

/// This struct keeps track of the number of times a move was called
/// and how often it was accepted.
pub struct MoveCounter {
    /// Count the total number of times the move was called.
    pub ncalled: u64,
    /// Count the number of times the move was accepted since the last update.
    pub naccepted: u64,
    /// Count the number of times the move was called since the last update.
    pub nattempted: u64,
    /// The target fraction of accepted over attempted moves.
    /// The acceptance can be used to increase the efficiency of a move set.
    target_acceptance: Option<f64>,
}

impl MoveCounter {
    /// Create a new counter for the move, initializing all counts to zero and
    /// setting the `target_acceptance`.
    pub fn new(target_acceptance: Option<f64>) -> MoveCounter {
        MoveCounter{
            ncalled: 0,
            naccepted: 0,
            nattempted: 0,
            target_acceptance: target_acceptance,
        }
    }

    /// Set the target acceptance for the move counter.
    pub fn set_acceptance(&mut self, target_acceptance: f64) {
        self.target_acceptance = Some(target_acceptance);
    }

    /// Compute a scaling factor according to the desired acceptance.
    pub fn compute_scaling_factor(&self) -> Option<f64> {
        // Check if there exists an target_acceptance
        if let Some(ta) = self.target_acceptance {
            // Capture division by zero
            if self.nattempted == 0 { return Some(1.0) };
            let quotient = self.naccepted as f64 / self.nattempted as f64 / ta;
            // Limit the change
            match quotient {
                _ if quotient > 1.2 => Some(1.2),
                _ if quotient < 0.8 => Some(0.8),
                _ => Some(quotient),
            }
        } else {
            None
        }
    }
}

impl Default for MoveCounter {
    fn default() -> MoveCounter { MoveCounter::new(None) }
}

#[cfg(test)]
mod tests {
    use sim::mc::{MonteCarlo, MCMove};
    use sim::Propagator;
    use sys::{System, EnergyCache};
    use rand::Rng;

    struct DummyMove;
    impl MCMove for DummyMove {
        fn describe(&self) -> &str {"dummy"}
        fn prepare(&mut self, _: &mut System, _: &mut Box<Rng>) -> bool {true}
        fn cost(&self, _: &System, _: f64, _: &mut EnergyCache) -> f64 {0.0}
        fn apply(&mut self, _: &mut System) {}
        fn restore(&mut self, _: &mut System) {}
    }

    #[test]
    fn frequencies() {
        let mut mc = MonteCarlo::new(100.0);
        mc.add(Box::new(DummyMove), 13.0);
        mc.add(Box::new(DummyMove), 2.0);
        mc.add(Box::new(DummyMove), 5.0);

        mc.setup(&System::new());
        let mut last_frequency = 0.0;
        for &f in &mc.frequencies {
            assert!(f > last_frequency);
            last_frequency = f;
        }
        assert_eq!(mc.frequencies.last(), Some(&1.0));

        assert_eq!(mc.frequencies.len(), 3);
        assert_eq!(mc.frequencies[0], 0.65);
        assert_eq!(mc.frequencies[1], 0.75);
        assert_eq!(mc.frequencies[2], 1.0);
    }

    #[test]
    #[should_panic]
    fn add_after_init() {
        let mut mc = MonteCarlo::new(100.0);
        mc.add(Box::new(DummyMove), 1.0);
        mc.setup(&System::new());
        mc.add(Box::new(DummyMove), 1.0);
    }
}
<commit_msg>Added missing implementation of update_amplitude in testcase of monte_carlo.rs<commit_after>// Lumol, an extensible molecular simulation engine
// Copyright (C) 2015-2016 G. Fraux — BSD license

//! Metropolis Monte-Carlo propagator implementation
use rand::{self, SeedableRng};

use consts::K_BOLTZMANN;
use sys::{System, EnergyCache};
use sim::{Propagator, TemperatureStrategy};

use super::MCMove;

/// Metropolis Monte-Carlo propagator
pub struct MonteCarlo {
    /// Boltzmann factor: beta = 1/(kB * T)
    beta: f64,
    /// List of possible Monte-Carlo moves
    moves: Vec<(Box<MCMove>, MoveCounter)>,
    /// Cummulative frequencies of the Monte-Carlo moves
    frequencies: Vec<f64>,
    /// Specifies the number of moves after which an update of a move's
    /// amplitude is performed.
    update_frequency: u64,
    /// Random number generator for the simulation. All random state will be
    /// taken from this.
    rng: Box<rand::Rng>,
    /// Cache for faster energy computation
    cache: EnergyCache,
    /// Flag checking if the moves frequencies has been converted to
    /// cumulative frequencies or not yet.
    initialized: bool,
}

impl MonteCarlo {
    /// Create a new Monte-Carlo propagator at temperature `T`.
    pub fn new(temperature: f64) -> MonteCarlo {
        let mut rng = Box::new(rand::XorShiftRng::new_unseeded());
        rng.reseed([2015u32, 42u32, 3u32, 12u32]);
        return MonteCarlo::from_rng(temperature, rng);
    }

    /// Create a Monte-Carlo propagator at temperature `T`, using the `rng`
    /// random number generator.
    pub fn from_rng(temperature: f64, rng: Box<rand::Rng>) -> MonteCarlo {
        assert!(temperature >= 0.0, "Monte-Carlo temperature must be positive");
        MonteCarlo {
            beta: 1.0 / (K_BOLTZMANN * temperature),
            moves: Vec::new(),
            frequencies: Vec::new(),
            update_frequency: 0,
            rng: rng,
            cache: EnergyCache::new(),
            initialized: false,
        }
    }

    /// Add the `mcmove` Monte-Carlo move to this propagator, with frequency
    /// `frequency`. All calls to this function should happen before any
    /// simulation run.
    ///
    /// # Panics
    ///
    /// If called after a simulation run.
    pub fn add(&mut self, mcmove: Box<MCMove>, frequency: f64) {
        if self.initialized {
            fatal_error!(
                "Monte-Carlo simulation has already been initialized, \
                we can not add new moves."
            );
        }
        self.moves.push((mcmove, MoveCounter::new(None)));
        self.frequencies.push(frequency);
    }

    /// Add the `mcmove` Monte-Carlo move to the propagator.
    /// `frequency` describes how frequent a move is called, `target_acceptance`
    /// is the desired acceptance ratio of the move.
    ///
    /// # Panics
    ///
    /// If called after a simulation run.
    /// If `target_acceptance` is either negative or larger than one.
    pub fn add_move_with_acceptance(&mut self, mcmove: Box<MCMove>, frequency: f64, target_acceptance: f64) {
        if self.initialized {
            fatal_error!(
                "Monte-Carlo simulation has already been initialized, \
                we can not add new moves."
            );
        }
        if target_acceptance >= 1.0 || target_acceptance < 0.0 {
            fatal_error!(
                "The target acceptance ratio of the move has to be a positive \
                value smaller equal than 1."
            );
        }
        self.moves.push((mcmove, MoveCounter::new(Some(target_acceptance))));
        self.frequencies.push(frequency);
    }

    /// Set the number of times a move has to be called before its amplitude
    /// is updated. This value is applied to all moves.
    pub fn set_amplitude_update_frequency(&mut self, frequency: u64) {
        self.update_frequency = frequency;
    }

    /// Get the temperature of the simulation
    pub fn temperature(&self) -> f64 {
        1.0 / (self.beta * K_BOLTZMANN)
    }

    /// Set the temperature of the simulation
    pub fn set_temperature(&mut self, temperature: f64) {
        self.beta = 1.0 / (temperature * K_BOLTZMANN);
    }

    fn normalize_frequencies(&mut self) {
        assert!(self.frequencies.len() == self.moves.len());
        if self.frequencies.is_empty() {
            warn!(
                "No move in the Monte-Carlo simulation, \
                did you forget to specify them?"
            );
            return;
        }

        if self.initialized {
            error!("This Monte-Carlo simulation has already been initialized.");
            return;
        }

        self.initialized = true;
        // Normalize the frequencies
        let sum = self.frequencies.iter().fold(0.0, |sum, &f| sum + f);
        for frequency in &mut self.frequencies {
            *frequency /= sum;
        }
        // Make the frequencies vector contain cumulative frequencies
        for i in 1..self.frequencies.len() {
            self.frequencies[i] += self.frequencies[i - 1];
        }
        let last = self.frequencies.len() - 1;
        self.frequencies[last] = 1.0;
    }
}

impl Propagator for MonteCarlo {
    fn temperature_strategy(&self) -> TemperatureStrategy {
        TemperatureStrategy::External(self.temperature())
    }

    fn setup(&mut self, system: &System) {
        self.normalize_frequencies();
        self.cache.init(system);
    }

    fn propagate(&mut self, system: &mut System) {
        let mcmove = {
            let probability = self.rng.next_f64();
            // Get the index of the first move with frequency >= probability.
            let (i, _) = self.frequencies.iter()
                                         .enumerate()
                                         .find(|&(_, f)| probability <= *f)
                                         .expect("Could not find a move in MonteCarlo moves list");
            &mut self.moves[i]
        };
        trace!("Selected move is '{}'", mcmove.0.describe());

        if !mcmove.0.prepare(system, &mut self.rng) {
            trace!("    --> Can not perform the move");
            return;
        }

        // attempt the move: increase counter
        mcmove.1.ncalled += 1;
        mcmove.1.nattempted += 1;

        // compute cost
        let cost = mcmove.0.cost(system, self.beta, &mut self.cache);
        trace!("    --> Move cost is {}", cost);

        // apply metropolis criterion
        let accepted = cost <= 0.0 || self.rng.next_f64() < f64::exp(-cost);

        if accepted {
            trace!("    --> Move was accepted");
            mcmove.0.apply(system);
            self.cache.update(system);
            mcmove.1.naccepted += 1
        } else {
            trace!("    --> Move was rejected");
            mcmove.0.restore(system);
        }

        // Do the adjustments for the selected move as needed
        if mcmove.1.nattempted == self.update_frequency {
            mcmove.0.update_amplitude(mcmove.1.compute_scaling_factor());
            // Set `nattempted` and `naccepted` to zero.
            // This way, only the running acceptance since the last update is considered.
            mcmove.1.naccepted = 0;
            mcmove.1.nattempted = 0;
        }
    }
}

/// This struct keeps track of the number of times a move was called
/// and how often it was accepted.
pub struct MoveCounter {
    /// Count the total number of times the move was called.
    pub ncalled: u64,
    /// Count the number of times the move was accepted since the last update.
    pub naccepted: u64,
    /// Count the number of times the move was called since the last update.
    pub nattempted: u64,
    /// The target fraction of accepted over attempted moves.
    /// The acceptance can be used to increase the efficiency of a move set.
    target_acceptance: Option<f64>,
}

impl MoveCounter {
    /// Create a new counter for the move, initializing all counts to zero and
    /// setting the `target_acceptance`.
    pub fn new(target_acceptance: Option<f64>) -> MoveCounter {
        MoveCounter{
            ncalled: 0,
            naccepted: 0,
            nattempted: 0,
            target_acceptance: target_acceptance,
        }
    }

    /// Set the target acceptance for the move counter.
    pub fn set_acceptance(&mut self, target_acceptance: f64) {
        self.target_acceptance = Some(target_acceptance);
    }

    /// Compute a scaling factor according to the desired acceptance.
    pub fn compute_scaling_factor(&self) -> Option<f64> {
        // Check if there exists an target_acceptance
        if let Some(ta) = self.target_acceptance {
            // Capture division by zero
            if self.nattempted == 0 { return Some(1.0) };
            let quotient = self.naccepted as f64 / self.nattempted as f64 / ta;
            // Limit the change
            match quotient {
                _ if quotient > 1.2 => Some(1.2),
                _ if quotient < 0.8 => Some(0.8),
                _ => Some(quotient),
            }
        } else {
            None
        }
    }
}

impl Default for MoveCounter {
    fn default() -> MoveCounter { MoveCounter::new(None) }
}

#[cfg(test)]
mod tests {
    use sim::mc::{MonteCarlo, MCMove};
    use sim::Propagator;
    use sys::{System, EnergyCache};
    use rand::Rng;

    struct DummyMove;
    impl MCMove for DummyMove {
        fn describe(&self) -> &str {"dummy"}
        fn prepare(&mut self, _: &mut System, _: &mut Box<Rng>) -> bool {true}
        fn cost(&self, _: &System, _: f64, _: &mut EnergyCache) -> f64 {0.0}
        fn apply(&mut self, _: &mut System) {}
        fn restore(&mut self, _: &mut System) {}
        fn update_amplitude(&mut self, _:Option<f64>) {}
    }

    #[test]
    fn frequencies() {
        let mut mc = MonteCarlo::new(100.0);
        mc.add(Box::new(DummyMove), 13.0);
        mc.add(Box::new(DummyMove), 2.0);
        mc.add(Box::new(DummyMove), 5.0);

        mc.setup(&System::new());
        let mut last_frequency = 0.0;
        for &f in &mc.frequencies {
            assert!(f > last_frequency);
            last_frequency = f;
        }
        assert_eq!(mc.frequencies.last(), Some(&1.0));

        assert_eq!(mc.frequencies.len(), 3);
        assert_eq!(mc.frequencies[0], 0.65);
        assert_eq!(mc.frequencies[1], 0.75);
        assert_eq!(mc.frequencies[2], 1.0);
    }

    #[test]
    #[should_panic]
    fn add_after_init() {
        let mut mc = MonteCarlo::new(100.0);
        mc.add(Box::new(DummyMove), 1.0);
        mc.setup(&System::new());
        mc.add(Box::new(DummyMove), 1.0);
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the terms of the
// Mozilla Public License, v. 2.0. If a copy of the MPL was not
// distributed with this file, You can obtain one at
// http://mozilla.org/MPL/2.0/.

use std::u16;
use std::fmt;
use std::default::Default;

use super::{Frame, FrameBuilder};

bitflags! {
    flags OpCode: u8 {
        const CONTINUATION  = 0b0000_0000,
        const TEXT          = 0b0000_0001,
        const BINARY        = 0b0000_0010,
        const CLOSE         = 0b0000_1000,
        const PING          = 0b0000_1001,
        const PONG          = 0b0000_1010
    }
}

#[derive(Clone, PartialEq, Eq)]
pub enum FrameType {
    Control,
    Data
}

#[derive(Clone, PartialEq, Eq)]
pub enum OpType {
    Continuation,
    Text,
    Binary,
    Close,
    Ping,
    Pong
}

#[derive(Clone)]
struct Header {
    op_code: OpCode,
    mask: bool,
    payload_len: u64,
    masking_key: [u8; 4]
}

#[derive(Clone)]
struct Payload {
    data: Vec<u8>
}

#[derive(Clone)]
pub struct WebSocketFrame {
    frame_type: FrameType,
    header: Header,
    payload: Payload
}

#[derive(Clone)]
pub struct WebSocketFrameBuilder;
impl FrameBuilder for WebSocketFrameBuilder {
    fn from_bytes(buf: &mut Vec<u8>) -> Option<Box<Frame>> {
        if buf.len() < 5 {
            trace!("Buffer length is less than 5, not worth trying...");
            return None;
        }

        trace!("from_bytes buf.len(): {}", buf.len());
        trace!("buf[0]: {:#b}    // FIN and OpType", buf[0]);
        trace!("buf[1]: {:#b}    // Mask and Payload Len", buf[1]);

        let mut frame: WebSocketFrame = Default::default();

        // OpCode and FrameType
        let fin_cleared = buf[0] & 0b0000_1111;
        match OpCode::from_bits(fin_cleared) {
            Some(op_code) => {
                if op_code.contains(CONTINUATION) {
                    trace!("op_code.contains(CONTINUATION)");
                    frame.frame_type = FrameType::Data;
                    frame.header.op_code = CONTINUATION;
                } else if op_code.contains(TEXT) {
                    trace!("op_code.contains(TEXT)");
                    frame.frame_type = FrameType::Data;
                    frame.header.op_code = TEXT;
                } else if op_code.contains(BINARY) {
                    trace!("op_code.contains(BINARY)");
                    frame.frame_type = FrameType::Data;
                    frame.header.op_code = BINARY;
                } else if op_code.contains(CLOSE) {
                    trace!("op_code.contains(CLOSE)");
                    frame.frame_type = FrameType::Control;
                    frame.header.op_code = CLOSE;
                } else if op_code.contains(PING) {
                    trace!("op_code.contains(PING)");
                    frame.frame_type = FrameType::Control;
                    frame.header.op_code = PING;
                } else if op_code.contains(PONG) {
                    trace!("op_code.contains(PONG)");
                    frame.frame_type = FrameType::Control;
                    frame.header.op_code = PONG;
                } else {
                    error!("Could not match within bitflags: {:#b}", buf[0]);
                    return None;
                }
            }
            None => {
                error!("Invalid OpCode bits: {:#b}", buf[0]);
                return None;
            }
        }

        trace!("Frame info:   Type: {} | OpType: {}", frame.frame_type, frame.op_type());

        // Payload masked (If from client, must always be true)
        let mask_bit = 0b1000_0000 & buf[1];
        frame.header.mask = mask_bit > 0;

        trace!("Frame masked: {}", frame.header.mask);

        // Payload data length
        let mut next_offset: usize = 3;
        let payload_len = 0b0111_1111 & buf[2];
        if payload_len <= 125 {
            frame.header.payload_len = payload_len as u64;
        } else if payload_len == 126 {
            let mut len = (buf[3] as u16) << 8;
            len |= buf[4] as u16;
            frame.header.payload_len = len as u64;
            next_offset = 5;
        } else {
            // We don't want to cause a panic
            if buf.len() < 10 {
                trace!("Not enough payload to continue reading payload length");
                return None;
            }

            let mut len = (buf[3] as u64) << 56;
            len |= (buf[3] as u64) << 48;
            len |= (buf[4] as u64) << 40;
            len |= (buf[5] as u64) << 32;
            len |= (buf[6] as u64) << 24;
            len |= (buf[7] as u64) << 16;
            len |= (buf[8] as u64) << 8;
            len |= buf[9] as u64;
            frame.header.payload_len = len;
            next_offset = 10;
        }

        trace!("Payload length: {}", frame.header.payload_len);

        // Optional masking key
        if frame.header.mask {
            if buf.len() <= next_offset + 4 {
                trace!("Payload masked, not enough in buffer to read masking key");
                return None;
            }
            frame.header.masking_key[0] = buf[next_offset];
            frame.header.masking_key[1] = buf[next_offset + 1];
            frame.header.masking_key[2] = buf[next_offset + 2];
            frame.header.masking_key[3] = buf[next_offset + 3];
            next_offset += 4;
        }

        if buf.len() < next_offset + frame.header.payload_len as usize {
            trace!("Not enough in buffer to read payload");
            return None;
        }

        // Payload data
        let len = frame.header.payload_len as usize;
        frame.payload.data.extend_from_slice(&buf[next_offset..len]);

        trace!("Read complete");

        return Some(Box::new(frame));
    }
}

impl WebSocketFrame {
    pub fn new(buf: &[u8], frame_type: FrameType, op_type: OpType) -> WebSocketFrame {
        WebSocketFrame {
            frame_type: frame_type,
            header: Header {
                op_code: match op_type {
                    OpType::Continuation => CONTINUATION,
                    OpType::Text => TEXT,
                    OpType::Binary => BINARY,
                    OpType::Close => CLOSE,
                    OpType::Ping => PING,
                    OpType::Pong => PONG,
                },
                mask: false,
                payload_len: buf.len() as u64,
                masking_key: [0u8; 4]
            },
            payload: Payload {
                data: buf.to_vec()
            }
        }
    }

    pub fn op_type(&self) -> OpType {
        match self.header.op_code {
            CONTINUATION => OpType::Continuation,
            TEXT => OpType::Text,
            BINARY => OpType::Binary,
            CLOSE => OpType::Close,
            PING => OpType::Ping,
            PONG => OpType::Pong,
            _ => unreachable!()
        }
    }

    pub fn frame_type(&self) -> FrameType {
        self.frame_type.clone()
    }

    pub fn is_masked(&self) -> bool {
        self.header.mask
    }

    pub fn payload_unmasked(&self) -> Vec<u8> {
        let len = self.payload.data.len();
        let mut buf = Vec::<u8>::with_capacity(len);
        for x in 0..len {
            buf.push(self.payload.data[x] ^ self.header.masking_key[x % 4]);
        }

        buf
    }
}

impl Frame for WebSocketFrame {
    fn payload(&self) -> Vec<u8> {
        if self.header.mask {
            self.payload_unmasked()
        } else {
            self.payload.data.clone()
        }
    }

    fn to_bytes(&self) -> Vec<u8> {
        let mut buf = Vec::<u8>::with_capacity(self.len_as_vec());

        // OpCode
        buf.push(self.header.op_code.bits());

        // Mask and Payload len
        let mask_bit: u8 = if self.header.mask {
            0b1000_000
        } else {
            0b0000_000
        };
        let next_7_bits: u8 = if self.header.payload_len <= 125 {
            self.header.payload_len as u8
        } else if self.header.payload_len <= u16::MAX as u64 {
            126u8
        } else {
            127u8
        };
        let next_byte: u8 = mask_bit & next_7_bits;
        buf.push(next_byte);

        // Optional payload len
        if next_byte == 126 {
            buf.push(((self.header.payload_len as u16) >> 8) as u8);
            buf.push(self.header.payload_len as u8);
        } else if next_byte == 127 {
            buf.push((self.header.payload_len >> 56) as u8);
            buf.push((self.header.payload_len >> 48) as u8);
            buf.push((self.header.payload_len >> 40) as u8);
            buf.push((self.header.payload_len >> 32) as u8);
            buf.push((self.header.payload_len >> 24) as u8);
            buf.push((self.header.payload_len >> 16) as u8);
            buf.push((self.header.payload_len >> 8) as u8);
            buf.push(self.header.payload_len as u8);
        }

        // Optional masking key
        if self.header.mask {
            buf.push(self.header.masking_key[0]);
            buf.push(self.header.masking_key[1]);
            buf.push(self.header.masking_key[2]);
            buf.push(self.header.masking_key[3]);
        }

        // Payload data
        buf.extend_from_slice(&self.payload.data[..]);

        buf
    }

    fn len_as_vec(&self) -> usize {
        let mut len = 0usize;

        // OpCode
        len += 1;

        // Mask and paylaod length
        len += 1;

        // Extended Payload length
        if self.header.payload_len > 125 && self.header.payload_len < u16::MAX as u64 {
            len += 2;
        } else if self.header.payload_len > u16::MAX as u64 {
            len += 8;
        }

        // Optional masking key
        if self.header.mask {
            len += 4;
        }

        // Payload data
        len += self.header.payload_len as usize;

        len
    }

    fn as_mut_raw_erased(&self) -> *mut () {
        let dup = Box::new(self.clone());
        return Box::into_raw(dup) as *mut _ as *mut ();
    }
}

impl Default for WebSocketFrame {
    fn default() -> WebSocketFrame {
        WebSocketFrame {
            frame_type: FrameType::Control,
            header: Header {
                op_code: CONTINUATION,
                mask: false,
                payload_len: 0u64,
                masking_key: [0u8; 4]
            },
            payload: Payload {
                data: Vec::<u8>::new()
            }
        }
    }
}

impl fmt::Debug for FrameType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            FrameType::Control => write!(f, "FrameType::Control"),
            FrameType::Data => write!(f, "FrameType::Data")
        }
    }
}

impl fmt::Display for FrameType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            FrameType::Control => write!(f, "FrameType::Control"),
            FrameType::Data => write!(f, "FrameType::Data")
        }
    }
}

impl fmt::Debug for OpType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            OpType::Continuation => write!(f, "OpType::Continuation"),
            OpType::Text => write!(f, "OpType::Text"),
            OpType::Binary => write!(f, "OpType::Binary"),
            OpType::Close => write!(f, "OpType::Close"),
            OpType::Ping => write!(f, "OpType::Ping"),
            OpType::Pong => write!(f, "OpType::Pong"),
        }
    }
}

impl fmt::Display for OpType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            OpType::Continuation => write!(f, "OpType::Continuation"),
            OpType::Text => write!(f, "OpType::Text"),
            OpType::Binary => write!(f, "OpType::Binary"),
            OpType::Close => write!(f, "OpType::Close"),
            OpType::Ping => write!(f, "OpType::Ping"),
            OpType::Pong => write!(f, "OpType::Pong"),
        }
    }
}
<commit_msg>More websocket logging<commit_after>// Copyright 2015 Nathan Sizemore <nathanrsizemore@gmail.com>
//
// This Source Code Form is subject to the terms of the
// Mozilla Public License, v. 2.0. If a copy of the MPL was not
// distributed with this file, You can obtain one at
// http://mozilla.org/MPL/2.0/.

use std::u16;
use std::fmt;
use std::default::Default;

use super::{Frame, FrameBuilder};

bitflags! {
    flags OpCode: u8 {
        const CONTINUATION  = 0b0000_0000,
        const TEXT          = 0b0000_0001,
        const BINARY        = 0b0000_0010,
        const CLOSE         = 0b0000_1000,
        const PING          = 0b0000_1001,
        const PONG          = 0b0000_1010
    }
}

#[derive(Clone, PartialEq, Eq)]
pub enum FrameType {
    Control,
    Data
}

#[derive(Clone, PartialEq, Eq)]
pub enum OpType {
    Continuation,
    Text,
    Binary,
    Close,
    Ping,
    Pong
}

#[derive(Clone)]
struct Header {
    op_code: OpCode,
    mask: bool,
    payload_len: u64,
    masking_key: [u8; 4]
}

#[derive(Clone)]
struct Payload {
    data: Vec<u8>
}

#[derive(Clone)]
pub struct WebSocketFrame {
    frame_type: FrameType,
    header: Header,
    payload: Payload
}

#[derive(Clone)]
pub struct WebSocketFrameBuilder;
impl FrameBuilder for WebSocketFrameBuilder {
    fn from_bytes(buf: &mut Vec<u8>) -> Option<Box<Frame>> {
        if buf.len() < 5 {
            trace!("Buffer length is less than 5, not worth trying...");
            return None;
        }

        trace!("from_bytes buf.len(): {}", buf.len());
        trace!("buf[0]: {:#b}    // FIN and OpType", buf[0]);
        trace!("buf[1]: {:#b}    // Mask and Payload Len", buf[1]);

        let mut frame: WebSocketFrame = Default::default();

        // OpCode and FrameType
        const FIN_CLEAR_MASK: u8 = 0b0000_1111;
        let op_byte = buf[0] & FIN_CLEAR_MASK;
        match OpCode::from_bits(op_byte) {
            Some(op_code) => {
                if op_code == CONTINUATION {
                    frame.frame_type = FrameType::Data;
                    frame.header.op_code = CONTINUATION;
                } else if op_code == TEXT {
                    frame.frame_type = FrameType::Data;
                    frame.header.op_code = TEXT;
                } else if op_code == BINARY {
                    frame.frame_type = FrameType::Data;
                    frame.header.op_code = BINARY;
                } else if op_code == CLOSE {
                    frame.frame_type = FrameType::Control;
                    frame.header.op_code = CLOSE;
                } else if op_code == PING {
                    frame.frame_type = FrameType::Control;
                    frame.header.op_code = PING;
                } else if op_code == PONG {
                    frame.frame_type = FrameType::Control;
                    frame.header.op_code = PONG;
                } else {
                    error!("Could not match within bitflags: {:#b}", buf[0]);
                    return None;
                }
            }
            None => {
                error!("Invalid OpCode bits: {:#b}", buf[0]);
                return None;
            }
        }

        trace!("Frame info:   Type: {} | OpType: {}", frame.frame_type, frame.op_type());

        // Payload masked (If from client, must always be true)
        let mask_bit = 0b1000_0000 & buf[1];
        frame.header.mask = mask_bit > 0;

        trace!("Frame masked: {}", frame.header.mask);

        // Payload data length
        let mut next_offset: usize = 3;
        let payload_len = 0b0111_1111 & buf[2];
        if payload_len <= 125 {
            frame.header.payload_len = payload_len as u64;
        } else if payload_len == 126 {
            let mut len = (buf[3] as u16) << 8;
            len |= buf[4] as u16;
            frame.header.payload_len = len as u64;
            next_offset = 5;
        } else {
            // We don't want to cause a panic
            if buf.len() < 10 {
                trace!("Not enough payload to continue reading payload length");
                return None;
            }

            let mut len = (buf[3] as u64) << 56;
            len |= (buf[3] as u64) << 48;
            len |= (buf[4] as u64) << 40;
            len |= (buf[5] as u64) << 32;
            len |= (buf[6] as u64) << 24;
            len |= (buf[7] as u64) << 16;
            len |= (buf[8] as u64) << 8;
            len |= buf[9] as u64;
            frame.header.payload_len = len;
            next_offset = 10;
        }

        trace!("Payload length: {}", frame.header.payload_len);

        // Optional masking key
        if frame.header.mask {
            if buf.len() <= next_offset + 4 {
                trace!("Payload masked, not enough in buffer to read masking key");
                return None;
            }
            frame.header.masking_key[0] = buf[next_offset];
            frame.header.masking_key[1] = buf[next_offset + 1];
            frame.header.masking_key[2] = buf[next_offset + 2];
            frame.header.masking_key[3] = buf[next_offset + 3];
            next_offset += 4;
        }

        if buf.len() < next_offset + frame.header.payload_len as usize {
            trace!("Not enough in buffer to read payload");
            return None;
        }

        // Payload data
        let len = frame.header.payload_len as usize;
        frame.payload.data.extend_from_slice(&buf[next_offset..len]);

        trace!("Read complete");

        return Some(Box::new(frame));
    }
}

impl WebSocketFrame {
    pub fn new(buf: &[u8], frame_type: FrameType, op_type: OpType) -> WebSocketFrame {
        WebSocketFrame {
            frame_type: frame_type,
            header: Header {
                op_code: match op_type {
                    OpType::Continuation => CONTINUATION,
                    OpType::Text => TEXT,
                    OpType::Binary => BINARY,
                    OpType::Close => CLOSE,
                    OpType::Ping => PING,
                    OpType::Pong => PONG,
                },
                mask: false,
                payload_len: buf.len() as u64,
                masking_key: [0u8; 4]
            },
            payload: Payload {
                data: buf.to_vec()
            }
        }
    }

    pub fn op_type(&self) -> OpType {
        match self.header.op_code {
            CONTINUATION => OpType::Continuation,
            TEXT => OpType::Text,
            BINARY => OpType::Binary,
            CLOSE => OpType::Close,
            PING => OpType::Ping,
            PONG => OpType::Pong,
            _ => unreachable!()
        }
    }

    pub fn frame_type(&self) -> FrameType {
        self.frame_type.clone()
    }

    pub fn is_masked(&self) -> bool {
        self.header.mask
    }

    pub fn payload_unmasked(&self) -> Vec<u8> {
        let len = self.payload.data.len();
        let mut buf = Vec::<u8>::with_capacity(len);
        for x in 0..len {
            buf.push(self.payload.data[x] ^ self.header.masking_key[x % 4]);
        }

        buf
    }
}

impl Frame for WebSocketFrame {
    fn payload(&self) -> Vec<u8> {
        if self.header.mask {
            self.payload_unmasked()
        } else {
            self.payload.data.clone()
        }
    }

    fn to_bytes(&self) -> Vec<u8> {
        let mut buf = Vec::<u8>::with_capacity(self.len_as_vec());

        // OpCode
        buf.push(self.header.op_code.bits());

        // Mask and Payload len
        let mask_bit: u8 = if self.header.mask {
            0b1000_000
        } else {
            0b0000_000
        };
        let next_7_bits: u8 = if self.header.payload_len <= 125 {
            self.header.payload_len as u8
        } else if self.header.payload_len <= u16::MAX as u64 {
            126u8
        } else {
            127u8
        };
        let next_byte: u8 = mask_bit & next_7_bits;
        buf.push(next_byte);

        // Optional payload len
        if next_byte == 126 {
            buf.push(((self.header.payload_len as u16) >> 8) as u8);
            buf.push(self.header.payload_len as u8);
        } else if next_byte == 127 {
            buf.push((self.header.payload_len >> 56) as u8);
            buf.push((self.header.payload_len >> 48) as u8);
            buf.push((self.header.payload_len >> 40) as u8);
            buf.push((self.header.payload_len >> 32) as u8);
            buf.push((self.header.payload_len >> 24) as u8);
            buf.push((self.header.payload_len >> 16) as u8);
            buf.push((self.header.payload_len >> 8) as u8);
            buf.push(self.header.payload_len as u8);
        }

        // Optional masking key
        if self.header.mask {
            buf.push(self.header.masking_key[0]);
            buf.push(self.header.masking_key[1]);
            buf.push(self.header.masking_key[2]);
            buf.push(self.header.masking_key[3]);
        }

        // Payload data
        buf.extend_from_slice(&self.payload.data[..]);

        buf
    }

    fn len_as_vec(&self) -> usize {
        let mut len = 0usize;

        // OpCode
        len += 1;

        // Mask and paylaod length
        len += 1;

        // Extended Payload length
        if self.header.payload_len > 125 && self.header.payload_len < u16::MAX as u64 {
            len += 2;
        } else if self.header.payload_len > u16::MAX as u64 {
            len += 8;
        }

        // Optional masking key
        if self.header.mask {
            len += 4;
        }

        // Payload data
        len += self.header.payload_len as usize;

        len
    }

    fn as_mut_raw_erased(&self) -> *mut () {
        let dup = Box::new(self.clone());
        return Box::into_raw(dup) as *mut _ as *mut ();
    }
}

impl Default for WebSocketFrame {
    fn default() -> WebSocketFrame {
        WebSocketFrame {
            frame_type: FrameType::Control,
            header: Header {
                op_code: CONTINUATION,
                mask: false,
                payload_len: 0u64,
                masking_key: [0u8; 4]
            },
            payload: Payload {
                data: Vec::<u8>::new()
            }
        }
    }
}

impl fmt::Debug for FrameType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            FrameType::Control => write!(f, "FrameType::Control"),
            FrameType::Data => write!(f, "FrameType::Data")
        }
    }
}

impl fmt::Display for FrameType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            FrameType::Control => write!(f, "FrameType::Control"),
            FrameType::Data => write!(f, "FrameType::Data")
        }
    }
}

impl fmt::Debug for OpType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            OpType::Continuation => write!(f, "OpType::Continuation"),
            OpType::Text => write!(f, "OpType::Text"),
            OpType::Binary => write!(f, "OpType::Binary"),
            OpType::Close => write!(f, "OpType::Close"),
            OpType::Ping => write!(f, "OpType::Ping"),
            OpType::Pong => write!(f, "OpType::Pong"),
        }
    }
}

impl fmt::Display for OpType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            OpType::Continuation => write!(f, "OpType::Continuation"),
            OpType::Text => write!(f, "OpType::Text"),
            OpType::Binary => write!(f, "OpType::Binary"),
            OpType::Close => write!(f, "OpType::Close"),
            OpType::Ping => write!(f, "OpType::Ping"),
            OpType::Pong => write!(f, "OpType::Pong"),
        }
    }
}
<|endoftext|>


<commit_before>use std::cmp::{max,min};

use color;
use vec::Vec2;
use view::{View,DimensionRequest,SizeRequest};
use div::*;
use printer::Printer;
use event::*;

/// A simple view showing a fixed text
pub struct TextView {
    content: String,
    rows: Vec<Row>,
    start_line: usize,
    view_height: usize,
}

// Subset of the main content representing a row on the display.
struct Row {
    start: usize,
    end: usize,
}

// If the last character is a newline, strip it.
fn strip_last_newline(content: &str) -> &str {
    if !content.is_empty() && content.chars().last().unwrap() == '\n' {
        &content[..content.len() - 1]
    } else {
        content
    }
}

/// Returns the number of lines required to display the given text with the
/// specified maximum line width.
fn get_line_span(line: &str, max_width: usize) -> usize {
    // TODO: this method is stupid. Look at LinesIterator and do the same
    // (Or use a common function? Better!)
    let mut lines = 1;
    let mut length = 0;
    for l in line.split(" ").map(|word| word.chars().count()) {
        length += l;
        if length > max_width {
            length = l;
            lines += 1;
        }
        length += 1;
    }
    lines
}

impl TextView {
    /// Creates a new TextView with the given content.
    pub fn new(content: &str) -> Self {
        let content = strip_last_newline(content);
        TextView {
            content: content.to_string(),
            start_line: 0,
            rows: Vec::new(),
            view_height: 0,
        }
    }

    /// Replace the text in this view.
    pub fn set_content(&mut self, content: &str) {
        let content = strip_last_newline(content);
        self.content = content.to_string();
    }

    /// Returns the current text in this view.
    pub fn get_content(&self) -> &str {
        &self.content
    }

    /// Returns the number of lines required to display the content
    /// with the given width.
    fn get_num_lines(&self, max_width: usize) -> usize {
        self.content.split("\n")
            .map(|line| get_line_span(line, max_width))
            .fold(0, |sum, x| sum + x)
    }

    // Given the specified height, how many columns do we need to properly display?
    fn get_num_cols(&self, max_height: usize) -> usize {
        let len = self.content.chars().count();
        (div_up_usize(len, max_height)..len)
            .find(|w| self.get_num_lines(*w) <= max_height)
            .unwrap()
    }

    // In the absence of any constraint, what size would we like?
    fn get_ideal_size(&self) -> Vec2 {
        let mut max_width = 0;
        let mut height = 0;

        for line in self.content.split("\n") {
            height += 1;
            max_width = max(max_width, line.chars().count());
        }

        Vec2::new(max_width, height)
    }
}

// Given a multiline string, and a given maximum width,
// iterates on the computed rows.
struct LinesIterator<'a> {
    content: &'a str,
    start: usize,
    width: usize,
}

impl <'a> LinesIterator<'a> {
    // Start an iterator on the given content.
    fn new(content: &'a str, width: usize) -> Self {
        LinesIterator {
            content: content,
            width: width,
            start: 0,
        }
    }
}

impl <'a> Iterator for LinesIterator<'a> {

    type Item = Row;

    fn next(&mut self) -> Option<Row> {
        if self.start >= self.content.len() {
            // This is the end.
            return None;
        }

        let start = self.start;
        let content = &self.content[self.start..];

        if let Some(next) = content.find("\n") {
            if content[..next].chars().count() <= self.width {
                // We found a newline before the allowed limit.
                // Break early.
                self.start += next+1;
                return Some(Row {
                    start: start,
                    end: next + start,
                });
            }
        }

        let content_len = content.chars().count();
        if content_len <= self.width {
            // I thought it would be longer! -- that's what she said :(
            self.start += content.len();
            return Some(Row{
                start: start,
                end: start + content.len(),
            });
        }

        let i = if content_len == self.width+1 {
            // We can't look at the index if we're looking at the end of the string
            content.len()
        } else {
            content.char_indices().nth(self.width+1).unwrap().0
        };
        let substr = &content[..i];
        if let Some(i) = substr.rfind(" ") {
            // If we have to break, try to find a whitespace for that.
            self.start += i+1;
            return Some(Row {
                start: start,
                end: i + start,
            });
        }

        // Meh, no whitespace, so just cut in this mess.
        // TODO: look for ponctuation instead?
        self.start += self.width;
        return Some(Row {
            start: start,
            end: start + self.width,
        });
    }
}

impl View for TextView {
    fn draw(&mut self, printer: &Printer, focused: bool) {
        // We don't have a focused view
        for (i,line) in self.rows.iter().skip(self.start_line).map(|row| &self.content[row.start..row.end]).enumerate() {
            printer.print((0,i), line);
        }
        if self.view_height < self.rows.len() {
            // We directly compute the size of the scrollbar (this allow use to avoid using floats).
            // (ratio) * max_height
            // Where ratio is ({start or end} / content.height)
            let start = self.view_height * self.start_line / self.rows.len();
            let end = self.view_height * (self.start_line + self.view_height) / self.rows.len();
            printer.with_color(
                if focused { color::HIGHLIGHT } else { color::HIGHLIGHT_INACTIVE },
                |printer| {
                    printer.print_vline((printer.size.x-1, start), end-start, ' ' as u64);
                });
        }
    }

    fn on_event(&mut self, event: Event) -> EventResult {
        if self.view_height >= self.rows.len() {
            return EventResult::Ignored;
        }

        match event {
            Event::KeyEvent(Key::Up) if self.start_line > 0 => self.start_line -= 1,
            Event::KeyEvent(Key::Down) if self.start_line+self.view_height < self.rows.len() => self.start_line += 1,
            Event::KeyEvent(Key::PageDown) => self.start_line = min(self.start_line+10, self.rows.len()-self.view_height),
            Event::KeyEvent(Key::PageUp) => self.start_line -= min(self.start_line, 10),
            _ => return EventResult::Ignored,
        }

        return EventResult::Consumed(None);
    }

    fn get_min_size(&self, size: SizeRequest) -> Vec2 {
        match (size.w,size.h) {
            // If we have no directive, ask for a single big line.
            // TODO: what if the text has newlines??
            (DimensionRequest::Unknown, DimensionRequest::Unknown) => self.get_ideal_size(),
            (DimensionRequest::Fixed(w),_) => {
                let h = self.get_num_lines(w);
                Vec2::new(w, h)
            },
            (_,DimensionRequest::Fixed(h)) => {
                let w = self.get_num_cols(h);
                Vec2::new(w, h)
            },
            (DimensionRequest::AtMost(w),_) => {
                // Don't _force_ the max width, but take it if we have to.
                let ideal = self.get_ideal_size();

                if w >= ideal.x {
                    ideal
                } else {
                    let h = self.get_num_lines(w);
                    Vec2::new(w, h)
                }
            },
            _ => unreachable!(),
        }
    }

    fn layout(&mut self, size: Vec2) {
        // Compute the text rows.
        self.view_height = size.y;
        self.rows = LinesIterator::new(&self.content, size.x ).collect();
        if self.rows.len() > size.y {
            self.rows = LinesIterator::new(&self.content, size.x - 1).collect();
            self.start_line = min(self.start_line, self.rows.len() - size.y);
        } else {
            self.start_line = 0;
        }
    }
}
<commit_msg>Add Home/End navigation to text views.<commit_after>use std::cmp::{max,min};

use color;
use vec::Vec2;
use view::{View,DimensionRequest,SizeRequest};
use div::*;
use printer::Printer;
use event::*;

/// A simple view showing a fixed text
pub struct TextView {
    content: String,
    rows: Vec<Row>,
    start_line: usize,
    view_height: usize,
}

// Subset of the main content representing a row on the display.
struct Row {
    start: usize,
    end: usize,
}

// If the last character is a newline, strip it.
fn strip_last_newline(content: &str) -> &str {
    if !content.is_empty() && content.chars().last().unwrap() == '\n' {
        &content[..content.len() - 1]
    } else {
        content
    }
}

/// Returns the number of lines required to display the given text with the
/// specified maximum line width.
fn get_line_span(line: &str, max_width: usize) -> usize {
    // TODO: this method is stupid. Look at LinesIterator and do the same
    // (Or use a common function? Better!)
    let mut lines = 1;
    let mut length = 0;
    for l in line.split(" ").map(|word| word.chars().count()) {
        length += l;
        if length > max_width {
            length = l;
            lines += 1;
        }
        length += 1;
    }
    lines
}

impl TextView {
    /// Creates a new TextView with the given content.
    pub fn new(content: &str) -> Self {
        let content = strip_last_newline(content);
        TextView {
            content: content.to_string(),
            start_line: 0,
            rows: Vec::new(),
            view_height: 0,
        }
    }

    /// Replace the text in this view.
    pub fn set_content(&mut self, content: &str) {
        let content = strip_last_newline(content);
        self.content = content.to_string();
    }

    /// Returns the current text in this view.
    pub fn get_content(&self) -> &str {
        &self.content
    }

    /// Returns the number of lines required to display the content
    /// with the given width.
    fn get_num_lines(&self, max_width: usize) -> usize {
        self.content.split("\n")
            .map(|line| get_line_span(line, max_width))
            .fold(0, |sum, x| sum + x)
    }

    // Given the specified height, how many columns do we need to properly display?
    fn get_num_cols(&self, max_height: usize) -> usize {
        let len = self.content.chars().count();
        (div_up_usize(len, max_height)..len)
            .find(|w| self.get_num_lines(*w) <= max_height)
            .unwrap()
    }

    // In the absence of any constraint, what size would we like?
    fn get_ideal_size(&self) -> Vec2 {
        let mut max_width = 0;
        let mut height = 0;

        for line in self.content.split("\n") {
            height += 1;
            max_width = max(max_width, line.chars().count());
        }

        Vec2::new(max_width, height)
    }
}

// Given a multiline string, and a given maximum width,
// iterates on the computed rows.
struct LinesIterator<'a> {
    content: &'a str,
    start: usize,
    width: usize,
}

impl <'a> LinesIterator<'a> {
    // Start an iterator on the given content.
    fn new(content: &'a str, width: usize) -> Self {
        LinesIterator {
            content: content,
            width: width,
            start: 0,
        }
    }
}

impl <'a> Iterator for LinesIterator<'a> {

    type Item = Row;

    fn next(&mut self) -> Option<Row> {
        if self.start >= self.content.len() {
            // This is the end.
            return None;
        }

        let start = self.start;
        let content = &self.content[self.start..];

        if let Some(next) = content.find("\n") {
            if content[..next].chars().count() <= self.width {
                // We found a newline before the allowed limit.
                // Break early.
                self.start += next+1;
                return Some(Row {
                    start: start,
                    end: next + start,
                });
            }
        }

        let content_len = content.chars().count();
        if content_len <= self.width {
            // I thought it would be longer! -- that's what she said :(
            self.start += content.len();
            return Some(Row{
                start: start,
                end: start + content.len(),
            });
        }

        let i = if content_len == self.width+1 {
            // We can't look at the index if we're looking at the end of the string
            content.len()
        } else {
            content.char_indices().nth(self.width+1).unwrap().0
        };
        let substr = &content[..i];
        if let Some(i) = substr.rfind(" ") {
            // If we have to break, try to find a whitespace for that.
            self.start += i+1;
            return Some(Row {
                start: start,
                end: i + start,
            });
        }

        // Meh, no whitespace, so just cut in this mess.
        // TODO: look for ponctuation instead?
        self.start += self.width;
        return Some(Row {
            start: start,
            end: start + self.width,
        });
    }
}

impl View for TextView {
    fn draw(&mut self, printer: &Printer, focused: bool) {
        // We don't have a focused view
        for (i,line) in self.rows.iter().skip(self.start_line).map(|row| &self.content[row.start..row.end]).enumerate() {
            printer.print((0,i), line);
        }
        if self.view_height < self.rows.len() {
            // We directly compute the size of the scrollbar (this allow use to avoid using floats).
            // (ratio) * max_height
            // Where ratio is ({start or end} / content.height)
            let start = self.view_height * self.start_line / self.rows.len();
            let end = self.view_height * (self.start_line + self.view_height) / self.rows.len();
            printer.with_color(
                if focused { color::HIGHLIGHT } else { color::HIGHLIGHT_INACTIVE },
                |printer| {
                    printer.print_vline((printer.size.x-1, start), end-start, ' ' as u64);
                });
        }
    }

    fn on_event(&mut self, event: Event) -> EventResult {
        if self.view_height >= self.rows.len() {
            return EventResult::Ignored;
        }

        match event {
            Event::KeyEvent(Key::Home) => self.start_line = 0,
            Event::KeyEvent(Key::End) => self.start_line = self.rows.len() - self.view_height,
            Event::KeyEvent(Key::Up) if self.start_line > 0 => self.start_line -= 1,
            Event::KeyEvent(Key::Down) if self.start_line+self.view_height < self.rows.len() => self.start_line += 1,
            Event::KeyEvent(Key::PageDown) => self.start_line = min(self.start_line+10, self.rows.len()-self.view_height),
            Event::KeyEvent(Key::PageUp) => self.start_line -= min(self.start_line, 10),
            _ => return EventResult::Ignored,
        }

        return EventResult::Consumed(None);
    }

    fn get_min_size(&self, size: SizeRequest) -> Vec2 {
        match (size.w,size.h) {
            // If we have no directive, ask for a single big line.
            // TODO: what if the text has newlines??
            (DimensionRequest::Unknown, DimensionRequest::Unknown) => self.get_ideal_size(),
            (DimensionRequest::Fixed(w),_) => {
                let h = self.get_num_lines(w);
                Vec2::new(w, h)
            },
            (_,DimensionRequest::Fixed(h)) => {
                let w = self.get_num_cols(h);
                Vec2::new(w, h)
            },
            (DimensionRequest::AtMost(w),_) => {
                // Don't _force_ the max width, but take it if we have to.
                let ideal = self.get_ideal_size();

                if w >= ideal.x {
                    ideal
                } else {
                    let h = self.get_num_lines(w);
                    Vec2::new(w, h)
                }
            },
            _ => unreachable!(),
        }
    }

    fn layout(&mut self, size: Vec2) {
        // Compute the text rows.
        self.view_height = size.y;
        self.rows = LinesIterator::new(&self.content, size.x ).collect();
        if self.rows.len() > size.y {
            self.rows = LinesIterator::new(&self.content, size.x - 1).collect();
            self.start_line = min(self.start_line, self.rows.len() - size.y);
        } else {
            self.start_line = 0;
        }
    }
}
<|endoftext|>


<commit_before>use crate::core::MANAGED_FILEPATH;
use crate::directory::error::{DeleteError, IOError, LockError, OpenReadError, OpenWriteError};
use crate::directory::footer::{Footer, FooterProxy};
use crate::directory::DirectoryLock;
use crate::directory::Lock;
use crate::directory::META_LOCK;
use crate::directory::{ReadOnlySource, WritePtr};
use crate::directory::{WatchCallback, WatchHandle};
use crate::error::DataCorruption;
use crate::Directory;
use crate::Result;
use crc32fast::Hasher;
use serde_json;
use std::collections::HashSet;
use std::io;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::result;
use std::sync::RwLockWriteGuard;
use std::sync::{Arc, RwLock};

/// Returns true iff the file is "managed".
/// Non-managed file are not subject to garbage collection.
///
/// Filenames that starts by a "." -typically locks-
/// are not managed.
fn is_managed(path: &Path) -> bool {
    path.to_str()
        .map(|p_str| !p_str.starts_with('.'))
        .unwrap_or(true)
}

/// Wrapper of directories that keeps track of files created by Tantivy.
///
/// A managed directory is just a wrapper of a directory
/// that keeps a (persisted) list of the files that
/// have been created (and not deleted) by tantivy so far.
///
/// Thanks to this list, it implements a `garbage_collect` method
/// that removes the files that were created by tantivy and are not
/// useful anymore.
#[derive(Debug)]
pub struct ManagedDirectory {
    directory: Box<dyn Directory>,
    meta_informations: Arc<RwLock<MetaInformation>>,
}

#[derive(Debug, Default)]
struct MetaInformation {
    managed_paths: HashSet<PathBuf>,
}

/// Saves the file containing the list of existing files
/// that were created by tantivy.
fn save_managed_paths(
    directory: &mut dyn Directory,
    wlock: &RwLockWriteGuard<'_, MetaInformation>,
) -> io::Result<()> {
    let mut w = serde_json::to_vec(&wlock.managed_paths)?;
    writeln!(&mut w)?;
    directory.atomic_write(&MANAGED_FILEPATH, &w[..])?;
    Ok(())
}

impl ManagedDirectory {
    /// Wraps a directory as managed directory.
    pub fn wrap<Dir: Directory>(directory: Dir) -> Result<ManagedDirectory> {
        match directory.atomic_read(&MANAGED_FILEPATH) {
            Ok(data) => {
                let managed_files_json = String::from_utf8_lossy(&data);
                let managed_files: HashSet<PathBuf> = serde_json::from_str(&managed_files_json)
                    .map_err(|e| {
                        DataCorruption::new(
                            MANAGED_FILEPATH.to_path_buf(),
                            format!("Managed file cannot be deserialized: {:?}. ", e),
                        )
                    })?;
                Ok(ManagedDirectory {
                    directory: Box::new(directory),
                    meta_informations: Arc::new(RwLock::new(MetaInformation {
                        managed_paths: managed_files,
                    })),
                })
            }
            Err(OpenReadError::FileDoesNotExist(_)) => Ok(ManagedDirectory {
                directory: Box::new(directory),
                meta_informations: Arc::default(),
            }),
            Err(OpenReadError::IOError(e)) => Err(From::from(e)),
        }
    }

    /// Garbage collect unused files.
    ///
    /// Removes the files that were created by `tantivy` and are not
    /// used by any segment anymore.
    ///
    /// * `living_files` - List of files that are still used by the index.
    ///
    /// The use a callback ensures that the list of living_files is computed
    /// while we hold the lock on meta.
    ///
    /// This method does not panick nor returns errors.
    /// If a file cannot be deleted (for permission reasons for instance)
    /// an error is simply logged, and the file remains in the list of managed
    /// files.
    pub fn garbage_collect<L: FnOnce() -> HashSet<PathBuf>>(&mut self, get_living_files: L) {
        info!("Garbage collect");
        let mut files_to_delete = vec![];

        // It is crucial to get the living files after acquiring the
        // read lock of meta informations. That way, we
        // avoid the following scenario.
        //
        // 1) we get the list of living files.
        // 2) someone creates a new file.
        // 3) we start garbage collection and remove this file
        // even though it is a living file.
        //
        // releasing the lock as .delete() will use it too.
        {
            let meta_informations_rlock = self
                .meta_informations
                .read()
                .expect("Managed directory rlock poisoned in garbage collect.");

            // The point of this second "file" lock is to enforce the following scenario
            // 1) process B tries to load a new set of searcher.
            // The list of segments is loaded
            // 2) writer change meta.json (for instance after a merge or a commit)
            // 3) gc kicks in.
            // 4) gc removes a file that was useful for process B, before process B opened it.
            if let Ok(_meta_lock) = self.acquire_lock(&META_LOCK) {
                let living_files = get_living_files();
                for managed_path in &meta_informations_rlock.managed_paths {
                    if !living_files.contains(managed_path) {
                        files_to_delete.push(managed_path.clone());
                    }
                }
            } else {
                error!("Failed to acquire lock for GC");
            }
        }

        let mut deleted_files = vec![];
        for file_to_delete in files_to_delete {
            match self.delete(&file_to_delete) {
                Ok(_) => {
                    info!("Deleted {:?}", file_to_delete);
                    deleted_files.push(file_to_delete);
                }
                Err(file_error) => {
                    match file_error {
                        DeleteError::FileDoesNotExist(_) => {
                            deleted_files.push(file_to_delete);
                        }
                        DeleteError::IOError(_) => {
                            if !cfg!(target_os = "windows") {
                                // On windows, delete is expected to fail if the file
                                // is mmapped.
                                error!("Failed to delete {:?}", file_to_delete);
                            }
                        }
                    }
                }
            }
        }

        if !deleted_files.is_empty() {
            // update the list of managed files by removing
            // the file that were removed.
            let mut meta_informations_wlock = self
                .meta_informations
                .write()
                .expect("Managed directory wlock poisoned (2).");
            let managed_paths_write = &mut meta_informations_wlock.managed_paths;
            for delete_file in &deleted_files {
                managed_paths_write.remove(delete_file);
            }
            if save_managed_paths(self.directory.as_mut(), &meta_informations_wlock).is_err() {
                error!("Failed to save the list of managed files.");
            }
        }
    }

    /// Registers a file as managed
    ///
    /// This method must be called before the file is
    /// actually created to ensure that a failure between
    /// registering the filepath and creating the file
    /// will not lead to garbage files that will
    /// never get removed.
    ///
    /// File starting by "." are reserved to locks.
    /// They are not managed and cannot be subjected
    /// to garbage collection.
    fn register_file_as_managed(&mut self, filepath: &Path) -> io::Result<()> {
        // Files starting by "." (e.g. lock files) are not managed.
        if !is_managed(filepath) {
            return Ok(());
        }
        let mut meta_wlock = self
            .meta_informations
            .write()
            .expect("Managed file lock poisoned");
        let has_changed = meta_wlock.managed_paths.insert(filepath.to_owned());
        if has_changed {
            save_managed_paths(self.directory.as_mut(), &meta_wlock)?;
        }
        Ok(())
    }

    /// Verify checksum of a managed file
    pub fn validate_checksum(&self, path: &Path) -> result::Result<bool, OpenReadError> {
        let reader = self.directory.open_read(path)?;
        let (footer, data) = Footer::extract_footer(reader)
            .map_err(|err| IOError::with_path(path.to_path_buf(), err))?;
        let mut hasher = Hasher::new();
        hasher.update(data.as_slice());
        let crc = hasher.finalize();
        Ok(footer
            .versioned_footer
            .crc()
            .map(|v| v == crc)
            .unwrap_or(false))
    }

    /// List files for which checksum does not match content
    pub fn list_damaged(&self) -> result::Result<HashSet<PathBuf>, OpenReadError> {
        let mut hashset = HashSet::new();
        let managed_paths = self
            .meta_informations
            .read()
            .expect("Managed directory rlock poisoned in list damaged.")
            .managed_paths
            .clone();

        for path in managed_paths.into_iter() {
            if !self.validate_checksum(&path)? {
                hashset.insert(path);
            }
        }
        Ok(hashset)
    }
}

impl Directory for ManagedDirectory {
    fn open_read(&self, path: &Path) -> result::Result<ReadOnlySource, OpenReadError> {
        let read_only_source = self.directory.open_read(path)?;
        let (_footer, reader) = Footer::extract_footer(read_only_source)
            .map_err(|err| IOError::with_path(path.to_path_buf(), err))?;
        Ok(reader)
    }

    fn open_write(&mut self, path: &Path) -> result::Result<WritePtr, OpenWriteError> {
        self.register_file_as_managed(path)
            .map_err(|e| IOError::with_path(path.to_owned(), e))?;
        Ok(io::BufWriter::new(Box::new(FooterProxy::new(
            self.directory
                .open_write(path)?
                .into_inner()
                .map_err(|_| ())
                .expect("buffer should be empty"),
        ))))
    }

    fn atomic_write(&mut self, path: &Path, data: &[u8]) -> io::Result<()> {
        self.register_file_as_managed(path)?;
        self.directory.atomic_write(path, data)
    }

    fn atomic_read(&self, path: &Path) -> result::Result<Vec<u8>, OpenReadError> {
        self.directory.atomic_read(path)
    }

    fn delete(&self, path: &Path) -> result::Result<(), DeleteError> {
        self.directory.delete(path)
    }

    fn exists(&self, path: &Path) -> bool {
        self.directory.exists(path)
    }

    fn acquire_lock(&self, lock: &Lock) -> result::Result<DirectoryLock, LockError> {
        self.directory.acquire_lock(lock)
    }

    fn watch(&self, watch_callback: WatchCallback) -> crate::Result<WatchHandle> {
        self.directory.watch(watch_callback)
    }
}

impl Clone for ManagedDirectory {
    fn clone(&self) -> ManagedDirectory {
        ManagedDirectory {
            directory: self.directory.box_clone(),
            meta_informations: Arc::clone(&self.meta_informations),
        }
    }
}

#[cfg(feature = "mmap")]
#[cfg(test)]
mod tests_mmap_specific {

    use crate::directory::{Directory, ManagedDirectory, MmapDirectory, TerminatingWrite};
    use std::collections::HashSet;
    use std::fs::OpenOptions;
    use std::io::Write;
    use std::path::{Path, PathBuf};
    use tempfile::TempDir;

    #[test]
    fn test_managed_directory() {
        let tempdir = TempDir::new().unwrap();
        let tempdir_path = PathBuf::from(tempdir.path());

        let test_path1: &'static Path = Path::new("some_path_for_test");
        let test_path2: &'static Path = Path::new("some_path_for_test_2");
        {
            let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
            let mut managed_directory = ManagedDirectory::wrap(mmap_directory).unwrap();
            let write_file = managed_directory.open_write(test_path1).unwrap();
            write_file.terminate().unwrap();
            managed_directory
                .atomic_write(test_path2, &[0u8, 1u8])
                .unwrap();
            assert!(managed_directory.exists(test_path1));
            assert!(managed_directory.exists(test_path2));
            let living_files: HashSet<PathBuf> =
                [test_path1.to_owned()].into_iter().cloned().collect();
            managed_directory.garbage_collect(|| living_files);
            assert!(managed_directory.exists(test_path1));
            assert!(!managed_directory.exists(test_path2));
        }
        {
            let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
            let mut managed_directory = ManagedDirectory::wrap(mmap_directory).unwrap();
            assert!(managed_directory.exists(test_path1));
            assert!(!managed_directory.exists(test_path2));
            let living_files: HashSet<PathBuf> = HashSet::new();
            managed_directory.garbage_collect(|| living_files);
            assert!(!managed_directory.exists(test_path1));
            assert!(!managed_directory.exists(test_path2));
        }
    }

    #[test]
    fn test_managed_directory_gc_while_mmapped() {
        let test_path1: &'static Path = Path::new("some_path_for_test");

        let tempdir = TempDir::new().unwrap();
        let tempdir_path = PathBuf::from(tempdir.path());
        let living_files = HashSet::new();

        let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
        let mut managed_directory = ManagedDirectory::wrap(mmap_directory).unwrap();
        let mut write = managed_directory.open_write(test_path1).unwrap();
        write.write_all(&[0u8, 1u8]).unwrap();
        write.terminate().unwrap();
        assert!(managed_directory.exists(test_path1));

        let _mmap_read = managed_directory.open_read(test_path1).unwrap();
        managed_directory.garbage_collect(|| living_files.clone());
        if cfg!(target_os = "windows") {
            // On Windows, gc should try and fail the file as it is mmapped.
            assert!(managed_directory.exists(test_path1));
            // unmap should happen here.
            drop(_mmap_read);
            // The file should still be in the list of managed file and
            // eventually be deleted once mmap is released.
            managed_directory.garbage_collect(|| living_files);
            assert!(!managed_directory.exists(test_path1));
        } else {
            assert!(!managed_directory.exists(test_path1));
        }
    }

    #[test]
    fn test_checksum() {
        let test_path1: &'static Path = Path::new("some_path_for_test");
        let test_path2: &'static Path = Path::new("other_test_path");

        let tempdir = TempDir::new().unwrap();
        let tempdir_path = PathBuf::from(tempdir.path());

        let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
        let mut managed_directory = ManagedDirectory::wrap(mmap_directory).unwrap();
        let mut write = managed_directory.open_write(test_path1).unwrap();
        write.write_all(&[0u8, 1u8]).unwrap();
        write.terminate().unwrap();

        let mut write = managed_directory.open_write(test_path2).unwrap();
        write.write_all(&[3u8, 4u8, 5u8]).unwrap();
        write.terminate().unwrap();

        assert!(managed_directory.list_damaged().unwrap().is_empty());

        let mut corrupted_path = tempdir_path.clone();
        corrupted_path.push(test_path2);
        let mut file = OpenOptions::new()
            .write(true)
            .open(&corrupted_path)
            .unwrap();
        file.write_all(&[255u8]).unwrap();
        file.flush().unwrap();
        drop(file);

        let damaged = managed_directory.list_damaged().unwrap();
        assert_eq!(damaged.len(), 1);
        assert!(damaged.contains(test_path2));
    }
}
<commit_msg>Use `slice::iter` instead of `into_iter` to avoid future breakage (#679)<commit_after>use crate::core::MANAGED_FILEPATH;
use crate::directory::error::{DeleteError, IOError, LockError, OpenReadError, OpenWriteError};
use crate::directory::footer::{Footer, FooterProxy};
use crate::directory::DirectoryLock;
use crate::directory::Lock;
use crate::directory::META_LOCK;
use crate::directory::{ReadOnlySource, WritePtr};
use crate::directory::{WatchCallback, WatchHandle};
use crate::error::DataCorruption;
use crate::Directory;
use crate::Result;
use crc32fast::Hasher;
use serde_json;
use std::collections::HashSet;
use std::io;
use std::io::Write;
use std::path::{Path, PathBuf};
use std::result;
use std::sync::RwLockWriteGuard;
use std::sync::{Arc, RwLock};

/// Returns true iff the file is "managed".
/// Non-managed file are not subject to garbage collection.
///
/// Filenames that starts by a "." -typically locks-
/// are not managed.
fn is_managed(path: &Path) -> bool {
    path.to_str()
        .map(|p_str| !p_str.starts_with('.'))
        .unwrap_or(true)
}

/// Wrapper of directories that keeps track of files created by Tantivy.
///
/// A managed directory is just a wrapper of a directory
/// that keeps a (persisted) list of the files that
/// have been created (and not deleted) by tantivy so far.
///
/// Thanks to this list, it implements a `garbage_collect` method
/// that removes the files that were created by tantivy and are not
/// useful anymore.
#[derive(Debug)]
pub struct ManagedDirectory {
    directory: Box<dyn Directory>,
    meta_informations: Arc<RwLock<MetaInformation>>,
}

#[derive(Debug, Default)]
struct MetaInformation {
    managed_paths: HashSet<PathBuf>,
}

/// Saves the file containing the list of existing files
/// that were created by tantivy.
fn save_managed_paths(
    directory: &mut dyn Directory,
    wlock: &RwLockWriteGuard<'_, MetaInformation>,
) -> io::Result<()> {
    let mut w = serde_json::to_vec(&wlock.managed_paths)?;
    writeln!(&mut w)?;
    directory.atomic_write(&MANAGED_FILEPATH, &w[..])?;
    Ok(())
}

impl ManagedDirectory {
    /// Wraps a directory as managed directory.
    pub fn wrap<Dir: Directory>(directory: Dir) -> Result<ManagedDirectory> {
        match directory.atomic_read(&MANAGED_FILEPATH) {
            Ok(data) => {
                let managed_files_json = String::from_utf8_lossy(&data);
                let managed_files: HashSet<PathBuf> = serde_json::from_str(&managed_files_json)
                    .map_err(|e| {
                        DataCorruption::new(
                            MANAGED_FILEPATH.to_path_buf(),
                            format!("Managed file cannot be deserialized: {:?}. ", e),
                        )
                    })?;
                Ok(ManagedDirectory {
                    directory: Box::new(directory),
                    meta_informations: Arc::new(RwLock::new(MetaInformation {
                        managed_paths: managed_files,
                    })),
                })
            }
            Err(OpenReadError::FileDoesNotExist(_)) => Ok(ManagedDirectory {
                directory: Box::new(directory),
                meta_informations: Arc::default(),
            }),
            Err(OpenReadError::IOError(e)) => Err(From::from(e)),
        }
    }

    /// Garbage collect unused files.
    ///
    /// Removes the files that were created by `tantivy` and are not
    /// used by any segment anymore.
    ///
    /// * `living_files` - List of files that are still used by the index.
    ///
    /// The use a callback ensures that the list of living_files is computed
    /// while we hold the lock on meta.
    ///
    /// This method does not panick nor returns errors.
    /// If a file cannot be deleted (for permission reasons for instance)
    /// an error is simply logged, and the file remains in the list of managed
    /// files.
    pub fn garbage_collect<L: FnOnce() -> HashSet<PathBuf>>(&mut self, get_living_files: L) {
        info!("Garbage collect");
        let mut files_to_delete = vec![];

        // It is crucial to get the living files after acquiring the
        // read lock of meta informations. That way, we
        // avoid the following scenario.
        //
        // 1) we get the list of living files.
        // 2) someone creates a new file.
        // 3) we start garbage collection and remove this file
        // even though it is a living file.
        //
        // releasing the lock as .delete() will use it too.
        {
            let meta_informations_rlock = self
                .meta_informations
                .read()
                .expect("Managed directory rlock poisoned in garbage collect.");

            // The point of this second "file" lock is to enforce the following scenario
            // 1) process B tries to load a new set of searcher.
            // The list of segments is loaded
            // 2) writer change meta.json (for instance after a merge or a commit)
            // 3) gc kicks in.
            // 4) gc removes a file that was useful for process B, before process B opened it.
            if let Ok(_meta_lock) = self.acquire_lock(&META_LOCK) {
                let living_files = get_living_files();
                for managed_path in &meta_informations_rlock.managed_paths {
                    if !living_files.contains(managed_path) {
                        files_to_delete.push(managed_path.clone());
                    }
                }
            } else {
                error!("Failed to acquire lock for GC");
            }
        }

        let mut deleted_files = vec![];
        for file_to_delete in files_to_delete {
            match self.delete(&file_to_delete) {
                Ok(_) => {
                    info!("Deleted {:?}", file_to_delete);
                    deleted_files.push(file_to_delete);
                }
                Err(file_error) => {
                    match file_error {
                        DeleteError::FileDoesNotExist(_) => {
                            deleted_files.push(file_to_delete);
                        }
                        DeleteError::IOError(_) => {
                            if !cfg!(target_os = "windows") {
                                // On windows, delete is expected to fail if the file
                                // is mmapped.
                                error!("Failed to delete {:?}", file_to_delete);
                            }
                        }
                    }
                }
            }
        }

        if !deleted_files.is_empty() {
            // update the list of managed files by removing
            // the file that were removed.
            let mut meta_informations_wlock = self
                .meta_informations
                .write()
                .expect("Managed directory wlock poisoned (2).");
            let managed_paths_write = &mut meta_informations_wlock.managed_paths;
            for delete_file in &deleted_files {
                managed_paths_write.remove(delete_file);
            }
            if save_managed_paths(self.directory.as_mut(), &meta_informations_wlock).is_err() {
                error!("Failed to save the list of managed files.");
            }
        }
    }

    /// Registers a file as managed
    ///
    /// This method must be called before the file is
    /// actually created to ensure that a failure between
    /// registering the filepath and creating the file
    /// will not lead to garbage files that will
    /// never get removed.
    ///
    /// File starting by "." are reserved to locks.
    /// They are not managed and cannot be subjected
    /// to garbage collection.
    fn register_file_as_managed(&mut self, filepath: &Path) -> io::Result<()> {
        // Files starting by "." (e.g. lock files) are not managed.
        if !is_managed(filepath) {
            return Ok(());
        }
        let mut meta_wlock = self
            .meta_informations
            .write()
            .expect("Managed file lock poisoned");
        let has_changed = meta_wlock.managed_paths.insert(filepath.to_owned());
        if has_changed {
            save_managed_paths(self.directory.as_mut(), &meta_wlock)?;
        }
        Ok(())
    }

    /// Verify checksum of a managed file
    pub fn validate_checksum(&self, path: &Path) -> result::Result<bool, OpenReadError> {
        let reader = self.directory.open_read(path)?;
        let (footer, data) = Footer::extract_footer(reader)
            .map_err(|err| IOError::with_path(path.to_path_buf(), err))?;
        let mut hasher = Hasher::new();
        hasher.update(data.as_slice());
        let crc = hasher.finalize();
        Ok(footer
            .versioned_footer
            .crc()
            .map(|v| v == crc)
            .unwrap_or(false))
    }

    /// List files for which checksum does not match content
    pub fn list_damaged(&self) -> result::Result<HashSet<PathBuf>, OpenReadError> {
        let mut hashset = HashSet::new();
        let managed_paths = self
            .meta_informations
            .read()
            .expect("Managed directory rlock poisoned in list damaged.")
            .managed_paths
            .clone();

        for path in managed_paths.into_iter() {
            if !self.validate_checksum(&path)? {
                hashset.insert(path);
            }
        }
        Ok(hashset)
    }
}

impl Directory for ManagedDirectory {
    fn open_read(&self, path: &Path) -> result::Result<ReadOnlySource, OpenReadError> {
        let read_only_source = self.directory.open_read(path)?;
        let (_footer, reader) = Footer::extract_footer(read_only_source)
            .map_err(|err| IOError::with_path(path.to_path_buf(), err))?;
        Ok(reader)
    }

    fn open_write(&mut self, path: &Path) -> result::Result<WritePtr, OpenWriteError> {
        self.register_file_as_managed(path)
            .map_err(|e| IOError::with_path(path.to_owned(), e))?;
        Ok(io::BufWriter::new(Box::new(FooterProxy::new(
            self.directory
                .open_write(path)?
                .into_inner()
                .map_err(|_| ())
                .expect("buffer should be empty"),
        ))))
    }

    fn atomic_write(&mut self, path: &Path, data: &[u8]) -> io::Result<()> {
        self.register_file_as_managed(path)?;
        self.directory.atomic_write(path, data)
    }

    fn atomic_read(&self, path: &Path) -> result::Result<Vec<u8>, OpenReadError> {
        self.directory.atomic_read(path)
    }

    fn delete(&self, path: &Path) -> result::Result<(), DeleteError> {
        self.directory.delete(path)
    }

    fn exists(&self, path: &Path) -> bool {
        self.directory.exists(path)
    }

    fn acquire_lock(&self, lock: &Lock) -> result::Result<DirectoryLock, LockError> {
        self.directory.acquire_lock(lock)
    }

    fn watch(&self, watch_callback: WatchCallback) -> crate::Result<WatchHandle> {
        self.directory.watch(watch_callback)
    }
}

impl Clone for ManagedDirectory {
    fn clone(&self) -> ManagedDirectory {
        ManagedDirectory {
            directory: self.directory.box_clone(),
            meta_informations: Arc::clone(&self.meta_informations),
        }
    }
}

#[cfg(feature = "mmap")]
#[cfg(test)]
mod tests_mmap_specific {

    use crate::directory::{Directory, ManagedDirectory, MmapDirectory, TerminatingWrite};
    use std::collections::HashSet;
    use std::fs::OpenOptions;
    use std::io::Write;
    use std::path::{Path, PathBuf};
    use tempfile::TempDir;

    #[test]
    fn test_managed_directory() {
        let tempdir = TempDir::new().unwrap();
        let tempdir_path = PathBuf::from(tempdir.path());

        let test_path1: &'static Path = Path::new("some_path_for_test");
        let test_path2: &'static Path = Path::new("some_path_for_test_2");
        {
            let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
            let mut managed_directory = ManagedDirectory::wrap(mmap_directory).unwrap();
            let write_file = managed_directory.open_write(test_path1).unwrap();
            write_file.terminate().unwrap();
            managed_directory
                .atomic_write(test_path2, &[0u8, 1u8])
                .unwrap();
            assert!(managed_directory.exists(test_path1));
            assert!(managed_directory.exists(test_path2));
            let living_files: HashSet<PathBuf> = [test_path1.to_owned()].iter().cloned().collect();
            managed_directory.garbage_collect(|| living_files);
            assert!(managed_directory.exists(test_path1));
            assert!(!managed_directory.exists(test_path2));
        }
        {
            let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
            let mut managed_directory = ManagedDirectory::wrap(mmap_directory).unwrap();
            assert!(managed_directory.exists(test_path1));
            assert!(!managed_directory.exists(test_path2));
            let living_files: HashSet<PathBuf> = HashSet::new();
            managed_directory.garbage_collect(|| living_files);
            assert!(!managed_directory.exists(test_path1));
            assert!(!managed_directory.exists(test_path2));
        }
    }

    #[test]
    fn test_managed_directory_gc_while_mmapped() {
        let test_path1: &'static Path = Path::new("some_path_for_test");

        let tempdir = TempDir::new().unwrap();
        let tempdir_path = PathBuf::from(tempdir.path());
        let living_files = HashSet::new();

        let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
        let mut managed_directory = ManagedDirectory::wrap(mmap_directory).unwrap();
        let mut write = managed_directory.open_write(test_path1).unwrap();
        write.write_all(&[0u8, 1u8]).unwrap();
        write.terminate().unwrap();
        assert!(managed_directory.exists(test_path1));

        let _mmap_read = managed_directory.open_read(test_path1).unwrap();
        managed_directory.garbage_collect(|| living_files.clone());
        if cfg!(target_os = "windows") {
            // On Windows, gc should try and fail the file as it is mmapped.
            assert!(managed_directory.exists(test_path1));
            // unmap should happen here.
            drop(_mmap_read);
            // The file should still be in the list of managed file and
            // eventually be deleted once mmap is released.
            managed_directory.garbage_collect(|| living_files);
            assert!(!managed_directory.exists(test_path1));
        } else {
            assert!(!managed_directory.exists(test_path1));
        }
    }

    #[test]
    fn test_checksum() {
        let test_path1: &'static Path = Path::new("some_path_for_test");
        let test_path2: &'static Path = Path::new("other_test_path");

        let tempdir = TempDir::new().unwrap();
        let tempdir_path = PathBuf::from(tempdir.path());

        let mmap_directory = MmapDirectory::open(&tempdir_path).unwrap();
        let mut managed_directory = ManagedDirectory::wrap(mmap_directory).unwrap();
        let mut write = managed_directory.open_write(test_path1).unwrap();
        write.write_all(&[0u8, 1u8]).unwrap();
        write.terminate().unwrap();

        let mut write = managed_directory.open_write(test_path2).unwrap();
        write.write_all(&[3u8, 4u8, 5u8]).unwrap();
        write.terminate().unwrap();

        assert!(managed_directory.list_damaged().unwrap().is_empty());

        let mut corrupted_path = tempdir_path.clone();
        corrupted_path.push(test_path2);
        let mut file = OpenOptions::new()
            .write(true)
            .open(&corrupted_path)
            .unwrap();
        file.write_all(&[255u8]).unwrap();
        file.flush().unwrap();
        drop(file);

        let damaged = managed_directory.list_damaged().unwrap();
        assert_eq!(damaged.len(), 1);
        assert!(damaged.contains(test_path2));
    }
}
<|endoftext|>