<commit_before>use uci_engine::UciEngine;
use uci;
use std::sync::Arc;
use std::sync::Mutex;
use uci_engine::UciOption;
use uci::TimeRestriction;
use uci::EngineComm;
use search_algorithms::board::Board;
use uci::UciInfo;
use search_algorithms::board::Color::*;
use search_algorithms::board::GameResult::*;
use search_algorithms::board::GameResult;
use std::ops::Add;
use time::Duration;

extern crate ordered_float;
use ordered_float::OrderedFloat;
use std::ops::Not;
use search_algorithms::alpha_beta;
use std::fmt::Debug;
use search_algorithms::alpha_beta::AlphaBeta;
use search_algorithms::board::ExtendedBoard;
use pgn::PgnBoard;
use std::hash::Hash;
use uci_engine::UciOptionType;
use uci::TimeInfo;
use std::time::Instant;

const GAME_TIME: Duration = Duration::from_millis(30000);
const INC_MS: Duration = Duration::from_millis(300);

const EVAL_DEPENDENCE : f32 = 10.0;

pub struct MonteCarlo<B: Board> {
    root: MonteCarloTree<B>,
    time_limit: uci::TimeRestriction,
    start_time: Instant,
    engine_comm : Arc<Mutex<EngineComm>>,
}

impl<B> UciEngine<B> for MonteCarlo<B>
where B: ExtendedBoard + PgnBoard + Debug + Hash + Eq + 'static {
    fn init() -> Self {
        MonteCarlo {
            root: MonteCarloTree::new_root(&mut B::start_board()),
            time_limit: TimeRestriction::Infinite,
            engine_comm: Arc::new(Mutex::new(EngineComm::new())),
            start_time: Instant::now(),
        }
    }

    fn uci_options(&mut self) -> Vec<UciOption> {
        vec![]
    }

    fn set_uci_option(&mut self, _: UciOption) {}

    fn search(mut self, mut board: B, time_limit: TimeRestriction, engine_comm: Arc<Mutex<EngineComm>>,
              _: Option<Vec<<B as Board>::Move>>) -> Box<Iterator<Item=UciInfo<B>>> {

        self.root = MonteCarloTree::new_root(&mut board);
        self.time_limit = time_limit;
        self.engine_comm = engine_comm;
        self.start_time = Instant::now();

        Box::new(DepthIterator { monte_carlo: self, root_board: board })
    }
}

struct DepthIterator<B: Board> {
    monte_carlo: MonteCarlo<B>,
    root_board: B,
}

impl<B> Iterator for DepthIterator<B>
where B: ExtendedBoard + PgnBoard + Debug + Hash + Eq + 'static {
    type Item = UciInfo<B>;

    fn next(&mut self) -> Option<<Self as Iterator>::Item> {
        for _ in 0..10 {
            let mut game = vec![];
            let result = self.monte_carlo.root.select(self.root_board.clone(), &mut game);

            let mut board = self.root_board.clone();
            for (i, mv) in game.into_iter().enumerate() {
                if i == 0 && board.side_to_move() == Black {
                    print!("{}... {} ", 1, board.move_to_san(&mv));
                }
                else if board.side_to_move() == White {
                    print!("{}. {} ", (i + 1) / 2 + 1, board.move_to_san(&mv));
                }
                else {
                    print!("{} ", board.move_to_san(&mv));
                }
                board.do_move(mv);
            }
            println!(" {:?}", result);
        }

        let root = &self.monte_carlo.root;
        println!("Root: {}, static winrate: {}, true winrate: {}, score: {:?}, exploration: {}",
                 root.static_eval, eval_to_win_pct(root.static_eval),
                 root.score.win_rate(),
                 root.score, root.expl_value(root.searches()));

        let (mv, _) = root.children.as_ref().unwrap().iter()
            .min_by_key(|&(_, ref child)| OrderedFloat(child.score.win_rate()))
            .unwrap();

        for (mv, child) in root.children.as_ref().unwrap() {
            println!("\tChild {:?}: {}, static winrate: {}, true winrate: {}, score: {:?}, exploration: {}",
                     mv,
                     child.static_eval, eval_to_win_pct(child.static_eval),
                     child.score.win_rate(),
                     child.score, child.expl_value(root.searches()));
        }
        println!();

        Some(UciInfo {
            color: self.root_board.side_to_move(),
            depth: 0,
            seldepth: 0,
            time: 0,
            nodes: root.searches(),
            hashfull: 0.0,
            pvs: vec![(alpha_beta::Score::Val(win_pct_to_eval(root.score.win_rate())), vec![mv.clone()])]
        })
    }
}

#[derive(Debug)]
struct MonteCarloTree<B: Board> {
    children: Option<Vec<(B::Move, MonteCarloTree<B>)>>,
    score: Score,
    static_eval: f32,
}

impl<B> MonteCarloTree<B>
where B: ExtendedBoard + PgnBoard + Debug + Hash + Eq + 'static {
    fn new_root(board: &mut B) -> MonteCarloTree<B> {
        MonteCarloTree {
            children: None,
            score: Score::new(),
            static_eval: board.static_eval(),
        }
    }

    fn searches(&self) -> u64 {
        self.score.wins + self.score.losses + self.score.draws
    }

    fn expl_value(self: &MonteCarloTree<B>, parent_searches: u64) -> f32 {

        let extra_losses = EVAL_DEPENDENCE - eval_to_win_pct(self.static_eval) * EVAL_DEPENDENCE;

        let exploitation = (self.score.draws as f32 * 0.5 + self.score.losses as f32 + extra_losses)
            / (self.searches() as f32 + EVAL_DEPENDENCE);

        let exploration = 1.42 * f32::sqrt(
            f32::ln(parent_searches as f32 + 1.0) / (self.searches() as f32 + EVAL_DEPENDENCE));
        exploitation + exploration
    }

    fn select(&mut self, mut board: B, game: &mut Vec<B::Move>) -> Score {
        let searches = self.searches();
        if let Some(result) = board.game_result() {
            let mut score = Score::new();
            match (result, board.side_to_move()) {
                (WhiteWin, White) | (BlackWin, Black) => score.wins += 1,
                (WhiteWin, Black) | (BlackWin, White) => score.losses += 1,
                (Draw, _) => score.draws += 1,
            }
            return score;
        }
        if let Some(children) = self.children.as_mut() {

            let child_index = (0..children.len())
                .max_by_key(|&i| OrderedFloat(children[i].1.expl_value(searches)))
                .unwrap().clone();

            let (ref mut mv, ref mut child) = children.get_mut(child_index).unwrap();
            board.do_move(mv.clone());
            game.push(mv.clone());

            let result = !child.select(board, game);

            self.score = self.score + result;
            return result;
        }

        return self.expand(board, game);
    }

    fn expand(&mut self, mut board: B, game: &mut Vec<B::Move>) -> Score {
        let mut moves = vec![];
        board.generate_moves(&mut moves);

        let mut engine: AlphaBeta<B> = AlphaBeta::init();

        engine.set_uci_option(UciOption {
            name: "multipv".to_string(),
            option_type: UciOptionType::Spin(moves.len() as i64, 1, 1)
        });

        engine.set_uci_option(UciOption {
            name: "hash".to_string(),
            option_type: UciOptionType::Spin(32, 1, 1)
        });

        let results = engine.best_moves_multipv(
            board.clone(),
            TimeRestriction::MoveTime(INC_MS * f32::sqrt(moves.len() as f32).ceil() as u32),
        None).unwrap();

        let children = results.into_iter()
            .map(|(score, mv)| {
                let child = MonteCarloTree {
                    children: None,
                    score: Score::new(),
                    static_eval: match score {
                        alpha_beta::Score::Val(val) => - val,
                        alpha_beta::Score::Draw(_) => 0.0,
                        alpha_beta::Score::Win(_) => -100.0,
                        alpha_beta::Score::Loss(_) => 100.0,
                    }
                };
                (mv, child)
            })
            .collect::<Vec<_>>();
        {
            let (mv, _) = children.get(0).unwrap();
            board.do_move(mv.clone());
            game.push(mv.clone());
        }

        self.children = Some(children);
        let result = !simulate(board, game);
        self.score = self.score + result;

        result
    }
}

fn simulate<B>(mut board: B, game: &mut Vec<B::Move>) -> Score
    where B: ExtendedBoard + PgnBoard + Debug + Hash + Eq + 'static {

    let start_color = board.side_to_move();

    let mut white_time = GAME_TIME;
    let mut black_time = GAME_TIME;

    loop {
        if let Some(result) = board.game_result() {
            let mut score = Score::new();
            match (result, board.side_to_move()) {
                (WhiteWin, White) | (BlackWin, Black) => score.wins += 1,
                (WhiteWin, Black) | (BlackWin, White) => score.losses += 1,
                (Draw, _) => score.draws += 1,
            }
            println!("{:?} ({} to move) on the board", result, board.side_to_move());
            return if board.side_to_move() == start_color { score } else { !score };
        }

        let turn_start = Instant::now();

        let mut engine: AlphaBeta<B> = AlphaBeta::init();

        engine.set_uci_option(UciOption {
            name: "hash".to_string(),
            option_type: UciOptionType::Spin(16, 1, 1)
        });

        let time_info = TimeInfo {
            white_time,
            black_time,
            white_inc: INC_MS,
            black_inc: INC_MS,
            moves_to_go: None
        };

        let (_, mv) = engine.best_move(
            board.clone(),
            TimeRestriction::GameTime(time_info),
            None)
            .unwrap(); // TODO: handle errors

        let time_elapsed = Instant::now() - turn_start;

        match board.side_to_move() {
            White if white_time < time_elapsed => {
                println!("White wins on time");
                return Score::from_game_result(&BlackWin)
            },
            Black if black_time < time_elapsed => {
                println!("Black wins on time");
                return Score::from_game_result(&WhiteWin)
            },
            White => {
                white_time -= time_elapsed;
                white_time += INC_MS;
            },
            Black => {
                black_time -= time_elapsed;
                black_time += INC_MS;
            },
        }

        game.push(mv.clone());
        board.do_move(mv);
    }
}

pub fn win_pct_to_eval(pct: f32) -> f32 {
    2.90680623072 * f32::tan(3.096181612 * (pct - 0.5))
}

pub fn eval_to_win_pct(cp: f32) -> f32 {
    (f32::atan(cp / 2.90680623072) / 3.096181612 + 0.5)
}

#[derive(PartialEq, Clone, Debug, Copy)]
pub struct Score {
    pub wins: u64,
    pub losses: u64,
    pub draws : u64,
}

impl Score {
    pub fn new() -> Self {
        Score { wins: 0, losses: 0, draws: 0 }
    }

