<commit_before>pub mod search;

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

// use rand;
// use rand::distributions::{Sample, Range};
// let mut rng = rand::thread_rng();
// let mut between = Range::new(0, 10);
// let x = between.sample(&mut rng);

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

/// Implements `UciEngine` trait.
pub struct Engine {
    position: Position,
    pondering_is_allowed: bool,
    is_thinking: bool,
    is_pondering: bool,
    stop_when: TimeManagement,
    replies: Vec<EngineReply>,
    tt: Arc<TranspositionTable>,
    commands: Sender<search::Command>,
    reports: Receiver<search::Report>,
    thinking_since: SystemTime,
    searched_nodes: NodeCount,
    searched_time: u64,  // milliseconds
    searched_depth: u8,
}

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 tt = Arc::new(tt);
        let (commands_tx, commands_rx) = channel();
        let (reports_tx, reports_rx) = channel();
        let tt2 = tt.clone();
        thread::spawn(move || {
            search::run_deepening(tt2, commands_rx, reports_tx);
        });

        Engine {
            position: Position::from_fen("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 \
                                          1")
                          .ok()
                          .unwrap(),
            pondering_is_allowed: false,
            is_thinking: false,
            is_pondering: false,
            stop_when: TimeManagement::Infinite,
            replies: vec![],
            tt: tt,
            commands: commands_tx,
            reports: reports_rx,
            thinking_since: SystemTime::now(),
            searched_nodes: 0,
            searched_time: 0,
            searched_depth: 0,
        }
    }
}

impl UciEngine for Engine {
    fn set_option(&mut self, name: &str, value: &str) {
        match name {
            // Tells the engine that it will be allowed to ponder.
            // This option is needed because the engine might change
            // its time management algorithm when pondering is
            // allowed.
            "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 change anything,
        // 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 do not support the "searchmoves" parameter.

            let (time, inc) = if self.position.board().to_move() == WHITE {
                (wtime, winc.unwrap_or(0))
            } else {
                (btime, binc.unwrap_or(0))
            };
            self.is_thinking = true;
            self.is_pondering = ponder;
            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(min(depth.unwrap(), MAX_DEPTH as u64) as u8)
            } else {
                // TODO: Robert Hyatt proposes this:
                //
                // nMoves =  min( numberOfMovesOutOfBook, 10 )
                // factor = 2 - nMoves / 10
                // target = timeLeft / numberOfMovesUntilNextTimeControl
                // time   = factor * target
                let time = time.unwrap_or(0);
                let movestogo = movestogo.unwrap_or(40);
                let movetime = (time + inc * movestogo) / movestogo;
                TimeManagement::MoveTimeHint(min(movetime, time / 2))
            };
            self.thinking_since = SystemTime::now();
            self.searched_nodes = 0;
            self.searched_time = 0;
            self.searched_depth = 0;

            self.commands
                .send(search::Command::Search {
                    search_id: 0,
                    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 best_move = self.get_best_move();
            self.replies.push(EngineReply::BestMove {
                best_move: best_move,
                ponder_move: None,
            });
            self.is_thinking = false;
        }
    }

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

    fn get_reply(&mut self) -> Option<EngineReply> {
        // TODO: implement panic mode for MoveTimeHint management.
        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.searched_depth >= d,
                TimeManagement::Infinite => false,
            } {
                self.stop();
            }
        }

        // Empty the reports queue.
        while let Ok(report) = self.reports.try_recv() {
            // let rep_type = match report {
            //     search::Report::Progress { searched_nodes, .. } => format!("progress {}", searched_nodes),
            //     search::Report::Done { .. } => "done".to_string(),
            // };
            // self.replies.push(EngineReply::Info(vec![
            //     ("string".to_string(), rep_type),
            // ]));
            if self.is_thinking {
                match report {
                    search::Report::Progress { searched_nodes, depth, .. } => {
                        self.searched_nodes = searched_nodes;
                        let duration = self.thinking_since.elapsed().unwrap();
                        self.searched_time = 1000 * duration.as_secs() +
                                             (duration.subsec_nanos() / 1000000) as u64;
                        if self.searched_depth < depth {
                            self.searched_depth = depth;
                            let mut p = self.position.clone();
                            let mut v = MoveStack::new();
                            let mut pv = String::from("");
                            let mut pv_length = 0;
                            let mut value = None;
                            'depthloop: for _ in 0..depth {
                                let m = match self.tt.probe(p.hash()) {
                                    Some(entry) => {
                                        if value.is_none() {
                                            value = Some(entry.value());
                                        }
                                        entry.move16()
                                    }
                                    None => 0,
                                };
                                if m != 0 {
                                    p.generate_moves(&mut v);
                                    while let Some(x) = v.pop() {
                                        if x.digest() == m && p.do_move(x) {
                                            pv_length += 1;
                                            pv.push_str(&x.notation());
                                            pv.push(' ');
                                            v.clear();
                                            continue 'depthloop;
                                        }
                                    }
                                }
                                break;
                            }
                            if searched_nodes > 0 && pv_length >= depth {
                                self.replies.push(EngineReply::Info(vec![
                                    ("depth".to_string(), format!("{}", depth)),
                                    ("score".to_string(), format!("cp {}", value.unwrap_or(666))),
                                    ("nodes".to_string(), format!("{}", searched_nodes)),
                                    ("time".to_string(), format!("{}", self.searched_time)),
                                    ("pv".to_string(), format!("{}", pv)),
                                ]));
                            }
                        }
                    }
                    _ => (),
                }
            }
        }

        if self.replies.len() > 0 {
            Some(self.replies.remove(0))
        } else {
            None
        }
    }
}

impl Engine {
    fn get_best_move(&mut self) -> String {
        let mut m = match self.tt.probe(self.position.hash()) {
            Some(entry) => entry.move16(),
            None => 0,
        };
        if m == 0 {
            // Pick the first legal move.
            let mut first_legal_move = Move::invalid();
            let mut v = MoveStack::new();
            self.position.generate_moves(&mut v);
            while let Some(x) = v.pop() {
                if self.position.do_move(x) {
                    self.position.undo_move();
                    first_legal_move = x;
                    break;
                }
            }
            m = first_legal_move.digest();
        }
        if m != 0 {
            let move_type = move_type(m);
            let orig_square = orig_square(m);
            let dest_square = dest_square(m);
            let promoted_piece = match aux_data(m) {
                0 => "q",
                1 => "r",
                2 => "b",
                3 => "n",
                _ => panic!("invalid promoted piece code"),
            };
            format!("{}{}{}",
                    notation(orig_square),
                    notation(dest_square),
                    if move_type == MOVE_PROMOTION {
                        promoted_piece
                    } else {
                        ""
                    })

        } else {
            "0000".to_string()
        }
    }
}

/// 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))
    }
}

enum TimeManagement {
    MoveTime(u64),
    MoveTimeHint(u64),
    Nodes(NodeCount),
    Depth(u8),
    Infinite,
}
<commit_msg>beautification<commit_after>pub mod search;

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

// use rand;
// use rand::distributions::{Sample, Range};
// let mut rng = rand::thread_rng();
// let mut between = Range::new(0, 10);
// let x = between.sample(&mut rng);

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

/// Implements `UciEngine` trait.
pub struct Engine {
    position: Position,
    pondering_is_allowed: bool,
    is_thinking: bool,
    is_pondering: bool,
    replies: Vec<EngineReply>,
    tt: Arc<TranspositionTable>,
    commands: Sender<search::Command>,
    reports: Receiver<search::Report>,
    thinking_since: SystemTime,
    searched_nodes: NodeCount,
    searched_time: u64,  // milliseconds
    searched_depth: u8,
    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 tt = Arc::new(tt);
        let (commands_tx, commands_rx) = channel();
        let (reports_tx, reports_rx) = channel();
        let tt2 = tt.clone();
        thread::spawn(move || {
            search::run_deepening(tt2, commands_rx, reports_tx);
        });

        Engine {
            position: Position::from_fen("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 \
                                          1")
                          .ok()
                          .unwrap(),
            pondering_is_allowed: false,
            is_thinking: false,
            is_pondering: false,
            replies: vec![],
            tt: tt,
            commands: commands_tx,
            reports: reports_rx,
            thinking_since: SystemTime::now(),
            searched_nodes: 0,
            searched_time: 0,
            searched_depth: 0,
            stop_when: TimeManagement::Infinite,
        }
    }
}

impl UciEngine for Engine {
    fn set_option(&mut self, name: &str, value: &str) {
        match name {
            // Tells the engine that it will be allowed to ponder.
            // This option is needed because the engine might change
            // its time management algorithm when pondering is
            // allowed.
            "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 change anything,
        // 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 do not support the "searchmoves" parameter.

            let (time, inc) = if self.position.board().to_move() == WHITE {
                (wtime, winc.unwrap_or(0))
            } else {
                (btime, binc.unwrap_or(0))
            };
            self.is_thinking = true;
            self.is_pondering = ponder;
            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(min(depth.unwrap(), MAX_DEPTH as u64) as u8)
            } else {
                // TODO: Robert Hyatt proposes this:
                //
                // nMoves =  min( numberOfMovesOutOfBook, 10 )
                // factor = 2 - nMoves / 10
                // target = timeLeft / numberOfMovesUntilNextTimeControl
                // time   = factor * target
                let time = time.unwrap_or(0);
                let movestogo = movestogo.unwrap_or(40);
                let movetime = (time + inc * movestogo) / movestogo;
                TimeManagement::MoveTimeHint(min(movetime, time / 2))
            };
            self.thinking_since = SystemTime::now();
            self.searched_nodes = 0;
            self.searched_time = 0;
            self.searched_depth = 0;

            self.commands
                .send(search::Command::Search {
                    search_id: 0,
                    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 best_move = self.get_best_move();
            self.replies.push(EngineReply::BestMove {
                best_move: best_move,
                ponder_move: None,
            });
            self.is_thinking = false;
        }
    }

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

    fn get_reply(&mut self) -> Option<EngineReply> {
        // TODO: implement panic mode for MoveTimeHint management.
        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.searched_depth >= d,
                TimeManagement::Infinite => false,
            } {
                self.stop();
            }
        }

