<commit_before>pub mod search;

use std::thread;
use std::cmp::min;
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver};
use std::time::SystemTime;
use std::num::Wrapping;
use basetypes::*;
use uci::{UciEngine, UciEngineFactory, EngineReply, OptionName, OptionDescription};
use position::Position;
use chess_move::*;
use tt::*;

const VERSION: &'static str = "0.1";
const MAX_DEPTH: u8 = 126;

// The number of nodes that can be searched without reporting search
// progress. If this value is too small the engine may become slow, if
// this value is too big the engine may become unresponsive.
const NODE_COUNT_REPORT_INTERVAL: NodeCount = 50000;

/// Implements `UciEngine` trait.
pub struct Engine {
    tt: Arc<TranspositionTable>,
    reply_queue: VecDeque<EngineReply>,
    position: Position,
    is_thinking: bool,
    is_pondering: bool,

    // Tells the engine if it will be allowed to ponder. This option
    // is needed because the engine might change its time management
    // algorithm when pondering is allowed.
    pondering_is_allowed: bool,

    // A channel for sending commands to the search thread.
    commands: Sender<search::Command>,

    // A channel for receiving reports from the search thread.
    reports: Receiver<search::Report>,

    // Search status info
    search_id: usize,
    thinking_since: SystemTime,
    silent_since: SystemTime,
    current_depth: u8,
    searched_nodes: NodeCount,
    searched_time: u64, // milliseconds
    nps: u64, // nodes per second
    mangled_pv: bool, // `true` if the primary variation is imperfect
    stop_when: TimeManagement,
}

impl Engine {
    /// Creates a new instance.
    ///
    /// `tt_size_mb` is the preferred size of the transposition
    /// table in Mbytes.
    pub fn new(tt_size_mb: usize) -> Engine {
        let mut tt = TranspositionTable::new();
        tt.resize(tt_size_mb);
        let tt1 = Arc::new(tt);
        let tt2 = tt1.clone();
        let (commands_tx, commands_rx) = channel();
        let (reports_tx, reports_rx) = channel();
        thread::spawn(move || {
            search::run_deepening(tt2, commands_rx, reports_tx);
        });

        Engine {
            tt: tt1,
            reply_queue: VecDeque::new(),
            position: Position::from_fen("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 \
                                          1")
                          .ok()
                          .unwrap(),
            is_thinking: false,
            is_pondering: false,
            pondering_is_allowed: false,
            commands: commands_tx,
            reports: reports_rx,
            search_id: 0,
            thinking_since: SystemTime::now(),
            silent_since: SystemTime::now(),
            current_depth: 0,
            searched_nodes: 0,
            searched_time: 0,
            nps: 0,
            mangled_pv: false,
            stop_when: TimeManagement::Infinite,
        }
    }

    // A helper method. It updates the search status info and makes
    // sure that a new PV is sent to the GUI for every newly reached
    // depth.
    fn register_progress(&mut self, depth: u8, searched_nodes: NodeCount) {
        if searched_nodes < NODE_COUNT_REPORT_INTERVAL || depth == 0 {
            return;
        }
        let thinking_duration = self.thinking_since.elapsed().unwrap();
        self.searched_time = 1000 * thinking_duration.as_secs() +
                             (thinking_duration.subsec_nanos() / 1000000) as u64;
        self.searched_nodes = searched_nodes;
        self.nps = 1000 * (self.nps + self.searched_nodes) / (1000 + self.searched_time);
        if self.current_depth < depth {
            self.current_depth = depth;
            self.report_pv(depth);
        }
    }

    // A helper method. It extracts the primary variation (PV) from
    // the transposition table (TT) and sends it to the GUI.
    fn report_pv(&mut self, depth: u8) {
        let mut prev_move = None;
        let mut p = self.position.clone();

        // Extract the PV, the value, and the bound from the TT.
        let mut pv = Vec::new();
        let mut value = -30000;
        let mut bound = BOUND_LOWER;
        while let Some(entry) = self.tt.probe(p.hash()) {
            if entry.bound() != BOUND_NONE {
                if let Some(m) = prev_move {
                    // Extend the PV.
                    pv.push(m);
                }

                // Get the value and the bound type. In half of the
                // cases the value stored in `entry` is from other
                // side's perspective.
                if pv.len() & 1 == 0 {
                    value = entry.value();
                    bound = entry.bound();
                } else {
                    value = -entry.value();
                    bound = match entry.bound() {
                        BOUND_UPPER => BOUND_LOWER,
                        BOUND_LOWER => BOUND_UPPER,
                        x => x,
                    };
                }

                // The values under -19999 and over 19999 may carry
                // information in how many moves is the inevitable
                // checkmate. We choose to not show this to the user,
                // because it would complicate unnecessarily the PV
                // extraction procedure.
                if value >= 20000 && bound != BOUND_UPPER {
                    value = 19999;
                    bound = BOUND_EXACT
                }
                if value <= -20000 && bound != BOUND_LOWER {
                    value = -19999;
                    bound = BOUND_EXACT
                }

                // Try to extend the PV.
                if bound == BOUND_EXACT {
                    if let Some(m) = p.try_move_digest(entry.move16()) {
                        if p.do_move(m) && !p.is_repeated() {
                            prev_move = Some(m);
                            continue;
                        }
                    }
                }
            }
            break;
        }

        // Check if the extracted PV is imperfect.
        self.mangled_pv = bound != BOUND_EXACT || pv.len() < depth as usize;

        // Send the newly extracted PV to the GUI.
        let value_suffix = match bound {
            BOUND_EXACT => "",
            BOUND_UPPER => " upperbound",
            BOUND_LOWER => " lowerbound",
            _ => panic!("unexpected bound type"),
        };
        let mut pv_string = String::new();
        for m in &pv {
            pv_string.push_str(&m.notation());
            pv_string.push(' ');
        }
        self.reply_queue.push_back(EngineReply::Info(vec![
            ("depth".to_string(), format!("{}", depth)),
            ("score".to_string(), format!("cp {}{}", value, value_suffix)),
            ("time".to_string(), format!("{}", self.searched_time)),
            ("nodes".to_string(), format!("{}", self.searched_nodes)),
            ("nps".to_string(), format!("{}", self.nps)),
            ("pv".to_string(), pv_string),
        ]));
        self.silent_since = SystemTime::now();
    }

    // A helper method. It reports the current depth, the searched
    // node count, and the nodes per second to the GUI.
    fn report_progress(&mut self) {
        self.reply_queue.push_back(EngineReply::Info(vec![
            ("depth".to_string(), format!("{}", self.current_depth)),
            ("nodes".to_string(), format!("{}", self.searched_nodes)),
            ("nps".to_string(), format!("{}", self.nps)),
        ]));
        self.silent_since = SystemTime::now();
    }

    // A helper method for `Engine::stop`. It probes the TT for the
    // best move in the position `p`, plays it, and returns it. If the
    // TT gives no move this function will play and return the first
    // legal move. `None` is returned only when there are no legal
    // moves.
    fn do_best_move(&self, p: &mut Position) -> Option<Move> {
        // Try to get best move from the TT.
        let move16 = self.tt.probe(p.hash()).map_or(0, |entry| entry.move16());
        if let Some(m) = p.try_move_digest(move16) {
            if p.do_move(m) {
                return Some(m);
            }
        }

        // Otherwise, pick the first legal move.
        let mut s = MoveStack::new();
        p.generate_moves(&mut s);
        while let Some(m) = s.pop() {
            if p.do_move(m) {
                return Some(m);
            }
        }

        // No legal moves.
        None
    }
}

impl UciEngine for Engine {
    fn set_option(&mut self, name: &str, value: &str) {
        match name {
            "Ponder" => {
                self.pondering_is_allowed = value == "true";
            }

            // An invalid option. Notice that we do not support
            // re-sizing of the transposition table once the engine
            // had started.
            _ => (),
        }
    }

    fn new_game(&mut self) {
        // Clearing the transposition table would not do any good, so
        // we do nothing.
    }

    fn position(&mut self, fen: &str, moves: &mut Iterator<Item = &str>) {
        if !self.is_thinking {
            if let Ok(p) = Position::from_history(fen, moves) {
                self.position = p;
                self.tt.new_search();
            }
        }
    }

    #[allow(unused_variables)]
    fn go(&mut self,
          searchmoves: Option<Vec<String>>,
          ponder: bool,
          wtime: Option<u64>,
          btime: Option<u64>,
          winc: Option<u64>,
          binc: Option<u64>,
          movestogo: Option<u64>,
          depth: Option<u64>,
          nodes: Option<u64>,
          mate: Option<u64>,
          movetime: Option<u64>,
          infinite: bool) {
        if !self.is_thinking {
            // Note: We ignore "searchmoves" and "mate" parameters.

            self.is_thinking = true;
            self.is_pondering = ponder;
            self.search_id = (Wrapping(self.search_id) + Wrapping(1)).0;
            self.thinking_since = SystemTime::now();
            self.silent_since = SystemTime::now();
            self.current_depth = 0;
            self.searched_nodes = 0;
            self.searched_time = 0;
            self.mangled_pv = false;

            // Figure out when we should stop thinking.
            //
            // TODO: Implement smarter time management.
            let (time, inc) = if self.position.board().to_move() == WHITE {
                (wtime, winc.unwrap_or(0))
            } else {
                (btime, binc.unwrap_or(0))
            };
            self.stop_when = if infinite {
                TimeManagement::Infinite
            } else if movetime.is_some() {
                TimeManagement::MoveTime(movetime.unwrap())
            } else if nodes.is_some() {
                TimeManagement::Nodes(nodes.unwrap())
            } else if depth.is_some() {
                TimeManagement::Depth(depth.unwrap() as u8)
            } else {
                let time = time.unwrap_or(0);
                let movestogo = movestogo.unwrap_or(40);
                let movetime = (time + inc * movestogo) / movestogo;
                TimeManagement::MoveTimeHint(min(movetime, time / 2))
            };

            // Start deepening search to the maximum possible depth.
            self.commands
                .send(search::Command::Search {
                    search_id: self.search_id,
                    position: self.position.clone(),
                    depth: MAX_DEPTH,
                    lower_bound: -20000,
                    upper_bound: 20000,
                })
                .unwrap();
        }
    }

    fn ponder_hit(&mut self) {
        if self.is_thinking && self.is_pondering {
            self.is_pondering = false;
        }
    }

    fn stop(&mut self) {
        if self.is_thinking {
            self.commands.send(search::Command::Stop).unwrap();
            let mut p = self.position.clone();
            let (best_move, ponder_move) = if let Some(m) = self.do_best_move(&mut p) {
                (m.notation(), self.do_best_move(&mut p).map(|m| m.notation()))
            } else {
                ("0000".to_string(), None)
            };
            self.reply_queue.push_back(EngineReply::BestMove {
                best_move: best_move,
                ponder_move: ponder_move,
            });
            self.is_thinking = false;
        }
    }

    fn is_thinking(&self) -> bool {
        self.is_thinking
    }