    /// Creates a score object containing a single result
    pub fn from_game_result(result: &GameResult) -> Self {
        let mut score = Score::new();
        match *result {
            WhiteWin => score.wins += 1,
            BlackWin => score.losses += 1,
            Draw => score.draws += 1,
        }
        score
    }

    pub fn win_rate(&self) -> f32 {
        (self.wins as f32 + self.draws as f32 * 0.5) / (self.wins + self.draws + self.losses) as f32
    }
}

impl Add for Score {
    type Output = Score;

    fn add(self, rhs: Score) -> <Self as Add<Score>>::Output {
        Score {
            wins: self.wins + rhs.wins,
            losses: self.losses + rhs.losses,
            draws: self.draws + rhs.draws,
        }
    }
}

impl Not for Score {
    type Output = Self;

    fn not(self) -> <Self as Not>::Output {
        Score {
            wins: self.losses,
            losses: self.wins,
            draws: self.draws,
        }
    }
}<commit_msg>When expanding, add score to the first child node searched<commit_after>use uci_engine::UciEngine;
use uci;
use std::sync::Arc;
use std::sync::Mutex;
use uci_engine::UciOption;
use uci::TimeRestriction;
use uci::EngineComm;
use search_algorithms::board::Board;
use uci::UciInfo;
use search_algorithms::board::Color::*;
use search_algorithms::board::GameResult::*;
use search_algorithms::board::GameResult;
use std::ops::Add;
use time::Duration;

extern crate ordered_float;
use ordered_float::OrderedFloat;
use std::ops::Not;
use search_algorithms::alpha_beta;
use std::fmt::Debug;
use search_algorithms::alpha_beta::AlphaBeta;
use search_algorithms::board::ExtendedBoard;
use pgn::PgnBoard;
use std::hash::Hash;
use uci_engine::UciOptionType;
use uci::TimeInfo;
use std::time::Instant;

const GAME_TIME: Duration = Duration::from_millis(30000);
const INC_MS: Duration = Duration::from_millis(300);

const EVAL_DEPENDENCE : f32 = 10.0;

pub struct MonteCarlo<B: Board> {
    root: MonteCarloTree<B>,
    time_limit: uci::TimeRestriction,
    start_time: Instant,
    engine_comm : Arc<Mutex<EngineComm>>,
}

impl<B> UciEngine<B> for MonteCarlo<B>
where B: ExtendedBoard + PgnBoard + Debug + Hash + Eq + 'static {
    fn init() -> Self {
        MonteCarlo {
            root: MonteCarloTree::new_root(&mut B::start_board()),
            time_limit: TimeRestriction::Infinite,
            engine_comm: Arc::new(Mutex::new(EngineComm::new())),
            start_time: Instant::now(),
        }
    }

    fn uci_options(&mut self) -> Vec<UciOption> {
        vec![]
    }

    fn set_uci_option(&mut self, _: UciOption) {}

    fn search(mut self, mut board: B, time_limit: TimeRestriction, engine_comm: Arc<Mutex<EngineComm>>,
              _: Option<Vec<<B as Board>::Move>>) -> Box<Iterator<Item=UciInfo<B>>> {

        self.root = MonteCarloTree::new_root(&mut board);
        self.time_limit = time_limit;
        self.engine_comm = engine_comm;
        self.start_time = Instant::now();

        Box::new(DepthIterator { monte_carlo: self, root_board: board })
    }
}

struct DepthIterator<B: Board> {
    monte_carlo: MonteCarlo<B>,
    root_board: B,
}

impl<B> Iterator for DepthIterator<B>
where B: ExtendedBoard + PgnBoard + Debug + Hash + Eq + 'static {
    type Item = UciInfo<B>;

    fn next(&mut self) -> Option<<Self as Iterator>::Item> {
        for _ in 0..10 {
            let mut game = vec![];
            let result = self.monte_carlo.root.select(self.root_board.clone(), &mut game);

            let mut board = self.root_board.clone();
            for (i, mv) in game.into_iter().enumerate() {
                if i == 0 && board.side_to_move() == Black {
                    print!("{}... {} ", 1, board.move_to_san(&mv));
                }
                else if board.side_to_move() == White {
                    print!("{}. {} ", (i + 1) / 2 + 1, board.move_to_san(&mv));
                }
                else {
                    print!("{} ", board.move_to_san(&mv));
                }
                board.do_move(mv);
            }
            println!(" {:?}", result);
        }

        let root = &self.monte_carlo.root;
        println!("Root: {}, static winrate: {}, true winrate: {}, score: {:?}, exploration: {}",
                 root.static_eval, eval_to_win_pct(root.static_eval),
                 root.score.win_rate(),
                 root.score, root.expl_value(root.searches()));

        let (mv, _) = root.children.as_ref().unwrap().iter()
            .min_by_key(|&(_, ref child)| OrderedFloat(child.score.win_rate()))
            .unwrap();

        for (mv, child) in root.children.as_ref().unwrap() {
            println!("\tChild {:?}: {}, static winrate: {}, true winrate: {}, score: {:?}, exploration: {}",
                     mv,
                     child.static_eval, eval_to_win_pct(child.static_eval),
                     child.score.win_rate(),
                     child.score, child.expl_value(root.searches()));
        }
        println!();

        Some(UciInfo {
            color: self.root_board.side_to_move(),
            depth: 0,
            seldepth: 0,
            time: 0,
            nodes: root.searches(),
            hashfull: 0.0,
            pvs: vec![(alpha_beta::Score::Val(win_pct_to_eval(root.score.win_rate())), vec![mv.clone()])]
        })
    }
}

#[derive(Debug)]
struct MonteCarloTree<B: Board> {
    children: Option<Vec<(B::Move, MonteCarloTree<B>)>>,
    score: Score,
    static_eval: f32,
}

impl<B> MonteCarloTree<B>
where B: ExtendedBoard + PgnBoard + Debug + Hash + Eq + 'static {
    fn new_root(board: &mut B) -> MonteCarloTree<B> {
        MonteCarloTree {
            children: None,
            score: Score::new(),
            static_eval: board.static_eval(),
        }
    }

    fn searches(&self) -> u64 {
        self.score.wins + self.score.losses + self.score.draws
    }

    fn expl_value(self: &MonteCarloTree<B>, parent_searches: u64) -> f32 {

        let extra_losses = EVAL_DEPENDENCE - eval_to_win_pct(self.static_eval) * EVAL_DEPENDENCE;

        let exploitation = (self.score.draws as f32 * 0.5 + self.score.losses as f32 + extra_losses)
            / (self.searches() as f32 + EVAL_DEPENDENCE);

        let exploration = 1.42 * f32::sqrt(
            f32::ln(parent_searches as f32 + 1.0) / (self.searches() as f32 + EVAL_DEPENDENCE));
        exploitation + exploration
    }

    fn select(&mut self, mut board: B, game: &mut Vec<B::Move>) -> Score {
        let searches = self.searches();
        if let Some(result) = board.game_result() {
            let mut score = Score::new();
            match (result, board.side_to_move()) {
                (WhiteWin, White) | (BlackWin, Black) => score.wins += 1,
                (WhiteWin, Black) | (BlackWin, White) => score.losses += 1,
                (Draw, _) => score.draws += 1,
            }
            return score;
        }
        if let Some(children) = self.children.as_mut() {

            let child_index = (0..children.len())
                .max_by_key(|&i| OrderedFloat(children[i].1.expl_value(searches)))
                .unwrap().clone();

            let (ref mut mv, ref mut child) = children.get_mut(child_index).unwrap();
            board.do_move(mv.clone());
            game.push(mv.clone());

            let result = !child.select(board, game);

            self.score = self.score + result;
            return result;
        }

        return self.expand(board, game);
    }

    fn expand(&mut self, mut board: B, game: &mut Vec<B::Move>) -> Score {
        let mut moves = vec![];
        board.generate_moves(&mut moves);

        let mut engine: AlphaBeta<B> = AlphaBeta::init();

        engine.set_uci_option(UciOption {
            name: "multipv".to_string(),
            option_type: UciOptionType::Spin(moves.len() as i64, 1, 1)
        });

        engine.set_uci_option(UciOption {
            name: "hash".to_string(),
            option_type: UciOptionType::Spin(32, 1, 1)
        });

        let results = engine.best_moves_multipv(
            board.clone(),
            TimeRestriction::MoveTime(INC_MS * f32::sqrt(moves.len() as f32).ceil() as u32),
        None).unwrap();

        let children = results.into_iter()
            .map(|(score, mv)| {
                let child = MonteCarloTree {
                    children: None,
                    score: Score::new(),
                    static_eval: match score {
                        alpha_beta::Score::Val(val) => - val,
                        alpha_beta::Score::Draw(_) => 0.0,
                        alpha_beta::Score::Win(_) => -100.0,
                        alpha_beta::Score::Loss(_) => 100.0,
                    }
                };
                (mv, child)
            })
            .collect::<Vec<_>>();

        self.children = Some(children);

        let (mv, child) = self.children.as_mut().unwrap().get_mut(0).unwrap();
        board.do_move(mv.clone());
        game.push(mv.clone());
        let result = !simulate(board, game);
        self.score = self.score + result;
        child.score = child.score + !result;

        result
    }
}

fn simulate<B>(mut board: B, game: &mut Vec<B::Move>) -> Score
    where B: ExtendedBoard + PgnBoard + Debug + Hash + Eq + 'static {

    let start_color = board.side_to_move();

    let mut white_time = GAME_TIME;
    let mut black_time = GAME_TIME;

    loop {
        if let Some(result) = board.game_result() {
            let mut score = Score::new();
            match (result, board.side_to_move()) {
                (WhiteWin, White) | (BlackWin, Black) => score.wins += 1,
                (WhiteWin, Black) | (BlackWin, White) => score.losses += 1,
                (Draw, _) => score.draws += 1,
            }
            println!("{:?} ({} to move) on the board", result, board.side_to_move());
            return if board.side_to_move() == start_color { score } else { !score };
        }

        let turn_start = Instant::now();

        let mut engine: AlphaBeta<B> = AlphaBeta::init();

        engine.set_uci_option(UciOption {
            name: "hash".to_string(),
            option_type: UciOptionType::Spin(16, 1, 1)
        });

        let time_info = TimeInfo {
            white_time,
            black_time,
            white_inc: INC_MS,
            black_inc: INC_MS,
            moves_to_go: None
        };

        let (_, mv) = engine.best_move(
            board.clone(),
            TimeRestriction::GameTime(time_info),
            None)
            .unwrap(); // TODO: handle errors

        let time_elapsed = Instant::now() - turn_start;

        match board.side_to_move() {
            White if white_time < time_elapsed => {
                println!("White wins on time");
                return Score::from_game_result(&BlackWin)
            },
            Black if black_time < time_elapsed => {
                println!("Black wins on time");
                return Score::from_game_result(&WhiteWin)
            },
            White => {
                white_time -= time_elapsed;
                white_time += INC_MS;
            },
            Black => {
                black_time -= time_elapsed;
                black_time += INC_MS;
            },
        }

        game.push(mv.clone());
        board.do_move(mv);
    }
}