        // Empty the reports queue.
        while let Ok(report) = self.reports.try_recv() {
            // let rep_type = match report {
            //     search::Report::Progress { searched_nodes, .. } => format!("progress {}", searched_nodes),
            //     search::Report::Done { .. } => "done".to_string(),
            // };
            // self.replies.push(EngineReply::Info(vec![
            //     ("string".to_string(), rep_type),
            // ]));
            if self.is_thinking {
                match report {
                    search::Report::Progress { searched_nodes, depth, .. } => {
                        self.searched_nodes = searched_nodes;
                        let duration = self.thinking_since.elapsed().unwrap();
                        self.searched_time = 1000 * duration.as_secs() +
                                             (duration.subsec_nanos() / 1000000) as u64;
                        if self.searched_depth < depth {
                            self.searched_depth = depth;
                            let mut p = self.position.clone();
                            let mut v = MoveStack::new();
                            let mut pv = String::from("");
                            let mut pv_length = 0;
                            let mut value = None;
                            'depthloop: for _ in 0..depth {
                                let m = match self.tt.probe(p.hash()) {
                                    Some(entry) => {
                                        if value.is_none() {
                                            value = Some(entry.value());
                                        }
                                        entry.move16()
                                    }
                                    None => 0,
                                };
                                if m != 0 {
                                    p.generate_moves(&mut v);
                                    while let Some(x) = v.pop() {
                                        if x.digest() == m && p.do_move(x) {
                                            pv_length += 1;
                                            pv.push_str(&x.notation());
                                            pv.push(' ');
                                            v.clear();
                                            continue 'depthloop;
                                        }
                                    }
                                }
                                break;
                            }
                            if searched_nodes > 0 && pv_length >= depth {
                                self.replies.push(EngineReply::Info(vec![
                                    ("depth".to_string(), format!("{}", depth)),
                                    ("score".to_string(), format!("cp {}", value.unwrap_or(666))),
                                    ("nodes".to_string(), format!("{}", searched_nodes)),
                                    ("time".to_string(), format!("{}", self.searched_time)),
                                    ("pv".to_string(), format!("{}", pv)),
                                ]));
                            }
                        }
                    }
                    _ => (),
                }
            }
        }

        if self.replies.len() > 0 {
            Some(self.replies.remove(0))
        } else {
            None
        }
    }
}

impl Engine {
    fn get_best_move(&mut self) -> String {
        let mut m = match self.tt.probe(self.position.hash()) {
            Some(entry) => entry.move16(),
            None => 0,
        };
        if m == 0 {
            // Pick the first legal move.
            let mut first_legal_move = Move::invalid();
            let mut v = MoveStack::new();
            self.position.generate_moves(&mut v);
            while let Some(x) = v.pop() {
                if self.position.do_move(x) {
                    self.position.undo_move();
                    first_legal_move = x;
                    break;
                }
            }
            m = first_legal_move.digest();
        }
        if m != 0 {
            let move_type = move_type(m);
            let orig_square = orig_square(m);
            let dest_square = dest_square(m);
            let promoted_piece = match aux_data(m) {
                0 => "q",
                1 => "r",
                2 => "b",
                3 => "n",
                _ => panic!("invalid promoted piece code"),
            };
            format!("{}{}{}",
                    notation(orig_square),
                    notation(dest_square),
                    if move_type == MOVE_PROMOTION {
                        promoted_piece
                    } else {
                        ""
                    })

        } else {
            "0000".to_string()
        }
    }
}

/// 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))
    }
}

enum TimeManagement {
    MoveTime(u64),
    MoveTimeHint(u64),
    Nodes(NodeCount),
    Depth(u8),
    Infinite,
}
<|endoftext|>


<commit_before>use crate::{
    core::{make_shared, Shared},
    fa::{StateActionFunction, FiniteStateActionFunction, Parameterised},
    geometry::{MatrixView, MatrixViewMut},
};
use std::marker::PhantomData;

pub struct MockQ {
    output: Option<Vec<f64>>,
}

impl MockQ {
    pub fn new(output: Option<Vec<f64>>) -> Self { MockQ { output: output } }

    pub fn new_shared(output: Option<Vec<f64>>) -> Shared<Self> {
        make_shared(Self::new(output))
    }

    #[allow(dead_code)]
    pub fn set_output(&mut self, output: Vec<f64>) { self.output = Some(output) }

    #[allow(dead_code)]
    pub fn clear_output(&mut self) { self.output = None }
}

impl Parameterised for MockQ {
    fn weights_view(&self) -> MatrixView<f64> { unimplemented!() }

    fn weights_view_mut(&mut self) -> MatrixViewMut<f64> { unimplemented!() }
}

impl StateActionFunction<Vec<f64>, usize> for MockQ {
    type Output = f64;

    fn evaluate(&self, x: &Vec<f64>, action: &usize) -> f64 {
        match self.output {
            Some(ref out) => out[*action].clone(),
            None => x[*action].clone(),
        }
    }

    fn update(&mut self, _: &Vec<f64>, _: &usize, _: f64) {}
}

impl FiniteStateActionFunction<Vec<f64>> for MockQ {
    fn n_actions(&self) -> usize {
        match self.output {
            Some(ref out) => out.len(),
            None => unimplemented!(),
        }
    }

    fn evaluate_all(&self, x: &Vec<f64>) -> Vec<f64> {
        match self.output {
            Some(ref out) => out.clone(),
            None => x.clone(),
        }
    }

    fn update_all(&mut self, _: &Vec<f64>, _: Vec<f64>) {}
}
<commit_msg>fix(fa): correct reference to deprecated trait in fa/mocking.rs<commit_after>use crate::{
    core::{make_shared, Shared},
    fa::{StateActionFunction, FiniteActionFunction, Parameterised},
    geometry::{MatrixView, MatrixViewMut},
};
use std::marker::PhantomData;

pub struct MockQ {
    output: Option<Vec<f64>>,
}

impl MockQ {
    pub fn new(output: Option<Vec<f64>>) -> Self { MockQ { output: output } }

    pub fn new_shared(output: Option<Vec<f64>>) -> Shared<Self> {
        make_shared(Self::new(output))
    }

    #[allow(dead_code)]
    pub fn set_output(&mut self, output: Vec<f64>) { self.output = Some(output) }

    #[allow(dead_code)]
    pub fn clear_output(&mut self) { self.output = None }
}

impl Parameterised for MockQ {
    fn weights_view(&self) -> MatrixView<f64> { unimplemented!() }

    fn weights_view_mut(&mut self) -> MatrixViewMut<f64> { unimplemented!() }
}

impl StateActionFunction<Vec<f64>, usize> for MockQ {
    type Output = f64;

    fn evaluate(&self, x: &Vec<f64>, action: &usize) -> f64 {
        match self.output {
            Some(ref out) => out[*action].clone(),
            None => x[*action].clone(),
        }
    }

    fn update(&mut self, _: &Vec<f64>, _: &usize, _: f64) {}
}

impl FiniteActionFunction<Vec<f64>> for MockQ {
    fn n_actions(&self) -> usize {
        match self.output {
            Some(ref out) => out.len(),
            None => unimplemented!(),
        }
    }

    fn evaluate_all(&self, x: &Vec<f64>) -> Vec<f64> {
        match self.output {
            Some(ref out) => out.clone(),
            None => x.clone(),
        }
    }

    fn update_all(&mut self, _: &Vec<f64>, _: Vec<f64>) {}
}
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::ast_discoverer::Discoveries;
use crate::CppCodegenOptions;
use crate::{
    cxxbridge::CxxBridge, Error as EngineError, GeneratedCpp, IncludeCppEngine,
    RebuildDependencyRecorder,
};
use autocxx_parser::directives::SUBCLASS;
use autocxx_parser::{RustPath, Subclass, SubclassAttrs};
use proc_macro2::{Span, TokenStream};
use quote::ToTokens;
use std::{collections::HashSet, fmt::Display, io::Read, path::PathBuf};
use std::{panic::UnwindSafe, path::Path, rc::Rc};
use syn::{token::Brace, Item, ItemMod, LitStr};

/// Errors which may occur when parsing a Rust source file to discover
/// and interpret include_cxx macros.
#[derive(Debug)]
pub enum ParseError {
    /// Unable to open the source file
    FileOpen(std::io::Error),
    /// The .rs file couldn't be read.
    FileRead(std::io::Error),
    /// The .rs file couldn't be parsed.
    Syntax(syn::Error),
    /// The include CPP macro could not be expanded into
    /// Rust bindings to C++, because of some problem during the conversion
    /// process. This could be anything from a C++ parsing error to some
    /// C++ feature that autocxx can't yet handle and isn't able to skip
    /// over. It could also cover errors in your syntax of the `include_cpp`
    /// macro or the directives inside.
    AutocxxCodegenError(EngineError),
    /// There are two or more `include_cpp` macros with the same
    /// mod name.
    ConflictingModNames,
    ZeroModsForDynamicDiscovery,
    MultipleModsForDynamicDiscovery,
    DiscoveredRustItemsWhenNotInAutoDiscover,
}

impl Display for ParseError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ParseError::FileOpen(err) => write!(f, "Unable to open file: {}", err)?,
            ParseError::FileRead(err) => write!(f, "Unable to read file: {}", err)?,
            ParseError::Syntax(err) => write!(f, "Syntax error parsing Rust file: {}", err)?,
            ParseError::AutocxxCodegenError(err) =>
                write!(f, "Unable to parse include_cpp! macro: {}", err)?,
            ParseError::ConflictingModNames =>
                write!(f, "There are two or more include_cpp! macros with the same output mod name. Use name!")?,
            ParseError::ZeroModsForDynamicDiscovery =>
                write!(f, "This file contains extra information to append to an include_cpp! but no such include_cpp! was found in this file.")?,
            ParseError::MultipleModsForDynamicDiscovery =>
                write!(f, "This file contains extra information to append to an include_cpp! but multiple such include_cpp! declarations were found in this file.")?,
            ParseError::DiscoveredRustItemsWhenNotInAutoDiscover =>
                write!(f, "This file contains extra information to append to an \"extern Rust\" but auto-discover was switched off.")?,
        }
        Ok(())
    }
}

/// Parse a Rust file, and spot any include_cpp macros within it.
pub fn parse_file<P1: AsRef<Path>>(
    rs_file: P1,
    auto_allowlist: bool,
) -> Result<ParsedFile, ParseError> {
    let mut source = String::new();
    let mut file = std::fs::File::open(rs_file).map_err(ParseError::FileOpen)?;
    file.read_to_string(&mut source)
        .map_err(ParseError::FileRead)?;
    proc_macro2::fallback::force();
    let source = syn::parse_file(&source).map_err(ParseError::Syntax)?;
    parse_file_contents(source, auto_allowlist)
}