    fn get_reply(&mut self) -> Option<EngineReply> {
        // Check if we should stop thinking.
        if self.is_thinking && !self.is_pondering {
            if match self.stop_when {
                TimeManagement::MoveTime(t) => self.searched_time >= t,
                TimeManagement::MoveTimeHint(t) => self.searched_time >= t,
                TimeManagement::Nodes(n) => self.searched_nodes >= n,
                TimeManagement::Depth(d) => self.current_depth > d,
                TimeManagement::Infinite => false,
            } {
                self.stop();
            }
        }

        // Check the reports queue.
        while let Ok(report) = self.reports.try_recv() {
            if self.is_thinking {
                match report {
                    search::Report::Progress { search_id, depth, searched_nodes }
                        if search_id == self.search_id => {
                        // Register search progress.
                        self.register_progress(depth, searched_nodes);
                        if self.silent_since.elapsed().unwrap().as_secs() > 20 {
                            if self.mangled_pv {
                                self.report_pv(depth);
                            } else {
                                self.report_progress();
                            }
                        }
                    }
                    search::Report::Done { search_id, .. } if search_id == self.search_id => {
                        // Terminate the search if not infinite.
                        self.stop_when = match self.stop_when {
                            TimeManagement::Infinite => TimeManagement::Infinite,
                            _ => TimeManagement::MoveTime(0),
                        };
                    }
                    // We may still receive stale reports from already
                    // stopped searches.
                    _ => (),
                }
            }
        }

        self.reply_queue.pop_front()
    }
}

enum TimeManagement {
    MoveTime(u64),
    MoveTimeHint(u64),
    Nodes(NodeCount),
    Depth(u8),
    Infinite,
}

/// Implements `UciEngineFactory` trait.
pub struct EngineFactory;

impl UciEngineFactory<Engine> for EngineFactory {
    fn name(&self) -> String {
        format!("Alcibiades {}", VERSION)
    }

    fn author(&self) -> String {
        "Evgeni Pandurski".to_string()
    }

    fn options(&self) -> Vec<(OptionName, OptionDescription)> {
        vec![
            // TODO: Calculate a sane limit for the hash size.
            ("Hash".to_string(), OptionDescription::Spin { min: 1, max: 2048, default: 16 }),
            ("Ponder".to_string(), OptionDescription::Check { default: false }),
        ]
    }

    fn create(&self, hash_size_mb: Option<usize>) -> Engine {
        Engine::new(hash_size_mb.unwrap_or(16))
    }
}
<commit_msg>limit the PV length to depth again<commit_after>pub mod search;

use std::thread;
use std::cmp::min;
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver};
use std::time::SystemTime;
use std::num::Wrapping;
use basetypes::*;
use uci::{UciEngine, UciEngineFactory, EngineReply, OptionName, OptionDescription};
use position::Position;
use chess_move::*;
use tt::*;

const VERSION: &'static str = "0.1";
const MAX_DEPTH: u8 = 126;

// The number of nodes that can be searched without reporting search
// progress. If this value is too small the engine may become slow, if
// this value is too big the engine may become unresponsive.
const NODE_COUNT_REPORT_INTERVAL: NodeCount = 50000;

/// Implements `UciEngine` trait.
pub struct Engine {
    tt: Arc<TranspositionTable>,
    reply_queue: VecDeque<EngineReply>,
    position: Position,
    is_thinking: bool,
    is_pondering: bool,

    // Tells the engine if it will be allowed to ponder. This option
    // is needed because the engine might change its time management
    // algorithm when pondering is allowed.
    pondering_is_allowed: bool,

    // A channel for sending commands to the search thread.
    commands: Sender<search::Command>,

    // A channel for receiving reports from the search thread.
    reports: Receiver<search::Report>,

    // Search status info
    search_id: usize,
    thinking_since: SystemTime,
    silent_since: SystemTime,
    current_depth: u8,
    searched_nodes: NodeCount,
    searched_time: u64, // milliseconds
    nps: u64, // nodes per second
    mangled_pv: bool, // `true` if the primary variation is imperfect
    stop_when: TimeManagement,
}

impl Engine {
    /// Creates a new instance.
    ///
    /// `tt_size_mb` is the preferred size of the transposition
    /// table in Mbytes.
    pub fn new(tt_size_mb: usize) -> Engine {
        let mut tt = TranspositionTable::new();
        tt.resize(tt_size_mb);
        let tt1 = Arc::new(tt);
        let tt2 = tt1.clone();
        let (commands_tx, commands_rx) = channel();
        let (reports_tx, reports_rx) = channel();
        thread::spawn(move || {
            search::run_deepening(tt2, commands_rx, reports_tx);
        });

        Engine {
            tt: tt1,
            reply_queue: VecDeque::new(),
            position: Position::from_fen("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 \
                                          1")
                          .ok()
                          .unwrap(),
            is_thinking: false,
            is_pondering: false,
            pondering_is_allowed: false,
            commands: commands_tx,
            reports: reports_rx,
            search_id: 0,
            thinking_since: SystemTime::now(),
            silent_since: SystemTime::now(),
            current_depth: 0,
            searched_nodes: 0,
            searched_time: 0,
            nps: 0,
            mangled_pv: false,
            stop_when: TimeManagement::Infinite,
        }
    }

    // A helper method. It updates the search status info and makes
    // sure that a new PV is sent to the GUI for every newly reached
    // depth.
    fn register_progress(&mut self, depth: u8, searched_nodes: NodeCount) {
        if searched_nodes < NODE_COUNT_REPORT_INTERVAL || depth == 0 {
            return;
        }
        let thinking_duration = self.thinking_since.elapsed().unwrap();
        self.searched_time = 1000 * thinking_duration.as_secs() +
                             (thinking_duration.subsec_nanos() / 1000000) as u64;
        self.searched_nodes = searched_nodes;
        self.nps = 1000 * (self.nps + self.searched_nodes) / (1000 + self.searched_time);
        if self.current_depth < depth {
            self.current_depth = depth;
            self.report_pv(depth);
        }
    }

    // A helper method. It extracts the primary variation (PV) from
    // the transposition table (TT) and sends it to the GUI.
    fn report_pv(&mut self, depth: u8) {
        let mut prev_move = None;
        let mut p = self.position.clone();

        // Extract the PV, the value, and the bound from the TT.
        let mut pv = Vec::new();
        let mut value = -30000;
        let mut bound = BOUND_LOWER;
        while let Some(entry) = self.tt.probe(p.hash()) {
            if entry.bound() != BOUND_NONE {
                if let Some(m) = prev_move {
                    // Extend the PV.
                    pv.push(m);
                }

                // Get the value and the bound type. In half of the
                // cases the value stored in `entry` is from other
                // side's perspective.
                if pv.len() & 1 == 0 {
                    value = entry.value();
                    bound = entry.bound();
                } else {
                    value = -entry.value();
                    bound = match entry.bound() {
                        BOUND_UPPER => BOUND_LOWER,
                        BOUND_LOWER => BOUND_UPPER,
                        x => x,
                    };
                }

                // The values under -19999 and over 19999 may carry
                // information in how many moves is the inevitable
                // checkmate. We choose to not show this to the user,
                // because it would complicate unnecessarily the PV
                // extraction procedure.
                if value >= 20000 && bound != BOUND_UPPER {
                    value = 19999;
                    bound = BOUND_EXACT
                }
                if value <= -20000 && bound != BOUND_LOWER {
                    value = -19999;
                    bound = BOUND_EXACT
                }

                // Try to extend the PV.
                if pv.len() < depth as usize && bound == BOUND_EXACT {
                    if let Some(m) = p.try_move_digest(entry.move16()) {
                        if p.do_move(m) && !p.is_repeated() {
                            prev_move = Some(m);
                            continue;
                        }
                    }
                }
            }
            break;
        }

        // Check if the extracted PV is imperfect.
        self.mangled_pv = bound != BOUND_EXACT || pv.len() < depth as usize;

        // Send the newly extracted PV to the GUI.
        let value_suffix = match bound {
            BOUND_EXACT => "",
            BOUND_UPPER => " upperbound",
            BOUND_LOWER => " lowerbound",
            _ => panic!("unexpected bound type"),
        };
        let mut pv_string = String::new();
        for m in &pv {
            pv_string.push_str(&m.notation());
            pv_string.push(' ');
        }
        self.reply_queue.push_back(EngineReply::Info(vec![
            ("depth".to_string(), format!("{}", depth)),
            ("score".to_string(), format!("cp {}{}", value, value_suffix)),
            ("time".to_string(), format!("{}", self.searched_time)),
            ("nodes".to_string(), format!("{}", self.searched_nodes)),
            ("nps".to_string(), format!("{}", self.nps)),
            ("pv".to_string(), pv_string),
        ]));
        self.silent_since = SystemTime::now();
    }

    // A helper method. It reports the current depth, the searched
    // node count, and the nodes per second to the GUI.
    fn report_progress(&mut self) {
        self.reply_queue.push_back(EngineReply::Info(vec![
            ("depth".to_string(), format!("{}", self.current_depth)),
            ("nodes".to_string(), format!("{}", self.searched_nodes)),
            ("nps".to_string(), format!("{}", self.nps)),
        ]));
        self.silent_since = SystemTime::now();
    }

    // A helper method for `Engine::stop`. It probes the TT for the
    // best move in the position `p`, plays it, and returns it. If the
    // TT gives no move this function will play and return the first
    // legal move. `None` is returned only when there are no legal
    // moves.
    fn do_best_move(&self, p: &mut Position) -> Option<Move> {
        // Try to get best move from the TT.
        let move16 = self.tt.probe(p.hash()).map_or(0, |entry| entry.move16());
        if let Some(m) = p.try_move_digest(move16) {
            if p.do_move(m) {
                return Some(m);
            }
        }

        // Otherwise, pick the first legal move.
        let mut s = MoveStack::new();
        p.generate_moves(&mut s);
        while let Some(m) = s.pop() {
            if p.do_move(m) {
                return Some(m);
            }
        }

        // No legal moves.
        None
    }
}

impl UciEngine for Engine {
    fn set_option(&mut self, name: &str, value: &str) {
        match name {
            "Ponder" => {
                self.pondering_is_allowed = value == "true";
            }

            // An invalid option. Notice that we do not support
            // re-sizing of the transposition table once the engine
            // had started.
            _ => (),
        }
    }

    fn new_game(&mut self) {
        // Clearing the transposition table would not do any good, so
        // we do nothing.
    }

    fn position(&mut self, fen: &str, moves: &mut Iterator<Item = &str>) {
        if !self.is_thinking {
            if let Ok(p) = Position::from_history(fen, moves) {
                self.position = p;
                self.tt.new_search();
            }
        }
    }

    #[allow(unused_variables)]
    fn go(&mut self,
          searchmoves: Option<Vec<String>>,
          ponder: bool,
          wtime: Option<u64>,
          btime: Option<u64>,
          winc: Option<u64>,
          binc: Option<u64>,
          movestogo: Option<u64>,
          depth: Option<u64>,
          nodes: Option<u64>,
          mate: Option<u64>,
          movetime: Option<u64>,
          infinite: bool) {
        if !self.is_thinking {
            // Note: We ignore "searchmoves" and "mate" parameters.

            self.is_thinking = true;
            self.is_pondering = ponder;
            self.search_id = (Wrapping(self.search_id) + Wrapping(1)).0;
            self.thinking_since = SystemTime::now();
            self.silent_since = SystemTime::now();
            self.current_depth = 0;
            self.searched_nodes = 0;
            self.searched_time = 0;
            self.mangled_pv = false;

            // Figure out when we should stop thinking.
            //
            // TODO: Implement smarter time management.
            let (time, inc) = if self.position.board().to_move() == WHITE {
                (wtime, winc.unwrap_or(0))
            } else {
                (btime, binc.unwrap_or(0))
            };
            self.stop_when = if infinite {
                TimeManagement::Infinite
            } else if movetime.is_some() {
                TimeManagement::MoveTime(movetime.unwrap())
            } else if nodes.is_some() {
                TimeManagement::Nodes(nodes.unwrap())
            } else if depth.is_some() {
                TimeManagement::Depth(depth.unwrap() as u8)
            } else {
                let time = time.unwrap_or(0);
                let movestogo = movestogo.unwrap_or(40);
                let movetime = (time + inc * movestogo) / movestogo;
                TimeManagement::MoveTimeHint(min(movetime, time / 2))
            };

            // Start deepening search to the maximum possible depth.
            self.commands
                .send(search::Command::Search {
                    search_id: self.search_id,
                    position: self.position.clone(),
                    depth: MAX_DEPTH,
                    lower_bound: -20000,
                    upper_bound: 20000,
                })
                .unwrap();
        }
    }