pub fn win_pct_to_eval(pct: f32) -> f32 {
    2.90680623072 * f32::tan(3.096181612 * (pct - 0.5))
}

pub fn eval_to_win_pct(cp: f32) -> f32 {
    (f32::atan(cp / 2.90680623072) / 3.096181612 + 0.5)
}

#[derive(PartialEq, Clone, Debug, Copy)]
pub struct Score {
    pub wins: u64,
    pub losses: u64,
    pub draws : u64,
}

impl Score {
    pub fn new() -> Self {
        Score { wins: 0, losses: 0, draws: 0 }
    }

    /// Creates a score object containing a single result
    pub fn from_game_result(result: &GameResult) -> Self {
        let mut score = Score::new();
        match *result {
            WhiteWin => score.wins += 1,
            BlackWin => score.losses += 1,
            Draw => score.draws += 1,
        }
        score
    }

    pub fn win_rate(&self) -> f32 {
        (self.wins as f32 + self.draws as f32 * 0.5) / (self.wins + self.draws + self.losses) as f32
    }
}

impl Add for Score {
    type Output = Score;

    fn add(self, rhs: Score) -> <Self as Add<Score>>::Output {
        Score {
            wins: self.wins + rhs.wins,
            losses: self.losses + rhs.losses,
            draws: self.draws + rhs.draws,
        }
    }
}

impl Not for Score {
    type Output = Self;

    fn not(self) -> <Self as Not>::Output {
        Score {
            wins: self.losses,
            losses: self.wins,
            draws: self.draws,
        }
    }
}<|endoftext|>


<commit_before>use std::collections::HashSet;

use protobuf::descriptor::*;
use protobuf::descriptorx::*;

use super::code_writer::*;
use super::customize::Customize;

#[derive(Clone)]
pub struct EnumValueGen {
    proto: EnumValueDescriptorProto,
    enum_rust_name: String,
}

impl EnumValueGen {
    fn parse(proto: &EnumValueDescriptorProto, enum_rust_name: &str) -> EnumValueGen {
        EnumValueGen {
            proto: proto.clone(),
            enum_rust_name: enum_rust_name.to_string(),
        }
    }

    // enum value
    fn number(&self) -> i32 {
        self.proto.get_number()
    }

    // name of enum variant in generated rust code
    fn rust_name_inner(&self) -> String {
        self.proto.rust_name()
    }

    pub fn rust_name_outer(&self) -> String {
        let mut r = String::new();
        r.push_str(&self.enum_rust_name);
        r.push_str("::");
        r.push_str(&self.rust_name_inner());
        r
    }
}

pub struct EnumGen<'a> {
    enum_with_scope: &'a EnumWithScope<'a>,
    type_name: String,
    lite_runtime: bool,
    customize: Customize
}

impl<'a> EnumGen<'a> {
    pub fn new(
        enum_with_scope: &'a EnumWithScope<'a>,
        current_file: &FileDescriptorProto,
        customize: &Customize
    ) -> EnumGen<'a> {
        let rust_name = if enum_with_scope.get_scope().get_file_descriptor().get_name() ==
            current_file.get_name()
        {
            // field type is a message or enum declared in the same file
            enum_with_scope.rust_name()
        } else {
            format!(
                "super::{}::{}",
                proto_path_to_rust_mod(
                    enum_with_scope.get_scope().get_file_descriptor().get_name(),
                ),
                enum_with_scope.rust_name()
            )
        };
        EnumGen {
            enum_with_scope: enum_with_scope,
            type_name: rust_name,
            lite_runtime: enum_with_scope
                .get_scope()
                .get_file_descriptor()
                .get_options()
                .get_optimize_for() ==
                FileOptions_OptimizeMode::LITE_RUNTIME,
            customize: customize.clone()
        }
    }

    fn allow_alias(&self) -> bool {
        self.enum_with_scope.en.get_options().get_allow_alias()
    }

    fn values_all(&self) -> Vec<EnumValueGen> {
        let mut r = Vec::new();
        for p in self.enum_with_scope.values() {
            r.push(EnumValueGen::parse(p, &self.type_name));
        }
        r
    }

    pub fn values_unique(&self) -> Vec<EnumValueGen> {
        let mut used = HashSet::new();
        let mut r = Vec::new();
        for p in self.enum_with_scope.values() {
            // skipping non-unique enums
            // TODO: should support it
            if !used.insert(p.get_number()) {
                continue;
            }
            r.push(EnumValueGen::parse(p, &self.type_name));
        }
        r
    }

    // find enum value by name
    pub fn value_by_name(&'a self, name: &str) -> EnumValueGen {
        EnumValueGen::parse(self.enum_with_scope.value_by_name(name), &self.type_name)
    }

    pub fn write(&self, w: &mut CodeWriter) {
        self.write_struct(w);
        if self.allow_alias() {
            w.write_line("");
            self.write_impl_eq(w);
            w.write_line("");
            self.write_impl_hash(w);
        }
        w.write_line("");
        self.write_impl_enum(w);
        if self.enum_with_scope.scope.file_scope.syntax() == Syntax::PROTO3 {
            w.write_line("");
            self.write_impl_default(w);
        }
        w.write_line("");
        self.write_impl_value(w);
    }

    fn write_struct(&self, w: &mut CodeWriter) {
        let mut derive = Vec::new();
        derive.push("Clone");
        derive.push("Copy");
        if !self.allow_alias() {
            derive.push("PartialEq");
        }
        derive.push("Eq");
        derive.push("Debug");
        if !self.allow_alias() {
            derive.push("Hash");
        } else {
            w.comment(
                "Note: you cannot use pattern matching for enums with allow_alias option",
            );
        }
        if self.customize.serde_derive.unwrap_or(false) {
            derive.extend(&["Serialize", "Deserialize"]);
        }
        w.derive(&derive);
        let ref type_name = self.type_name;
        w.expr_block(&format!("pub enum {}", type_name), |w| {
            for value in self.values_all() {
                if self.allow_alias() {
                    w.write_line(&format!(
                        "{}, // {}",
                        value.rust_name_inner(),
                        value.number()
                    ));
                } else {
                    w.write_line(&format!(
                        "{} = {},",
                        value.rust_name_inner(),
                        value.number()
                    ));
                }
            }
        });
    }

    fn write_fn_value(&self, w: &mut CodeWriter) {
        w.def_fn("value(&self) -> i32", |w| if self.allow_alias() {
            w.match_expr("*self", |w| for value in self.values_all() {
                w.case_expr(value.rust_name_outer(), format!("{}", value.number()));
            });
        } else {
            w.write_line("*self as i32")
        });
    }

    fn write_impl_enum(&self, w: &mut CodeWriter) {
        let ref type_name = self.type_name;
        w.impl_for_block("::protobuf::ProtobufEnum", &type_name, |w| {
            self.write_fn_value(w);

            w.write_line("");
            let ref type_name = self.type_name;
            w.def_fn(&format!("from_i32(value: i32) -> ::std::option::Option<{}>", type_name), |w| {
                w.match_expr("value", |w| {
                    let values = self.values_unique();
                    for value in values {
                        w.write_line(&format!("{} => ::std::option::Option::Some({}),",
                            value.number(), value.rust_name_outer()));
                    }
                    w.write_line(&format!("_ => ::std::option::Option::None"));
                });
            });

            w.write_line("");
            w.def_fn(&format!("values() -> &'static [Self]"), |w| {
                w.write_line(&format!("static values: &'static [{}] = &[", type_name));
                w.indented(|w| {
                    for value in self.values_all() {
                        w.write_line(&format!("{},", value.rust_name_outer()));
                    }
                });
                w.write_line("];");
                w.write_line("values");
            });

            if !self.lite_runtime {
                w.write_line("");
                w.def_fn(&format!("enum_descriptor_static() -> &'static ::protobuf::reflect::EnumDescriptor"), |w| {
                    w.lazy_static_decl_get("descriptor", "::protobuf::reflect::EnumDescriptor", |w| {
                        let ref type_name = self.type_name;
                        w.write_line(&format!(
                            "::protobuf::reflect::EnumDescriptor::new::<{}>(\"{}\", file_descriptor_proto())",
                            type_name,
                            type_name));
                    });
                });
            }
        });
    }

    fn write_impl_value(&self, w: &mut CodeWriter) {
        w.impl_for_block("::protobuf::reflect::ProtobufValue", &self.type_name, |_w| {
        })
    }

    fn write_impl_eq(&self, w: &mut CodeWriter) {
        assert!(self.allow_alias());
        w.impl_for_block("::std::cmp::PartialEq", &self.type_name, |w| {
            w.def_fn("eq(&self, other: &Self) -> bool", |w| {
                w.write_line("self.value() == other.value()");
            });
        });
    }

    fn write_impl_hash(&self, w: &mut CodeWriter) {
        assert!(self.allow_alias());
        w.impl_for_block("::std::hash::Hash", &self.type_name, |w| {
            w.def_fn("hash<H : ::std::hash::Hasher>(&self, state: &mut H)", |w| {
                w.write_line("state.write_i32(self.value())");
            });
        });
    }

    fn write_impl_default(&self, w: &mut CodeWriter) {
        assert!(self.enum_with_scope.scope.file_scope.syntax() == Syntax::PROTO3);
        w.impl_for_block("::std::default::Default", &self.type_name, |w| {
            w.def_fn("default() -> Self", |w| {
                w.write_line(&format!(
                    "{}::{}",
                    &self.type_name,
                    &self.enum_with_scope.values()[0].rust_name()
                ))
            });
        });
    }
}
<commit_msg>Implement `Default` for enums in proto2<commit_after>use std::collections::HashSet;

use protobuf::descriptor::*;
use protobuf::descriptorx::*;

use super::code_writer::*;
use super::customize::Customize;

#[derive(Clone)]
pub struct EnumValueGen {
    proto: EnumValueDescriptorProto,
    enum_rust_name: String,
}

impl EnumValueGen {
    fn parse(proto: &EnumValueDescriptorProto, enum_rust_name: &str) -> EnumValueGen {
        EnumValueGen {
            proto: proto.clone(),
            enum_rust_name: enum_rust_name.to_string(),
        }
    }