fn parse_file_contents(source: syn::File, auto_allowlist: bool) -> Result<ParsedFile, ParseError> {
    #[derive(Default)]
    struct State {
        auto_allowlist: bool,
        results: Vec<Segment>,
        extra_superclasses: Vec<Subclass>,
        discoveries: Discoveries,
    }
    impl State {
        fn parse_item(&mut self, item: Item, mod_path: Option<RustPath>) -> Result<(), ParseError> {
            let result = match item {
                Item::Macro(mac)
                    if mac
                        .mac
                        .path
                        .segments
                        .last()
                        .map(|s| s.ident == "include_cpp")
                        .unwrap_or(false) =>
                {
                    Segment::Autocxx(
                        crate::IncludeCppEngine::new_from_syn(mac.mac.clone())
                            .map_err(ParseError::AutocxxCodegenError)?,
                    )
                }
                Item::Mod(itm)
                    if itm
                        .attrs
                        .iter()
                        .any(|attr| attr.path.to_token_stream().to_string() == "cxx :: bridge") =>
                {
                    Segment::Cxx(CxxBridge::from(itm))
                }
                Item::Mod(itm) => {
                    if let Some((brace, items)) = itm.content {
                        let mut mod_state = State {
                            auto_allowlist: self.auto_allowlist,
                            ..Default::default()
                        };
                        let mod_path = match &mod_path {
                            None => RustPath::new_from_ident(itm.ident.clone()),
                            Some(mod_path) => mod_path.append(itm.ident.clone()),
                        };
                        for item in items {
                            mod_state.parse_item(item, Some(mod_path.clone()))?
                        }
                        self.extra_superclasses.extend(mod_state.extra_superclasses);
                        self.discoveries.extend(mod_state.discoveries);
                        Segment::Mod(
                            mod_state.results,
                            (
                                brace,
                                ItemMod {
                                    content: None,
                                    ..itm
                                },
                            ),
                        )
                    } else {
                        Segment::Other(Item::Mod(itm))
                    }
                }
                Item::Struct(ref its) if self.auto_allowlist => {
                    let attrs = &its.attrs;
                    let is_superclass_attr = attrs.iter().find(|attr| {
                        attr.path
                            .segments
                            .last()
                            .map(|seg| seg.ident == "is_subclass" || seg.ident == SUBCLASS)
                            .unwrap_or(false)
                    });
                    if let Some(is_superclass_attr) = is_superclass_attr {
                        if !is_superclass_attr.tokens.is_empty() {
                            let subclass = its.ident.clone();
                            let args: SubclassAttrs = is_superclass_attr
                                .parse_args()
                                .map_err(ParseError::Syntax)?;
                            if let Some(superclass) = args.superclass {
                                self.extra_superclasses.push(Subclass {
                                    superclass,
                                    subclass,
                                })
                            }
                        }
                    }
                    self.discoveries.search_item(&item, mod_path);
                    Segment::Other(item)
                }
                _ => {
                    self.discoveries.search_item(&item, mod_path);
                    Segment::Other(item)
                }
            };
            self.results.push(result);
            Ok(())
        }
    }
    let mut state = State {
        auto_allowlist,
        ..Default::default()
    };
    for item in source.items {
        state.parse_item(item, None)?
    }
    let State {
        auto_allowlist,
        mut results,
        mut extra_superclasses,
        mut discoveries,
    } = state;
    if !auto_allowlist
        && (!discoveries.extern_rust_types.is_empty() || !discoveries.extern_rust_funs.is_empty())
    {
        return Err(ParseError::DiscoveredRustItemsWhenNotInAutoDiscover);
    }
    if !extra_superclasses.is_empty() || (auto_allowlist && !discoveries.is_empty()) {
        let mut autocxx_seg_iterator = results.iter_mut().filter_map(|seg| match seg {
            Segment::Autocxx(engine) => Some(engine),
            _ => None,
        });
        let our_seg = autocxx_seg_iterator.next();
        match our_seg {
            None => return Err(ParseError::ZeroModsForDynamicDiscovery),
            Some(engine) => {
                engine
                    .config_mut()
                    .subclasses
                    .append(&mut extra_superclasses);
                if auto_allowlist {
                    for cpp in discoveries.cpp_list {
                        engine
                            .config_mut()
                            .allowlist
                            .push(LitStr::new(&cpp, Span::call_site()))
                            .map_err(ParseError::Syntax)?;
                    }
                }
                engine
                    .config_mut()
                    .extern_rust_funs
                    .append(&mut discoveries.extern_rust_funs);
                engine
                    .config_mut()
                    .rust_types
                    .append(&mut discoveries.extern_rust_types);
            }
        }
        if autocxx_seg_iterator.next().is_some() {
            return Err(ParseError::MultipleModsForDynamicDiscovery);
        }
    }
    let autocxx_seg_iterator = results.iter_mut().filter_map(|seg| match seg {
        Segment::Autocxx(engine) => Some(engine),
        _ => None,
    });
    for seg in autocxx_seg_iterator {
        seg.config
            .confirm_complete(auto_allowlist)
            .map_err(ParseError::Syntax)?;
    }
    Ok(ParsedFile(results))
}

/// A Rust file parsed by autocxx. May contain zero or more autocxx 'engines',
/// i.e. the `IncludeCpp` class, corresponding to zero or more include_cpp
/// macros within this file. Also contains `syn::Item` structures for all
/// the rest of the Rust code, such that it can be reconstituted if necessary.
pub struct ParsedFile(Vec<Segment>);

#[allow(clippy::large_enum_variant)]
enum Segment {
    Autocxx(IncludeCppEngine),
    Cxx(CxxBridge),
    Mod(Vec<Segment>, (Brace, ItemMod)),
    Other(Item),
}

pub trait CppBuildable {
    fn generate_h_and_cxx(
        &self,
        cpp_codegen_options: &CppCodegenOptions,
    ) -> Result<GeneratedCpp, cxx_gen::Error>;
}

impl ParsedFile {
    /// Get all the autocxxes in this parsed file.
    pub fn get_rs_buildables(&self) -> impl Iterator<Item = &IncludeCppEngine> {
        fn do_get_rs_buildables(
            segments: &Vec<Segment>,
        ) -> impl Iterator<Item = &IncludeCppEngine> {
            segments
                .iter()
                .flat_map(|s| -> Box<dyn Iterator<Item = &IncludeCppEngine>> {
                    match s {
                        Segment::Autocxx(includecpp) => Box::new(std::iter::once(includecpp)),
                        Segment::Mod(segments, _) => Box::new(do_get_rs_buildables(segments)),
                        _ => Box::new(std::iter::empty()),
                    }
                })
        }

        do_get_rs_buildables(&self.0)
    }

    /// Get all items which can result in C++ code
    pub fn get_cpp_buildables(&self) -> impl Iterator<Item = &dyn CppBuildable> {
        fn do_get_cpp_buildables(
            segments: &Vec<Segment>,
        ) -> impl Iterator<Item = &dyn CppBuildable> {
            segments
                .iter()
                .flat_map(|s| -> Box<dyn Iterator<Item = &dyn CppBuildable>> {
                    match s {
                        Segment::Autocxx(includecpp) => {
                            Box::new(std::iter::once(includecpp as &dyn CppBuildable))
                        }
                        Segment::Cxx(cxxbridge) => {
                            Box::new(std::iter::once(cxxbridge as &dyn CppBuildable))
                        }
                        Segment::Mod(segments, _) => Box::new(do_get_cpp_buildables(segments)),
                        _ => Box::new(std::iter::empty()),
                    }
                })
        }

        do_get_cpp_buildables(&self.0)
    }

    fn get_autocxxes_mut(&mut self) -> impl Iterator<Item = &mut IncludeCppEngine> {
        fn do_get_autocxxes_mut(
            segments: &mut Vec<Segment>,
        ) -> impl Iterator<Item = &mut IncludeCppEngine> {
            segments
                .iter_mut()
                .flat_map(|s| -> Box<dyn Iterator<Item = &mut IncludeCppEngine>> {
                    match s {
                        Segment::Autocxx(includecpp) => Box::new(std::iter::once(includecpp)),
                        Segment::Mod(segments, _) => Box::new(do_get_autocxxes_mut(segments)),
                        _ => Box::new(std::iter::empty()),
                    }
                })
        }

        do_get_autocxxes_mut(&mut self.0)
    }

    pub fn include_dirs(&self) -> impl Iterator<Item = &PathBuf> {
        fn do_get_include_dirs(segments: &Vec<Segment>) -> impl Iterator<Item = &PathBuf> {
            segments
                .iter()
                .flat_map(|s| -> Box<dyn Iterator<Item = &PathBuf>> {
                    match s {
                        Segment::Autocxx(includecpp) => Box::new(includecpp.include_dirs()),
                        Segment::Mod(segments, _) => Box::new(do_get_include_dirs(segments)),
                        _ => Box::new(std::iter::empty()),
                    }
                })
        }

        do_get_include_dirs(&self.0)
    }

    pub fn resolve_all(
        &mut self,
        autocxx_inc: Vec<PathBuf>,
        extra_clang_args: &[&str],
        dep_recorder: Option<Box<dyn RebuildDependencyRecorder>>,
        cpp_codegen_options: &CppCodegenOptions,
    ) -> Result<(), ParseError> {
        let mut mods_found = HashSet::new();
        let inner_dep_recorder: Option<Rc<dyn RebuildDependencyRecorder>> =
            dep_recorder.map(Rc::from);
        for include_cpp in self.get_autocxxes_mut() {
            #[allow(clippy::manual_map)] // because of dyn shenanigans
            let dep_recorder: Option<Box<dyn RebuildDependencyRecorder>> = match &inner_dep_recorder
            {
                None => None,
                Some(inner_dep_recorder) => Some(Box::new(CompositeDepRecorder::new(
                    inner_dep_recorder.clone(),
                ))),
            };
            if !mods_found.insert(include_cpp.get_mod_name()) {
                return Err(ParseError::ConflictingModNames);
            }
            include_cpp
                .generate(
                    autocxx_inc.clone(),
                    extra_clang_args,
                    dep_recorder,
                    cpp_codegen_options,
                )
                .map_err(ParseError::AutocxxCodegenError)?
        }
        Ok(())
    }
}

impl ToTokens for ParsedFile {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        for seg in &self.0 {
            seg.to_tokens(tokens)
        }
    }
}

impl ToTokens for Segment {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        match self {
            Segment::Other(item) => item.to_tokens(tokens),
            Segment::Autocxx(autocxx) => {
                let these_tokens = autocxx.generate_rs();
                tokens.extend(these_tokens);
            }
            Segment::Cxx(cxxbridge) => cxxbridge.to_tokens(tokens),
            Segment::Mod(segments, (brace, itemmod)) => {
                let mut mod_items = Vec::with_capacity(segments.len());
                for segment in segments.iter() {
                    mod_items.push(syn::parse2::<Item>(segment.to_token_stream()).unwrap());
                }
                let itemmod = ItemMod {
                    content: Some((*brace, mod_items)),
                    ..itemmod.clone()
                };
                Item::Mod(itemmod).to_tokens(tokens)
            }
        }
    }
}

/// Shenanigans required to share the same RebuildDependencyRecorder
/// with all of the include_cpp instances in this one file.
#[derive(Debug, Clone)]
struct CompositeDepRecorder(Rc<dyn RebuildDependencyRecorder>);

impl CompositeDepRecorder {
    fn new(inner: Rc<dyn RebuildDependencyRecorder>) -> Self {
        CompositeDepRecorder(inner)
    }
}

impl UnwindSafe for CompositeDepRecorder {}

impl RebuildDependencyRecorder for CompositeDepRecorder {
    fn record_header_file_dependency(&self, filename: &str) {
        self.0.record_header_file_dependency(filename);
    }
}
<commit_msg>Code review feedback<commit_after>// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::ast_discoverer::Discoveries;
use crate::CppCodegenOptions;
use crate::{
    cxxbridge::CxxBridge, Error as EngineError, GeneratedCpp, IncludeCppEngine,
    RebuildDependencyRecorder,
};
use autocxx_parser::directives::SUBCLASS;
use autocxx_parser::{RustPath, Subclass, SubclassAttrs};
use proc_macro2::{Span, TokenStream};
use quote::ToTokens;
use std::{collections::HashSet, fmt::Display, io::Read, path::PathBuf};
use std::{panic::UnwindSafe, path::Path, rc::Rc};
use syn::{token::Brace, Item, ItemMod, LitStr};