    fn ponder_hit(&mut self) {
        if self.is_thinking && self.is_pondering {
            self.is_pondering = false;
        }
    }

    fn stop(&mut self) {
        if self.is_thinking {
            self.commands.send(search::Command::Stop).unwrap();
            let mut p = self.position.clone();
            let (best_move, ponder_move) = if let Some(m) = self.do_best_move(&mut p) {
                (m.notation(), self.do_best_move(&mut p).map(|m| m.notation()))
            } else {
                ("0000".to_string(), None)
            };
            self.reply_queue.push_back(EngineReply::BestMove {
                best_move: best_move,
                ponder_move: ponder_move,
            });
            self.is_thinking = false;
        }
    }

    fn is_thinking(&self) -> bool {
        self.is_thinking
    }

    fn get_reply(&mut self) -> Option<EngineReply> {
        // Check if we should stop thinking.
        if self.is_thinking && !self.is_pondering {
            if match self.stop_when {
                TimeManagement::MoveTime(t) => self.searched_time >= t,
                TimeManagement::MoveTimeHint(t) => self.searched_time >= t,
                TimeManagement::Nodes(n) => self.searched_nodes >= n,
                TimeManagement::Depth(d) => self.current_depth > d,
                TimeManagement::Infinite => false,
            } {
                self.stop();
            }
        }

        // Check the reports queue.
        while let Ok(report) = self.reports.try_recv() {
            if self.is_thinking {
                match report {
                    search::Report::Progress { search_id, depth, searched_nodes }
                        if search_id == self.search_id => {
                        // Register search progress.
                        self.register_progress(depth, searched_nodes);
                        if self.silent_since.elapsed().unwrap().as_secs() > 20 {
                            if self.mangled_pv {
                                self.report_pv(depth);
                            } else {
                                self.report_progress();
                            }
                        }
                    }
                    search::Report::Done { search_id, .. } if search_id == self.search_id => {
                        // Terminate the search if not infinite.
                        self.stop_when = match self.stop_when {
                            TimeManagement::Infinite => TimeManagement::Infinite,
                            _ => TimeManagement::MoveTime(0),
                        };
                    }
                    // We may still receive stale reports from already
                    // stopped searches.
                    _ => (),
                }
            }
        }

        self.reply_queue.pop_front()
    }
}

enum TimeManagement {
    MoveTime(u64),
    MoveTimeHint(u64),
    Nodes(NodeCount),
    Depth(u8),
    Infinite,
}

/// Implements `UciEngineFactory` trait.
pub struct EngineFactory;

impl UciEngineFactory<Engine> for EngineFactory {
    fn name(&self) -> String {
        format!("Alcibiades {}", VERSION)
    }

    fn author(&self) -> String {
        "Evgeni Pandurski".to_string()
    }

    fn options(&self) -> Vec<(OptionName, OptionDescription)> {
        vec![
            // TODO: Calculate a sane limit for the hash size.
            ("Hash".to_string(), OptionDescription::Spin { min: 1, max: 2048, default: 16 }),
            ("Ponder".to_string(), OptionDescription::Check { default: false }),
        ]
    }

    fn create(&self, hash_size_mb: Option<usize>) -> Engine {
        Engine::new(hash_size_mb.unwrap_or(16))
    }
}
<|endoftext|>


<commit_before>// Copyright 2022 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use audio_streams::shm_streams::NullShmStreamSource;
use audio_streams::shm_streams::ShmStreamSource;
#[cfg(feature = "audio_cras")]
use libcras::CrasClient;
#[cfg(feature = "audio_cras")]
use libcras::CrasClientType;
#[cfg(feature = "audio_cras")]
use libcras::CrasSocketType;
use serde::Deserialize;
use serde::Serialize;
use vm_memory::GuestMemory;

use crate::pci::ac97::Ac97Dev;
use crate::pci::ac97::Ac97Error;
use crate::pci::ac97::Ac97Parameters;
#[cfg(feature = "audio_cras")]
use crate::pci::pci_device;
use crate::pci::pci_device::Result;

pub(in crate::pci::ac97) type AudioStreamSource = Box<dyn ShmStreamSource<base::Error>>;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Ac97Backend {
    #[cfg(feature = "audio_cras")]
    Cras,
}

impl Ac97Dev {
    pub(in crate::pci::ac97) fn initialize_backend(
        ac97_backend: &Ac97Backend,
        #[allow(unused_variables)] mem: GuestMemory,
        #[allow(unused_variables)] param: &Ac97Parameters,
    ) -> Result<Self> {
        match *ac97_backend {
            #[cfg(feature = "audio_cras")]
            Ac97Backend::Cras => Self::create_cras_audio_device(param, mem.clone()).or_else(|e| {
                error!(
                    "Ac97Dev: create_cras_audio_device: {}. Fallback to null audio device",
                    e
                );
                Ok(Self::create_null_audio_device(mem))
            }),
        }
    }

    #[cfg(feature = "audio_cras")]
    fn create_cras_audio_device(params: &Ac97Parameters, mem: GuestMemory) -> Result<Self> {
        let mut server = Box::new(
            CrasClient::with_type(params.socket_type.unwrap_or(CrasSocketType::Unified))
                .map_err(pci_device::Error::CreateCrasClientFailed)?,
        );
        server.set_client_type(
            params
                .client_type
                .unwrap_or(CrasClientType::CRAS_CLIENT_TYPE_CROSVM),
        );
        if params.capture {
            server.enable_cras_capture();
        }

        let cras_audio = Self::new(
            mem,
            crate::pci::ac97::Ac97Backend::System(Ac97Backend::Cras),
            server,
        );
        Ok(cras_audio)
    }

    /// Return the minijail policy file path for the current Ac97Dev.
    pub fn minijail_policy(&self) -> &'static str {
        match &self.backend {
            crate::pci::ac97::Ac97Backend::System(backend) => match *backend {
                #[cfg(feature = "audio_cras")]
                Ac97Backend::Cras => "cras_audio_device",
            },
            crate::pci::ac97::Ac97Backend::NULL => "null_audio_device",
        }
    }
}

pub(in crate::pci::ac97) fn ac97_backend_from_str(
    s: &str,
) -> std::result::Result<crate::pci::ac97::Ac97Backend, Ac97Error> {
    match s {
        #[cfg(feature = "audio_cras")]
        "cras" => Ok(crate::pci::ac97::Ac97Backend::System(Ac97Backend::Cras)),
        _ => Err(Ac97Error::InvalidBackend),
    }
}

pub(in crate::pci::ac97) fn create_null_server() -> AudioStreamSource {
    Box::new(NullShmStreamSource::new())
}

#[cfg(test)]
pub(in crate::pci::ac97) mod tests {
    use audio_streams::shm_streams::MockShmStreamSource;

    use super::*;

    pub(in crate::pci::ac97) fn create_ac97_device(
        mem: GuestMemory,
        backend: crate::pci::ac97::Ac97Backend,
    ) -> Ac97Dev {
        Ac97Dev::new(mem, backend, Box::new(MockShmStreamSource::new()))
    }
}
<commit_msg>devices: ac97: fix missing import with audio_cras<commit_after>// Copyright 2022 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use audio_streams::shm_streams::NullShmStreamSource;
use audio_streams::shm_streams::ShmStreamSource;
#[cfg(feature = "audio_cras")]
use base::error;
#[cfg(feature = "audio_cras")]
use libcras::CrasClient;
#[cfg(feature = "audio_cras")]
use libcras::CrasClientType;
#[cfg(feature = "audio_cras")]
use libcras::CrasSocketType;
use serde::Deserialize;
use serde::Serialize;
use vm_memory::GuestMemory;

use crate::pci::ac97::Ac97Dev;
use crate::pci::ac97::Ac97Error;
use crate::pci::ac97::Ac97Parameters;
#[cfg(feature = "audio_cras")]
use crate::pci::pci_device;
use crate::pci::pci_device::Result;

pub(in crate::pci::ac97) type AudioStreamSource = Box<dyn ShmStreamSource<base::Error>>;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Ac97Backend {
    #[cfg(feature = "audio_cras")]
    Cras,
}

impl Ac97Dev {
    pub(in crate::pci::ac97) fn initialize_backend(
        ac97_backend: &Ac97Backend,
        #[allow(unused_variables)] mem: GuestMemory,
        #[allow(unused_variables)] param: &Ac97Parameters,
    ) -> Result<Self> {
        match *ac97_backend {
            #[cfg(feature = "audio_cras")]
            Ac97Backend::Cras => Self::create_cras_audio_device(param, mem.clone()).or_else(|e| {
                error!(
                    "Ac97Dev: create_cras_audio_device: {}. Fallback to null audio device",
                    e
                );
                Ok(Self::create_null_audio_device(mem))
            }),
        }
    }

    #[cfg(feature = "audio_cras")]
    fn create_cras_audio_device(params: &Ac97Parameters, mem: GuestMemory) -> Result<Self> {
        let mut server = Box::new(
            CrasClient::with_type(params.socket_type.unwrap_or(CrasSocketType::Unified))
                .map_err(pci_device::Error::CreateCrasClientFailed)?,
        );
        server.set_client_type(
            params
                .client_type
                .unwrap_or(CrasClientType::CRAS_CLIENT_TYPE_CROSVM),
        );
        if params.capture {
            server.enable_cras_capture();
        }

        let cras_audio = Self::new(
            mem,
            crate::pci::ac97::Ac97Backend::System(Ac97Backend::Cras),
            server,
        );
        Ok(cras_audio)
    }

    /// Return the minijail policy file path for the current Ac97Dev.
    pub fn minijail_policy(&self) -> &'static str {
        match &self.backend {
            crate::pci::ac97::Ac97Backend::System(backend) => match *backend {
                #[cfg(feature = "audio_cras")]
                Ac97Backend::Cras => "cras_audio_device",
            },
            crate::pci::ac97::Ac97Backend::NULL => "null_audio_device",
        }
    }
}

pub(in crate::pci::ac97) fn ac97_backend_from_str(
    s: &str,
) -> std::result::Result<crate::pci::ac97::Ac97Backend, Ac97Error> {
    match s {
        #[cfg(feature = "audio_cras")]
        "cras" => Ok(crate::pci::ac97::Ac97Backend::System(Ac97Backend::Cras)),
        _ => Err(Ac97Error::InvalidBackend),
    }
}

pub(in crate::pci::ac97) fn create_null_server() -> AudioStreamSource {
    Box::new(NullShmStreamSource::new())
}

#[cfg(test)]
pub(in crate::pci::ac97) mod tests {
    use audio_streams::shm_streams::MockShmStreamSource;

    use super::*;

    pub(in crate::pci::ac97) fn create_ac97_device(
        mem: GuestMemory,
        backend: crate::pci::ac97::Ac97Backend,
    ) -> Ac97Dev {
        Ac97Dev::new(mem, backend, Box::new(MockShmStreamSource::new()))
    }
}
<|endoftext|>


<commit_before>use support::{database, project};