    // enum value
    fn number(&self) -> i32 {
        self.proto.get_number()
    }

    // name of enum variant in generated rust code
    fn rust_name_inner(&self) -> String {
        self.proto.rust_name()
    }

    pub fn rust_name_outer(&self) -> String {
        let mut r = String::new();
        r.push_str(&self.enum_rust_name);
        r.push_str("::");
        r.push_str(&self.rust_name_inner());
        r
    }
}

pub struct EnumGen<'a> {
    enum_with_scope: &'a EnumWithScope<'a>,
    type_name: String,
    lite_runtime: bool,
    customize: Customize
}

impl<'a> EnumGen<'a> {
    pub fn new(
        enum_with_scope: &'a EnumWithScope<'a>,
        current_file: &FileDescriptorProto,
        customize: &Customize
    ) -> EnumGen<'a> {
        let rust_name = if enum_with_scope.get_scope().get_file_descriptor().get_name() ==
            current_file.get_name()
        {
            // field type is a message or enum declared in the same file
            enum_with_scope.rust_name()
        } else {
            format!(
                "super::{}::{}",
                proto_path_to_rust_mod(
                    enum_with_scope.get_scope().get_file_descriptor().get_name(),
                ),
                enum_with_scope.rust_name()
            )
        };
        EnumGen {
            enum_with_scope: enum_with_scope,
            type_name: rust_name,
            lite_runtime: enum_with_scope
                .get_scope()
                .get_file_descriptor()
                .get_options()
                .get_optimize_for() ==
                FileOptions_OptimizeMode::LITE_RUNTIME,
            customize: customize.clone()
        }
    }

    fn allow_alias(&self) -> bool {
        self.enum_with_scope.en.get_options().get_allow_alias()
    }

    fn values_all(&self) -> Vec<EnumValueGen> {
        let mut r = Vec::new();
        for p in self.enum_with_scope.values() {
            r.push(EnumValueGen::parse(p, &self.type_name));
        }
        r
    }

    pub fn values_unique(&self) -> Vec<EnumValueGen> {
        let mut used = HashSet::new();
        let mut r = Vec::new();
        for p in self.enum_with_scope.values() {
            // skipping non-unique enums
            // TODO: should support it
            if !used.insert(p.get_number()) {
                continue;
            }
            r.push(EnumValueGen::parse(p, &self.type_name));
        }
        r
    }

    // find enum value by name
    pub fn value_by_name(&'a self, name: &str) -> EnumValueGen {
        EnumValueGen::parse(self.enum_with_scope.value_by_name(name), &self.type_name)
    }

    pub fn write(&self, w: &mut CodeWriter) {
        self.write_struct(w);
        if self.allow_alias() {
            w.write_line("");
            self.write_impl_eq(w);
            w.write_line("");
            self.write_impl_hash(w);
        }
        w.write_line("");
        self.write_impl_enum(w);
        w.write_line("");
        self.write_impl_default(w);
        w.write_line("");
        self.write_impl_value(w);
    }

    fn write_struct(&self, w: &mut CodeWriter) {
        let mut derive = Vec::new();
        derive.push("Clone");
        derive.push("Copy");
        if !self.allow_alias() {
            derive.push("PartialEq");
        }
        derive.push("Eq");
        derive.push("Debug");
        if !self.allow_alias() {
            derive.push("Hash");
        } else {
            w.comment(
                "Note: you cannot use pattern matching for enums with allow_alias option",
            );
        }
        if self.customize.serde_derive.unwrap_or(false) {
            derive.extend(&["Serialize", "Deserialize"]);
        }
        w.derive(&derive);
        let ref type_name = self.type_name;
        w.expr_block(&format!("pub enum {}", type_name), |w| {
            for value in self.values_all() {
                if self.allow_alias() {
                    w.write_line(&format!(
                        "{}, // {}",
                        value.rust_name_inner(),
                        value.number()
                    ));
                } else {
                    w.write_line(&format!(
                        "{} = {},",
                        value.rust_name_inner(),
                        value.number()
                    ));
                }
            }
        });
    }

    fn write_fn_value(&self, w: &mut CodeWriter) {
        w.def_fn("value(&self) -> i32", |w| if self.allow_alias() {
            w.match_expr("*self", |w| for value in self.values_all() {
                w.case_expr(value.rust_name_outer(), format!("{}", value.number()));
            });
        } else {
            w.write_line("*self as i32")
        });
    }

    fn write_impl_enum(&self, w: &mut CodeWriter) {
        let ref type_name = self.type_name;
        w.impl_for_block("::protobuf::ProtobufEnum", &type_name, |w| {
            self.write_fn_value(w);

            w.write_line("");
            let ref type_name = self.type_name;
            w.def_fn(&format!("from_i32(value: i32) -> ::std::option::Option<{}>", type_name), |w| {
                w.match_expr("value", |w| {
                    let values = self.values_unique();
                    for value in values {
                        w.write_line(&format!("{} => ::std::option::Option::Some({}),",
                            value.number(), value.rust_name_outer()));
                    }
                    w.write_line(&format!("_ => ::std::option::Option::None"));
                });
            });

            w.write_line("");
            w.def_fn(&format!("values() -> &'static [Self]"), |w| {
                w.write_line(&format!("static values: &'static [{}] = &[", type_name));
                w.indented(|w| {
                    for value in self.values_all() {
                        w.write_line(&format!("{},", value.rust_name_outer()));
                    }
                });
                w.write_line("];");
                w.write_line("values");
            });

            if !self.lite_runtime {
                w.write_line("");
                w.def_fn(&format!("enum_descriptor_static() -> &'static ::protobuf::reflect::EnumDescriptor"), |w| {
                    w.lazy_static_decl_get("descriptor", "::protobuf::reflect::EnumDescriptor", |w| {
                        let ref type_name = self.type_name;
                        w.write_line(&format!(
                            "::protobuf::reflect::EnumDescriptor::new::<{}>(\"{}\", file_descriptor_proto())",
                            type_name,
                            type_name));
                    });
                });
            }
        });
    }

    fn write_impl_value(&self, w: &mut CodeWriter) {
        w.impl_for_block("::protobuf::reflect::ProtobufValue", &self.type_name, |_w| {
        })
    }

    fn write_impl_eq(&self, w: &mut CodeWriter) {
        assert!(self.allow_alias());
        w.impl_for_block("::std::cmp::PartialEq", &self.type_name, |w| {
            w.def_fn("eq(&self, other: &Self) -> bool", |w| {
                w.write_line("self.value() == other.value()");
            });
        });
    }

    fn write_impl_hash(&self, w: &mut CodeWriter) {
        assert!(self.allow_alias());
        w.impl_for_block("::std::hash::Hash", &self.type_name, |w| {
            w.def_fn("hash<H : ::std::hash::Hasher>(&self, state: &mut H)", |w| {
                w.write_line("state.write_i32(self.value())");
            });
        });
    }

    fn write_impl_default(&self, w: &mut CodeWriter) {
        let first_value = &self.enum_with_scope.values()[0];
        if first_value.get_number() != 0 {
            // This warning is emitted only for proto2
            // (because in proto3 first enum variant number is always 0).
            // `Default` implemented unconditionally to simplify certain
            // generic operations, e. g. reading a map.
            // Also, note that even in proto2 some operations fallback to
            // first enum value, e. g. `get_xxx` for unset field,
            // so this implementation is not completely unreasonable.
            w.comment("Note, `Default` is implemented although default value is not 0");
        }
        w.impl_for_block("::std::default::Default", &self.type_name, |w| {
            w.def_fn("default() -> Self", |w| {
                w.write_line(&format!(
                    "{}::{}",
                    &self.type_name,
                    &first_value.rust_name()
                ))
            });
        });
    }
}
<|endoftext|>


<commit_before>use crate::{
    frame::*,
    protocol::{
        *,
        basic::parse_properties,
    },
    types::parsing::*,
};

use nom::{
    bytes::streaming::{tag, take},
    combinator::{all_consuming, flat_map, map, map_opt, map_res},
    error::context,
};

/// Parse a channel id
pub fn parse_channel(i: &[u8]) -> ParserResult<'_, AMQPChannel> {
    context("parse_channel", map(parse_id, From::from))(i)
}

/// Parse the protocol header
pub fn parse_protocol_header(i: &[u8]) -> ParserResult<'_, ()> {
    context("parse_protocol_header", map(flat_map(tag(metadata::NAME.as_bytes()), |_| flat_map(tag(&[0]), |_| tag(&[metadata::MAJOR_VERSION, metadata::MINOR_VERSION, metadata::REVISION]))), |_| ()))(i)
}

/// Parse the frame type
pub fn parse_frame_type(i: &[u8]) -> ParserResult<'_, AMQPFrameType> {
    context("parse_frame_type", map_opt(parse_short_short_uint, |method| match method {
        constants::FRAME_METHOD    => Some(AMQPFrameType::Method),
        constants::FRAME_HEADER    => Some(AMQPFrameType::Header),
        constants::FRAME_BODY      => Some(AMQPFrameType::Body),
        constants::FRAME_HEARTBEAT => Some(AMQPFrameType::Heartbeat),
        _                          => None,
    }))(i)
}

/// Parse a full AMQP Frame (with contents)
pub fn parse_frame(i: &[u8]) -> ParserResult<'_, AMQPFrame> {
    context("parse_frame", map_res(parse_raw_frame, |raw| match raw.frame_type {
        AMQPFrameType::Method    => all_consuming(parse_class)(raw.payload).map(|(_, m)| AMQPFrame::Method(raw.channel_id, m)),
        AMQPFrameType::Header    => all_consuming(parse_content_header)(raw.payload).map(|(_, h)| AMQPFrame::Header(raw.channel_id, h.class_id, Box::new(h))),
        AMQPFrameType::Body      => Ok(AMQPFrame::Body(raw.channel_id, Vec::from(raw.payload))),
        AMQPFrameType::Heartbeat => Ok(AMQPFrame::Heartbeat(raw.channel_id)),
    }))(i)
}

/// Parse a raw AMQP frame
pub fn parse_raw_frame(i: &[u8]) -> ParserResult<'_, AMQPRawFrame<'_>> {
    context("parse_raw_frame", flat_map(parse_frame_type, |frame_type| flat_map(parse_id, move |channel_id| flat_map(parse_long_uint, move |size| flat_map(take(size), move |payload| map(tag(&[constants::FRAME_END]), move |_| AMQPRawFrame { frame_type, channel_id, size, payload }))))))(i)
}