/// Errors which may occur when parsing a Rust source file to discover
/// and interpret include_cxx macros.
#[derive(Debug)]
pub enum ParseError {
    /// Unable to open the source file
    FileOpen(std::io::Error),
    /// The .rs file couldn't be read.
    FileRead(std::io::Error),
    /// The .rs file couldn't be parsed.
    Syntax(syn::Error),
    /// The include CPP macro could not be expanded into
    /// Rust bindings to C++, because of some problem during the conversion
    /// process. This could be anything from a C++ parsing error to some
    /// C++ feature that autocxx can't yet handle and isn't able to skip
    /// over. It could also cover errors in your syntax of the `include_cpp`
    /// macro or the directives inside.
    AutocxxCodegenError(EngineError),
    /// There are two or more `include_cpp` macros with the same
    /// mod name.
    ConflictingModNames,
    ZeroModsForDynamicDiscovery,
    MultipleModsForDynamicDiscovery,
    DiscoveredRustItemsWhenNotInAutoDiscover,
}

impl Display for ParseError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ParseError::FileOpen(err) => write!(f, "Unable to open file: {}", err)?,
            ParseError::FileRead(err) => write!(f, "Unable to read file: {}", err)?,
            ParseError::Syntax(err) => write!(f, "Syntax error parsing Rust file: {}", err)?,
            ParseError::AutocxxCodegenError(err) =>
                write!(f, "Unable to parse include_cpp! macro: {}", err)?,
            ParseError::ConflictingModNames =>
                write!(f, "There are two or more include_cpp! macros with the same output mod name. Use name!")?,
            ParseError::ZeroModsForDynamicDiscovery =>
                write!(f, "This file contains extra information to append to an include_cpp! but no such include_cpp! was found in this file.")?,
            ParseError::MultipleModsForDynamicDiscovery =>
                write!(f, "This file contains extra information to append to an include_cpp! but multiple such include_cpp! declarations were found in this file.")?,
            ParseError::DiscoveredRustItemsWhenNotInAutoDiscover =>
                write!(f, "This file contains extra information to append to an \"extern Rust\" but auto-discover was switched off.")?,
        }
        Ok(())
    }
}

/// Parse a Rust file, and spot any include_cpp macros within it.
pub fn parse_file<P1: AsRef<Path>>(
    rs_file: P1,
    auto_allowlist: bool,
) -> Result<ParsedFile, ParseError> {
    let mut source = String::new();
    let mut file = std::fs::File::open(rs_file).map_err(ParseError::FileOpen)?;
    file.read_to_string(&mut source)
        .map_err(ParseError::FileRead)?;
    proc_macro2::fallback::force();
    let source = syn::parse_file(&source).map_err(ParseError::Syntax)?;
    parse_file_contents(source, auto_allowlist)
}

fn parse_file_contents(source: syn::File, auto_allowlist: bool) -> Result<ParsedFile, ParseError> {
    #[derive(Default)]
    struct State {
        auto_allowlist: bool,
        results: Vec<Segment>,
        extra_superclasses: Vec<Subclass>,
        discoveries: Discoveries,
    }
    impl State {
        fn parse_item(&mut self, item: Item, mod_path: Option<RustPath>) -> Result<(), ParseError> {
            let result = match item {
                Item::Macro(mac)
                    if mac
                        .mac
                        .path
                        .segments
                        .last()
                        .map(|s| s.ident == "include_cpp")
                        .unwrap_or(false) =>
                {
                    Segment::Autocxx(
                        crate::IncludeCppEngine::new_from_syn(mac.mac.clone())
                            .map_err(ParseError::AutocxxCodegenError)?,
                    )
                }
                Item::Mod(itm)
                    if itm
                        .attrs
                        .iter()
                        .any(|attr| attr.path.to_token_stream().to_string() == "cxx :: bridge") =>
                {
                    Segment::Cxx(CxxBridge::from(itm))
                }
                Item::Mod(itm) => {
                    if let Some((brace, items)) = itm.content {
                        let mut mod_state = State {
                            auto_allowlist: self.auto_allowlist,
                            ..Default::default()
                        };
                        let mod_path = match &mod_path {
                            None => RustPath::new_from_ident(itm.ident.clone()),
                            Some(mod_path) => mod_path.append(itm.ident.clone()),
                        };
                        for item in items {
                            mod_state.parse_item(item, Some(mod_path.clone()))?
                        }
                        self.extra_superclasses.extend(mod_state.extra_superclasses);
                        self.discoveries.extend(mod_state.discoveries);
                        Segment::Mod(
                            mod_state.results,
                            (
                                brace,
                                ItemMod {
                                    content: None,
                                    ..itm
                                },
                            ),
                        )
                    } else {
                        Segment::Other(Item::Mod(itm))
                    }
                }
                Item::Struct(ref its) if self.auto_allowlist => {
                    let attrs = &its.attrs;
                    let is_superclass_attr = attrs.iter().find(|attr| {
                        attr.path
                            .segments
                            .last()
                            .map(|seg| seg.ident == "is_subclass" || seg.ident == SUBCLASS)
                            .unwrap_or(false)
                    });
                    if let Some(is_superclass_attr) = is_superclass_attr {
                        if !is_superclass_attr.tokens.is_empty() {
                            let subclass = its.ident.clone();
                            let args: SubclassAttrs = is_superclass_attr
                                .parse_args()
                                .map_err(ParseError::Syntax)?;
                            if let Some(superclass) = args.superclass {
                                self.extra_superclasses.push(Subclass {
                                    superclass,
                                    subclass,
                                })
                            }
                        }
                    }
                    self.discoveries.search_item(&item, mod_path);
                    Segment::Other(item)
                }
                _ => {
                    self.discoveries.search_item(&item, mod_path);
                    Segment::Other(item)
                }
            };
            self.results.push(result);
            Ok(())
        }
    }
    let mut state = State {
        auto_allowlist,
        ..Default::default()
    };
    for item in source.items {
        state.parse_item(item, None)?
    }
    let State {
        auto_allowlist,
        mut results,
        mut extra_superclasses,
        mut discoveries,
    } = state;
    if !auto_allowlist
        && (!discoveries.extern_rust_types.is_empty() || !discoveries.extern_rust_funs.is_empty())
    {
        return Err(ParseError::DiscoveredRustItemsWhenNotInAutoDiscover);
    }
    if !extra_superclasses.is_empty() || (auto_allowlist && !discoveries.is_empty()) {
        let mut autocxx_seg_iterator = results.iter_mut().filter_map(|seg| match seg {
            Segment::Autocxx(engine) => Some(engine),
            _ => None,
        });
        let our_seg = autocxx_seg_iterator.next();
        match our_seg {
            None => return Err(ParseError::ZeroModsForDynamicDiscovery),
            Some(engine) => {
                engine
                    .config_mut()
                    .subclasses
                    .append(&mut extra_superclasses);
                if auto_allowlist {
                    for cpp in discoveries.cpp_list {
                        engine
                            .config_mut()
                            .allowlist
                            .push(LitStr::new(&cpp, Span::call_site()))
                            .map_err(ParseError::Syntax)?;
                    }
                }
                engine
                    .config_mut()
                    .extern_rust_funs
                    .append(&mut discoveries.extern_rust_funs);
                engine
                    .config_mut()
                    .rust_types
                    .append(&mut discoveries.extern_rust_types);
            }
        }
        if autocxx_seg_iterator.next().is_some() {
            return Err(ParseError::MultipleModsForDynamicDiscovery);
        }
    }
    let autocxx_seg_iterator = results.iter_mut().filter_map(|seg| match seg {
        Segment::Autocxx(engine) => Some(engine),
        _ => None,
    });
    for seg in autocxx_seg_iterator {
        seg.config
            .confirm_complete(auto_allowlist)
            .map_err(ParseError::Syntax)?;
    }
    Ok(ParsedFile(results))
}

/// A Rust file parsed by autocxx. May contain zero or more autocxx 'engines',
/// i.e. the `IncludeCpp` class, corresponding to zero or more include_cpp
/// macros within this file. Also contains `syn::Item` structures for all
/// the rest of the Rust code, such that it can be reconstituted if necessary.
pub struct ParsedFile(Vec<Segment>);

#[allow(clippy::large_enum_variant)]
enum Segment {
    Autocxx(IncludeCppEngine),
    Cxx(CxxBridge),
    Mod(Vec<Segment>, (Brace, ItemMod)),
    Other(Item),
}

pub trait CppBuildable {
    fn generate_h_and_cxx(
        &self,
        cpp_codegen_options: &CppCodegenOptions,
    ) -> Result<GeneratedCpp, cxx_gen::Error>;
}

impl ParsedFile {
    /// Get all the autocxxes in this parsed file.
    pub fn get_rs_buildables(&self) -> impl Iterator<Item = &IncludeCppEngine> {
        fn do_get_rs_buildables(
            segments: &Vec<Segment>,
        ) -> impl Iterator<Item = &IncludeCppEngine> {
            segments
                .iter()
                .flat_map(|s| -> Box<dyn Iterator<Item = &IncludeCppEngine>> {
                    match s {
                        Segment::Autocxx(includecpp) => Box::new(std::iter::once(includecpp)),
                        Segment::Mod(segments, _) => Box::new(do_get_rs_buildables(segments)),
                        _ => Box::new(std::iter::empty()),
                    }
                })
        }

        do_get_rs_buildables(&self.0)
    }

    /// Get all items which can result in C++ code
    pub fn get_cpp_buildables(&self) -> impl Iterator<Item = &dyn CppBuildable> {
        fn do_get_cpp_buildables(
            segments: &Vec<Segment>,
        ) -> impl Iterator<Item = &dyn CppBuildable> {
            segments
                .iter()
                .flat_map(|s| -> Box<dyn Iterator<Item = &dyn CppBuildable>> {
                    match s {
                        Segment::Autocxx(includecpp) => {
                            Box::new(std::iter::once(includecpp as &dyn CppBuildable))
                        }
                        Segment::Cxx(cxxbridge) => {
                            Box::new(std::iter::once(cxxbridge as &dyn CppBuildable))
                        }
                        Segment::Mod(segments, _) => Box::new(do_get_cpp_buildables(segments)),
                        _ => Box::new(std::iter::empty()),
                    }
                })
        }

        do_get_cpp_buildables(&self.0)
    }

    fn get_autocxxes_mut(&mut self) -> impl Iterator<Item = &mut IncludeCppEngine> {
        fn do_get_autocxxes_mut(
            segments: &mut Vec<Segment>,
        ) -> impl Iterator<Item = &mut IncludeCppEngine> {
            segments
                .iter_mut()
                .flat_map(|s| -> Box<dyn Iterator<Item = &mut IncludeCppEngine>> {
                    match s {
                        Segment::Autocxx(includecpp) => Box::new(std::iter::once(includecpp)),
                        Segment::Mod(segments, _) => Box::new(do_get_autocxxes_mut(segments)),
                        _ => Box::new(std::iter::empty()),
                    }
                })
        }

        do_get_autocxxes_mut(&mut self.0)
    }