#[test]
fn run_infer_schema() {
    let p = project("print-schema").build();
    let db = database(&p.database_url());

    // Make sure the project is setup
    p.command("setup").run();

    db.execute("CREATE TABLE users1 (id INTEGER PRIMARY KEY);");
    db.execute("CREATE TABLE users2 (id INTEGER PRIMARY KEY);");

    assert!(db.table_exists("users1"));
    assert!(db.table_exists("users2"));

    let result = p.command("print-schema").run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);
    assert_eq!(result.stdout(),
r"mod infer_users1 {
    table! {
        users1(id) {
            id -> Nullable<Integer>,
        }
    }
}
pub use self::infer_users1::*;
mod infer_users2 {
    table! {
        users2(id) {
            id -> Nullable<Integer>,
        }
    }
}
pub use self::infer_users2::*;

");
}

#[test]
fn run_infer_schema_whitelist() {
    let p = project("print-schema").build();
    let db = database(&p.database_url());

    // Make sure the project is setup
    p.command("setup").run();

    db.execute("CREATE TABLE users1 (id INTEGER PRIMARY KEY);");
    db.execute("CREATE TABLE users2 (id INTEGER PRIMARY KEY);");

    assert!(db.table_exists("users1"));
    assert!(db.table_exists("users2"));

    let result = p.command("print-schema").arg("users1").arg("-w").run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);
    assert_eq!(result.stdout(),
r"mod infer_users1 {
    table! {
        users1(id) {
            id -> Nullable<Integer>,
        }
    }
}
pub use self::infer_users1::*;

");
}

#[test]
fn run_infer_schema_blacklist() {
    let p = project("print-schema").build();
    let db = database(&p.database_url());

    // Make sure the project is setup
    p.command("setup").run();

    db.execute("CREATE TABLE users1 (id INTEGER PRIMARY KEY);");
    db.execute("CREATE TABLE users2 (id INTEGER PRIMARY KEY);");

    assert!(db.table_exists("users1"));
    assert!(db.table_exists("users2"));

    let result = p.command("print-schema").arg("users1").arg("-b").run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);
    assert_eq!(result.stdout(),
r"mod infer_users2 {
    table! {
        users2(id) {
            id -> Nullable<Integer>,
        }
    }
}
pub use self::infer_users2::*;

");
}
<commit_msg>Fix postgres cli-tests (postgres did not like tablenames containing '-')<commit_after>use support::{database, project};

#[test]
fn run_infer_schema() {
    let p = project("print_schema").build();
    let db = database(&p.database_url());

    // Make sure the project is setup
    p.command("setup").run();

    db.execute("CREATE TABLE users1 (id INTEGER PRIMARY KEY);");
    db.execute("CREATE TABLE users2 (id INTEGER PRIMARY KEY);");

    assert!(db.table_exists("users1"));
    assert!(db.table_exists("users2"));

    let result = p.command("print-schema").run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);
    assert_eq!(result.stdout(),
r"mod infer_users1 {
    table! {
        users1(id) {
            id -> Nullable<Integer>,
        }
    }
}
pub use self::infer_users1::*;
mod infer_users2 {
    table! {
        users2(id) {
            id -> Nullable<Integer>,
        }
    }
}
pub use self::infer_users2::*;

");
}

#[test]
fn run_infer_schema_whitelist() {
    let p = project("print_schema").build();
    let db = database(&p.database_url());

    // Make sure the project is setup
    p.command("setup").run();

    db.execute("CREATE TABLE users1 (id INTEGER PRIMARY KEY);");
    db.execute("CREATE TABLE users2 (id INTEGER PRIMARY KEY);");

    assert!(db.table_exists("users1"));
    assert!(db.table_exists("users2"));

    let result = p.command("print-schema").arg("users1").arg("-w").run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);
    assert_eq!(result.stdout(),
r"mod infer_users1 {
    table! {
        users1(id) {
            id -> Nullable<Integer>,
        }
    }
}
pub use self::infer_users1::*;

");
}

#[test]
fn run_infer_schema_blacklist() {
    let p = project("print_schema").build();
    let db = database(&p.database_url());

    // Make sure the project is setup
    p.command("setup").run();

    db.execute("CREATE TABLE users1 (id INTEGER PRIMARY KEY);");
    db.execute("CREATE TABLE users2 (id INTEGER PRIMARY KEY);");

    assert!(db.table_exists("users1"));
    assert!(db.table_exists("users2"));

    let result = p.command("print-schema").arg("users1").arg("-b").run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);
    assert_eq!(result.stdout(),
r"mod infer_users2 {
    table! {
        users2(id) {
            id -> Nullable<Integer>,
        }
    }
}
pub use self::infer_users2::*;

");
}
<|endoftext|>


<commit_before>// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Tidy check to ensure that there are no binaries checked into the source tree
//! by accident.
//!
//! In the past we've accidentally checked in test binaries and such which add a
//! huge amount of bloat to the git history, so it's good to just ensure we
//! don't do that again :)

use std::path::Path;

// All files are executable on Windows, so just check on Unix
#[cfg(windows)]
pub fn check(_path: &Path, _bad: &mut bool) {}

#[cfg(unix)]
pub fn check(path: &Path, bad: &mut bool) {
    use std::fs;
    use std::os::unix::prelude::*;

    super::walk(path,
                &mut |path| super::filter_dirs(path) || path.ends_with("src/etc"),
                &mut |file| {
        let filename = file.file_name().unwrap().to_string_lossy();
        let extensions = [".py", ".sh"];
        if extensions.iter().any(|e| filename.ends_with(e)) {
            return
        }

        let metadata = t!(fs::metadata(&file), &file);
        if metadata.mode() & 0o111 != 0 {
            println!("binary checked into source: {}", file.display());
            *bad = true;
        }
    })
}

<commit_msg>Rollup merge of #33560 - eddyb:symtidy, r=alexcrichton<commit_after>// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Tidy check to ensure that there are no binaries checked into the source tree
//! by accident.
//!
//! In the past we've accidentally checked in test binaries and such which add a
//! huge amount of bloat to the git history, so it's good to just ensure we
//! don't do that again :)

use std::path::Path;

// All files are executable on Windows, so just check on Unix
#[cfg(windows)]
pub fn check(_path: &Path, _bad: &mut bool) {}

#[cfg(unix)]
pub fn check(path: &Path, bad: &mut bool) {
    use std::fs;
    use std::os::unix::prelude::*;

    super::walk(path,
                &mut |path| super::filter_dirs(path) || path.ends_with("src/etc"),
                &mut |file| {
        let filename = file.file_name().unwrap().to_string_lossy();
        let extensions = [".py", ".sh"];
        if extensions.iter().any(|e| filename.ends_with(e)) {
            return
        }

        let metadata = t!(fs::symlink_metadata(&file), &file);
        if metadata.mode() & 0o111 != 0 {
            println!("binary checked into source: {}", file.display());
            *bad = true;
        }
    })
}

<|endoftext|>


<commit_before>use std::num::Zero;
use traits::division_ring::DivisionRing;
use traits::workarounds::scalar_op::{ScalarMul, ScalarDiv};

/// Trait of elements of a vector space. A vector space is an algebraic
/// structure, the elements of which have addition, substraction, negation,
/// scalar multiplication (the scalar being a element of a `DivisionRing`), and
/// has a distinct element (`Zero`) neutral wrt the addition.
pub trait VectorSpace<N: DivisionRing>
: Sub<Self, Self> + Add<Self, Self> + Neg<Self> + Zero +
  ScalarMul<N> + ScalarDiv<N>
{ }
<commit_msg>Add automatic implementation of VectorSpace.<commit_after>use std::num::Zero;
use traits::division_ring::DivisionRing;
use traits::workarounds::scalar_op::{ScalarMul, ScalarDiv};

/// Trait of elements of a vector space. A vector space is an algebraic
/// structure, the elements of which have addition, substraction, negation,
/// scalar multiplication (the scalar being a element of a `DivisionRing`), and
/// has a distinct element (`Zero`) neutral wrt the addition.
pub trait VectorSpace<N>
: Sub<Self, Self> + Add<Self, Self> + Neg<Self> + Zero +
  ScalarMul<N> + ScalarDiv<N>
{ }

impl<V: Sub<V, V> + Add<V, V> + Neg<V> + Zero +
        ScalarMul<N> + ScalarDiv<N>,
     N: DivisionRing> VectorSpace<N> for V;
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::collections::Vec;
use std::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};
use std::cell::Cell;

use super::variable::{TVar};
use super::res::*;

static GLOBAL_SEQ_LOCK: AtomicUsize = ATOMIC_USIZE_INIT;

pub struct Transaction {
    usize: snapshot
    writeSet: HashMap<&mut TVar<T>, u32>
    readSet: Vec<(&mut TVar<T>, u32)>
}

impl Transaction {
    fn new(ss: usize) -> Transaction {
        Transaction {
            snapShot: ss
            writeSet: BTreeMap::new()
            readSet: Vec::new()
        }
    }

    pub fn run<T, F>(f: F) -> T
    where F: Fn(&mut Transaction) -> StmResult<T>,
    {
        let mut ss = GLOBAL_SEQ_LOCK.load(Ordering::SeqCst);
        while ((ss & 1) != 0) {
            ss = GLOBAL_SEQ_LOCK.load(Ordering::SeqCst);
        }
        let mut transaction = Transaction::new(ss);
        loop {
            match f(&mut transaction) {
                Ok(t) => {
                    if transaction.commit() {
                        return t;
                    }
                }
                Err(_) => { }
            }
            transaction.clear();
        }
    }

    pub fn read<T>(&mut self, var: &mut TVar<T>) -> StmResult<T> {
        match self.writeSet.get(&var) {
            Some &val => return Ok(val)
            None => {}
        };
        let mut val = (*var).value;
        while (self.snapshot != GLOBAL_SEQ_LOCK.load(Ordering::SeqCst)) {
            self.snapshot = self.validate()
            val = (*var).value;
            if (self.snapshot < 0) {
                return Err(Retry);
            }
        }
        self.readset.push((addr, val));
        Ok(val)
    }

    pub fn write<T>(&mut self, var: &mut TVar<T>, value: T) -> StmResult<()> {
        self.writeSet.insert(var, value);
        Ok(())

    }

    fn clear(&mut self) {
        self.readSet.clear();
        self.writeSet.clear();
    }

    fn validate(&mut self) -> <usize> {
        loop {
            let mut usize time = GLOBAL_SEQ_LOCK.load(Ordering::SeqCst);
            if ((time & 1) != 0) {
                continue;
            }
            for (addr, val) in self.readSet {
                if ((*addr).value != val) {
                    return -1;
                }
            }
            if (time == lock.load(Ordering::SeqCst)) {
                return time;
            }
        }
    }

    fn commit(&mut self) -> bool {
        if (self.writeSet.len() == 0) {
            return true;
        }
        while (GLOBAL_SEQ_LOCK.compare_and_swap.(self.snapshot, self.snapshot +  1, Ordering::SeqCst) != self.snapshot) {
            self.snapshot = self.validate()
            if (self.snapshot < 0) {
                return false;
            }
        }
        for (addr, val) in self.writeSet.iter() {
            (*addr).value = val;
        }
        GLOBAL_SEQ_LOCK.store(self.snapshot + 2, Ordering::SeqCst);
        return true;
    }
}<commit_msg>fix obvious stuff<commit_after>use std::collections::HashMap;
use std::collections::Vec;
use std::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};
use std::cell::Cell;

use super::variable::{TVar};
use super::res::*;

static GLOBAL_SEQ_LOCK: AtomicUsize = ATOMIC_USIZE_INIT;