/// Parse a content header frame
pub fn parse_content_header(i: &[u8]) -> ParserResult<'_, AMQPContentHeader> {
    context("parse_content_header", flat_map(parse_id, |class_id| flat_map(parse_short_uint, move |weight| flat_map(parse_long_long_uint, move |body_size| map(context("parse_propertes", parse_properties), move |properties| AMQPContentHeader { class_id, weight, body_size, properties })))))(i)
}
<commit_msg>protocol: simplify frames parsing<commit_after>use crate::{
    frame::*,
    protocol::{
        *,
        basic::parse_properties,
    },
    types::parsing::*,
};

use nom::{
    bytes::streaming::{tag, take},
    combinator::{all_consuming, flat_map, map, map_opt, map_res},
    error::context,
    sequence::{pair, tuple},
};

/// Parse a channel id
pub fn parse_channel(i: &[u8]) -> ParserResult<'_, AMQPChannel> {
    context("parse_channel", map(parse_id, From::from))(i)
}

/// Parse the protocol header
pub fn parse_protocol_header(i: &[u8]) -> ParserResult<'_, ()> {
    context("parse_protocol_header", map(tuple((tag(metadata::NAME.as_bytes()), tag(&[0]), tag(&[metadata::MAJOR_VERSION, metadata::MINOR_VERSION, metadata::REVISION]))), |_| ()))(i)
}

/// Parse the frame type
pub fn parse_frame_type(i: &[u8]) -> ParserResult<'_, AMQPFrameType> {
    context("parse_frame_type", map_opt(parse_short_short_uint, |method| match method {
        constants::FRAME_METHOD    => Some(AMQPFrameType::Method),
        constants::FRAME_HEADER    => Some(AMQPFrameType::Header),
        constants::FRAME_BODY      => Some(AMQPFrameType::Body),
        constants::FRAME_HEARTBEAT => Some(AMQPFrameType::Heartbeat),
        _                          => None,
    }))(i)
}

/// Parse a full AMQP Frame (with contents)
pub fn parse_frame(i: &[u8]) -> ParserResult<'_, AMQPFrame> {
    context("parse_frame", map_res(parse_raw_frame, |raw| match raw.frame_type {
        AMQPFrameType::Method    => all_consuming(parse_class)(raw.payload).map(|(_, m)| AMQPFrame::Method(raw.channel_id, m)),
        AMQPFrameType::Header    => all_consuming(parse_content_header)(raw.payload).map(|(_, h)| AMQPFrame::Header(raw.channel_id, h.class_id, Box::new(h))),
        AMQPFrameType::Body      => Ok(AMQPFrame::Body(raw.channel_id, Vec::from(raw.payload))),
        AMQPFrameType::Heartbeat => Ok(AMQPFrame::Heartbeat(raw.channel_id)),
    }))(i)
}

/// Parse a raw AMQP frame
pub fn parse_raw_frame(i: &[u8]) -> ParserResult<'_, AMQPRawFrame<'_>> {
    context("parse_raw_frame", flat_map(tuple((parse_frame_type, parse_id, parse_long_uint)), move |(frame_type, channel_id, size)| map(pair(take(size), tag(&[constants::FRAME_END])), move |(payload, _)| AMQPRawFrame { frame_type, channel_id, size, payload })))(i)
}

/// Parse a content header frame
pub fn parse_content_header(i: &[u8]) -> ParserResult<'_, AMQPContentHeader> {
    context("parse_content_header", map(tuple((parse_id, parse_short_uint, parse_long_long_uint, context("parse_propertes", parse_properties))), |(class_id, weight, body_size, properties)| AMQPContentHeader { class_id, weight, body_size, properties }))(i)
}
<|endoftext|>


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

// run-pass

const FOO: &[u8] = b"foo";
const BAR: &[u8] = &[1, 2, 3];

const BOO: &i32 = &42;

fn main() {
    match &[1u8, 2, 3] as &[u8] {
        FOO => panic!("a"),
        BAR => println!("b"),
        _ => panic!("c"),
    }

    match b"foo" as &[u8] {
        FOO => println!("a"),
        BAR => panic!("b"),
        _ => panic!("c"),
    }

    match &43 {
        &42 => panic!(),
        BOO => panic!(),
        _ => println!("d"),
    }
}
<commit_msg>Diagnostics improvement!<commit_after>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// run-pass

const FOO: &[u8] = b"foo";
const BAR: &[u8] = &[1, 2, 3];

const BOO: &i32 = &42;

fn main() {
    match &[1u8, 2, 3] as &[u8] {
        FOO => panic!("a"),
        BAR => println!("b"),
        _ => panic!("c"),
    }

    match b"foo" as &[u8] {
        FOO => println!("a"),
        BAR => panic!("b"),
        _ => panic!("c"),
    }

    #[allow(unreachable_patterns)]
    match &43 {
        &42 => panic!(),
        BOO => panic!(), // pattern is unreachable
        _ => println!("d"),
    }
}
<|endoftext|>


<commit_before>// This file is part of the uutils coreutils package.
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.

//! Common functions to manage permissions

use crate::error::strip_errno;
use crate::error::UResult;
use crate::error::USimpleError;
pub use crate::features::entries;
use crate::fs::resolve_relative_path;
use crate::show_error;
use clap::App;
use clap::Arg;
use clap::ArgMatches;
use libc::{self, gid_t, uid_t};
use walkdir::WalkDir;

use std::io::Error as IOError;
use std::io::Result as IOResult;

use std::ffi::CString;
use std::fs::Metadata;
use std::os::unix::fs::MetadataExt;

use std::os::unix::ffi::OsStrExt;
use std::path::Path;

/// The various level of verbosity
#[derive(PartialEq, Clone, Debug)]
pub enum VerbosityLevel {
    Silent,
    Changes,
    Verbose,
    Normal,
}
#[derive(PartialEq, Clone, Debug)]
pub struct Verbosity {
    pub groups_only: bool,
    pub level: VerbosityLevel,
}

/// Actually perform the change of owner on a path
fn chown<P: AsRef<Path>>(path: P, uid: uid_t, gid: gid_t, follow: bool) -> IOResult<()> {
    let path = path.as_ref();
    let s = CString::new(path.as_os_str().as_bytes()).unwrap();
    let ret = unsafe {
        if follow {
            libc::chown(s.as_ptr(), uid, gid)
        } else {
            libc::lchown(s.as_ptr(), uid, gid)
        }
    };
    if ret == 0 {
        Ok(())
    } else {
        Err(IOError::last_os_error())
    }
}

/// Perform the change of owner on a path
/// with the various options
/// and error messages management
pub fn wrap_chown<P: AsRef<Path>>(
    path: P,
    meta: &Metadata,
    dest_uid: Option<u32>,
    dest_gid: Option<u32>,
    follow: bool,
    verbosity: Verbosity,
) -> Result<String, String> {
    let dest_uid = dest_uid.unwrap_or_else(|| meta.uid());
    let dest_gid = dest_gid.unwrap_or_else(|| meta.gid());
    let path = path.as_ref();
    let mut out: String = String::new();

    if let Err(e) = chown(path, dest_uid, dest_gid, follow) {
        match verbosity.level {
            VerbosityLevel::Silent => (),
            level => {
                out = format!(
                    "changing {} of '{}': {}",
                    if verbosity.groups_only {
                        "group"
                    } else {
                        "ownership"
                    },
                    path.display(),
                    e
                );
                if level == VerbosityLevel::Verbose {
                    out = if verbosity.groups_only {
                        format!(
                            "{}\nfailed to change group of '{}' from {} to {}",
                            out,
                            path.display(),
                            entries::gid2grp(meta.gid()).unwrap(),
                            entries::gid2grp(dest_gid).unwrap()
                        )
                    } else {
                        format!(
                            "{}\nfailed to change ownership of '{}' from {}:{} to {}:{}",
                            out,
                            path.display(),
                            entries::uid2usr(meta.uid()).unwrap(),
                            entries::gid2grp(meta.gid()).unwrap(),
                            entries::uid2usr(dest_uid).unwrap(),
                            entries::gid2grp(dest_gid).unwrap()
                        )
                    };
                };
            }
        }
        return Err(out);
    } else {
        let changed = dest_uid != meta.uid() || dest_gid != meta.gid();
        if changed {
            match verbosity.level {
                VerbosityLevel::Changes | VerbosityLevel::Verbose => {
                    out = if verbosity.groups_only {
                        format!(
                            "changed group of '{}' from {} to {}",
                            path.display(),
                            entries::gid2grp(meta.gid()).unwrap(),
                            entries::gid2grp(dest_gid).unwrap()
                        )
                    } else {
                        format!(
                            "changed ownership of '{}' from {}:{} to {}:{}",
                            path.display(),
                            entries::uid2usr(meta.uid()).unwrap(),
                            entries::gid2grp(meta.gid()).unwrap(),
                            entries::uid2usr(dest_uid).unwrap(),
                            entries::gid2grp(dest_gid).unwrap()
                        )
                    };
                }
                _ => (),
            };
        } else if verbosity.level == VerbosityLevel::Verbose {
            out = if verbosity.groups_only {
                format!(
                    "group of '{}' retained as {}",
                    path.display(),
                    entries::gid2grp(dest_gid).unwrap_or_default()
                )
            } else {
                format!(
                    "ownership of '{}' retained as {}:{}",
                    path.display(),
                    entries::uid2usr(dest_uid).unwrap(),
                    entries::gid2grp(dest_gid).unwrap()
                )
            };
        }
    }
    Ok(out)
}

pub enum IfFrom {
    All,
    User(u32),
    Group(u32),
    UserGroup(u32, u32),
}

#[derive(PartialEq, Eq)]
pub enum TraverseSymlinks {
    None,
    First,
    All,
}

pub struct ChownExecutor {
    pub dest_uid: Option<u32>,
    pub dest_gid: Option<u32>,
    pub traverse_symlinks: TraverseSymlinks,
    pub verbosity: Verbosity,
    pub filter: IfFrom,
    pub files: Vec<String>,
    pub recursive: bool,
    pub preserve_root: bool,
    // Must be true if traverse_symlinks is not None
    pub dereference: bool,
}

impl ChownExecutor {
    pub fn exec(&self) -> UResult<()> {
        let mut ret = 0;
        for f in &self.files {
            ret |= self.traverse(f);
        }
        if ret != 0 {
            return Err(ret.into());
        }
        Ok(())
    }