    pub fn include_dirs(&self) -> impl Iterator<Item = &PathBuf> {
        fn do_get_include_dirs(segments: &Vec<Segment>) -> impl Iterator<Item = &PathBuf> {
            segments
                .iter()
                .flat_map(|s| -> Box<dyn Iterator<Item = &PathBuf>> {
                    match s {
                        Segment::Autocxx(includecpp) => Box::new(includecpp.include_dirs()),
                        Segment::Mod(segments, _) => Box::new(do_get_include_dirs(segments)),
                        _ => Box::new(std::iter::empty()),
                    }
                })
        }

        do_get_include_dirs(&self.0)
    }

    pub fn resolve_all(
        &mut self,
        autocxx_inc: Vec<PathBuf>,
        extra_clang_args: &[&str],
        dep_recorder: Option<Box<dyn RebuildDependencyRecorder>>,
        cpp_codegen_options: &CppCodegenOptions,
    ) -> Result<(), ParseError> {
        let mut mods_found = HashSet::new();
        let inner_dep_recorder: Option<Rc<dyn RebuildDependencyRecorder>> =
            dep_recorder.map(Rc::from);
        for include_cpp in self.get_autocxxes_mut() {
            #[allow(clippy::manual_map)] // because of dyn shenanigans
            let dep_recorder: Option<Box<dyn RebuildDependencyRecorder>> = match &inner_dep_recorder
            {
                None => None,
                Some(inner_dep_recorder) => Some(Box::new(CompositeDepRecorder::new(
                    inner_dep_recorder.clone(),
                ))),
            };
            if !mods_found.insert(include_cpp.get_mod_name()) {
                return Err(ParseError::ConflictingModNames);
            }
            include_cpp
                .generate(
                    autocxx_inc.clone(),
                    extra_clang_args,
                    dep_recorder,
                    cpp_codegen_options,
                )
                .map_err(ParseError::AutocxxCodegenError)?
        }
        Ok(())
    }
}

impl ToTokens for ParsedFile {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        for seg in &self.0 {
            seg.to_tokens(tokens)
        }
    }
}

impl ToTokens for Segment {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        match self {
            Segment::Other(item) => item.to_tokens(tokens),
            Segment::Autocxx(autocxx) => {
                let these_tokens = autocxx.generate_rs();
                tokens.extend(these_tokens);
            }
            Segment::Cxx(cxxbridge) => cxxbridge.to_tokens(tokens),
            Segment::Mod(segments, (brace, itemmod)) => {
                let mod_items = segments
                    .iter()
                    .map(|segment| syn::parse2::<Item>(segment.to_token_stream()).unwrap())
                    .collect();
                let itemmod = ItemMod {
                    content: Some((*brace, mod_items)),
                    ..itemmod.clone()
                };
                Item::Mod(itemmod).to_tokens(tokens)
            }
        }
    }
}

/// Shenanigans required to share the same RebuildDependencyRecorder
/// with all of the include_cpp instances in this one file.
#[derive(Debug, Clone)]
struct CompositeDepRecorder(Rc<dyn RebuildDependencyRecorder>);

impl CompositeDepRecorder {
    fn new(inner: Rc<dyn RebuildDependencyRecorder>) -> Self {
        CompositeDepRecorder(inner)
    }
}

impl UnwindSafe for CompositeDepRecorder {}

impl RebuildDependencyRecorder for CompositeDepRecorder {
    fn record_header_file_dependency(&self, filename: &str) {
        self.0.record_header_file_dependency(filename);
    }
}
<|endoftext|>


<commit_before>use dispatch_table::{DispatchTable, DispatchFn};
use spin::RwLock;
use isr;

#[derive(Copy, Clone, Default)]
#[repr(packed)]
struct IDTEntry {
    base_low: u16,
    selector: u16,
    always0: u8,
    flags: u8,
    base_high: u16,
    base_highest: u32,
    _reserved: u32,
}

pub struct IDT([IDTEntry; 256]);

lazy_static! {
    static ref INTERRUPT_TABLE: RwLock<DispatchTable<u64, InterruptCPUState>> =
        RwLock::new(DispatchTable::new(256));
    static ref DESCRIPTOR_TABLE: RwLock<IDT> = RwLock::new(IDT([Default::default(); 256]));
}

impl IDT {
    fn set_entry(&mut self, num: usize, base: u64, selector: u16, flags: u8) {
        self.0[num].base_low = (base & 0xffff) as u16;
        self.0[num].base_high = ((base >> 16) & 0xffff) as u16;
        self.0[num].base_highest = (base >> 32) as u32;
        self.0[num].always0 = 0;
        self.0[num].selector = selector;
        self.0[num].flags = flags;
    }
}

#[repr(C)]
pub struct InterruptCPUState {
    ds: u64,
    r15: u64,
    r14: u64,
    r13: u64,
    r12: u64,
    r11: u64,
    r10: u64,
    r9: u64,
    r8: u64,
    rdi: u64,
    rsi: u64,
    rbp: u64,
    rdx: u64,
    rcx: u64,
    rbx: u64,
    rax: u64,
    interrupt_number: u64,
    error_code: u64,
    rip: u64,
    cs: u64,
    rflags: u64,
    rsp: u64,
    ss: u64,
}

pub fn register_interrupt_handler(interrupt: u64,
                                  stub: isr::InterruptServiceRoutine,
                                  handler: DispatchFn<u64, InterruptCPUState>) {
    INTERRUPT_TABLE.write().register(interrupt, handler);
    DESCRIPTOR_TABLE.write().set_entry(interrupt as usize, stub as u64, 0x08, 0x8e);
}

pub fn unregister_interrupt_handler(interrupt: u64) {
    INTERRUPT_TABLE.write().unregister(interrupt);
}

#[no_mangle]
pub extern "C" fn dispatch_interrupt(state: &mut InterruptCPUState) {
    INTERRUPT_TABLE.read().dispatch(state.interrupt_number, state);
}

fn default_handler(_interrupt_number: u64, _regs: &mut InterruptCPUState) -> bool {
    true
}

pub fn init_interrupt_handlers() {
    for i in 0..20 {
        register_interrupt_handler(i, isr::ISR[i as usize], default_handler);
    }
}

#[cfg(feature = "runtime_tests")]
pub mod runtime_tests {

    use interrupts;
    use vmx;
    use cli;
    use isr;
    use core::mem;

    extern "C" {
        fn _test_division_by_zero_routine() -> bool;

        // DO NOT CALL THIS FUNCTION. YOU WILL MESS UP THE STACK.
        fn _after_division() -> bool;
    }

    fn new_host_idt_descriptor() -> vmx::CPUTableDescriptor {
        let base = (*interrupts::DESCRIPTOR_TABLE.write()).0.as_ptr() as u64;
        vmx::CPUTableDescriptor {
            base: base,
            limit: (mem::size_of::<interrupts::IDT>() - 1) as u16,
        }
    }

    pub fn run() {
        test_load_and_restore_idt();
        test_divide_by_zero_interrupt();
    }

    fn test_load_and_restore_idt() {
        cli::ClearLocalInterruptsGuard::new();
        let idt_desc = new_host_idt_descriptor();
        let mut orig_idt_desc: vmx::CPUTableDescriptor = Default::default();
        vmx::sidt(&mut orig_idt_desc);
        vmx::lidt(&idt_desc);
        vmx::lidt(&orig_idt_desc);
    }

    fn division_by_zero_handler(interrupt_number: u64,
                                regs: &mut interrupts::InterruptCPUState)
                                -> bool {
        assert_eq!(interrupt_number, 0);
        info!("Handling division by zero interrupt.");
        regs.rip = _after_division as u64;
        true
    }

    #[allow(unused_variables)]
    fn test_divide_by_zero_interrupt() {
        let mut orig_idt_desc: vmx::CPUTableDescriptor = Default::default();
        interrupts::register_interrupt_handler(0, isr::_isr0, division_by_zero_handler);

        {
            cli::ClearLocalInterruptsGuard::new();
            let idt_desc = new_host_idt_descriptor();
            vmx::sidt(&mut orig_idt_desc);
            vmx::lidt(&idt_desc);
        }

        info!("Here we go!");

        unsafe {
            assert!(_test_division_by_zero_routine());
        }

        info!("Successfully returned from a division by zero interrupt.");

        {
            cli::ClearLocalInterruptsGuard::new();
            vmx::lidt(&orig_idt_desc);
        }

    }
}
<commit_msg>Remove division by zero tests<commit_after>use dispatch_table::{DispatchTable, DispatchFn};
use spin::RwLock;
use isr;

#[derive(Copy, Clone, Default)]
#[repr(packed)]
struct IDTEntry {
    base_low: u16,
    selector: u16,
    always0: u8,
    flags: u8,
    base_high: u16,
    base_highest: u32,
    _reserved: u32,
}

pub struct IDT([IDTEntry; 256]);

lazy_static! {
    static ref INTERRUPT_TABLE: RwLock<DispatchTable<u64, InterruptCPUState>> =
        RwLock::new(DispatchTable::new(256));
    static ref DESCRIPTOR_TABLE: RwLock<IDT> = RwLock::new(IDT([Default::default(); 256]));
}

impl IDT {
    fn set_entry(&mut self, num: usize, base: u64, selector: u16, flags: u8) {
        self.0[num].base_low = (base & 0xffff) as u16;
        self.0[num].base_high = ((base >> 16) & 0xffff) as u16;
        self.0[num].base_highest = (base >> 32) as u32;
        self.0[num].always0 = 0;
        self.0[num].selector = selector;
        self.0[num].flags = flags;
    }
}

#[repr(C)]
pub struct InterruptCPUState {
    ds: u64,
    r15: u64,
    r14: u64,
    r13: u64,
    r12: u64,
    r11: u64,
    r10: u64,
    r9: u64,
    r8: u64,
    rdi: u64,
    rsi: u64,
    rbp: u64,
    rdx: u64,
    rcx: u64,
    rbx: u64,
    rax: u64,
    interrupt_number: u64,
    error_code: u64,
    rip: u64,
    cs: u64,
    rflags: u64,
    rsp: u64,
    ss: u64,
}

pub fn register_interrupt_handler(interrupt: u64,
                                  stub: isr::InterruptServiceRoutine,
                                  handler: DispatchFn<u64, InterruptCPUState>) {
    INTERRUPT_TABLE.write().register(interrupt, handler);
    DESCRIPTOR_TABLE.write().set_entry(interrupt as usize, stub as u64, 0x08, 0x8e);
}

pub fn unregister_interrupt_handler(interrupt: u64) {
    INTERRUPT_TABLE.write().unregister(interrupt);
}

#[no_mangle]
pub extern "C" fn dispatch_interrupt(state: &mut InterruptCPUState) {
    INTERRUPT_TABLE.read().dispatch(state.interrupt_number, state);
}

fn default_handler(_interrupt_number: u64, _regs: &mut InterruptCPUState) -> bool {
    true
}

pub fn init_interrupt_handlers() {
    for i in 0..20 {
        register_interrupt_handler(i, isr::ISR[i as usize], default_handler);
    }
}