pub struct Transaction {
    usize: snapshot,
    writeSet: HashMap<&mut TVar<T>, u32>,
    readSet: Vec<(&mut TVar<T>, u32)>,
}

impl Transaction {
    fn new(ss: usize) -> Transaction {
        Transaction {
            snapShot: ss,
            writeSet: BTreeMap::new(),
            readSet: Vec::new(),
        }
    }

    pub fn run<T, F>(f: F) -> T
    where F: Fn(&mut Transaction) -> StmResult<T>,
    {
        let mut ss = GLOBAL_SEQ_LOCK.load(Ordering::SeqCst);
        while (ss & 1) != 0 {
            ss = GLOBAL_SEQ_LOCK.load(Ordering::SeqCst);
        }
        let mut transaction = Transaction::new(ss);
        loop {
            match f(&mut transaction) {
                Ok(t) => {
                    if transaction.commit() {
                        return t;
                    }
                }
                Err(_) => { }
            }
            transaction.clear();
        }
    }

    pub fn read<T>(&mut self, var: &mut TVar<T>) -> StmResult<T> {
        match self.writeSet.get(&var) {
            Some(&val) => { return Ok(val) }
            None => {}
        };
        let mut val = (*var).value;
        while self.snapshot != GLOBAL_SEQ_LOCK.load(Ordering::SeqCst) {
            self.snapshot = self.validate();
            val = (*var).value;
            if self.snapshot < 0 {
                return Err(Retry);
            }
        }
        self.readset.push((addr, val));
        Ok(val)
    }

    pub fn write<T>(&mut self, var: &mut TVar<T>, value: T) -> StmResult<()> {
        self.writeSet.insert(var, value);
        Ok(())

    }

    fn clear(&mut self) {
        self.readSet.clear();
        self.writeSet.clear();
    }

    fn validate(&mut self) -> usize {
        loop {
            let mut time : usize = GLOBAL_SEQ_LOCK.load(Ordering::SeqCst);
            if (time & 1) != 0 {
                continue;
            }
            for (addr, val) in self.readSet {
                if (*addr).value != val {
                    return -1;
                }
            }
            if time == lock.load(Ordering::SeqCst) {
                return time;
            }
        }
    }

    fn commit(&mut self) -> bool {
        if self.writeSet.len() == 0 {
            return true;
        }
        while GLOBAL_SEQ_LOCK.compare_and_swap(self.snapshot, self.snapshot+1, Ordering::SeqCst) != self.snapshot {
            self.snapshot = self.validate();
            if self.snapshot < 0 {
                return false;
            }
        }
        for (addr, val) in self.writeSet.iter() {
            (*addr).value = val;
        }
        GLOBAL_SEQ_LOCK.store(self.snapshot + 2, Ordering::SeqCst);
        return true;
    }
}

<|endoftext|>


<commit_before>// Copyright 2016 Amanieu d'Antras
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

use crate::raw_fair_mutex::RawFairMutex;
use lock_api;

/// A mutual exclusive primitive that is always fair, useful for protecting shared data
///
/// This mutex will block threads waiting for the lock to become available. The
/// mutex can also be statically initialized or created via a `new`
/// constructor. Each mutex has a type parameter which represents the data that
/// it is protecting. The data can only be accessed through the RAII guards
/// returned from `lock` and `try_lock`, which guarantees that the data is only
/// ever accessed when the mutex is locked.
///
/// The regular mutex provided by `parking_lot` uses eventual locking fairness (after some
/// time it will default to the fair algorithm), but eventual fairness does not provide the same
/// garantees a always fair method would. Fair mutexes are generally slower, but sometimes needed.
/// This wrapper was created to avoid using a unfair protocol when it's forbidden by mistake.
///
/// In a fair mutex the lock is provided to whichever thread asked first. Always following the
/// first-in first-out order. By not allowing other threads to steal the lock even if it would mean
/// a faster execution.
///
/// # Differences from the standard library `Mutex`
///
/// - No poisoning, the lock is released normally on panic.
/// - Only requires 1 byte of space, whereas the standard library boxes the
///   `FairMutex` due to platform limitations.
/// - Can be statically constructed (requires the `const_fn` nightly feature).
/// - Does not require any drop glue when dropped.
/// - Inline fast path for the uncontended case.
/// - Efficient handling of micro-contention using adaptive spinning.
/// - Allows raw locking & unlocking without a guard.
///
/// # Examples
///
/// ```
/// use parking_lot::FairMutex;
/// use std::sync::{Arc, mpsc::channel};
/// use std::thread;
///
/// const N: usize = 10;
///
/// // Spawn a few threads to increment a shared variable (non-atomically), and
/// // let the main thread know once all increments are done.
/// //
/// // Here we're using an Arc to share memory among threads, and the data inside
/// // the Arc is protected with a mutex.
/// let data = Arc::new(FairMutex::new(0));
///
/// let (tx, rx) = channel();
/// for _ in 0..10 {
///     let (data, tx) = (Arc::clone(&data), tx.clone());
///     thread::spawn(move || {
///         // The shared state can only be accessed once the lock is held.
///         // Our non-atomic increment is safe because we're the only thread
///         // which can access the shared state when the lock is held.
///         let mut data = data.lock();
///         *data += 1;
///         if *data == N {
///             tx.send(()).unwrap();
///         }
///         // the lock is unlocked here when `data` goes out of scope.
///     });
/// }
///
/// rx.recv().unwrap();
/// ```
pub type FairMutex<T> = lock_api::Mutex<RawFairMutex, T>;

/// An RAII implementation of a "scoped lock" of a mutex. When this structure is
/// dropped (falls out of scope), the lock will be unlocked.
///
/// The data protected by the mutex can be accessed through this guard via its
/// `Deref` and `DerefMut` implementations.
pub type FairMutexGuard<'a, T> = lock_api::MutexGuard<'a, RawFairMutex, T>;

/// An RAII mutex guard returned by `FairMutexGuard::map`, which can point to a
/// subfield of the protected data.
///
/// The main difference between `MappedFairMutexGuard` and `FairMutexGuard` is that the
/// former doesn't support temporarily unlocking and re-locking, since that
/// could introduce soundness issues if the locked object is modified by another
/// thread.
pub type MappedFairMutexGuard<'a, T> = lock_api::MappedMutexGuard<'a, RawFairMutex, T>;

#[cfg(test)]
mod tests {
    use crate::FairMutex;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::mpsc::channel;
    use std::sync::Arc;
    use std::thread;

    #[cfg(feature = "serde")]
    use bincode::{deserialize, serialize};

    #[derive(Eq, PartialEq, Debug)]
    struct NonCopy(i32);

    #[test]
    fn smoke() {
        let m = FairMutex::new(());
        drop(m.lock());
        drop(m.lock());
    }

    #[test]
    fn lots_and_lots() {
        const J: u32 = 1000;
        const K: u32 = 3;

        let m = Arc::new(FairMutex::new(0));

        fn inc(m: &FairMutex<u32>) {
            for _ in 0..J {
                *m.lock() += 1;
            }
        }

        let (tx, rx) = channel();
        for _ in 0..K {
            let tx2 = tx.clone();
            let m2 = m.clone();
            thread::spawn(move || {
                inc(&m2);
                tx2.send(()).unwrap();
            });
            let tx2 = tx.clone();
            let m2 = m.clone();
            thread::spawn(move || {
                inc(&m2);
                tx2.send(()).unwrap();
            });
        }

        drop(tx);
        for _ in 0..2 * K {
            rx.recv().unwrap();
        }
        assert_eq!(*m.lock(), J * K * 2);
    }

    #[test]
    fn try_lock() {
        let m = FairMutex::new(());
        *m.try_lock().unwrap() = ();
    }

    #[test]
    fn test_into_inner() {
        let m = FairMutex::new(NonCopy(10));
        assert_eq!(m.into_inner(), NonCopy(10));
    }