    fn traverse<P: AsRef<Path>>(&self, root: P) -> i32 {
        let path = root.as_ref();
        let meta = match self.obtain_meta(path, self.dereference) {
            Some(m) => m,
            _ => return 1,
        };

        // Prohibit only if:
        // (--preserve-root and -R present) &&
        // (
        //     (argument is not symlink && resolved to be '/') ||
        //     (argument is symlink && should follow argument && resolved to be '/')
        // )
        if self.recursive && self.preserve_root {
            let may_exist = if self.dereference {
                path.canonicalize().ok()
            } else {
                let real = resolve_relative_path(path);
                if real.is_dir() {
                    Some(real.canonicalize().expect("failed to get real path"))
                } else {
                    Some(real.into_owned())
                }
            };

            if let Some(p) = may_exist {
                if p.parent().is_none() {
                    show_error!("it is dangerous to operate recursively on '/'");
                    show_error!("use --no-preserve-root to override this failsafe");
                    return 1;
                }
            }
        }

        let ret = if self.matched(meta.uid(), meta.gid()) {
            match wrap_chown(
                path,
                &meta,
                self.dest_uid,
                self.dest_gid,
                self.dereference,
                self.verbosity.clone(),
            ) {
                Ok(n) => {
                    if !n.is_empty() {
                        show_error!("{}", n);
                    }
                    0
                }
                Err(e) => {
                    if self.verbosity.level != VerbosityLevel::Silent {
                        show_error!("{}", e);
                    }
                    1
                }
            }
        } else {
            0
        };

        if !self.recursive {
            ret
        } else {
            ret | self.dive_into(&root)
        }
    }

    fn dive_into<P: AsRef<Path>>(&self, root: P) -> i32 {
        let root = root.as_ref();

        // walkdir always dereferences the root directory, so we have to check it ourselves
        // TODO: replace with `root.is_symlink()` once it is stable
        if self.traverse_symlinks == TraverseSymlinks::None
            && std::fs::symlink_metadata(root)
                .map(|m| m.file_type().is_symlink())
                .unwrap_or(false)
        {
            return 0;
        }

        let mut ret = 0;
        let mut iterator = WalkDir::new(root)
            .follow_links(self.traverse_symlinks == TraverseSymlinks::All)
            .min_depth(1)
            .into_iter();
        // We can't use a for loop because we need to manipulate the iterator inside the loop.
        while let Some(entry) = iterator.next() {
            let entry = match entry {
                Err(e) => {
                    ret = 1;
                    if let Some(path) = e.path() {
                        show_error!(
                            "cannot access '{}': {}",
                            path.display(),
                            if let Some(error) = e.io_error() {
                                strip_errno(error)
                            } else {
                                "Too many levels of symbolic links".into()
                            }
                        )
                    } else {
                        show_error!("{}", e)
                    }
                    continue;
                }
                Ok(entry) => entry,
            };
            let path = entry.path();
            let meta = match self.obtain_meta(path, self.dereference) {
                Some(m) => m,
                _ => {
                    ret = 1;
                    if entry.file_type().is_dir() {
                        // Instruct walkdir to skip this directory to avoid getting another error
                        // when walkdir tries to query the children of this directory.
                        iterator.skip_current_dir();
                    }
                    continue;
                }
            };

            if !self.matched(meta.uid(), meta.gid()) {
                continue;
            }

            ret = match wrap_chown(
                path,
                &meta,
                self.dest_uid,
                self.dest_gid,
                self.dereference,
                self.verbosity.clone(),
            ) {
                Ok(n) => {
                    if !n.is_empty() {
                        show_error!("{}", n);
                    }
                    0
                }
                Err(e) => {
                    if self.verbosity.level != VerbosityLevel::Silent {
                        show_error!("{}", e);
                    }
                    1
                }
            }
        }
        ret
    }

    fn obtain_meta<P: AsRef<Path>>(&self, path: P, follow: bool) -> Option<Metadata> {
        let path = path.as_ref();
        let meta = if follow {
            path.metadata()
        } else {
            path.symlink_metadata()
        };
        match meta {
            Err(e) => {
                match self.verbosity.level {
                    VerbosityLevel::Silent => (),
                    _ => show_error!("cannot access '{}': {}", path.display(), strip_errno(&e)),
                }
                None
            }
            Ok(meta) => Some(meta),
        }
    }

    #[inline]
    fn matched(&self, uid: uid_t, gid: gid_t) -> bool {
        match self.filter {
            IfFrom::All => true,
            IfFrom::User(u) => u == uid,
            IfFrom::Group(g) => g == gid,
            IfFrom::UserGroup(u, g) => u == uid && g == gid,
        }
    }
}

pub mod options {
    pub mod verbosity {
        pub const CHANGES: &str = "changes";
        pub const QUIET: &str = "quiet";
        pub const SILENT: &str = "silent";
        pub const VERBOSE: &str = "verbose";
    }
    pub mod preserve_root {
        pub const PRESERVE: &str = "preserve-root";
        pub const NO_PRESERVE: &str = "no-preserve-root";
    }
    pub mod dereference {
        pub const DEREFERENCE: &str = "dereference";
        pub const NO_DEREFERENCE: &str = "no-dereference";
    }
    pub const FROM: &str = "from";
    pub const RECURSIVE: &str = "recursive";
    pub mod traverse {
        pub const TRAVERSE: &str = "H";
        pub const NO_TRAVERSE: &str = "P";
        pub const EVERY: &str = "L";
    }
    pub const REFERENCE: &str = "reference";
    pub const ARG_OWNER: &str = "OWNER";
    pub const ARG_GROUP: &str = "GROUP";
    pub const ARG_FILES: &str = "FILE";
}

type GidUidFilterParser<'a> = fn(&ArgMatches<'a>) -> UResult<(Option<u32>, Option<u32>, IfFrom)>;

/// Base implementation for `chgrp` and `chown`.
///
/// An argument called `add_arg_if_not_reference` will be added to `app` if
/// `args` does not contain the `--reference` option.
/// `parse_gid_uid_and_filter` will be called to obtain the target gid and uid, and the filter,
/// from `ArgMatches`.
/// `groups_only` determines whether verbose output will only mention the group.
pub fn chown_base<'a>(
    mut app: App<'a, 'a>,
    args: impl crate::Args,
    add_arg_if_not_reference: &'a str,
    parse_gid_uid_and_filter: GidUidFilterParser<'a>,
    groups_only: bool,
) -> UResult<()> {
    let args: Vec<_> = args.collect();
    let mut reference = false;
    let mut help = false;
    // stop processing options on --
    for arg in args.iter().take_while(|s| *s != "--") {
        if arg.to_string_lossy().starts_with("--reference=") || arg == "--reference" {
            reference = true;
        } else if arg == "--help" {
            // we stop processing once we see --help,
            // as it doesn't matter if we've seen reference or not
            help = true;
            break;
        }
    }

    if help || !reference {
        // add both positional arguments
        // arg_group is only required if
        app = app.arg(
            Arg::with_name(add_arg_if_not_reference)
                .value_name(add_arg_if_not_reference)
                .required(true)
                .takes_value(true)
                .multiple(false),
        )
    }
    app = app.arg(
        Arg::with_name(options::ARG_FILES)
            .value_name(options::ARG_FILES)
            .multiple(true)
            .takes_value(true)
            .required(true)
            .min_values(1),
    );
    let matches = app.get_matches_from(args);

    let files: Vec<String> = matches
        .values_of(options::ARG_FILES)
        .map(|v| v.map(ToString::to_string).collect())
        .unwrap_or_default();

    let preserve_root = matches.is_present(options::preserve_root::PRESERVE);

    let mut dereference = if matches.is_present(options::dereference::DEREFERENCE) {
        Some(true)
    } else if matches.is_present(options::dereference::NO_DEREFERENCE) {
        Some(false)
    } else {
        None
    };

    let mut traverse_symlinks = if matches.is_present(options::traverse::TRAVERSE) {
        TraverseSymlinks::First
    } else if matches.is_present(options::traverse::EVERY) {
        TraverseSymlinks::All
    } else {
        TraverseSymlinks::None
    };

    let recursive = matches.is_present(options::RECURSIVE);
    if recursive {
        if traverse_symlinks == TraverseSymlinks::None {
            if dereference == Some(true) {
                return Err(USimpleError::new(1, "-R --dereference requires -H or -L"));
            }
            dereference = Some(false);
        }
    } else {
        traverse_symlinks = TraverseSymlinks::None;
    }

    let verbosity_level = if matches.is_present(options::verbosity::CHANGES) {
        VerbosityLevel::Changes
    } else if matches.is_present(options::verbosity::SILENT)
        || matches.is_present(options::verbosity::QUIET)
    {
        VerbosityLevel::Silent
    } else if matches.is_present(options::verbosity::VERBOSE) {
        VerbosityLevel::Verbose
    } else {
        VerbosityLevel::Normal
    };
    let (dest_gid, dest_uid, filter) = parse_gid_uid_and_filter(&matches)?;

    let executor = ChownExecutor {
        traverse_symlinks,
        dest_gid,
        dest_uid,
        verbosity: Verbosity {
            groups_only,
            level: verbosity_level,
        },
        recursive,
        dereference: dereference.unwrap_or(true),
        preserve_root,
        files,
        filter,
    };
    executor.exec()
}
<commit_msg>uucore/perms: add more information to an error message<commit_after>// This file is part of the uutils coreutils package.
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.

//! Common functions to manage permissions

use crate::error::strip_errno;
use crate::error::UResult;
use crate::error::USimpleError;
pub use crate::features::entries;
use crate::fs::resolve_relative_path;
use crate::show_error;
use clap::App;
use clap::Arg;
use clap::ArgMatches;
use libc::{self, gid_t, uid_t};
use walkdir::WalkDir;

use std::io::Error as IOError;
use std::io::Result as IOResult;

use std::ffi::CString;
use std::fs::Metadata;
use std::os::unix::fs::MetadataExt;

use std::os::unix::ffi::OsStrExt;
use std::path::Path;

/// The various level of verbosity
#[derive(PartialEq, Clone, Debug)]
pub enum VerbosityLevel {
    Silent,
    Changes,
    Verbose,
    Normal,
}
#[derive(PartialEq, Clone, Debug)]
pub struct Verbosity {
    pub groups_only: bool,
    pub level: VerbosityLevel,
}

/// Actually perform the change of owner on a path
fn chown<P: AsRef<Path>>(path: P, uid: uid_t, gid: gid_t, follow: bool) -> IOResult<()> {
    let path = path.as_ref();
    let s = CString::new(path.as_os_str().as_bytes()).unwrap();
    let ret = unsafe {
        if follow {
            libc::chown(s.as_ptr(), uid, gid)
        } else {
            libc::lchown(s.as_ptr(), uid, gid)
        }
    };
    if ret == 0 {
        Ok(())
    } else {
        Err(IOError::last_os_error())
    }
}