#[cfg(feature = "runtime_tests")]
pub mod runtime_tests {

    use interrupts;
    use vmx;
    use cli;

    fn new_host_idt_descriptor() -> vmx::CPUTableDescriptor {
        let base = (*interrupts::DESCRIPTOR_TABLE.write()).0.as_ptr() as u64;
        vmx::CPUTableDescriptor {
            base: base,
            limit: (mem::size_of::<interrupts::IDT>() - 1) as u16,
        }
    }

    pub fn run() {
        test_load_and_restore_idt();
    }

    fn test_load_and_restore_idt() {
        cli::ClearLocalInterruptsGuard::new();
        let idt_desc = new_host_idt_descriptor();
        let mut orig_idt_desc: vmx::CPUTableDescriptor = Default::default();
        vmx::sidt(&mut orig_idt_desc);
        vmx::lidt(&idt_desc);
        vmx::lidt(&orig_idt_desc);
    }
}
<|endoftext|>


<commit_before>// Zinc, the bare metal stack for rust.
// Copyright 2014 Ben Gamari <bgamari@gmail.com>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Tests for ioreg! syntax extension

#![feature(plugin)]

#[plugin] extern crate ioreg;
#[plugin] extern crate shiny;
extern crate core;

#[path="../zinc/util/volatile_cell.rs"] mod volatile_cell;

#[cfg(test)]
mod test {
  use std::mem::{transmute, zeroed};
  use std::ptr::PtrExt;
  use volatile_cell::VolatileCell;

  fn get_value<'a, T>(v: &'a T, offset: uint) -> u32 {
    unsafe {
      let ptr: *const u32 = transmute(v);
      *(ptr.offset(offset as int))
    }
  }

  fn zeroed_safe<T: Copy>() -> T {
    unsafe {
      return zeroed();
    }
  }

  ioregs!(BASIC_TEST = {
    0x0 => reg32 reg1 {
      0      => field1,
      1..3   => field2,
      16..24 => field3,
      25     => field4: set_to_clear,
    }
    0x4 => reg32 reg2 {
      0      => field1,
    }
    0x8 => reg32 wo_reg {
      0..15  => field1: wo,
      16..31 => field2: wo,
    }
  });

  describe!(
    before_each {
      let test: BASIC_TEST = zeroed_safe();
    }

    it "can round_trip simple field values 1" {
      test.reg1.set_field1(true);
      assert_eq!(test.reg1.field1(), true);
      assert_eq!(get_value(&test, 0), 1);
      assert_eq!(get_value(&test, 1), 0);
    }

    it "can round trip simple field values 2" {
      test.reg1.set_field3(0xde);
      assert_eq!(test.reg1.field3(), 0xde);
      assert_eq!(get_value(&test, 0), 0xde<<16);
    }

    it "sets set_to_clear fields" {
      test.reg1.clear_field4();
      assert_eq!(get_value(&test, 0), 1<<25);
    }

    it "does not read from writeonly registers" {
      test.wo_reg.set_field1(0xdead);
      assert_eq!(get_value(&test, 2), 0xdead);
      test.wo_reg.set_field2(0xdead);
      assert_eq!(get_value(&test, 2), 0xdead<<16);
    }
  );

  ioregs!(GROUP_TEST = {
    0x0 => group regs[5] {
      0x0 => reg32 reg1 {
        0..31 => field1
      }
      0x4 => reg32 reg2 {
        0..31 => field2
      }
    }
  });

  describe!(
    before_each {
      let test: GROUP_TEST = zeroed_safe();
    }

    it "sets groups correctly" {
      test.regs[0].reg1.set_field1(0xdeadbeef);
      assert_eq!(test.regs[0].reg1.field1(), 0xdeadbeef);
      assert_eq!(get_value(&test, 0), 0xdeadbeef);
      for i in range(1, 10) {
        assert_eq!(get_value(&test, i), 0);
      }

      test.regs[2].reg2.set_field2(0xfeedbeef);
      assert_eq!(test.regs[2].reg2.field2(), 0xfeedbeef);
      assert_eq!(get_value(&test, 5), 0xfeedbeef);
    }
  );

  ioregs!(FIELD_ARRAY_TEST = {
    0x0 => reg32 reg1 {
      0..31 => field[16]
    }
  });

  describe!(
    before_each {
      let test: FIELD_ARRAY_TEST = zeroed_safe();
    }

    it "sets field arrays correctly" {
      test.reg1.set_field(0, 1);
      assert_eq!(test.reg1.field(0), 1);
      assert_eq!(get_value(&test, 0), 0x1);

      test.reg1.set_field(4, 3);
      assert_eq!(test.reg1.field(4), 3);
      assert_eq!(get_value(&test, 0), 0x1 | 0x3<<8);
    }
  );

  ioregs!(GAP_TEST = {
    0x0 => reg32 reg1 {
      0..31 => field,
    }
    0x10 => reg32 reg2 {
      0..31 => field,
    }
    0x14 => reg32 reg3 {
      0..31 => field,
    }
    0x20 => reg32 reg4 {
      0..31 => field,
    }
  });

  describe!(
    before_each {
      let test: GAP_TEST = zeroed_safe();
      let base = &test as *const GAP_TEST;
    }
    it "has zero base offset" {
      let addr = &test.reg1 as *const GAP_TEST_reg1;
      assert_eq!(addr as uint - base as uint, 0x0);
    }
    it "computes the correct first gap" {
      let addr = &test.reg2 as *const GAP_TEST_reg2;
      assert_eq!(addr as uint - base as uint, 0x10);
    }
    it "computes the correct second gap" {
      let addr = &test.reg4 as *const GAP_TEST_reg4;
      assert_eq!(addr as uint - base as uint, 0x20);
    }
  );
}
<commit_msg>ioreg::test: Fix uint and int references<commit_after>// Zinc, the bare metal stack for rust.
// Copyright 2014 Ben Gamari <bgamari@gmail.com>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Tests for ioreg! syntax extension

#![feature(plugin)]

#[plugin] extern crate ioreg;
#[plugin] extern crate shiny;
extern crate core;

#[path="../zinc/util/volatile_cell.rs"] mod volatile_cell;

#[cfg(test)]
mod test {
  use std::mem::{transmute, zeroed};
  use std::ptr::PtrExt;
  use volatile_cell::VolatileCell;

  fn get_value<'a, T>(v: &'a T, offset: usize) -> u32 {
    unsafe {
      let ptr: *const u32 = transmute(v);
      *(ptr.offset(offset as isize))
    }
  }

  fn zeroed_safe<T: Copy>() -> T {
    unsafe {
      return zeroed();
    }
  }

  ioregs!(BASIC_TEST = {
    0x0 => reg32 reg1 {
      0      => field1,
      1..3   => field2,
      16..24 => field3,
      25     => field4: set_to_clear,
    }
    0x4 => reg32 reg2 {
      0      => field1,
    }
    0x8 => reg32 wo_reg {
      0..15  => field1: wo,
      16..31 => field2: wo,
    }
  });

  describe!(
    before_each {
      let test: BASIC_TEST = zeroed_safe();
    }

    it "can round_trip simple field values 1" {
      test.reg1.set_field1(true);
      assert_eq!(test.reg1.field1(), true);
      assert_eq!(get_value(&test, 0), 1);
      assert_eq!(get_value(&test, 1), 0);
    }

    it "can round trip simple field values 2" {
      test.reg1.set_field3(0xde);
      assert_eq!(test.reg1.field3(), 0xde);
      assert_eq!(get_value(&test, 0), 0xde<<16);
    }

    it "sets set_to_clear fields" {
      test.reg1.clear_field4();
      assert_eq!(get_value(&test, 0), 1<<25);
    }

    it "does not read from writeonly registers" {
      test.wo_reg.set_field1(0xdead);
      assert_eq!(get_value(&test, 2), 0xdead);
      test.wo_reg.set_field2(0xdead);
      assert_eq!(get_value(&test, 2), 0xdead<<16);
    }
  );

  ioregs!(GROUP_TEST = {
    0x0 => group regs[5] {
      0x0 => reg32 reg1 {
        0..31 => field1
      }
      0x4 => reg32 reg2 {
        0..31 => field2
      }
    }
  });

  describe!(
    before_each {
      let test: GROUP_TEST = zeroed_safe();
    }

    it "sets groups correctly" {
      test.regs[0].reg1.set_field1(0xdeadbeef);
      assert_eq!(test.regs[0].reg1.field1(), 0xdeadbeef);
      assert_eq!(get_value(&test, 0), 0xdeadbeef);
      for i in range(1, 10) {
        assert_eq!(get_value(&test, i), 0);
      }

      test.regs[2].reg2.set_field2(0xfeedbeef);
      assert_eq!(test.regs[2].reg2.field2(), 0xfeedbeef);
      assert_eq!(get_value(&test, 5), 0xfeedbeef);
    }
  );

  ioregs!(FIELD_ARRAY_TEST = {
    0x0 => reg32 reg1 {
      0..31 => field[16]
    }
  });

  describe!(
    before_each {
      let test: FIELD_ARRAY_TEST = zeroed_safe();
    }

    it "sets field arrays correctly" {
      test.reg1.set_field(0, 1);
      assert_eq!(test.reg1.field(0), 1);
      assert_eq!(get_value(&test, 0), 0x1);

      test.reg1.set_field(4, 3);
      assert_eq!(test.reg1.field(4), 3);
      assert_eq!(get_value(&test, 0), 0x1 | 0x3<<8);
    }
  );

  ioregs!(GAP_TEST = {
    0x0 => reg32 reg1 {
      0..31 => field,
    }
    0x10 => reg32 reg2 {
      0..31 => field,
    }
    0x14 => reg32 reg3 {
      0..31 => field,
    }
    0x20 => reg32 reg4 {
      0..31 => field,
    }
  });

  describe!(
    before_each {
      let test: GAP_TEST = zeroed_safe();
      let base = &test as *const GAP_TEST;
    }
    it "has zero base offset" {
      let addr = &test.reg1 as *const GAP_TEST_reg1;
      assert_eq!(addr as usize - base as usize, 0x0);
    }
    it "computes the correct first gap" {
      let addr = &test.reg2 as *const GAP_TEST_reg2;
      assert_eq!(addr as usize - base as usize, 0x10);
    }
    it "computes the correct second gap" {
      let addr = &test.reg4 as *const GAP_TEST_reg4;
      assert_eq!(addr as usize - base as usize, 0x20);
    }
  );
}
<|endoftext|>


<commit_before>extern crate rustbox;

pub struct UIBuffer {
    width: uint,
    height: uint,
    rows: Vec<Vec<Cell>>
}

impl UIBuffer {
    pub fn new(width: uint, height: uint) -> UIBuffer {
        let rows = Cell::create_grid(width, height, ' ');

        UIBuffer {
            width: width,
            height: height,
            rows: rows,
        }
    }

    pub fn draw_range(&self, start: uint, stop: uint) {
        let rows = self.rows.slice(start, stop);
        for row in rows.iter() {
            for cell in row.iter() {
                rustbox::print_char(cell.x, cell.y, rustbox::Style::Normal, cell.fg, cell.bg, cell.ch);
            }
        }
    }