    #[test]
    fn test_into_inner_drop() {
        struct Foo(Arc<AtomicUsize>);
        impl Drop for Foo {
            fn drop(&mut self) {
                self.0.fetch_add(1, Ordering::SeqCst);
            }
        }
        let num_drops = Arc::new(AtomicUsize::new(0));
        let m = FairMutex::new(Foo(num_drops.clone()));
        assert_eq!(num_drops.load(Ordering::SeqCst), 0);
        {
            let _inner = m.into_inner();
            assert_eq!(num_drops.load(Ordering::SeqCst), 0);
        }
        assert_eq!(num_drops.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn test_get_mut() {
        let mut m = FairMutex::new(NonCopy(10));
        *m.get_mut() = NonCopy(20);
        assert_eq!(m.into_inner(), NonCopy(20));
    }

    #[test]
    fn test_mutex_arc_nested() {
        // Tests nested mutexes and access
        // to underlying data.
        let arc = Arc::new(FairMutex::new(1));
        let arc2 = Arc::new(FairMutex::new(arc));
        let (tx, rx) = channel();
        let _t = thread::spawn(move || {
            let lock = arc2.lock();
            let lock2 = lock.lock();
            assert_eq!(*lock2, 1);
            tx.send(()).unwrap();
        });
        rx.recv().unwrap();
    }

    #[test]
    fn test_mutex_arc_access_in_unwind() {
        let arc = Arc::new(FairMutex::new(1));
        let arc2 = arc.clone();
        let _ = thread::spawn(move || {
            struct Unwinder {
                i: Arc<FairMutex<i32>>,
            }
            impl Drop for Unwinder {
                fn drop(&mut self) {
                    *self.i.lock() += 1;
                }
            }
            let _u = Unwinder { i: arc2 };
            panic!();
        })
        .join();
        let lock = arc.lock();
        assert_eq!(*lock, 2);
    }

    #[test]
    fn test_mutex_unsized() {
        let mutex: &FairMutex<[i32]> = &FairMutex::new([1, 2, 3]);
        {
            let b = &mut *mutex.lock();
            b[0] = 4;
            b[2] = 5;
        }
        let comp: &[i32] = &[4, 2, 5];
        assert_eq!(&*mutex.lock(), comp);
    }

    #[test]
    fn test_mutexguard_sync() {
        fn sync<T: Sync>(_: T) {}

        let mutex = FairMutex::new(());
        sync(mutex.lock());
    }

    #[test]
    fn test_mutex_debug() {
        let mutex = FairMutex::new(vec![0u8, 10]);

        assert_eq!(format!("{:?}", mutex), "Mutex { data: [0, 10] }");
        let _lock = mutex.lock();
        assert_eq!(format!("{:?}", mutex), "Mutex { data: <locked> }");
    }

    #[cfg(feature = "serde")]
    #[test]
    fn test_serde() {
        let contents: Vec<u8> = vec![0, 1, 2];
        let mutex = FairMutex::new(contents.clone());

        let serialized = serialize(&mutex).unwrap();
        let deserialized: FairMutex<Vec<u8>> = deserialize(&serialized).unwrap();

        assert_eq!(*(mutex.lock()), *(deserialized.lock()));
        assert_eq!(contents, *(deserialized.lock()));
    }
}
<commit_msg>Update docs<commit_after>// Copyright 2016 Amanieu d'Antras
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

use crate::raw_fair_mutex::RawFairMutex;
use lock_api;

/// A mutual exclusive primitive that is always fair, useful for protecting shared data
///
/// This mutex will block threads waiting for the lock to become available. The
/// mutex can also be statically initialized or created via a `new`
/// constructor. Each mutex has a type parameter which represents the data that
/// it is protecting. The data can only be accessed through the RAII guards
/// returned from `lock` and `try_lock`, which guarantees that the data is only
/// ever accessed when the mutex is locked.
///
/// The regular mutex provided by `parking_lot` uses eventual locking fairness
/// (after some time it will default to the fair algorithm), but eventual
/// fairness does not provide the same garantees a always fair method would.
/// Fair mutexes are generally slower, but sometimes needed. This wrapper was
/// created to avoid using a unfair protocol when it's forbidden by mistake.
///
/// In a fair mutex the lock is provided to whichever thread asked first,
/// they form a queue and always follow the first-in first-out order. This
/// means some thread in the queue won't be able to steal the lock and use it fast
/// to increase throughput, at the cost of latency. Since the response time will grow
/// for some threads that are waiting for the lock and losing to faster but later ones,
/// but it may make sending more responses possible.
///
/// A fair mutex may not be interesting if threads have different priorities (this is known as
/// priority inversion).
///
/// # Differences from the standard library `Mutex`
///
/// - No poisoning, the lock is released normally on panic.
/// - Only requires 1 byte of space, whereas the standard library boxes the
///   `FairMutex` due to platform limitations.
/// - Can be statically constructed (requires the `const_fn` nightly feature).
/// - Does not require any drop glue when dropped.
/// - Inline fast path for the uncontended case.
/// - Efficient handling of micro-contention using adaptive spinning.
/// - Allows raw locking & unlocking without a guard.
///
/// # Examples
///
/// ```
/// use parking_lot::FairMutex;
/// use std::sync::{Arc, mpsc::channel};
/// use std::thread;
///
/// const N: usize = 10;
///
/// // Spawn a few threads to increment a shared variable (non-atomically), and
/// // let the main thread know once all increments are done.
/// //
/// // Here we're using an Arc to share memory among threads, and the data inside
/// // the Arc is protected with a mutex.
/// let data = Arc::new(FairMutex::new(0));
///
/// let (tx, rx) = channel();
/// for _ in 0..10 {
///     let (data, tx) = (Arc::clone(&data), tx.clone());
///     thread::spawn(move || {
///         // The shared state can only be accessed once the lock is held.
///         // Our non-atomic increment is safe because we're the only thread
///         // which can access the shared state when the lock is held.
///         let mut data = data.lock();
///         *data += 1;
///         if *data == N {
///             tx.send(()).unwrap();
///         }
///         // the lock is unlocked here when `data` goes out of scope.
///     });
/// }
///
/// rx.recv().unwrap();
/// ```
pub type FairMutex<T> = lock_api::Mutex<RawFairMutex, T>;

/// An RAII implementation of a "scoped lock" of a mutex. When this structure is
/// dropped (falls out of scope), the lock will be unlocked.
///
/// The data protected by the mutex can be accessed through this guard via its
/// `Deref` and `DerefMut` implementations.
pub type FairMutexGuard<'a, T> = lock_api::MutexGuard<'a, RawFairMutex, T>;

/// An RAII mutex guard returned by `FairMutexGuard::map`, which can point to a
/// subfield of the protected data.
///
/// The main difference between `MappedFairMutexGuard` and `FairMutexGuard` is that the
/// former doesn't support temporarily unlocking and re-locking, since that
/// could introduce soundness issues if the locked object is modified by another
/// thread.
pub type MappedFairMutexGuard<'a, T> = lock_api::MappedMutexGuard<'a, RawFairMutex, T>;

#[cfg(test)]
mod tests {
    use crate::FairMutex;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::mpsc::channel;
    use std::sync::Arc;
    use std::thread;

    #[cfg(feature = "serde")]
    use bincode::{deserialize, serialize};

    #[derive(Eq, PartialEq, Debug)]
    struct NonCopy(i32);

    #[test]
    fn smoke() {
        let m = FairMutex::new(());
        drop(m.lock());
        drop(m.lock());
    }

    #[test]
    fn lots_and_lots() {
        const J: u32 = 1000;
        const K: u32 = 3;

        let m = Arc::new(FairMutex::new(0));

        fn inc(m: &FairMutex<u32>) {
            for _ in 0..J {
                *m.lock() += 1;
            }
        }

        let (tx, rx) = channel();
        for _ in 0..K {
            let tx2 = tx.clone();
            let m2 = m.clone();
            thread::spawn(move || {
                inc(&m2);
                tx2.send(()).unwrap();
            });
            let tx2 = tx.clone();
            let m2 = m.clone();
            thread::spawn(move || {
                inc(&m2);
                tx2.send(()).unwrap();
            });
        }

        drop(tx);
        for _ in 0..2 * K {
            rx.recv().unwrap();
        }
        assert_eq!(*m.lock(), J * K * 2);
    }

    #[test]
    fn try_lock() {
        let m = FairMutex::new(());
        *m.try_lock().unwrap() = ();
    }

    #[test]
    fn test_into_inner() {
        let m = FairMutex::new(NonCopy(10));
        assert_eq!(m.into_inner(), NonCopy(10));
    }

    #[test]
    fn test_into_inner_drop() {
        struct Foo(Arc<AtomicUsize>);
        impl Drop for Foo {
            fn drop(&mut self) {
                self.0.fetch_add(1, Ordering::SeqCst);
            }
        }
        let num_drops = Arc::new(AtomicUsize::new(0));
        let m = FairMutex::new(Foo(num_drops.clone()));
        assert_eq!(num_drops.load(Ordering::SeqCst), 0);
        {
            let _inner = m.into_inner();
            assert_eq!(num_drops.load(Ordering::SeqCst), 0);
        }
        assert_eq!(num_drops.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn test_get_mut() {
        let mut m = FairMutex::new(NonCopy(10));
        *m.get_mut() = NonCopy(20);
        assert_eq!(m.into_inner(), NonCopy(20));
    }

    #[test]
    fn test_mutex_arc_nested() {
        // Tests nested mutexes and access
        // to underlying data.
        let arc = Arc::new(FairMutex::new(1));
        let arc2 = Arc::new(FairMutex::new(arc));
        let (tx, rx) = channel();
        let _t = thread::spawn(move || {
            let lock = arc2.lock();
            let lock2 = lock.lock();
            assert_eq!(*lock2, 1);
            tx.send(()).unwrap();
        });
        rx.recv().unwrap();
    }

    #[test]
    fn test_mutex_arc_access_in_unwind() {
        let arc = Arc::new(FairMutex::new(1));
        let arc2 = arc.clone();
        let _ = thread::spawn(move || {
            struct Unwinder {
                i: Arc<FairMutex<i32>>,
            }
            impl Drop for Unwinder {
                fn drop(&mut self) {
                    *self.i.lock() += 1;
                }
            }
            let _u = Unwinder { i: arc2 };
            panic!();
        })
        .join();
        let lock = arc.lock();
        assert_eq!(*lock, 2);
    }

    #[test]
    fn test_mutex_unsized() {
        let mutex: &FairMutex<[i32]> = &FairMutex::new([1, 2, 3]);
        {
            let b = &mut *mutex.lock();
            b[0] = 4;
            b[2] = 5;
        }
        let comp: &[i32] = &[4, 2, 5];
        assert_eq!(&*mutex.lock(), comp);
    }

    #[test]
    fn test_mutexguard_sync() {
        fn sync<T: Sync>(_: T) {}

        let mutex = FairMutex::new(());
        sync(mutex.lock());
    }

    #[test]
    fn test_mutex_debug() {
        let mutex = FairMutex::new(vec![0u8, 10]);

        assert_eq!(format!("{:?}", mutex), "Mutex { data: [0, 10] }");
        let _lock = mutex.lock();
        assert_eq!(format!("{:?}", mutex), "Mutex { data: <locked> }");
    }

    #[cfg(feature = "serde")]
    #[test]
    fn test_serde() {
        let contents: Vec<u8> = vec![0, 1, 2];
        let mutex = FairMutex::new(contents.clone());

        let serialized = serialize(&mutex).unwrap();
        let deserialized: FairMutex<Vec<u8>> = deserialize(&serialized).unwrap();

        assert_eq!(*(mutex.lock()), *(deserialized.lock()));
        assert_eq!(contents, *(deserialized.lock()));
    }
}
<|endoftext|>


<commit_before>use std::{io::ErrorKind, path::PathBuf};
use std::fs::File;
use std::error::Error;
use std::fs;
use std::fmt;
use std::io::{Seek, Read, SeekFrom};
use std::collections::HashMap;
use flate2::read::GzDecoder;
use bzip2::Compression;
use bzip2::read::{BzEncoder, BzDecoder};

#[derive(Debug)]
pub enum FsError {
    FileNotFound,
    InvalidDirectory,
    NoFileHandle,
    MalformedDataSequence,
    CorruptedData,
}
impl Error for FsError {
    fn description(&self) -> &str {
        match *self {
            FsError::FileNotFound => "the folder does not exist or cannot be read from",
            FsError::InvalidDirectory => "the specified directory is not a valid directory",
            FsError::NoFileHandle => "the filesystem did not load a file yet",
            FsError::MalformedDataSequence => "the data sequence did not complete correctly",
            FsError::CorruptedData => "the data was corrupt",
        }
    }
}
impl fmt::Display for FsError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            FsError::FileNotFound => write!(f, "the folder specified could not be found or read from"),
            FsError::InvalidDirectory => write!(f, "the specified directory is not a valid directory"),
            FsError::NoFileHandle => write!(f, "the filesystem did not load a file yet"),
            FsError::MalformedDataSequence => write!(f, "the data sequence did not complete correctly"),
            FsError::CorruptedData => write!(f, "the data was corrupt"),
        }
    }
}
impl From<FsError> for std::io::Error {
    fn from(other: FsError) -> std::io::Error {
        std::io::Error::new(ErrorKind::Other, other)
    }
}

#[derive(Debug,Clone)]
pub enum CompressionType {
    /// The archive is not compressed and the raw data is the real data.
    None,
    /// The archive is compressed with the Bzip2 compression algorithm.
    Bzip2,
    /// The archive is compressed with (a slightly modified, headerless) Gzip codec.
    Gzip,
    /// The archive is compressed with a modified LZMA2 variant (without size field).
    Lzma
}

#[derive(Debug)]
pub struct FileSystem {
    path: PathBuf,
    mainfile: MainFile,
    indices: HashMap<u32, IndexFile>
}

#[derive(Debug)]
pub struct MainFile {
    file: Option<File>
}

#[derive(Debug)]
pub struct IndexFile {
    id: u32,
    file: File
}

#[derive(Debug,Clone)]
pub struct IndexEntry {
    index: u8,
    id: u32,
    size: u32,
    offset: u64
}

#[derive(Debug,Clone)]
pub struct EntryHeader {
    raw_size: u32,
    real_size: u32,
    compression: CompressionType
}

impl EntryHeader {

    pub fn from_bytes(bytes: [u8; 9]) -> Result<EntryHeader, std::io::Error> {
        // Parse the 9 bytes of important info
        let compression_type = bytes[0];
        let raw_size: u32 = ((bytes[1] as u32) << 24) | ((bytes[2] as u32) << 16) | ((bytes[3] as u32) << 8) | (bytes[4] as u32);
        let real_size: u32 = ((bytes[5] as u32) << 24) | ((bytes[6] as u32) << 16) | ((bytes[7] as u32) << 8) | (bytes[8] as u32);

        // Return the new entry header
        return Ok(EntryHeader {
            raw_size: raw_size,
            real_size: real_size,
            compression: CompressionType::from_code(compression_type)
        });
    }

}

#[derive(Debug,Clone)]
pub struct BlockHeader {
    big: bool,
    entry_id: u32,
    index_id: u8,