/// Perform the change of owner on a path
/// with the various options
/// and error messages management
pub fn wrap_chown<P: AsRef<Path>>(
    path: P,
    meta: &Metadata,
    dest_uid: Option<u32>,
    dest_gid: Option<u32>,
    follow: bool,
    verbosity: Verbosity,
) -> Result<String, String> {
    let dest_uid = dest_uid.unwrap_or_else(|| meta.uid());
    let dest_gid = dest_gid.unwrap_or_else(|| meta.gid());
    let path = path.as_ref();
    let mut out: String = String::new();

    if let Err(e) = chown(path, dest_uid, dest_gid, follow) {
        match verbosity.level {
            VerbosityLevel::Silent => (),
            level => {
                out = format!(
                    "changing {} of '{}': {}",
                    if verbosity.groups_only {
                        "group"
                    } else {
                        "ownership"
                    },
                    path.display(),
                    e
                );
                if level == VerbosityLevel::Verbose {
                    out = if verbosity.groups_only {
                        format!(
                            "{}\nfailed to change group of '{}' from {} to {}",
                            out,
                            path.display(),
                            entries::gid2grp(meta.gid()).unwrap(),
                            entries::gid2grp(dest_gid).unwrap()
                        )
                    } else {
                        format!(
                            "{}\nfailed to change ownership of '{}' from {}:{} to {}:{}",
                            out,
                            path.display(),
                            entries::uid2usr(meta.uid()).unwrap(),
                            entries::gid2grp(meta.gid()).unwrap(),
                            entries::uid2usr(dest_uid).unwrap(),
                            entries::gid2grp(dest_gid).unwrap()
                        )
                    };
                };
            }
        }
        return Err(out);
    } else {
        let changed = dest_uid != meta.uid() || dest_gid != meta.gid();
        if changed {
            match verbosity.level {
                VerbosityLevel::Changes | VerbosityLevel::Verbose => {
                    out = if verbosity.groups_only {
                        format!(
                            "changed group of '{}' from {} to {}",
                            path.display(),
                            entries::gid2grp(meta.gid()).unwrap(),
                            entries::gid2grp(dest_gid).unwrap()
                        )
                    } else {
                        format!(
                            "changed ownership of '{}' from {}:{} to {}:{}",
                            path.display(),
                            entries::uid2usr(meta.uid()).unwrap(),
                            entries::gid2grp(meta.gid()).unwrap(),
                            entries::uid2usr(dest_uid).unwrap(),
                            entries::gid2grp(dest_gid).unwrap()
                        )
                    };
                }
                _ => (),
            };
        } else if verbosity.level == VerbosityLevel::Verbose {
            out = if verbosity.groups_only {
                format!(
                    "group of '{}' retained as {}",
                    path.display(),
                    entries::gid2grp(dest_gid).unwrap_or_default()
                )
            } else {
                format!(
                    "ownership of '{}' retained as {}:{}",
                    path.display(),
                    entries::uid2usr(dest_uid).unwrap(),
                    entries::gid2grp(dest_gid).unwrap()
                )
            };
        }
    }
    Ok(out)
}

pub enum IfFrom {
    All,
    User(u32),
    Group(u32),
    UserGroup(u32, u32),
}

#[derive(PartialEq, Eq)]
pub enum TraverseSymlinks {
    None,
    First,
    All,
}

pub struct ChownExecutor {
    pub dest_uid: Option<u32>,
    pub dest_gid: Option<u32>,
    pub traverse_symlinks: TraverseSymlinks,
    pub verbosity: Verbosity,
    pub filter: IfFrom,
    pub files: Vec<String>,
    pub recursive: bool,
    pub preserve_root: bool,
    // Must be true if traverse_symlinks is not None
    pub dereference: bool,
}

impl ChownExecutor {
    pub fn exec(&self) -> UResult<()> {
        let mut ret = 0;
        for f in &self.files {
            ret |= self.traverse(f);
        }
        if ret != 0 {
            return Err(ret.into());
        }
        Ok(())
    }

    fn traverse<P: AsRef<Path>>(&self, root: P) -> i32 {
        let path = root.as_ref();
        let meta = match self.obtain_meta(path, self.dereference) {
            Some(m) => m,
            _ => return 1,
        };

        // Prohibit only if:
        // (--preserve-root and -R present) &&
        // (
        //     (argument is not symlink && resolved to be '/') ||
        //     (argument is symlink && should follow argument && resolved to be '/')
        // )
        if self.recursive && self.preserve_root {
            let may_exist = if self.dereference {
                path.canonicalize().ok()
            } else {
                let real = resolve_relative_path(path);
                if real.is_dir() {
                    Some(real.canonicalize().expect("failed to get real path"))
                } else {
                    Some(real.into_owned())
                }
            };

            if let Some(p) = may_exist {
                if p.parent().is_none() {
                    show_error!("it is dangerous to operate recursively on '/'");
                    show_error!("use --no-preserve-root to override this failsafe");
                    return 1;
                }
            }
        }

        let ret = if self.matched(meta.uid(), meta.gid()) {
            match wrap_chown(
                path,
                &meta,
                self.dest_uid,
                self.dest_gid,
                self.dereference,
                self.verbosity.clone(),
            ) {
                Ok(n) => {
                    if !n.is_empty() {
                        show_error!("{}", n);
                    }
                    0
                }
                Err(e) => {
                    if self.verbosity.level != VerbosityLevel::Silent {
                        show_error!("{}", e);
                    }
                    1
                }
            }
        } else {
            0
        };

        if !self.recursive {
            ret
        } else {
            ret | self.dive_into(&root)
        }
    }

    fn dive_into<P: AsRef<Path>>(&self, root: P) -> i32 {
        let root = root.as_ref();

        // walkdir always dereferences the root directory, so we have to check it ourselves
        // TODO: replace with `root.is_symlink()` once it is stable
        if self.traverse_symlinks == TraverseSymlinks::None
            && std::fs::symlink_metadata(root)
                .map(|m| m.file_type().is_symlink())
                .unwrap_or(false)
        {
            return 0;
        }

        let mut ret = 0;
        let mut iterator = WalkDir::new(root)
            .follow_links(self.traverse_symlinks == TraverseSymlinks::All)
            .min_depth(1)
            .into_iter();
        // We can't use a for loop because we need to manipulate the iterator inside the loop.
        while let Some(entry) = iterator.next() {
            let entry = match entry {
                Err(e) => {
                    ret = 1;
                    if let Some(path) = e.path() {
                        show_error!(
                            "cannot access '{}': {}",
                            path.display(),
                            if let Some(error) = e.io_error() {
                                strip_errno(error)
                            } else {
                                "Too many levels of symbolic links".into()
                            }
                        )
                    } else {
                        show_error!("{}", e)
                    }
                    continue;
                }
                Ok(entry) => entry,
            };
            let path = entry.path();
            let meta = match self.obtain_meta(path, self.dereference) {
                Some(m) => m,
                _ => {
                    ret = 1;
                    if entry.file_type().is_dir() {
                        // Instruct walkdir to skip this directory to avoid getting another error
                        // when walkdir tries to query the children of this directory.
                        iterator.skip_current_dir();
                    }
                    continue;
                }
            };

            if !self.matched(meta.uid(), meta.gid()) {
                continue;
            }

            ret = match wrap_chown(
                path,
                &meta,
                self.dest_uid,
                self.dest_gid,
                self.dereference,
                self.verbosity.clone(),
            ) {
                Ok(n) => {
                    if !n.is_empty() {
                        show_error!("{}", n);
                    }
                    0
                }
                Err(e) => {
                    if self.verbosity.level != VerbosityLevel::Silent {
                        show_error!("{}", e);
                    }
                    1
                }
            }
        }
        ret
    }

    fn obtain_meta<P: AsRef<Path>>(&self, path: P, follow: bool) -> Option<Metadata> {
        let path = path.as_ref();
        let meta = if follow {
            path.metadata()
        } else {
            path.symlink_metadata()
        };
        match meta {
            Err(e) => {
                match self.verbosity.level {
                    VerbosityLevel::Silent => (),
                    _ => show_error!(
                        "cannot {} '{}': {}",
                        if follow { "dereference" } else { "access" },
                        path.display(),
                        strip_errno(&e)
                    ),
                }
                None
            }
            Ok(meta) => Some(meta),
        }
    }

    #[inline]
    fn matched(&self, uid: uid_t, gid: gid_t) -> bool {
        match self.filter {
            IfFrom::All => true,
            IfFrom::User(u) => u == uid,
            IfFrom::Group(g) => g == gid,
            IfFrom::UserGroup(u, g) => u == uid && g == gid,
        }
    }
}

pub mod options {
    pub mod verbosity {
        pub const CHANGES: &str = "changes";
        pub const QUIET: &str = "quiet";
        pub const SILENT: &str = "silent";
        pub const VERBOSE: &str = "verbose";
    }
    pub mod preserve_root {
        pub const PRESERVE: &str = "preserve-root";
        pub const NO_PRESERVE: &str = "no-preserve-root";
    }
    pub mod dereference {
        pub const DEREFERENCE: &str = "dereference";
        pub const NO_DEREFERENCE: &str = "no-dereference";
    }
    pub const FROM: &str = "from";
    pub const RECURSIVE: &str = "recursive";
    pub mod traverse {
        pub const TRAVERSE: &str = "H";
        pub const NO_TRAVERSE: &str = "P";
        pub const EVERY: &str = "L";
    }
    pub const REFERENCE: &str = "reference";
    pub const ARG_OWNER: &str = "OWNER";
    pub const ARG_GROUP: &str = "GROUP";
    pub const ARG_FILES: &str = "FILE";
}

type GidUidFilterParser<'a> = fn(&ArgMatches<'a>) -> UResult<(Option<u32>, Option<u32>, IfFrom)>;