    pub fn get_height(&self) -> uint {
        self.height
    }

    /// Recreated the entire grid, will cells containing `ch`.
    pub fn fill(&mut self, ch: char) {
        self.rows = Cell::create_grid(self.width, self.height, ch);
    }

    /// Update the `ch` attribute of an individual cell
    pub fn update_cell_content(&mut self, cell_num: uint, row_num: uint, ch: char) {
        self.rows[row_num][cell_num].ch = ch
    }

    /// Update the `ch`, `fg`, and `bg` attributes of an indivudual cell
    pub fn update_cell(&mut self, cell_num: uint, row_num: uint, ch: char, fg: rustbox::Color, bg: rustbox::Color) {
        let cell = &mut self.rows[row_num][cell_num];
        cell.ch = ch;
        cell.fg = fg;
        cell.bg = bg;
    }
}

pub struct Cell {
    pub bg: rustbox::Color,
    pub fg: rustbox::Color,
    pub ch: char,
    pub x: uint,
    pub y: uint,
}

impl Cell {
    pub fn new() -> Cell {
        Cell {
            bg: rustbox::Color::Default,
            fg: rustbox::Color::White,
            ch: ' ',
            x: 0,
            y: 0,
        }
    }

    pub fn create_grid(width: uint, height: uint, ch: char) -> Vec<Vec<Cell>> {
        let mut rows = Vec::new();
        for voffset in range(0, height) {
            let mut cells = Vec::new();
            for boffset in range(0, width) {
                let mut cell = Cell::new();
                cell.x = boffset;
                cell.y = voffset;
                cell.ch = ch;
                cells.push(cell);
            }
            rows.push(cells);
        }

        return rows
    }
}

#[cfg(test)]
mod tests {

    extern crate rustbox;

    use uibuf::UIBuffer;

    fn setup_uibuf() -> UIBuffer {
        UIBuffer::new(50, 50)
    }

    #[test]
    fn grid_is_created_correctly() {
        let uibuf = setup_uibuf();
        assert_eq!(uibuf.height, 50);
        assert_eq!(uibuf.width, 50);
        assert_eq!(uibuf.rows.len(), 50);
        assert_eq!(uibuf.rows[0][5].ch, ' ');
    }

    #[test]
    fn fill_updates_contents_of_all_cells() {
        let mut uibuf = setup_uibuf();
        uibuf.fill('x');

        // check some random cells
        assert_eq!(uibuf.rows[20][5].ch, 'x');
        assert_eq!(uibuf.rows[0][30].ch, 'x');
    }

    #[test]
    fn update_cell_content_updates_a_single_cell() {
        let mut uibuf = setup_uibuf();
        // update cell 10 in row 32
        uibuf.update_cell_content(10, 32, 'y');

        assert_eq!(uibuf.rows[32][10].ch, 'y');
    }

    #[test]
    fn update_cell_updates_all_attrs_of_cell() {
        let mut uibuf = setup_uibuf();
        let cell_num = 10u;
        let row_num = 0u;
        let ch = 'q';
        let fg = rustbox::Color::Red;
        let bg = rustbox::Color::Blue;

        uibuf.update_cell(cell_num, row_num, ch, fg, bg);

        let cell = &uibuf.rows[row_num][cell_num];
        assert_eq!(cell.ch, ch);

        // FIXME(greg): this fails cos '==' is not implemented for rustbox::Color.
        // assert_eq!(cell.fg, fg);
        // assert_eq!(cell.bg, bg);
    }

}
<commit_msg>Added UIBuffer::get_cell_mut method<commit_after>extern crate rustbox;

pub struct UIBuffer {
    width: uint,
    height: uint,
    rows: Vec<Vec<Cell>>
}

impl UIBuffer {
    pub fn new(width: uint, height: uint) -> UIBuffer {
        let rows = Cell::create_grid(width, height, ' ');

        UIBuffer {
            width: width,
            height: height,
            rows: rows,
        }
    }

    pub fn draw_range(&self, start: uint, stop: uint) {
        let rows = self.rows.slice(start, stop);
        for row in rows.iter() {
            for cell in row.iter() {
                rustbox::print_char(cell.x, cell.y, rustbox::Style::Normal, cell.fg, cell.bg, cell.ch);
            }
        }
    }

    pub fn get_height(&self) -> uint {
        self.height
    }

    /// Recreated the entire grid, will cells containing `ch`.
    pub fn fill(&mut self, ch: char) {
        self.rows = Cell::create_grid(self.width, self.height, ch);
    }

    /// Update the `ch` attribute of an individual cell
    pub fn update_cell_content(&mut self, cell_num: uint, row_num: uint, ch: char) {
        self.rows[row_num][cell_num].ch = ch
    }

    /// Update the `ch`, `fg`, and `bg` attributes of an indivudual cell
    pub fn update_cell(&mut self, cell_num: uint, row_num: uint, ch: char, fg: rustbox::Color, bg: rustbox::Color) {
        let cell = self.get_cell_mut(cell_num, row_num);
        cell.ch = ch;
        cell.fg = fg;
        cell.bg = bg;
    }

    pub fn get_cell_mut(&mut self, cell_num: uint, row_num: uint) -> &mut Cell {
        &mut self.rows[row_num][cell_num]
    }
}

pub struct Cell {
    pub bg: rustbox::Color,
    pub fg: rustbox::Color,
    pub ch: char,
    pub x: uint,
    pub y: uint,
}

impl Cell {
    pub fn new() -> Cell {
        Cell {
            bg: rustbox::Color::Default,
            fg: rustbox::Color::White,
            ch: ' ',
            x: 0,
            y: 0,
        }
    }

    pub fn create_grid(width: uint, height: uint, ch: char) -> Vec<Vec<Cell>> {
        let mut rows = Vec::new();
        for voffset in range(0, height) {
            let mut cells = Vec::new();
            for boffset in range(0, width) {
                let mut cell = Cell::new();
                cell.x = boffset;
                cell.y = voffset;
                cell.ch = ch;
                cells.push(cell);
            }
            rows.push(cells);
        }

        return rows
    }
}

#[cfg(test)]
mod tests {

    extern crate rustbox;

    use uibuf::UIBuffer;

    fn setup_uibuf() -> UIBuffer {
        UIBuffer::new(50, 50)
    }

    #[test]
    fn grid_is_created_correctly() {
        let uibuf = setup_uibuf();
        assert_eq!(uibuf.height, 50);
        assert_eq!(uibuf.width, 50);
        assert_eq!(uibuf.rows.len(), 50);
        assert_eq!(uibuf.rows[0][5].ch, ' ');
    }

    #[test]
    fn fill_updates_contents_of_all_cells() {
        let mut uibuf = setup_uibuf();
        uibuf.fill('x');

        // check some random cells
        assert_eq!(uibuf.rows[20][5].ch, 'x');
        assert_eq!(uibuf.rows[0][30].ch, 'x');
    }

    #[test]
    fn update_cell_content_updates_a_single_cell() {
        let mut uibuf = setup_uibuf();
        // update cell 10 in row 32
        uibuf.update_cell_content(10, 32, 'y');

        assert_eq!(uibuf.rows[32][10].ch, 'y');
    }

    #[test]
    fn update_cell_updates_all_attrs_of_cell() {
        let mut uibuf = setup_uibuf();
        let cell_num = 10u;
        let row_num = 0u;
        let ch = 'q';
        let fg = rustbox::Color::Red;
        let bg = rustbox::Color::Blue;

        uibuf.update_cell(cell_num, row_num, ch, fg, bg);

        let cell = &uibuf.rows[row_num][cell_num];
        assert_eq!(cell.ch, ch);

        // FIXME(greg): this fails cos '==' is not implemented for rustbox::Color.
        // assert_eq!(cell.fg, fg);
        // assert_eq!(cell.bg, bg);
    }

}
<|endoftext|>


<commit_before>//!Pack consecutive duplicates of list elements into sublists. (medium)
///```
///    use p99::p9::pack;
///assert_eq!(pack(&['a','a','a','a','b','c','c','a','a','d','d','e','e','e','e']),vec![vec!['a', 'a', 'a', 'a'], vec!['b'], vec!['c', 'c'], vec!['a', 'a'], vec!['d', 'd'], vec!['e', 'e', 'e', 'e']])
///```

pub fn pack(list: &[char]) ->Vec<Vec<char>> {
    let list = list.to_vec();
    let mut result  = vec![];
    let mut last = list[0];
    let mut sub_result = vec![last];
    for c in &list[1..] {
        if *c == last {
            sub_result.push(last);
        }
        else {
            last = *c;
            result.push(sub_result.clone());
            sub_result = vec![last];
        }
    }
    result.push(sub_result.clone());
    result
}

<commit_msg>not clone<commit_after>//!Pack consecutive duplicates of list elements into sublists. (medium)
///```
///    use p99::p9::pack;
///assert_eq!(pack(&['a','a','a','a','b','c','c','a','a','d','d','e','e','e','e']),vec![vec!['a', 'a', 'a', 'a'], vec!['b'], vec!['c', 'c'], vec!['a', 'a'], vec!['d', 'd'], vec!['e', 'e', 'e', 'e']])
///```

pub fn pack(list: &[char]) ->Vec<Vec<char>> {
    let list = list.to_vec();
    let mut result  = vec![];
    let mut last = list[0];
    let mut sub_result = vec![last];
    for c in &list[1..] {
        if *c == last {
            sub_result.push(last);
        }
        else {
            last = *c;
            result.push(sub_result);
            sub_result = vec![last];
        }
    }
    result.push(sub_result);
    result
}

<|endoftext|>


<commit_before>#![deny(missing_docs)]
#![warn(rust_2018_idioms)]
#![doc(html_root_url = "https://docs.rs/amq-protocol-tcp/7.0.0-alpha.1/")]

//! # AMQP URI TCP connection handling
//!
//! amq-protocol-tcp is a library aiming at providing tools to help
//! connecting to an AMQP URI

use amq_protocol_uri::{AMQPScheme, AMQPUri};
use std::time::Duration;
use tracing::trace;

/// Re-export TcpStream
pub use tcp_stream::{
    HandshakeError, HandshakeResult, Identity, MidHandshakeTlsStream, OwnedIdentity,
    OwnedTLSConfig, TLSConfig, TcpStream,
};

#[cfg(feature = "native-tls")]
pub use tcp_stream::NativeTlsConnector;

#[cfg(feature = "openssl")]
pub use tcp_stream::OpenSslConnector;

#[cfg(feature = "rustls-connector")]
pub use tcp_stream::{RustlsConnector, RustlsConnectorConfig};

/// Trait providing a method to connect to a TcpStream
pub trait AMQPUriTcpExt {
    /// connect to a TcpStream
    fn connect(&self) -> HandshakeResult
    where
        Self: Sized,
    {
        self.connect_with_config(TLSConfig::default())
    }