    next_seq: i32,
    next_block: u32
}

impl CompressionType {
    /// Fetches the appropriate type of compression based on the header field
    /// in the archive header.
    pub fn from_code(code: u8) -> CompressionType {
        match code {
            1 => CompressionType::Bzip2,
            2 => CompressionType::Gzip,
            3 => CompressionType::Lzma,
            _ => CompressionType::None
        }
    }
}

impl BlockHeader {
    pub fn from_block(big: bool, data: [u8; 520]) -> BlockHeader {
        match big {
            true => {
                BlockHeader {
                    big: true,
                    entry_id: ((data[0] as u32) << 24) | ((data[1] as u32) << 16) | ((data[2] as u32) << 8) | (data[3] as u32),
                    next_seq: (((data[4] as u32) << 8) | (data[5] as u32)) as i32,
                    next_block: ((data[6] as u32) << 16) | ((data[7] as u32) << 8) | (data[8] as u32),
                    index_id: data[9] as u8
                }
            },
            false => {
                BlockHeader {
                    big: false,
                    entry_id: ((data[0] as u32) << 8) | (data[1] as u32),
                    next_seq: (((data[2] as u32) << 8) | (data[3] as u32)) as i32,
                    next_block: ((data[4] as u32) << 16) | ((data[5] as u32) << 8) | (data[6] as u32),
                    index_id: data[7] as u8
                }
            }
        }
    }
}

impl IndexEntry {
    pub fn index(&self) -> u8 {
        self.index
    }

    /// Gets the id of this block.
    pub fn id(&self) -> u32 {
        self.id
    }

    /// Gets the absolute size of the entry data, not counting the 8-10
    /// byte header in the blocks.
    pub fn size(&self) -> u32 {
        self.size
    }

    /// Gets the absolute offset of the very first 520-byte block of this
    /// entry in the main data file.
    pub fn offset(&self) -> u64 {
        self.offset
    }

    pub fn block(&self) -> u32 {
        (self.offset / 520u64) as u32
    }
}

impl IndexFile {
    pub fn last_entry(&self) -> u64 {
         self.file.metadata().unwrap().len() / 6u64
    }

    pub fn entry(&mut self, id: u32) -> Option<IndexEntry> {
        let ref mut file = self.file;
        let mut tmp: [u8; 6] = [0; 6];

        // Seek to the proper position and read into the temp buffer
        let seek_offset = id as u64 * 6u64;
        let res1 = file.seek(SeekFrom::Start(seek_offset));
        let res2 = file.read(&mut tmp);

        // Check if the seek and read operation succeeded
        if res1.is_err() || res2.is_err() {
            return None;
        }

        // Check if the operations returned the expected results
        if res1.unwrap() != seek_offset || res2.unwrap() != 6 {
            return None;
        }

        // Decode the size and offset from the temp buffer
        let size: u32 = ((tmp[0] as u32) << 16) | ((tmp[1] as u32) << 8) | (tmp[2] as u32);
        let offset: u64 = ((tmp[3] as u64) << 16) | ((tmp[4] as u64) << 8) | (tmp[5] as u64);

        Some(IndexEntry {index: self.id as u8, id: id, size: size, offset: offset * 520u64})
    }
}

impl FileSystem {
    pub fn new(string: &'static str) -> Result<FileSystem, FsError> {
        // Declare some nice variables!!!
        let path = PathBuf::from(string);
        let metadata = fs::metadata(&path);

        // Make sure the folder exists
        if !metadata.is_ok() {
            return Err(FsError::FileNotFound);
        }

        // Make sure it is a directory
        if !metadata.unwrap().is_dir() {
            return Err(FsError::InvalidDirectory);
        }

        // Create mainfile path
        let mut mainfile_path = PathBuf::from(string);
        mainfile_path.push("main_file_cache.dat2");

        // Find all valid index files
        let mut indices: HashMap<u32, IndexFile> = HashMap::new();
        let entries = fs::read_dir(&path).unwrap();
        for entry in entries {
            let e = entry.unwrap();
            let fname = e.file_name().into_string().unwrap();

            // Is this an index?
            if fname.starts_with("main_file_cache.idx") {
                // Parse the index id into an integer
                let idx = fname[19..].parse::<u32>().unwrap();

                // Add the index file to our map with indices
                indices.insert(idx, IndexFile {id: idx, file: File::open(e.path()).unwrap()});
            }
        }

        // Create the filesystem object and return it
        let file = File::open(mainfile_path).ok();
        let mainfile = MainFile{file: file};

        Ok(FileSystem {path: path, mainfile: mainfile, indices: indices})
    }

    /// Gets the mainfile, that is, the main_file_cache.dat2 entry in the folder
    /// that holds the actual binary data of the filesystem entries.
    pub fn mainfile(&mut self) -> &mut MainFile {
        &mut self.mainfile
    }

    /// Gets an index with a specific id if it exists. The index can only exist if the file exists
    /// on the file system.
    pub fn index(&mut self, index: u32) -> Option<&mut IndexFile> {
        self.indices.get_mut(&index)
    }
}

impl MainFile {
    /// Checks if the file exists.
    pub fn exists(&self) -> bool {
        self.file.is_some()
    }

    /// Gets the backing file, if existant. Returns a new instance with a fresh seek pointer.
    pub fn file(&mut self) -> Option<&mut File> {
        self.file.as_mut()
    }

    /// Calculates the number of data blocks in the mainfile (if existant). This is done by
    /// taking the file size and dividing that by 520 (rouding up), because each block
    /// takes up 520 bytes of data.
    pub fn num_blocks(&self) -> Option<u64> {
        match self.file {
            Some(ref x) => Some((x.metadata().unwrap().len() + 519u64) / 520u64),
            None => None
        }
    }

    /// Reads a block of data, specified by the block id. The data is read at 520 * block_id
    /// and is exactly 520 bytes big. It is not guaranteed all 520 bytes are occupied if the
    /// block is the last one, thus possible to be trimmed.
    pub fn read_block(&mut self, block: u32) -> Option<[u8; 520]> {
        // Do we have a valid file?
        if self.file.is_none() {
            return None;
        }

        let mut data: [u8; 520] = [0; 520];
        let file = self.file().unwrap();

        // Seek to the right position and read the data
        file.seek(SeekFrom::Start(block as u64 * 520u64)).unwrap();
        file.read(&mut data).unwrap();

        return Some(data);
    }

    pub fn read_header(&mut self, entry: IndexEntry) -> Option<EntryHeader> {
        // Do we have a valid file?
        if self.file.is_none() {
            return None;
        }

        let mut hdr: [u8; 9] = [0; 9];
        let file = self.file().unwrap();

        // Seek to the right position and read the data, skipping the block header at start
        let block_header_len = if entry.id() > 0xFFFF { 10 } else { 8 };
        file.seek(SeekFrom::Start(entry.offset() + block_header_len)).unwrap();
        file.read(&mut hdr).unwrap();

        return Some(EntryHeader::from_bytes(hdr).unwrap());
    }

    pub fn read_entry(&mut self, entry: IndexEntry) -> Result<Vec<u8>, FsError> {
        // Do we have a valid file?
        if self.file.is_none() {
            return Err(FsError::NoFileHandle);
        }

        // Create a vec with what we assume is the size. If not, the vec will
        // perfectly resize itself, so it's only an estimation to help us speed up.
        let mut data: Vec<u8> = Vec::with_capacity(entry.size() as usize);

        let mut current_block = entry.block();
        let mut remaining = entry.size();
        let mut current_seq = 0; // We expect a next part to be '1'

        while remaining > 0 {
            let block_data = self.read_block(current_block).unwrap();
            let block_info = BlockHeader::from_block(entry.id() > 65535, block_data);

            let header_size = if block_info.big {10} else {8};
            let available_data = 520 - header_size;
            let consumable = if remaining > available_data {available_data} else {remaining};

            if block_info.big {
                data.extend(&block_data[10..10+consumable as usize]);
            } else {
                data.extend(&block_data[8..8+consumable as usize]);
            }

            remaining -= consumable;

            // Do some checks to validate this block.
            if remaining > 0 {
                if block_info.index_id != entry.index() || block_info.next_seq != current_seq {
                    return Err(FsError::MalformedDataSequence);
                }
            }

            current_block += 1;
            current_seq += 1;
        }

        Ok(data)
    }

    pub fn read_decompressed(&mut self, entry: IndexEntry) -> Result<Vec<u8>, FsError> {
        let mut data = self.read_entry(entry.clone())?;
        let header = self.read_header(entry).unwrap();

        match header.compression {
            CompressionType::None => {
                Ok(data[5usize..(header.raw_size+5) as usize].to_vec())
            }
            CompressionType::Gzip => {
                let mut cursor = std::io::Cursor::new(&mut data);
                cursor.seek(SeekFrom::Current(9)).unwrap();

                let mut decoder = GzDecoder::new(cursor);
                let mut out = Vec::<u8>::new();
                out.resize(header.real_size as usize, 0);

                match decoder.read_exact(out.as_mut_slice()) {
                    Err(_) => return Err(FsError::CorruptedData),
                    Ok(_) => return Ok(out),
                }
            }
            CompressionType::Lzma => {
                panic!("Lzma compression unsupported")
            }
            CompressionType::Bzip2 => {
                // Patch the data so that the prefix is present
                data[5] = b'B';
                data[6] = b'Z';
                data[7] = b'h';
                data[8] = b'1';

                let mut cursor = std::io::Cursor::new(&mut data);
                cursor.seek(SeekFrom::Current(5)).unwrap();

                let mut decoder = BzDecoder::new(cursor);
                let mut out = Vec::<u8>::new();
                out.resize(header.real_size as usize, 0);

                match decoder.read_exact(out.as_mut_slice()) {
                    Err(_) => return Err(FsError::CorruptedData),
                    Ok(_) => return Ok(out),
                }
            }
        }
    }

}
<commit_msg>Clean up<commit_after>use std::{io::ErrorKind, path::PathBuf};
use std::fs::File;
use std::error::Error;
use std::fs;
use std::fmt;
use std::io::{Seek, Read, SeekFrom};
use std::collections::HashMap;
use flate2::read::GzDecoder;
use bzip2::read::BzDecoder;

#[derive(Debug)]
pub enum FsError {
    FileNotFound,
    InvalidDirectory,
    NoFileHandle,
    MalformedDataSequence,
    CorruptedData,
}
impl Error for FsError {
    fn description(&self) -> &str {
        match *self {
            FsError::FileNotFound => "the folder does not exist or cannot be read from",
            FsError::InvalidDirectory => "the specified directory is not a valid directory",
            FsError::NoFileHandle => "the filesystem did not load a file yet",
            FsError::MalformedDataSequence => "the data sequence did not complete correctly",
            FsError::CorruptedData => "the data was corrupt",
        }
    }
}
impl fmt::Display for FsError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            FsError::FileNotFound => write!(f, "the folder specified could not be found or read from"),
            FsError::InvalidDirectory => write!(f, "the specified directory is not a valid directory"),
            FsError::NoFileHandle => write!(f, "the filesystem did not load a file yet"),
            FsError::MalformedDataSequence => write!(f, "the data sequence did not complete correctly"),
            FsError::CorruptedData => write!(f, "the data was corrupt"),
        }
    }
}
impl From<FsError> for std::io::Error {
    fn from(other: FsError) -> std::io::Error {
        std::io::Error::new(ErrorKind::Other, other)
    }
}