/// Base implementation for `chgrp` and `chown`.
///
/// An argument called `add_arg_if_not_reference` will be added to `app` if
/// `args` does not contain the `--reference` option.
/// `parse_gid_uid_and_filter` will be called to obtain the target gid and uid, and the filter,
/// from `ArgMatches`.
/// `groups_only` determines whether verbose output will only mention the group.
pub fn chown_base<'a>(
    mut app: App<'a, 'a>,
    args: impl crate::Args,
    add_arg_if_not_reference: &'a str,
    parse_gid_uid_and_filter: GidUidFilterParser<'a>,
    groups_only: bool,
) -> UResult<()> {
    let args: Vec<_> = args.collect();
    let mut reference = false;
    let mut help = false;
    // stop processing options on --
    for arg in args.iter().take_while(|s| *s != "--") {
        if arg.to_string_lossy().starts_with("--reference=") || arg == "--reference" {
            reference = true;
        } else if arg == "--help" {
            // we stop processing once we see --help,
            // as it doesn't matter if we've seen reference or not
            help = true;
            break;
        }
    }

    if help || !reference {
        // add both positional arguments
        // arg_group is only required if
        app = app.arg(
            Arg::with_name(add_arg_if_not_reference)
                .value_name(add_arg_if_not_reference)
                .required(true)
                .takes_value(true)
                .multiple(false),
        )
    }
    app = app.arg(
        Arg::with_name(options::ARG_FILES)
            .value_name(options::ARG_FILES)
            .multiple(true)
            .takes_value(true)
            .required(true)
            .min_values(1),
    );
    let matches = app.get_matches_from(args);

    let files: Vec<String> = matches
        .values_of(options::ARG_FILES)
        .map(|v| v.map(ToString::to_string).collect())
        .unwrap_or_default();

    let preserve_root = matches.is_present(options::preserve_root::PRESERVE);

    let mut dereference = if matches.is_present(options::dereference::DEREFERENCE) {
        Some(true)
    } else if matches.is_present(options::dereference::NO_DEREFERENCE) {
        Some(false)
    } else {
        None
    };

    let mut traverse_symlinks = if matches.is_present(options::traverse::TRAVERSE) {
        TraverseSymlinks::First
    } else if matches.is_present(options::traverse::EVERY) {
        TraverseSymlinks::All
    } else {
        TraverseSymlinks::None
    };

    let recursive = matches.is_present(options::RECURSIVE);
    if recursive {
        if traverse_symlinks == TraverseSymlinks::None {
            if dereference == Some(true) {
                return Err(USimpleError::new(1, "-R --dereference requires -H or -L"));
            }
            dereference = Some(false);
        }
    } else {
        traverse_symlinks = TraverseSymlinks::None;
    }

    let verbosity_level = if matches.is_present(options::verbosity::CHANGES) {
        VerbosityLevel::Changes
    } else if matches.is_present(options::verbosity::SILENT)
        || matches.is_present(options::verbosity::QUIET)
    {
        VerbosityLevel::Silent
    } else if matches.is_present(options::verbosity::VERBOSE) {
        VerbosityLevel::Verbose
    } else {
        VerbosityLevel::Normal
    };
    let (dest_gid, dest_uid, filter) = parse_gid_uid_and_filter(&matches)?;

    let executor = ChownExecutor {
        traverse_symlinks,
        dest_gid,
        dest_uid,
        verbosity: Verbosity {
            groups_only,
            level: verbosity_level,
        },
        recursive,
        dereference: dereference.unwrap_or(true),
        preserve_root,
        files,
        filter,
    };
    executor.exec()
}
<|endoftext|>


<commit_before>extern crate sourcemap;

use std::collections::HashMap;
use sourcemap::SourceMapIndex;

#[test]
fn test_basic_indexed_sourcemap() {
    let input: &[_] = br#"{
        "version": 3,
        "file": "min.js",
        "sections": [
            {
                "offset": {
                    "line": 0,
                    "column": 0
                },
                "map": {
                    "version":3,
                    "sources":["file1.js"],
                    "names":["add","a","b"],
                    "mappings":"AAAA,QAASA,KAAIC,EAAGC,GACf,YACA,OAAOD,GAAIC",
                    "file":"file1.min.js"
                }
            },
            {
                "offset": {
                    "line": 1,
                    "column": 0
                },
                "map": {
                    "version":3,
                    "sources":["file2.js"],
                    "names":["multiply","a","b","divide","add","c","e","Raven","captureException"],
                    "mappings":"AAAA,QAASA,UAASC,EAAGC,GACpB,YACA,OAAOD,GAAIC,EAEZ,QAASC,QAAOF,EAAGC,GAClB,YACA,KACC,MAAOF,UAASI,IAAIH,EAAGC,GAAID,EAAGC,GAAKG,EAClC,MAAOC,GACRC,MAAMC,iBAAiBF",
                    "file":"file2.min.js"
                }
            }
        ]
    }"#;

    let f1 : Vec<_> = "function add(a, b) {\n \"use strict\";\n \
                       return a + b; // fôo\n}".lines().collect();
    let f2 : Vec<_> = "function multiply(a, b) {\n \"use strict\";\n \
                       return a * b;\n}\nfunction divide(a, b) {\n \
                       \"use strict\";\n try {\n  return multiply(add(a, \
                       b), a, b) / c;\n } catch (e) {\n  \
                       Raven.captureException(e);\n }\n}".lines().collect();

    let mut files = HashMap::new();
    files.insert("file1.js", f1);
    files.insert("file2.js", f2);

    let ism = SourceMapIndex::from_reader(input).unwrap();
    let flat_map = ism.flatten().unwrap();

    for token in flat_map.tokens() {
        let src = files.get(token.get_source().unwrap()).unwrap();
        if let Some(name) = token.get_name() {
            let line = src[token.get_src_line() as usize];
            let idx = token.get_src_col() as usize;
            let span = &line[idx..idx + name.len()];
            assert_eq!(span, name);
        }
    }
}
<commit_msg>Added a test for source contents setting<commit_after>extern crate sourcemap;

use std::collections::HashMap;
use sourcemap::SourceMapIndex;

#[test]
fn test_basic_indexed_sourcemap() {
    let input: &[_] = br#"{
        "version": 3,
        "file": "min.js",
        "sections": [
            {
                "offset": {
                    "line": 0,
                    "column": 0
                },
                "map": {
                    "version":3,
                    "sources":["file1.js"],
                    "names":["add","a","b"],
                    "mappings":"AAAA,QAASA,KAAIC,EAAGC,GACf,YACA,OAAOD,GAAIC",
                    "file":"file1.min.js"
                }
            },
            {
                "offset": {
                    "line": 1,
                    "column": 0
                },
                "map": {
                    "version":3,
                    "sources":["file2.js"],
                    "names":["multiply","a","b","divide","add","c","e","Raven","captureException"],
                    "mappings":"AAAA,QAASA,UAASC,EAAGC,GACpB,YACA,OAAOD,GAAIC,EAEZ,QAASC,QAAOF,EAAGC,GAClB,YACA,KACC,MAAOF,UAASI,IAAIH,EAAGC,GAAID,EAAGC,GAAKG,EAClC,MAAOC,GACRC,MAAMC,iBAAiBF",
                    "file":"file2.min.js"
                }
            }
        ]
    }"#;

    let f1 = "function add(a, b) {\n \"use strict\";\n \
              return a + b; // fôo\n}";
    let f2 = "function multiply(a, b) {\n \"use strict\";\n \
              return a * b;\n}\nfunction divide(a, b) {\n \
              \"use strict\";\n try {\n  return multiply(add(a, \
              b), a, b) / c;\n } catch (e) {\n  \
              Raven.captureException(e);\n }\n}";

    let mut raw_files = HashMap::new();
    raw_files.insert("file1.js", f1);
    raw_files.insert("file2.js", f2);

    let mut files : HashMap<_, Vec<_>> = HashMap::new();
    files.insert("file1.js", f1.lines().collect());
    files.insert("file2.js", f2.lines().collect());

    let mut ism = SourceMapIndex::from_reader(input).unwrap();
    for section_id in 0..ism.get_section_count() {
        let mut section = ism.get_section_mut(section_id).unwrap();
        let mut map = section.get_sourcemap_mut().unwrap();
        let contents = {
            let filename = map.get_source(0).unwrap();
            *raw_files.get(filename).unwrap()
        };
        map.set_source_contents(0, Some(contents));
    }
    let flat_map = ism.flatten().unwrap();

    let mut out : Vec<u8> = vec![];
    flat_map.to_writer(&mut out);
    println!("{}", String::from_utf8(out).unwrap());

    for token in flat_map.tokens() {
        let src = files.get(token.get_source().unwrap()).unwrap();
        if let Some(name) = token.get_name() {
            let line = src[token.get_src_line() as usize];
            let idx = token.get_src_col() as usize;
            let span = &line[idx..idx + name.len()];
            assert_eq!(span, name);
        }
    }

    for (src_id, filename) in flat_map.sources().enumerate() {
        let ref_contents = files.get(filename).unwrap();
        let contents : Vec<_> = flat_map.get_source_contents(src_id as u32).expect(
            &format!("no source for {}", filename)).lines().collect();
        assert_eq!(&contents, ref_contents);
    }
}
<|endoftext|>


<commit_before>#![crate_name = "foo"]

pub trait Foo {}
pub trait Foo2 {}

pub struct Bar;

impl Foo for Bar {}
impl Foo2 for Bar {}

// @!has foo/fn.foo.html '//section[@id="main"]//pre' "x: &\'x impl Foo"
// @!has foo/fn.foo.html '//section[@id="main"]//pre' "-> &\'x impl Foo {"
pub fn foo<'x>(x: &'x impl Foo) -> &'x impl Foo {
    x
}

// @!has foo/fn.foo2.html '//section[@id="main"]//pre' "x: &\'x impl Foo"
// @!has foo/fn.foo2.html '//section[@id="main"]//pre' '-> impl Foo2 {'
pub fn foo2<'x>(_x: &'x impl Foo) -> impl Foo2 {
    Bar
}

// @!has foo/fn.foo_foo.html '//section[@id="main"]//pre' '-> impl Foo + Foo2 {'
pub fn foo_foo() -> impl Foo + Foo2 {
    Bar
}

// @!has foo/fn.foo2.html '//section[@id="main"]//pre' "x: &'x (impl Foo + Foo2)"
pub fn foo_foo_foo<'x>(_x: &'x (impl Foo + Foo2)) {
}
<commit_msg>(rustdoc) fix test for trait impl display<commit_after>#![crate_name = "foo"]

pub trait Foo {}
pub trait Foo2 {}

pub struct Bar;

impl Foo for Bar {}
impl Foo2 for Bar {}

// @has foo/fn.foo.html '//section[@id="main"]//pre' "x: &'x impl Foo"
// @has foo/fn.foo.html '//section[@id="main"]//pre' "-> &'x impl Foo"
pub fn foo<'x>(x: &'x impl Foo) -> &'x impl Foo {
    x
}

// @has foo/fn.foo2.html '//section[@id="main"]//pre' "x: &'x impl Foo"
// @has foo/fn.foo2.html '//section[@id="main"]//pre' '-> impl Foo2'
pub fn foo2<'x>(_x: &'x impl Foo) -> impl Foo2 {
    Bar
}

// @has foo/fn.foo_foo.html '//section[@id="main"]//pre' '-> impl Foo + Foo2'
pub fn foo_foo() -> impl Foo + Foo2 {
    Bar
}

// @has foo/fn.foo_foo_foo.html '//section[@id="main"]//pre' "x: &'x impl Foo + Foo2"
pub fn foo_foo_foo<'x>(_x: &'x (impl Foo + Foo2)) {
}
<|endoftext|>