    /// connect to a TcpStream with the given configuration
    fn connect_with_config(&self, config: TLSConfig<'_, '_, '_>) -> HandshakeResult;
}

impl AMQPUriTcpExt for AMQPUri {
    fn connect_with_config(&self, config: TLSConfig<'_, '_, '_>) -> HandshakeResult {
        let uri = format!("{}:{}", self.authority.host, self.authority.port);
        trace!(uri = %uri, "Connecting.");
        let stream = if let Some(timeout) = self.query.connection_timeout {
            TcpStream::connect_timeout(uri, Duration::from_millis(timeout))
        } else {
            TcpStream::connect(uri)
        }?;
        stream.set_nonblocking(true)?;

        match self.scheme {
            AMQPScheme::AMQP => Ok(stream),
            AMQPScheme::AMQPS => stream.into_tls(&self.authority.host, config),
        }
    }
}
<commit_msg>make socket nonblocking later<commit_after>#![deny(missing_docs)]
#![warn(rust_2018_idioms)]
#![doc(html_root_url = "https://docs.rs/amq-protocol-tcp/7.0.0-alpha.1/")]

//! # AMQP URI TCP connection handling
//!
//! amq-protocol-tcp is a library aiming at providing tools to help
//! connecting to an AMQP URI

use amq_protocol_uri::{AMQPScheme, AMQPUri};
use std::time::Duration;
use tracing::trace;

/// Re-export TcpStream
pub use tcp_stream::{
    HandshakeError, HandshakeResult, Identity, MidHandshakeTlsStream, OwnedIdentity,
    OwnedTLSConfig, TLSConfig, TcpStream,
};

#[cfg(feature = "native-tls")]
pub use tcp_stream::NativeTlsConnector;

#[cfg(feature = "openssl")]
pub use tcp_stream::OpenSslConnector;

#[cfg(feature = "rustls-connector")]
pub use tcp_stream::{RustlsConnector, RustlsConnectorConfig};

/// Trait providing a method to connect to a TcpStream
pub trait AMQPUriTcpExt {
    /// connect to a TcpStream
    fn connect(&self) -> HandshakeResult
    where
        Self: Sized,
    {
        self.connect_with_config(TLSConfig::default())
    }

    /// connect to a TcpStream with the given configuration
    fn connect_with_config(&self, config: TLSConfig<'_, '_, '_>) -> HandshakeResult;
}

impl AMQPUriTcpExt for AMQPUri {
    fn connect_with_config(&self, config: TLSConfig<'_, '_, '_>) -> HandshakeResult {
        let uri = format!("{}:{}", self.authority.host, self.authority.port);
        trace!(uri = %uri, "Connecting.");
        let stream = if let Some(timeout) = self.query.connection_timeout {
            TcpStream::connect_timeout(uri, Duration::from_millis(timeout))
        } else {
            TcpStream::connect(uri)
        }?;
        let stream = match self.scheme {
            AMQPScheme::AMQP => stream,
            AMQPScheme::AMQPS => stream.into_tls(&self.authority.host, config)?,
        };
        stream.set_nonblocking(true)?;
        Ok(stream)
    }
}
<|endoftext|>


<commit_before>//
// Sysinfo
//
// Copyright (c) 2017 Guillaume Gomez
//

use ::DiskExt;
use super::system::get_all_data;

use libc::statvfs;
use std::{mem, str};
use std::fmt::{Debug, Error, Formatter};
use std::path::{Path, PathBuf};
use std::ffi::{OsStr, OsString};
use std::os::unix::ffi::OsStrExt;

/// Enum containing the different handled disks types.
#[derive(Debug, PartialEq, Clone, Copy)]
pub enum DiskType {
    /// HDD type.
    HDD,
    /// SSD type.
    SSD,
    /// Unknown type.
    Unknown(isize),
}

impl From<isize> for DiskType {
    fn from(t: isize) -> DiskType {
        match t {
            0 => DiskType::SSD,
            1 => DiskType::HDD,
            id => DiskType::Unknown(id),
        }
    }
}

fn find_type_for_name(name: &OsStr) -> DiskType
{
    /* turn "sda1" into "sda": */
    let mut trimmed: &[u8] = name.as_bytes();
    while trimmed.len() > 1 && trimmed[trimmed.len()-1] >= b'0' && trimmed[trimmed.len()-1] <= b'9' {
	trimmed = &trimmed[..trimmed.len()-1]
    }
    let trimmed: &OsStr = OsStrExt::from_bytes(trimmed);

    let path = Path::new("/sys/block/").to_owned()
        .join(trimmed)
        .join("queue/rotational");
    let rotational_int = get_all_data(path).unwrap_or(String::new()).trim().parse();
    DiskType::from(rotational_int.unwrap_or(-1))
}

/* convert a path to a NUL-terminated Vec<u8> suitable for use with C functions */
fn to_cpath(path: &Path) -> Vec<u8>
{
    let path_os: &OsStr = path.as_ref();
    let mut cpath = path_os.as_bytes().to_vec();
    cpath.push(0);
    cpath
}

pub fn new(name: &OsStr, mount_point: &Path, file_system: &[u8]) -> Disk {
    #[allow(or_fun_call)]
    let mount_point_cpath = to_cpath(mount_point);
    let type_ = find_type_for_name(name);
    let mut total_space = 0;
    let mut available_space = 0;
    unsafe {
        let mut stat: statvfs = mem::zeroed();
        if statvfs(mount_point_cpath.as_ptr() as *const i8, &mut stat) == 0 {
            total_space = stat.f_bsize * stat.f_blocks;
            available_space = stat.f_bsize * stat.f_bavail;
        }
    }
    Disk {
        type_: type_,
        name: name.to_owned(),
        file_system: file_system.to_owned(),
        mount_point: mount_point.to_owned(),
        total_space: total_space,
        available_space: available_space,
    }
}

/// Struct containing a disk information.
pub struct Disk {
    type_: DiskType,
    name: OsString,
    file_system: Vec<u8>,
    mount_point: PathBuf,
    total_space: u64,
    available_space: u64,
}

impl Debug for Disk {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        write!(fmt,
               "Disk({:?})[FS: {:?}][Type: {:?}] mounted on {:?}: {}/{} B",
               self.get_name(), self.get_file_system(), self.get_type(), self.get_mount_point(),
               self.get_available_space(), self.get_total_space())
    }
}

impl DiskExt for Disk {
    fn get_type(&self) -> DiskType {
        self.type_
    }

    fn get_name(&self) -> &OsStr {
        &self.name
    }

    fn get_file_system(&self) -> &[u8] {
        &self.file_system
    }

    fn get_mount_point(&self) -> &Path {
        &self.mount_point
    }

    fn get_total_space(&self) -> u64 {
        self.total_space
    }

    fn get_available_space(&self) -> u64 {
        self.available_space
    }

    fn update(&mut self) -> bool {
        unsafe {
            let mut stat: statvfs = mem::zeroed();
            let mount_point_cpath = to_cpath(&self.mount_point);
            if statvfs(mount_point_cpath.as_ptr() as *const i8, &mut stat) == 0 {
                self.available_space = stat.f_bsize * stat.f_bavail;
                true
            } else {
                false
            }
        }
    }
}
<commit_msg>appease clippy in find_type_for_name<commit_after>//
// Sysinfo
//
// Copyright (c) 2017 Guillaume Gomez
//

use ::DiskExt;
use super::system::get_all_data;

use libc::statvfs;
use std::{mem, str};
use std::fmt::{Debug, Error, Formatter};
use std::path::{Path, PathBuf};
use std::ffi::{OsStr, OsString};
use std::os::unix::ffi::OsStrExt;

/// Enum containing the different handled disks types.
#[derive(Debug, PartialEq, Clone, Copy)]
pub enum DiskType {
    /// HDD type.
    HDD,
    /// SSD type.
    SSD,
    /// Unknown type.
    Unknown(isize),
}

impl From<isize> for DiskType {
    fn from(t: isize) -> DiskType {
        match t {
            0 => DiskType::SSD,
            1 => DiskType::HDD,
            id => DiskType::Unknown(id),
        }
    }
}

fn find_type_for_name(name: &OsStr) -> DiskType
{
    #![allow(or_fun_call)]
    /* turn "sda1" into "sda": */
    let mut trimmed: &[u8] = name.as_bytes();
    while trimmed.len() > 1 && trimmed[trimmed.len()-1] >= b'0' && trimmed[trimmed.len()-1] <= b'9' {
        trimmed = &trimmed[..trimmed.len()-1]
    }
    let trimmed: &OsStr = OsStrExt::from_bytes(trimmed);

    let path = Path::new("/sys/block/").to_owned()
        .join(trimmed)
        .join("queue/rotational");
    let rotational_int = get_all_data(path).unwrap_or(String::new()).trim().parse();
    DiskType::from(rotational_int.unwrap_or(-1))
}

/* convert a path to a NUL-terminated Vec<u8> suitable for use with C functions */
fn to_cpath(path: &Path) -> Vec<u8>
{
    let path_os: &OsStr = path.as_ref();
    let mut cpath = path_os.as_bytes().to_vec();
    cpath.push(0);
    cpath
}

pub fn new(name: &OsStr, mount_point: &Path, file_system: &[u8]) -> Disk {
    let mount_point_cpath = to_cpath(mount_point);
    let type_ = find_type_for_name(name);
    let mut total_space = 0;
    let mut available_space = 0;
    unsafe {
        let mut stat: statvfs = mem::zeroed();
        if statvfs(mount_point_cpath.as_ptr() as *const i8, &mut stat) == 0 {
            total_space = stat.f_bsize * stat.f_blocks;
            available_space = stat.f_bsize * stat.f_bavail;
        }
    }
    Disk {
        type_: type_,
        name: name.to_owned(),
        file_system: file_system.to_owned(),
        mount_point: mount_point.to_owned(),
        total_space: total_space,
        available_space: available_space,
    }
}

/// Struct containing a disk information.
pub struct Disk {
    type_: DiskType,
    name: OsString,
    file_system: Vec<u8>,
    mount_point: PathBuf,
    total_space: u64,
    available_space: u64,
}

impl Debug for Disk {
    fn fmt(&self, fmt: &mut Formatter) -> Result<(), Error> {
        write!(fmt,
               "Disk({:?})[FS: {:?}][Type: {:?}] mounted on {:?}: {}/{} B",
               self.get_name(), self.get_file_system(), self.get_type(), self.get_mount_point(),
               self.get_available_space(), self.get_total_space())
    }
}

impl DiskExt for Disk {
    fn get_type(&self) -> DiskType {
        self.type_
    }

    fn get_name(&self) -> &OsStr {
        &self.name
    }

    fn get_file_system(&self) -> &[u8] {
        &self.file_system
    }

    fn get_mount_point(&self) -> &Path {
        &self.mount_point
    }

    fn get_total_space(&self) -> u64 {
        self.total_space
    }

    fn get_available_space(&self) -> u64 {
        self.available_space
    }

    fn update(&mut self) -> bool {
        unsafe {
            let mut stat: statvfs = mem::zeroed();
            let mount_point_cpath = to_cpath(&self.mount_point);
            if statvfs(mount_point_cpath.as_ptr() as *const i8, &mut stat) == 0 {
                self.available_space = stat.f_bsize * stat.f_bavail;
                true
            } else {
                false
            }
        }
    }
}
<|endoftext|>