#[derive(Debug,Clone)]
pub enum CompressionType {
    /// The archive is not compressed and the raw data is the real data.
    None,
    /// The archive is compressed with the Bzip2 compression algorithm.
    Bzip2,
    /// The archive is compressed with (a slightly modified, headerless) Gzip codec.
    Gzip,
    /// The archive is compressed with a modified LZMA2 variant (without size field).
    Lzma
}

#[derive(Debug)]
pub struct FileSystem {
    path: PathBuf,
    mainfile: MainFile,
    indices: HashMap<u32, IndexFile>
}

#[derive(Debug)]
pub struct MainFile {
    file: Option<File>
}

#[derive(Debug)]
pub struct IndexFile {
    id: u32,
    file: File
}

#[derive(Debug,Clone)]
pub struct IndexEntry {
    index: u8,
    id: u32,
    size: u32,
    offset: u64
}

#[derive(Debug,Clone)]
pub struct EntryHeader {
    raw_size: u32,
    real_size: u32,
    compression: CompressionType
}

impl EntryHeader {

    pub fn from_bytes(bytes: [u8; 9]) -> Result<EntryHeader, std::io::Error> {
        // Parse the 9 bytes of important info
        let compression_type = bytes[0];
        let raw_size: u32 = ((bytes[1] as u32) << 24) | ((bytes[2] as u32) << 16) | ((bytes[3] as u32) << 8) | (bytes[4] as u32);
        let real_size: u32 = ((bytes[5] as u32) << 24) | ((bytes[6] as u32) << 16) | ((bytes[7] as u32) << 8) | (bytes[8] as u32);

        // Return the new entry header
        return Ok(EntryHeader {
            raw_size: raw_size,
            real_size: real_size,
            compression: CompressionType::from_code(compression_type)
        });
    }

}

#[derive(Debug,Clone)]
pub struct BlockHeader {
    big: bool,
    entry_id: u32,
    index_id: u8,

    next_seq: i32,
    next_block: u32
}

impl CompressionType {
    /// Fetches the appropriate type of compression based on the header field
    /// in the archive header.
    pub fn from_code(code: u8) -> CompressionType {
        match code {
            1 => CompressionType::Bzip2,
            2 => CompressionType::Gzip,
            3 => CompressionType::Lzma,
            _ => CompressionType::None
        }
    }
}

impl BlockHeader {
    pub fn from_block(big: bool, data: [u8; 520]) -> BlockHeader {
        match big {
            true => {
                BlockHeader {
                    big: true,
                    entry_id: ((data[0] as u32) << 24) | ((data[1] as u32) << 16) | ((data[2] as u32) << 8) | (data[3] as u32),
                    next_seq: (((data[4] as u32) << 8) | (data[5] as u32)) as i32,
                    next_block: ((data[6] as u32) << 16) | ((data[7] as u32) << 8) | (data[8] as u32),
                    index_id: data[9] as u8
                }
            },
            false => {
                BlockHeader {
                    big: false,
                    entry_id: ((data[0] as u32) << 8) | (data[1] as u32),
                    next_seq: (((data[2] as u32) << 8) | (data[3] as u32)) as i32,
                    next_block: ((data[4] as u32) << 16) | ((data[5] as u32) << 8) | (data[6] as u32),
                    index_id: data[7] as u8
                }
            }
        }
    }
}

impl IndexEntry {
    pub fn index(&self) -> u8 {
        self.index
    }

    /// Gets the id of this block.
    pub fn id(&self) -> u32 {
        self.id
    }

    /// Gets the absolute size of the entry data, not counting the 8-10
    /// byte header in the blocks.
    pub fn size(&self) -> u32 {
        self.size
    }

    /// Gets the absolute offset of the very first 520-byte block of this
    /// entry in the main data file.
    pub fn offset(&self) -> u64 {
        self.offset
    }

    pub fn block(&self) -> u32 {
        (self.offset / 520u64) as u32
    }
}

impl IndexFile {
    pub fn last_entry(&self) -> u64 {
         self.file.metadata().unwrap().len() / 6u64
    }

    pub fn entry(&mut self, id: u32) -> Option<IndexEntry> {
        let ref mut file = self.file;
        let mut tmp: [u8; 6] = [0; 6];

        // Seek to the proper position and read into the temp buffer
        let seek_offset = id as u64 * 6u64;
        let res1 = file.seek(SeekFrom::Start(seek_offset));
        let res2 = file.read(&mut tmp);

        // Check if the seek and read operation succeeded
        if res1.is_err() || res2.is_err() {
            return None;
        }

        // Check if the operations returned the expected results
        if res1.unwrap() != seek_offset || res2.unwrap() != 6 {
            return None;
        }

        // Decode the size and offset from the temp buffer
        let size: u32 = ((tmp[0] as u32) << 16) | ((tmp[1] as u32) << 8) | (tmp[2] as u32);
        let offset: u64 = ((tmp[3] as u64) << 16) | ((tmp[4] as u64) << 8) | (tmp[5] as u64);

        Some(IndexEntry {index: self.id as u8, id: id, size: size, offset: offset * 520u64})
    }
}

impl FileSystem {
    pub fn new(string: &'static str) -> Result<FileSystem, FsError> {
        // Declare some nice variables!!!
        let path = PathBuf::from(string);
        let metadata = fs::metadata(&path);

        // Make sure the folder exists
        if !metadata.is_ok() {
            return Err(FsError::FileNotFound);
        }

        // Make sure it is a directory
        if !metadata.unwrap().is_dir() {
            return Err(FsError::InvalidDirectory);
        }

        // Create mainfile path
        let mut mainfile_path = PathBuf::from(string);
        mainfile_path.push("main_file_cache.dat2");

        // Find all valid index files
        let mut indices: HashMap<u32, IndexFile> = HashMap::new();
        let entries = fs::read_dir(&path).unwrap();
        for entry in entries {
            let e = entry.unwrap();
            let fname = e.file_name().into_string().unwrap();

            // Is this an index?
            if fname.starts_with("main_file_cache.idx") {
                // Parse the index id into an integer
                let idx = fname[19..].parse::<u32>().unwrap();

                // Add the index file to our map with indices
                indices.insert(idx, IndexFile {id: idx, file: File::open(e.path()).unwrap()});
            }
        }

        // Create the filesystem object and return it
        let file = File::open(mainfile_path).ok();
        let mainfile = MainFile{file: file};

        Ok(FileSystem {path: path, mainfile: mainfile, indices: indices})
    }

    /// Gets the mainfile, that is, the main_file_cache.dat2 entry in the folder
    /// that holds the actual binary data of the filesystem entries.
    pub fn mainfile(&mut self) -> &mut MainFile {
        &mut self.mainfile
    }

    /// Gets an index with a specific id if it exists. The index can only exist if the file exists
    /// on the file system.
    pub fn index(&mut self, index: u32) -> Option<&mut IndexFile> {
        self.indices.get_mut(&index)
    }
}

impl MainFile {
    /// Checks if the file exists.
    pub fn exists(&self) -> bool {
        self.file.is_some()
    }

    /// Gets the backing file, if existant. Returns a new instance with a fresh seek pointer.
    pub fn file(&mut self) -> Option<&mut File> {
        self.file.as_mut()
    }

    /// Calculates the number of data blocks in the mainfile (if existant). This is done by
    /// taking the file size and dividing that by 520 (rouding up), because each block
    /// takes up 520 bytes of data.
    pub fn num_blocks(&self) -> Option<u64> {
        match self.file {
            Some(ref x) => Some((x.metadata().unwrap().len() + 519u64) / 520u64),
            None => None
        }
    }

    /// Reads a block of data, specified by the block id. The data is read at 520 * block_id
    /// and is exactly 520 bytes big. It is not guaranteed all 520 bytes are occupied if the
    /// block is the last one, thus possible to be trimmed.
    pub fn read_block(&mut self, block: u32) -> Option<[u8; 520]> {
        // Do we have a valid file?
        if self.file.is_none() {
            return None;
        }

        let mut data: [u8; 520] = [0; 520];
        let file = self.file().unwrap();

        // Seek to the right position and read the data
        file.seek(SeekFrom::Start(block as u64 * 520u64)).unwrap();
        file.read(&mut data).unwrap();

        return Some(data);
    }

    pub fn read_header(&mut self, entry: IndexEntry) -> Option<EntryHeader> {
        // Do we have a valid file?
        if self.file.is_none() {
            return None;
        }

        let mut hdr: [u8; 9] = [0; 9];
        let file = self.file().unwrap();

        // Seek to the right position and read the data, skipping the block header at start
        let block_header_len = if entry.id() > 0xFFFF { 10 } else { 8 };
        file.seek(SeekFrom::Start(entry.offset() + block_header_len)).unwrap();
        file.read(&mut hdr).unwrap();

        return Some(EntryHeader::from_bytes(hdr).unwrap());
    }

    pub fn read_entry(&mut self, entry: IndexEntry) -> Result<Vec<u8>, FsError> {
        // Do we have a valid file?
        if self.file.is_none() {
            return Err(FsError::NoFileHandle);
        }

        // Create a vec with what we assume is the size. If not, the vec will
        // perfectly resize itself, so it's only an estimation to help us speed up.
        let mut data: Vec<u8> = Vec::with_capacity(entry.size() as usize);

        let mut current_block = entry.block();
        let mut remaining = entry.size();
        let mut current_seq = 0; // We expect a next part to be '1'

        while remaining > 0 {
            let block_data = self.read_block(current_block).unwrap();
            let block_info = BlockHeader::from_block(entry.id() > 65535, block_data);

            let header_size = if block_info.big {10} else {8};
            let available_data = 520 - header_size;
            let consumable = if remaining > available_data {available_data} else {remaining};

            if block_info.big {
                data.extend(&block_data[10..10+consumable as usize]);
            } else {
                data.extend(&block_data[8..8+consumable as usize]);
            }

            remaining -= consumable;

            // Do some checks to validate this block.
            if remaining > 0 {
                if block_info.index_id != entry.index() || block_info.next_seq != current_seq {
                    return Err(FsError::MalformedDataSequence);
                }
            }

            current_block += 1;
            current_seq += 1;
        }

        Ok(data)
    }

    pub fn read_decompressed(&mut self, entry: IndexEntry) -> Result<Vec<u8>, FsError> {
        let mut data = self.read_entry(entry.clone())?;
        let header = self.read_header(entry).unwrap();

        match header.compression {
            CompressionType::None => {
                Ok(data[5usize..(header.raw_size+5) as usize].to_vec())
            }
            CompressionType::Gzip => {
                let mut cursor = std::io::Cursor::new(&mut data);
                cursor.seek(SeekFrom::Current(9)).unwrap();

                let mut decoder = GzDecoder::new(cursor);
                let mut out = Vec::<u8>::new();
                out.resize(header.real_size as usize, 0);

                match decoder.read_exact(out.as_mut_slice()) {
                    Err(_) => return Err(FsError::CorruptedData),
                    Ok(_) => return Ok(out),
                }
            }
            CompressionType::Lzma => {
                panic!("Lzma compression unsupported")
            }
            CompressionType::Bzip2 => {
                // Patch the data so that the prefix is present
                data[5] = b'B';
                data[6] = b'Z';
                data[7] = b'h';
                data[8] = b'1';

                let mut cursor = std::io::Cursor::new(&mut data);
                cursor.seek(SeekFrom::Current(5)).unwrap();

                let mut decoder = BzDecoder::new(cursor);
                let mut out = Vec::<u8>::new();
                out.resize(header.real_size as usize, 0);

                match decoder.read_exact(out.as_mut_slice()) {
                    Err(_) => return Err(FsError::CorruptedData),
                    Ok(_) => return Ok(out),
                }
            }
        }
    }

}
<|endoftext|>