<commit_before>use std::string::*;
use std::slice;
use std::fmt;
use std::ops::Index;

#[derive(Debug,Clone,Copy,PartialEq)]
pub enum RGBSubmodel {
    RGB,
    SRGB,
}

impl fmt::Display for RGBSubmodel {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = match *self {
            RGBSubmodel::RGB => "RGB".to_string(),
            RGBSubmodel::SRGB => "sRGB".to_string(),
        };
        write!(f, "{}", name)
    }
}

#[derive(Debug,Clone,Copy,PartialEq)]
pub enum YUVSubmodel {
    YCbCr,
    YIQ,
    YUVJ,
}

impl fmt::Display for YUVSubmodel {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = match *self {
            YUVSubmodel::YCbCr => "YCbCr".to_string(),
            YUVSubmodel::YIQ => "YIQ".to_string(),
            YUVSubmodel::YUVJ => "YUVJ".to_string(),
        };
        write!(f, "{}", name)
    }
}

#[derive(Debug, Clone,Copy,PartialEq)]
pub enum ColorModel {
    RGB(RGBSubmodel),
    YUV(YUVSubmodel),
    CMYK,
    HSV,
    LAB,
    XYZ,
}

impl ColorModel {
    pub fn get_default_components(&self) -> usize {
        match *self {
            ColorModel::CMYK => 4,
            _ => 3,
        }
    }
}

impl fmt::Display for ColorModel {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let name = match *self {
            ColorModel::RGB(fmt) => format!("RGB({})", fmt).to_string(),
            ColorModel::YUV(fmt) => format!("YUV({})", fmt).to_string(),
            ColorModel::CMYK => "CMYK".to_string(),
            ColorModel::HSV => "HSV".to_string(),
            ColorModel::LAB => "LAB".to_string(),
            ColorModel::XYZ => "XYZ".to_string(),
        };
        write!(f, "{}", name)
    }
}

#[derive(Clone,Copy,PartialEq)]
pub struct Chromaton {
    h_ss: u8,
    v_ss: u8,
    packed: bool,
    depth: u8,
    shift: u8,
    comp_offs: u8,
    next_elem: u8,
}

fn align(v: usize, a: usize) -> usize {
    (v + a - 1) & !(a - 1)
}

impl Chromaton {
    pub fn get_subsampling(&self) -> (u8, u8) {
        (self.h_ss, self.v_ss)
    }
    pub fn is_packed(&self) -> bool {
        self.packed
    }
    pub fn get_depth(&self) -> u8 {
        self.depth
    }
    pub fn get_shift(&self) -> u8 {
        self.shift
    }
    pub fn get_offset(&self) -> u8 {
        self.comp_offs
    }
    pub fn get_step(&self) -> u8 {
        self.next_elem
    }

    pub fn get_width(&self, width: usize) -> usize {
        (width + ((1 << self.h_ss) - 1)) >> self.h_ss
    }
    pub fn get_height(&self, height: usize) -> usize {
        (height + ((1 << self.v_ss) - 1)) >> self.v_ss
    }
    pub fn get_linesize(&self, width: usize, alignment: usize) -> usize {
        let d = self.depth as usize;
        align((self.get_width(width) * d + d - 1) >> 3, alignment)
    }
    pub fn get_data_size(&self, width: usize, height: usize, align: usize) -> usize {
        let nh = (height + ((1 << self.v_ss) - 1)) >> self.v_ss;
        self.get_linesize(width, align) * nh
    }
}

impl fmt::Display for Chromaton {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let pfmt = if self.packed {
            let mask = ((1 << self.depth) - 1) << self.shift;
            format!("packed(+{},{:X}, step {})",
                    self.comp_offs,
                    mask,
                    self.next_elem)
        } else {
            format!("planar({},{})", self.comp_offs, self.next_elem)
        };
        write!(f, "({}x{}, {})", self.h_ss, self.v_ss, pfmt)
    }
}

#[derive(Clone,Copy,PartialEq)]
pub struct Formaton {
    model: ColorModel,
    components: u8,
    comp_info: [Option<Chromaton>; 5],
    elem_size: u8,
    be: bool,
    alpha: bool,
    palette: bool,
}

bitflags! {
    pub flags Flags: u8 {
        const BE       = 0x01,
        const ALPHA    = 0x02,
        const PALETTE  = 0x04,
    }
}

impl Formaton {
    pub fn new(model: ColorModel, components: &[Chromaton], flags: Flags, elem_size: u8) -> Self {
        let be = flags.contains(BE);
        let alpha = flags.contains(ALPHA);
        let palette = flags.contains(PALETTE);
        let mut c: [Option<Chromaton>; 5] = [None; 5];

        if components.len() > 5 {
            panic!("too many components");
        }

        for (i, v) in components.iter().enumerate() {
            c[i] = Some(*v);
        }

        Formaton {
            model: model,
            components: components.len() as u8,
            comp_info: c,
            elem_size: elem_size,
            be: be,
            alpha: alpha,
            palette: palette,
        }
    }

    pub fn get_model(&self) -> ColorModel {
        self.model
    }
    pub fn get_num_comp(&self) -> usize {
        self.components as usize
    }
    pub fn get_chromaton(&self, idx: usize) -> Option<Chromaton> {
        if idx < self.comp_info.len() {
            return self.comp_info[idx];
        }
        None
    }
    pub fn is_be(&self) -> bool {
        self.be
    }
    pub fn has_alpha(&self) -> bool {
        self.alpha
    }
    pub fn is_paletted(&self) -> bool {
        self.palette
    }
    pub fn get_elem_size(&self) -> u8 {
        self.elem_size
    }
}

impl<'a> Index<usize> for &'a Formaton {
    type Output = Option<Chromaton>;

    fn index(&self, index: usize) -> &Self::Output {
        self.comp_info.index(index)
    }
}

impl<'a> IntoIterator for &'a Formaton {
    type Item = &'a Option<Chromaton>;
    type IntoIter = slice::Iter<'a, Option<Chromaton>>;

    fn into_iter(self) -> Self::IntoIter {
        self.comp_info.iter()
    }
}

impl fmt::Display for Formaton {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let end = if self.be {
            "BE"
        } else {
            "LE"
        };
        let palstr = if self.palette {
            "palette "
        } else {
            ""
        };
        let astr = if self.alpha {
            "alpha "
        } else {
            ""
        };
        let mut str = format!("Formaton for {} ({}{}elem {} size {}): ",
                              self.model,
                              palstr,
                              astr,
                              end,
                              self.elem_size);
        for &i in self.into_iter() {
            if let Some(chr) = i {
                str = format!("{} {}", str, chr);
            }
        }
        write!(f, "[{}]", str)
    }
}

macro_rules! chromaton {
    ($hs: expr, $vs: expr, $pck: expr, $d: expr, $sh: expr, $co: expr, $ne: expr) => ({
        Some(Chromaton {
                h_ss: $hs,
                v_ss: $vs,
                packed: $pck,
                depth: $d,
                shift: $sh,
                comp_offs: $co,
                next_elem: $ne })
    });
    (yuv8; $hs: expr, $vs: expr, $co: expr) => ({
        chromaton!($hs, $vs, false, 8, 0, $co, 1)
    });
    (packrgb; $d: expr, $s: expr, $co: expr, $ne: expr) => ({
        chromaton!(0, 0, true, $d, $s, $co, $ne)
    });
    (pal8; $co: expr) => ({
        chromaton!(0, 0, true, 8, 0, $co, 3)
    });
}

pub mod formats {
    use data::pixel::*;

    pub const YUV420: Formaton = Formaton {
        model: ColorModel::YUV(YUVSubmodel::YUVJ),
        components: 3,
        comp_info: [chromaton!(0, 0, false, 8, 0, 0, 1),
                    chromaton!(yuv8; 1, 1, 1),
                    chromaton!(yuv8; 1, 1, 2),
                    None,
                    None],
        elem_size: 0,
        be: false,
        alpha: false,
        palette: false,
    };

    pub const YUV410: Formaton = Formaton {
        model: ColorModel::YUV(YUVSubmodel::YUVJ),
        components: 3,
        comp_info: [chromaton!(0, 0, false, 8, 0, 0, 1),
                    chromaton!(yuv8; 2, 2, 1),
                    chromaton!(yuv8; 2, 2, 2),
                    None,
                    None],
        elem_size: 0,
        be: false,
        alpha: false,
        palette: false,
    };

    pub const PAL8: Formaton = Formaton {
        model: ColorModel::RGB(RGBSubmodel::RGB),
        components: 3,
        comp_info: [chromaton!(pal8; 0), chromaton!(pal8; 1), chromaton!(pal8; 2), None, None],
        elem_size: 3,
        be: false,
        alpha: false,
        palette: true,
    };

    pub const RGB565: Formaton = Formaton {
        model: ColorModel::RGB(RGBSubmodel::RGB),
        components: 3,
        comp_info: [chromaton!(packrgb; 5, 11, 0, 2),
                    chromaton!(packrgb; 6,  5, 0, 2),
                    chromaton!(packrgb; 5,  0, 0, 2),
                    None,
                    None],
        elem_size: 2,
        be: false,
        alpha: false,
        palette: false,
    };

    pub const RGB24: Formaton = Formaton {
        model: ColorModel::RGB(RGBSubmodel::RGB),
        components: 3,
        comp_info: [chromaton!(packrgb; 8, 0, 2, 3),
                    chromaton!(packrgb; 8, 0, 1, 3),
                    chromaton!(packrgb; 8, 0, 0, 3),
                    None,
                    None],
        elem_size: 3,
        be: false,
        alpha: false,
        palette: false,
    };
}

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

    #[test]
    fn fmt() {
        println!("formaton yuv- {}", formats::YUV420);
        println!("formaton pal- {}", formats::PAL8);
        println!("formaton rgb565- {}", formats::RGB565);
    }
}
<commit_msg>pixel: Make the ColorModel definitions more correct<commit_after>use std::slice;
use std::fmt;
use std::ops::Index;

#[derive(Debug,Clone,Copy,PartialEq)]
pub enum YUVRange {
    Limited,
    Full
}

impl fmt::Display for YUVRange {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            YUVRange::Limited => write!(f, "Limited range"),
            YUVRange::Full => write!(f, "Full range")
        }
    }
}

#[derive(Debug,Clone,Copy,PartialEq)]
pub enum YUVSystem {
    YCbCr(YUVRange),
    YCoCg,
    ICtCp
}

impl fmt::Display for YUVSystem {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::YUVSystem::*;
        match *self {
            YCbCr(range) => write!(f, "YCbCr ({})", range),
            YCoCg => write!(f, "YCbCg"),
            ICtCp => write!(f, "ICtCp"),
        }
    }
}

#[derive(Debug, Clone,Copy,PartialEq)]
pub enum TrichromaticEncodingSystem {
    RGB,
    YUV(YUVSystem),
    XYZ
}

impl fmt::Display for TrichromaticEncodingSystem {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::TrichromaticEncodingSystem::*;
        match *self {
            YUV(system) => write!(f, "{}", system),
            RGB => write!(f, "RGB"),
            XYZ => write!(f, "XYZ"),
        }
    }
}

#[derive(Debug, Clone,Copy,PartialEq)]
pub enum ColorModel {
    Trichromatic(TrichromaticEncodingSystem),
    CMYK,
    HSV,
    LAB,
}

impl fmt::Display for ColorModel {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::ColorModel::*;
        match *self {
            Trichromatic(system) => write!(f, "{}", system),
            CMYK => write!(f, "CMYK"),
            HSV => write!(f, "HSV"),
            LAB => write!(f, "LAB"),
        }
    }
}

impl ColorModel {
    pub fn get_default_components(&self) -> usize {
        match *self {
            ColorModel::CMYK => 4,
            _ => 3,
        }
    }
}

#[derive(Clone,Copy,PartialEq)]
pub struct Chromaton {
    h_ss: u8,
    v_ss: u8,
    packed: bool,
    depth: u8,
    shift: u8,
    comp_offs: u8,
    next_elem: u8,
}

fn align(v: usize, a: usize) -> usize {
    (v + a - 1) & !(a - 1)
}

impl Chromaton {
    pub fn get_subsampling(&self) -> (u8, u8) {
        (self.h_ss, self.v_ss)
    }
    pub fn is_packed(&self) -> bool {
        self.packed
    }
    pub fn get_depth(&self) -> u8 {
        self.depth
    }
    pub fn get_shift(&self) -> u8 {
        self.shift
    }
    pub fn get_offset(&self) -> u8 {
        self.comp_offs
    }
    pub fn get_step(&self) -> u8 {
        self.next_elem
    }

    pub fn get_width(&self, width: usize) -> usize {
        (width + ((1 << self.h_ss) - 1)) >> self.h_ss
    }
    pub fn get_height(&self, height: usize) -> usize {
        (height + ((1 << self.v_ss) - 1)) >> self.v_ss
    }
    pub fn get_linesize(&self, width: usize, alignment: usize) -> usize {
        let d = self.depth as usize;
        align((self.get_width(width) * d + d - 1) >> 3, alignment)
    }
    pub fn get_data_size(&self, width: usize, height: usize, align: usize) -> usize {
        let nh = (height + ((1 << self.v_ss) - 1)) >> self.v_ss;
        self.get_linesize(width, align) * nh
    }
}

impl fmt::Display for Chromaton {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let pfmt = if self.packed {
            let mask = ((1 << self.depth) - 1) << self.shift;
            format!("packed(+{},{:X}, step {})",
                    self.comp_offs,
                    mask,
                    self.next_elem)
        } else {
            format!("planar({},{})", self.comp_offs, self.next_elem)
        };
        write!(f, "({}x{}, {})", self.h_ss, self.v_ss, pfmt)
    }
}

#[derive(Clone,Copy,PartialEq)]
pub struct Formaton {
    model: ColorModel,
    components: u8,
    comp_info: [Option<Chromaton>; 5],
    elem_size: u8,
    be: bool,
    alpha: bool,
    palette: bool,
}

bitflags! {
    pub flags Flags: u8 {
        const BE       = 0x01,
        const ALPHA    = 0x02,
        const PALETTE  = 0x04,
    }
}

impl Formaton {
    pub fn new(model: ColorModel, components: &[Chromaton], flags: Flags, elem_size: u8) -> Self {
        let be = flags.contains(BE);
        let alpha = flags.contains(ALPHA);
        let palette = flags.contains(PALETTE);
        let mut c: [Option<Chromaton>; 5] = [None; 5];

        if components.len() > 5 {
            panic!("too many components");
        }

        for (i, v) in components.iter().enumerate() {
            c[i] = Some(*v);
        }

        Formaton {
            model: model,
            components: components.len() as u8,
            comp_info: c,
            elem_size: elem_size,
            be: be,
            alpha: alpha,
            palette: palette,
        }
    }

    pub fn get_model(&self) -> ColorModel {
        self.model
    }
    pub fn get_num_comp(&self) -> usize {
        self.components as usize
    }
    pub fn get_chromaton(&self, idx: usize) -> Option<Chromaton> {
        if idx < self.comp_info.len() {
            return self.comp_info[idx];
        }
        None
    }
    pub fn is_be(&self) -> bool {
        self.be
    }
    pub fn has_alpha(&self) -> bool {
        self.alpha
    }
    pub fn is_paletted(&self) -> bool {
        self.palette
    }
    pub fn get_elem_size(&self) -> u8 {
        self.elem_size
    }
}

impl<'a> Index<usize> for &'a Formaton {
    type Output = Option<Chromaton>;

    fn index(&self, index: usize) -> &Self::Output {
        self.comp_info.index(index)
    }
}

impl<'a> IntoIterator for &'a Formaton {
    type Item = &'a Option<Chromaton>;
    type IntoIter = slice::Iter<'a, Option<Chromaton>>;

    fn into_iter(self) -> Self::IntoIter {
        self.comp_info.iter()
    }
}

impl fmt::Display for Formaton {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let end = if self.be {
            "BE"
        } else {
            "LE"
        };
        let palstr = if self.palette {
            "palette "
        } else {
            ""
        };
        let astr = if self.alpha {
            "alpha "
        } else {
            ""
        };
        let mut str = format!("Formaton for {} ({}{}elem {} size {}): ",
                              self.model,
                              palstr,
                              astr,
                              end,
                              self.elem_size);
        for &i in self.into_iter() {
            if let Some(chr) = i {
                str = format!("{} {}", str, chr);
            }
        }
        write!(f, "[{}]", str)
    }
}

macro_rules! chromaton {
    ($hs: expr, $vs: expr, $pck: expr, $d: expr, $sh: expr, $co: expr, $ne: expr) => ({
        Some(Chromaton {
                h_ss: $hs,
                v_ss: $vs,
                packed: $pck,
                depth: $d,
                shift: $sh,
                comp_offs: $co,
                next_elem: $ne })
    });
    (yuv8; $hs: expr, $vs: expr, $co: expr) => ({
        chromaton!($hs, $vs, false, 8, 0, $co, 1)
    });
    (packrgb; $d: expr, $s: expr, $co: expr, $ne: expr) => ({
        chromaton!(0, 0, true, $d, $s, $co, $ne)
    });
    (pal8; $co: expr) => ({
        chromaton!(0, 0, true, 8, 0, $co, 3)
    });
}

pub mod formats {
    use data::pixel::*;
    use self::ColorModel::*;
    use self::TrichromaticEncodingSystem::*;
    use self::YUVSystem::*;
    use self::YUVRange::*;

    pub const YUV420: Formaton = Formaton {
        model: Trichromatic(YUV(YCbCr(Limited))),
        components: 3,
        comp_info: [chromaton!(0, 0, false, 8, 0, 0, 1),
                    chromaton!(yuv8; 1, 1, 1),
                    chromaton!(yuv8; 1, 1, 2),
                    None,
                    None],
        elem_size: 0,
        be: false,
        alpha: false,
        palette: false,
    };

    pub const YUV410: Formaton = Formaton {
        model: Trichromatic(YUV(YCbCr(Limited))),
        components: 3,
        comp_info: [chromaton!(0, 0, false, 8, 0, 0, 1),
                    chromaton!(yuv8; 2, 2, 1),
                    chromaton!(yuv8; 2, 2, 2),
                    None,
                    None],
        elem_size: 0,
        be: false,
        alpha: false,
        palette: false,
    };

    pub const PAL8: Formaton = Formaton {
        model: Trichromatic(RGB),
        components: 3,
        comp_info: [chromaton!(pal8; 0), chromaton!(pal8; 1), chromaton!(pal8; 2), None, None],
        elem_size: 3,
        be: false,
        alpha: false,
        palette: true,
    };

    pub const RGB565: Formaton = Formaton {
        model: Trichromatic(RGB),
        components: 3,
        comp_info: [chromaton!(packrgb; 5, 11, 0, 2),
                    chromaton!(packrgb; 6,  5, 0, 2),
                    chromaton!(packrgb; 5,  0, 0, 2),
                    None,
                    None],
        elem_size: 2,
        be: false,
        alpha: false,
        palette: false,
    };

    pub const RGB24: Formaton = Formaton {
        model: Trichromatic(RGB),
        components: 3,
        comp_info: [chromaton!(packrgb; 8, 0, 2, 3),
                    chromaton!(packrgb; 8, 0, 1, 3),
                    chromaton!(packrgb; 8, 0, 0, 3),
                    None,
                    None],
        elem_size: 3,
        be: false,
        alpha: false,
        palette: false,
    };
}

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

    #[test]
    fn fmt() {
        println!("formaton yuv- {}", formats::YUV420);
        println!("formaton pal- {}", formats::PAL8);
        println!("formaton rgb565- {}", formats::RGB565);
    }
}
<|endoftext|>


<commit_before>extern crate rand;
use rand::Rng;
use image::CONF_LINE_IDENTIFIER__IMAGE;

// Config file syntax error message
pub const CONF_SYNTAX_ERROR: &'static str = "

SYNTAX ERROR in config file!
Every line has to start with one of the following characters:
'#' (comment line), '-' (guessing string), '|' (ASCII-Art image) or ':' (game modifier).
Edit config file and start again.\n";

// comments in config file start with
pub const CONF_LINE_IDENTIFIER__COMMENT: char = '#';

// comands in config-file start with
pub const CONF_LINE_IDENTIFIER__CONTROL: char = ':';

// guessing strings in config-file start with
pub const CONF_LINE_IDENTIFIER__WORD: char = '-';

// a modifier tagging parts of the string to be visible from the start, e.g.
// "guess*-me*: will be shown as "_ _ _ _ _ - m e"
pub const CONF_LINE_WORD_MODIFIER__VISIBLE: char = '*';

#[derive(Debug)]
pub struct Dict {
    wordlist: Vec<String>
}

impl Dict {

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

    pub fn new(lines: &str) -> Dict {

        Dict{wordlist :
            lines
            .lines()
            .enumerate()
            .filter(|&(_,l)|!( l.trim().len() == 0 ||
                          l.starts_with(CONF_LINE_IDENTIFIER__COMMENT) ||
                          l.starts_with(CONF_LINE_IDENTIFIER__CONTROL) ||
                          l.starts_with(CONF_LINE_IDENTIFIER__IMAGE)
                        )
            )
        .map(|(n,l)| if l.starts_with(CONF_LINE_IDENTIFIER__WORD) {
                             l[1..].trim().to_string()
                         } else {
                             panic!("{}\nError in line: {}: \"{}\"\n\n",
                                     CONF_SYNTAX_ERROR, n+1, l)
                     })
            .collect()
        }
    }

    pub fn get_random_word(&self) -> String {
        (*rand::thread_rng().choose(&self.wordlist).unwrap()).to_string()
    }
}
<commit_msg>allow Unicode BOM in config file<commit_after>extern crate rand;
use rand::Rng;
use image::CONF_LINE_IDENTIFIER__IMAGE;

// Config file syntax error message
pub const CONF_SYNTAX_ERROR: &'static str = "

SYNTAX ERROR in config file!
Every line has to start with one of the following characters:
'#' (comment line), '-' (guessing string), '|' (ASCII-Art image) or ':' (game modifier).
Edit config file and start again.\n";

// comments in config file start with
pub const CONF_LINE_IDENTIFIER__COMMENT: char = '#';

// comands in config-file start with
pub const CONF_LINE_IDENTIFIER__CONTROL: char = ':';

// guessing strings in config-file start with
pub const CONF_LINE_IDENTIFIER__WORD: char = '-';

// a modifier tagging parts of the string to be visible from the start, e.g.
// "guess*-me*: will be shown as "_ _ _ _ _ - m e"
pub const CONF_LINE_WORD_MODIFIER__VISIBLE: char = '*';

#[derive(Debug)]
pub struct Dict {
    wordlist: Vec<String>
}

impl Dict {

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

    pub fn new(lines: &str) -> Dict {

        Dict{wordlist :
          // remove Unicode BOM if present (\u{feff} has in UTF8 3 bytes).
          if lines.starts_with('\u{feff}') { &lines[3..] } else { &lines[..] }
            .lines()
            .enumerate()
            .filter(|&(_,l)|!( l.trim().len() == 0 ||
                          l.starts_with(CONF_LINE_IDENTIFIER__COMMENT) ||
                          l.starts_with(CONF_LINE_IDENTIFIER__CONTROL) ||
                          l.starts_with(CONF_LINE_IDENTIFIER__IMAGE)
                        )
            )
            .map(|(n,l)| if l.starts_with(CONF_LINE_IDENTIFIER__WORD) {
                             l[1..].trim().to_string()
                         } else {
                             panic!("{}\nError in line: {}: \"{}\"\n\n",
                                     CONF_SYNTAX_ERROR, n+1, l)
                     })
            .collect()
        }
    }

    pub fn get_random_word(&self) -> String {
        (*rand::thread_rng().choose(&self.wordlist).unwrap()).to_string()
    }
}
<|endoftext|>


<commit_before>//! The JIT driver uses [`cranelift_simplejit`] to JIT execute programs without writing any object
//! files.

use std::cell::RefCell;
use std::ffi::CString;
use std::os::raw::{c_char, c_int};

use rustc_codegen_ssa::CrateInfo;
use rustc_middle::mir::mono::MonoItem;

use cranelift_jit::{JITBuilder, JITModule};

use crate::{prelude::*, BackendConfig};
use crate::{CodegenCx, CodegenMode};

thread_local! {
    pub static BACKEND_CONFIG: RefCell<Option<BackendConfig>> = RefCell::new(None);
    pub static CURRENT_MODULE: RefCell<Option<JITModule>> = RefCell::new(None);
}

pub(super) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
    if !tcx.sess.opts.output_types.should_codegen() {
        tcx.sess.fatal("JIT mode doesn't work with `cargo check`.");
    }

    let imported_symbols = load_imported_symbols_for_jit(tcx);

    let mut jit_builder =
        JITBuilder::with_isa(crate::build_isa(tcx.sess), cranelift_module::default_libcall_names());
    jit_builder.hotswap(matches!(backend_config.codegen_mode, CodegenMode::JitLazy));
    jit_builder.symbols(imported_symbols);
    let mut jit_module = JITModule::new(jit_builder);
    assert_eq!(pointer_ty(tcx), jit_module.target_config().pointer_type());

    let sig = Signature {
        params: vec![
            AbiParam::new(jit_module.target_config().pointer_type()),
            AbiParam::new(jit_module.target_config().pointer_type()),
        ],
        returns: vec![AbiParam::new(jit_module.target_config().pointer_type() /*isize*/)],
        call_conv: CallConv::triple_default(&crate::target_triple(tcx.sess)),
    };
    let main_func_id = jit_module.declare_function("main", Linkage::Import, &sig).unwrap();

    let (_, cgus) = tcx.collect_and_partition_mono_items(LOCAL_CRATE);
    let mono_items = cgus
        .iter()
        .map(|cgu| cgu.items_in_deterministic_order(tcx).into_iter())
        .flatten()
        .collect::<FxHashMap<_, (_, _)>>()
        .into_iter()
        .collect::<Vec<(_, (_, _))>>();

    let mut cx = crate::CodegenCx::new(tcx, backend_config, &mut jit_module, false);

    super::time(tcx, "codegen mono items", || {
        super::predefine_mono_items(&mut cx, &mono_items);
        for (mono_item, (linkage, visibility)) in mono_items {
            let linkage = crate::linkage::get_clif_linkage(mono_item, linkage, visibility);
            match mono_item {
                MonoItem::Fn(inst) => match backend_config.codegen_mode {
                    CodegenMode::Aot => unreachable!(),
                    CodegenMode::Jit => {
                        cx.tcx
                            .sess
                            .time("codegen fn", || crate::base::codegen_fn(&mut cx, inst, linkage));
                    }
                    CodegenMode::JitLazy => codegen_shim(&mut cx, inst),
                },
                MonoItem::Static(def_id) => {
                    crate::constant::codegen_static(&mut cx.constants_cx, def_id);
                }
                MonoItem::GlobalAsm(item_id) => {
                    let item = cx.tcx.hir().item(item_id);
                    tcx.sess.span_fatal(item.span, "Global asm is not supported in JIT mode");
                }
            }
        }
    });

    let (global_asm, _debug, mut unwind_context) =
        tcx.sess.time("finalize CodegenCx", || cx.finalize());
    jit_module.finalize_definitions();

    if !global_asm.is_empty() {
        tcx.sess.fatal("Inline asm is not supported in JIT mode");
    }

    crate::main_shim::maybe_create_entry_wrapper(tcx, &mut jit_module, &mut unwind_context);
    crate::allocator::codegen(tcx, &mut jit_module, &mut unwind_context);

    tcx.sess.abort_if_errors();

    jit_module.finalize_definitions();

    let _unwind_register_guard = unsafe { unwind_context.register_jit(&jit_module) };

    let finalized_main: *const u8 = jit_module.get_finalized_function(main_func_id);

    println!(
        "Rustc codegen cranelift will JIT run the executable, because -Cllvm-args=mode=jit was passed"
    );

    let f: extern "C" fn(c_int, *const *const c_char) -> c_int =
        unsafe { ::std::mem::transmute(finalized_main) };

    let args = ::std::env::var("CG_CLIF_JIT_ARGS").unwrap_or_else(|_| String::new());
    let args = std::iter::once(&*tcx.crate_name(LOCAL_CRATE).as_str().to_string())
        .chain(args.split(' '))
        .map(|arg| CString::new(arg).unwrap())
        .collect::<Vec<_>>();
    let mut argv = args.iter().map(|arg| arg.as_ptr()).collect::<Vec<_>>();

    // Push a null pointer as a terminating argument. This is required by POSIX and
    // useful as some dynamic linkers use it as a marker to jump over.
    argv.push(std::ptr::null());

    BACKEND_CONFIG.with(|tls_backend_config| {
        assert!(tls_backend_config.borrow_mut().replace(backend_config).is_none())
    });
    CURRENT_MODULE
        .with(|current_module| assert!(current_module.borrow_mut().replace(jit_module).is_none()));

    let ret = f(args.len() as c_int, argv.as_ptr());

    std::process::exit(ret);
}

#[no_mangle]
extern "C" fn __clif_jit_fn(instance_ptr: *const Instance<'static>) -> *const u8 {
    rustc_middle::ty::tls::with(|tcx| {
        // lift is used to ensure the correct lifetime for instance.
        let instance = tcx.lift(unsafe { *instance_ptr }).unwrap();

        CURRENT_MODULE.with(|jit_module| {
            let mut jit_module = jit_module.borrow_mut();
            let jit_module = jit_module.as_mut().unwrap();
            let backend_config =
                BACKEND_CONFIG.with(|backend_config| backend_config.borrow().clone().unwrap());

            let name = tcx.symbol_name(instance).name.to_string();
            let sig = crate::abi::get_function_sig(tcx, jit_module.isa().triple(), instance);
            let func_id = jit_module.declare_function(&name, Linkage::Export, &sig).unwrap();
            jit_module.prepare_for_function_redefine(func_id).unwrap();

            let mut cx = crate::CodegenCx::new(tcx, backend_config, jit_module, false);
            tcx.sess
                .time("codegen fn", || crate::base::codegen_fn(&mut cx, instance, Linkage::Export));

            let (global_asm, _debug_context, unwind_context) = cx.finalize();
            assert!(global_asm.is_empty());
            jit_module.finalize_definitions();
            std::mem::forget(unsafe { unwind_context.register_jit(&jit_module) });
            jit_module.get_finalized_function(func_id)
        })
    })
}

fn load_imported_symbols_for_jit(tcx: TyCtxt<'_>) -> Vec<(String, *const u8)> {
    use rustc_middle::middle::dependency_format::Linkage;

    let mut dylib_paths = Vec::new();

    let crate_info = CrateInfo::new(tcx);
    let formats = tcx.dependency_formats(LOCAL_CRATE);
    let data = &formats
        .iter()
        .find(|(crate_type, _data)| *crate_type == rustc_session::config::CrateType::Executable)
        .unwrap()
        .1;
    for &(cnum, _) in &crate_info.used_crates_dynamic {
        let src = &crate_info.used_crate_source[&cnum];
        match data[cnum.as_usize() - 1] {
            Linkage::NotLinked | Linkage::IncludedFromDylib => {}
            Linkage::Static => {
                let name = tcx.crate_name(cnum);
                let mut err =
                    tcx.sess.struct_err(&format!("Can't load static lib {}", name.as_str()));
                err.note("rustc_codegen_cranelift can only load dylibs in JIT mode.");
                err.emit();
            }
            Linkage::Dynamic => {
                dylib_paths.push(src.dylib.as_ref().unwrap().0.clone());
            }
        }
    }

    let mut imported_symbols = Vec::new();
    for path in dylib_paths {
        use object::{Object, ObjectSymbol};
        let lib = libloading::Library::new(&path).unwrap();
        let obj = std::fs::read(path).unwrap();
        let obj = object::File::parse(&obj).unwrap();
        imported_symbols.extend(obj.dynamic_symbols().filter_map(|symbol| {
            let name = symbol.name().unwrap().to_string();
            if name.is_empty() || !symbol.is_global() || symbol.is_undefined() {
                return None;
            }
            if name.starts_with("rust_metadata_") {
                // The metadata is part of a section that is not loaded by the dynamic linker in
                // case of cg_llvm.
                return None;
            }
            let dlsym_name = if cfg!(target_os = "macos") {
                // On macOS `dlsym` expects the name without leading `_`.
                assert!(name.starts_with('_'), "{:?}", name);
                &name[1..]
            } else {
                &name
            };
            let symbol: libloading::Symbol<'_, *const u8> =
                unsafe { lib.get(dlsym_name.as_bytes()) }.unwrap();
            Some((name, *symbol))
        }));
        std::mem::forget(lib)
    }

    tcx.sess.abort_if_errors();

    imported_symbols
}

pub(super) fn codegen_shim<'tcx>(cx: &mut CodegenCx<'_, 'tcx>, inst: Instance<'tcx>) {
    let tcx = cx.tcx;

    let pointer_type = cx.module.target_config().pointer_type();

    let name = tcx.symbol_name(inst).name.to_string();
    let sig = crate::abi::get_function_sig(tcx, cx.module.isa().triple(), inst);
    let func_id = cx.module.declare_function(&name, Linkage::Export, &sig).unwrap();

    let instance_ptr = Box::into_raw(Box::new(inst));

    let jit_fn = cx
        .module
        .declare_function(
            "__clif_jit_fn",
            Linkage::Import,
            &Signature {
                call_conv: cx.module.target_config().default_call_conv,
                params: vec![AbiParam::new(pointer_type)],
                returns: vec![AbiParam::new(pointer_type)],
            },
        )
        .unwrap();

    let mut trampoline = Function::with_name_signature(ExternalName::default(), sig.clone());
    let mut builder_ctx = FunctionBuilderContext::new();
    let mut trampoline_builder = FunctionBuilder::new(&mut trampoline, &mut builder_ctx);

    let jit_fn = cx.module.declare_func_in_func(jit_fn, trampoline_builder.func);
    let sig_ref = trampoline_builder.func.import_signature(sig);

    let entry_block = trampoline_builder.create_block();
    trampoline_builder.append_block_params_for_function_params(entry_block);
    let fn_args = trampoline_builder.func.dfg.block_params(entry_block).to_vec();

    trampoline_builder.switch_to_block(entry_block);
    let instance_ptr = trampoline_builder.ins().iconst(pointer_type, instance_ptr as u64 as i64);
    let jitted_fn = trampoline_builder.ins().call(jit_fn, &[instance_ptr]);
    let jitted_fn = trampoline_builder.func.dfg.inst_results(jitted_fn)[0];
    let call_inst = trampoline_builder.ins().call_indirect(sig_ref, jitted_fn, &fn_args);
    let ret_vals = trampoline_builder.func.dfg.inst_results(call_inst).to_vec();
    trampoline_builder.ins().return_(&ret_vals);

    cx.module
        .define_function(
            func_id,
            &mut Context::for_function(trampoline),
            &mut cranelift_codegen::binemit::NullTrapSink {},
        )
        .unwrap();
}
<commit_msg>Directly invoke the main function in JIT mode<commit_after>//! The JIT driver uses [`cranelift_simplejit`] to JIT execute programs without writing any object
//! files.

use std::cell::RefCell;
use std::ffi::CString;
use std::os::raw::{c_char, c_int};

use rustc_codegen_ssa::CrateInfo;
use rustc_middle::mir::mono::MonoItem;
use rustc_session::config::EntryFnType;

use cranelift_jit::{JITBuilder, JITModule};

use crate::{prelude::*, BackendConfig};
use crate::{CodegenCx, CodegenMode};

thread_local! {
    pub static BACKEND_CONFIG: RefCell<Option<BackendConfig>> = RefCell::new(None);
    pub static CURRENT_MODULE: RefCell<Option<JITModule>> = RefCell::new(None);
}

pub(super) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
    if !tcx.sess.opts.output_types.should_codegen() {
        tcx.sess.fatal("JIT mode doesn't work with `cargo check`.");
    }

    let imported_symbols = load_imported_symbols_for_jit(tcx);

    let mut jit_builder =
        JITBuilder::with_isa(crate::build_isa(tcx.sess), cranelift_module::default_libcall_names());
    jit_builder.hotswap(matches!(backend_config.codegen_mode, CodegenMode::JitLazy));
    jit_builder.symbols(imported_symbols);
    let mut jit_module = JITModule::new(jit_builder);
    assert_eq!(pointer_ty(tcx), jit_module.target_config().pointer_type());

    let (_, cgus) = tcx.collect_and_partition_mono_items(LOCAL_CRATE);
    let mono_items = cgus
        .iter()
        .map(|cgu| cgu.items_in_deterministic_order(tcx).into_iter())
        .flatten()
        .collect::<FxHashMap<_, (_, _)>>()
        .into_iter()
        .collect::<Vec<(_, (_, _))>>();

    let mut cx = crate::CodegenCx::new(tcx, backend_config, &mut jit_module, false);

    super::time(tcx, "codegen mono items", || {
        super::predefine_mono_items(&mut cx, &mono_items);
        for (mono_item, (linkage, visibility)) in mono_items {
            let linkage = crate::linkage::get_clif_linkage(mono_item, linkage, visibility);
            match mono_item {
                MonoItem::Fn(inst) => match backend_config.codegen_mode {
                    CodegenMode::Aot => unreachable!(),
                    CodegenMode::Jit => {
                        cx.tcx
                            .sess
                            .time("codegen fn", || crate::base::codegen_fn(&mut cx, inst, linkage));
                    }
                    CodegenMode::JitLazy => codegen_shim(&mut cx, inst),
                },
                MonoItem::Static(def_id) => {
                    crate::constant::codegen_static(&mut cx.constants_cx, def_id);
                }
                MonoItem::GlobalAsm(item_id) => {
                    let item = cx.tcx.hir().item(item_id);
                    tcx.sess.span_fatal(item.span, "Global asm is not supported in JIT mode");
                }
            }
        }
    });

    let (global_asm, _debug, mut unwind_context) =
        tcx.sess.time("finalize CodegenCx", || cx.finalize());
    jit_module.finalize_definitions();

    if !global_asm.is_empty() {
        tcx.sess.fatal("Inline asm is not supported in JIT mode");
    }

    crate::allocator::codegen(tcx, &mut jit_module, &mut unwind_context);

    tcx.sess.abort_if_errors();

    jit_module.finalize_definitions();
    let _unwind_register_guard = unsafe { unwind_context.register_jit(&jit_module) };

    println!(
        "Rustc codegen cranelift will JIT run the executable, because -Cllvm-args=mode=jit was passed"
    );

    let args = ::std::env::var("CG_CLIF_JIT_ARGS").unwrap_or_else(|_| String::new());
    let args = std::iter::once(&*tcx.crate_name(LOCAL_CRATE).as_str().to_string())
        .chain(args.split(' '))
        .map(|arg| CString::new(arg).unwrap())
        .collect::<Vec<_>>();
    let mut argv = args.iter().map(|arg| arg.as_ptr()).collect::<Vec<_>>();

    // Push a null pointer as a terminating argument. This is required by POSIX and
    // useful as some dynamic linkers use it as a marker to jump over.
    argv.push(std::ptr::null());

    BACKEND_CONFIG.with(|tls_backend_config| {
        assert!(tls_backend_config.borrow_mut().replace(backend_config).is_none())
    });

    let (main_def_id, entry_ty) = tcx.entry_fn(LOCAL_CRATE).unwrap();
    let instance = Instance::mono(tcx, main_def_id.to_def_id()).polymorphize(tcx);

    match entry_ty {
        EntryFnType::Main => {
            // FIXME set program arguments somehow

            let main_sig = Signature {
                params: vec![],
                returns: vec![],
                call_conv: CallConv::triple_default(&crate::target_triple(tcx.sess)),
            };
            let main_func_id = jit_module
                .declare_function(tcx.symbol_name(instance).name, Linkage::Import, &main_sig)
                .unwrap();
            let finalized_main: *const u8 = jit_module.get_finalized_function(main_func_id);

            CURRENT_MODULE.with(|current_module| {
                assert!(current_module.borrow_mut().replace(jit_module).is_none())
            });

            let f: extern "C" fn() = unsafe { ::std::mem::transmute(finalized_main) };
            f();
            std::process::exit(0);
        }
        EntryFnType::Start => {
            let start_sig = Signature {
                params: vec![
                    AbiParam::new(jit_module.target_config().pointer_type()),
                    AbiParam::new(jit_module.target_config().pointer_type()),
                ],
                returns: vec![AbiParam::new(
                    jit_module.target_config().pointer_type(), /*isize*/
                )],
                call_conv: CallConv::triple_default(&crate::target_triple(tcx.sess)),
            };
            let start_func_id = jit_module
                .declare_function(tcx.symbol_name(instance).name, Linkage::Import, &start_sig)
                .unwrap();
            let finalized_start: *const u8 = jit_module.get_finalized_function(start_func_id);

            CURRENT_MODULE.with(|current_module| {
                assert!(current_module.borrow_mut().replace(jit_module).is_none())
            });

            let f: extern "C" fn(c_int, *const *const c_char) -> c_int =
                unsafe { ::std::mem::transmute(finalized_start) };
            let ret = f(args.len() as c_int, argv.as_ptr());
            std::process::exit(ret);
        }
    }
}

#[no_mangle]
extern "C" fn __clif_jit_fn(instance_ptr: *const Instance<'static>) -> *const u8 {
    rustc_middle::ty::tls::with(|tcx| {
        // lift is used to ensure the correct lifetime for instance.
        let instance = tcx.lift(unsafe { *instance_ptr }).unwrap();

        CURRENT_MODULE.with(|jit_module| {
            let mut jit_module = jit_module.borrow_mut();
            let jit_module = jit_module.as_mut().unwrap();
            let backend_config =
                BACKEND_CONFIG.with(|backend_config| backend_config.borrow().clone().unwrap());

            let name = tcx.symbol_name(instance).name.to_string();
            let sig = crate::abi::get_function_sig(tcx, jit_module.isa().triple(), instance);
            let func_id = jit_module.declare_function(&name, Linkage::Export, &sig).unwrap();
            jit_module.prepare_for_function_redefine(func_id).unwrap();

            let mut cx = crate::CodegenCx::new(tcx, backend_config, jit_module, false);
            tcx.sess
                .time("codegen fn", || crate::base::codegen_fn(&mut cx, instance, Linkage::Export));

            let (global_asm, _debug_context, unwind_context) = cx.finalize();
            assert!(global_asm.is_empty());
            jit_module.finalize_definitions();
            std::mem::forget(unsafe { unwind_context.register_jit(&jit_module) });
            jit_module.get_finalized_function(func_id)
        })
    })
}

fn load_imported_symbols_for_jit(tcx: TyCtxt<'_>) -> Vec<(String, *const u8)> {
    use rustc_middle::middle::dependency_format::Linkage;

    let mut dylib_paths = Vec::new();

    let crate_info = CrateInfo::new(tcx);
    let formats = tcx.dependency_formats(LOCAL_CRATE);
    let data = &formats
        .iter()
        .find(|(crate_type, _data)| *crate_type == rustc_session::config::CrateType::Executable)
        .unwrap()
        .1;
    for &(cnum, _) in &crate_info.used_crates_dynamic {
        let src = &crate_info.used_crate_source[&cnum];
        match data[cnum.as_usize() - 1] {
            Linkage::NotLinked | Linkage::IncludedFromDylib => {}
            Linkage::Static => {
                let name = tcx.crate_name(cnum);
                let mut err =
                    tcx.sess.struct_err(&format!("Can't load static lib {}", name.as_str()));
                err.note("rustc_codegen_cranelift can only load dylibs in JIT mode.");
                err.emit();
            }
            Linkage::Dynamic => {
                dylib_paths.push(src.dylib.as_ref().unwrap().0.clone());
            }
        }
    }

    let mut imported_symbols = Vec::new();
    for path in dylib_paths {
        use object::{Object, ObjectSymbol};
        let lib = libloading::Library::new(&path).unwrap();
        let obj = std::fs::read(path).unwrap();
        let obj = object::File::parse(&obj).unwrap();
        imported_symbols.extend(obj.dynamic_symbols().filter_map(|symbol| {
            let name = symbol.name().unwrap().to_string();
            if name.is_empty() || !symbol.is_global() || symbol.is_undefined() {
                return None;
            }
            if name.starts_with("rust_metadata_") {
                // The metadata is part of a section that is not loaded by the dynamic linker in
                // case of cg_llvm.
                return None;
            }
            let dlsym_name = if cfg!(target_os = "macos") {
                // On macOS `dlsym` expects the name without leading `_`.
                assert!(name.starts_with('_'), "{:?}", name);
                &name[1..]
            } else {
                &name
            };
            let symbol: libloading::Symbol<'_, *const u8> =
                unsafe { lib.get(dlsym_name.as_bytes()) }.unwrap();
            Some((name, *symbol))
        }));
        std::mem::forget(lib)
    }

    tcx.sess.abort_if_errors();

    imported_symbols
}

fn codegen_shim<'tcx>(cx: &mut CodegenCx<'_, 'tcx>, inst: Instance<'tcx>) {
    let tcx = cx.tcx;

    let pointer_type = cx.module.target_config().pointer_type();

    let name = tcx.symbol_name(inst).name.to_string();
    let sig = crate::abi::get_function_sig(tcx, cx.module.isa().triple(), inst);
    let func_id = cx.module.declare_function(&name, Linkage::Export, &sig).unwrap();

    let instance_ptr = Box::into_raw(Box::new(inst));

    let jit_fn = cx
        .module
        .declare_function(
            "__clif_jit_fn",
            Linkage::Import,
            &Signature {
                call_conv: cx.module.target_config().default_call_conv,
                params: vec![AbiParam::new(pointer_type)],
                returns: vec![AbiParam::new(pointer_type)],
            },
        )
        .unwrap();

    let mut trampoline = Function::with_name_signature(ExternalName::default(), sig.clone());
    let mut builder_ctx = FunctionBuilderContext::new();
    let mut trampoline_builder = FunctionBuilder::new(&mut trampoline, &mut builder_ctx);

    let jit_fn = cx.module.declare_func_in_func(jit_fn, trampoline_builder.func);
    let sig_ref = trampoline_builder.func.import_signature(sig);

    let entry_block = trampoline_builder.create_block();
    trampoline_builder.append_block_params_for_function_params(entry_block);
    let fn_args = trampoline_builder.func.dfg.block_params(entry_block).to_vec();

    trampoline_builder.switch_to_block(entry_block);
    let instance_ptr = trampoline_builder.ins().iconst(pointer_type, instance_ptr as u64 as i64);
    let jitted_fn = trampoline_builder.ins().call(jit_fn, &[instance_ptr]);
    let jitted_fn = trampoline_builder.func.dfg.inst_results(jitted_fn)[0];
    let call_inst = trampoline_builder.ins().call_indirect(sig_ref, jitted_fn, &fn_args);
    let ret_vals = trampoline_builder.func.dfg.inst_results(call_inst).to_vec();
    trampoline_builder.ins().return_(&ret_vals);

    cx.module
        .define_function(
            func_id,
            &mut Context::for_function(trampoline),
            &mut cranelift_codegen::binemit::NullTrapSink {},
        )
        .unwrap();
}
<|endoftext|>


<commit_before>pub mod search;

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

const VERSION: &'static str = "0.1";

/// Implements `UciEngine` trait.
pub struct Engine {
    position: Position,
    pondering_is_allowed: bool,
    multi_pv: usize,
    is_thinking: bool,
    is_pondering: bool,
    replies: Vec<EngineReply>,
    infinite: bool,
    tt: Arc<TranspositionTable>,
    commands: Sender<search::Command>,
    reports: Receiver<search::Progress>,
    results: Receiver<search::Done>,
    searcher: thread::JoinHandle<()>,
}

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 (commands_tx, commands_rx) = channel();
        let (reports_tx, reports_rx) = channel();
        let (results_tx, results_rx) = channel();
        let mut tt = TranspositionTable::new();
        tt.resize(tt_size_mb);
        let tt = Arc::new(tt);
        Engine {
            position: Position::from_fen("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 \
                                          1")
                          .ok()
                          .unwrap(),
            pondering_is_allowed: false,
            multi_pv: 1,
            is_thinking: false,
            is_pondering: false,
            replies: vec![],
            infinite: false,
            tt: tt.clone(),
            commands: commands_tx,
            reports: reports_rx,
            results: results_rx,
            searcher: thread::spawn(move || {
                search::run(&tt, commands_rx, reports_tx, results_tx);
            }),
        }
    }
}

impl UciEngine for Engine {
    fn set_option(&mut self, name: &str, value: &str) {
        match name {
            // We do not support re-sizing of the transposition table
            // once the engine had started, so we do nothing.
            "Hash" => (),

            // 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";
            }

            // Tells the engine to output multiple best lines.
            // the default value is `1`.)
            "MultiPV" => {
                self.multi_pv = value.parse::<usize>().ok().unwrap_or(1);
            }

            // An invalid option.
            _ => {
                self.replies.push(EngineReply::Info(vec![("string".to_string(),
                                                          format!("Invalid option \"{}\"", name))]));
            }
        }
    }

    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;
            }
        }
    }

    #[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 {
            return;
        }

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

        // Start a new search
        self.commands
            .send(search::Command::Search(search::Parameters {
                id: 0,
                position: self.position.clone(),
                depth: 5,
                lower_bound: -20000,
                upper_bound: 20000,
            }))
            .unwrap();
        self.is_thinking = true;
        self.is_pondering = ponder;
        self.infinite = infinite;
    }

    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: do interactive deepening instead.
        if !self.is_pondering {
            if let Ok(_) = self.results.try_recv() {
                self.stop();
            }
        }

        self.replies.pop()
    }
}

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::from_u32(0);
            let mut v = Vec::with_capacity(128);
            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.move16();
        }
        if m != 0 {
            let move_type = ((m & 0b1100000000000000) >> 14) as MoveType;
            let orig_square = ((m & 0b0011111100000000) >> 8) as Square;
            let dest_square = ((m & 0b0000000011111100) >> 2) as Square;
            let promoted_piece = match m & 0b11 {
                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 }),
            ("MultiPV".to_string(), OptionDescription::Spin { min: 1, max: 32, default: 1 }),
        ]
    }

    fn create(&self, hash_size_mb: Option<usize>) -> Engine {
        Engine::new(hash_size_mb.unwrap_or(16))
    }
}
<commit_msg>move thing around<commit_after>pub mod search;

use std::thread;
use std::sync::Arc;
use std::sync::mpsc::{channel, Sender, Receiver};
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";

/// Implements `UciEngine` trait.
pub struct Engine {
    position: Position,
    pondering_is_allowed: bool,
    multi_pv: usize,
    is_thinking: bool,
    is_pondering: bool,
    replies: Vec<EngineReply>,
    infinite: bool,
    tt: Arc<TranspositionTable>,
    commands: Sender<search::Command>,
    reports: Receiver<search::Progress>,
    results: Receiver<search::Done>,
    searcher: thread::JoinHandle<()>,
}

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 (commands_tx, commands_rx) = channel();
        let (reports_tx, reports_rx) = channel();
        let (results_tx, results_rx) = channel();
        let mut tt = TranspositionTable::new();
        tt.resize(tt_size_mb);
        let tt = Arc::new(tt);
        Engine {
            position: Position::from_fen("rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w QKqk - 0 \
                                          1")
                          .ok()
                          .unwrap(),
            pondering_is_allowed: false,
            multi_pv: 1,
            is_thinking: false,
            is_pondering: false,
            replies: vec![],
            infinite: false,
            tt: tt.clone(),
            commands: commands_tx,
            reports: reports_rx,
            results: results_rx,
            searcher: thread::spawn(move || {
                search::run(&tt, commands_rx, reports_tx, results_tx);
            }),
        }
    }
}

impl UciEngine for Engine {
    fn set_option(&mut self, name: &str, value: &str) {
        match name {
            // We do not support re-sizing of the transposition table
            // once the engine had started, so we do nothing.
            "Hash" => (),

            // 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";
            }

            // Tells the engine to output multiple best lines.
            // the default value is `1`.)
            "MultiPV" => {
                self.multi_pv = value.parse::<usize>().ok().unwrap_or(1);
            }

            // An invalid option.
            _ => {
                self.replies.push(EngineReply::Info(vec![("string".to_string(),
                                                          format!("Invalid option \"{}\"", name))]));
            }
        }
    }

    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;
            }
        }
    }

    #[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 {
            self.commands
                .send(search::Command::Search(search::Parameters {
                    id: 0,
                    position: self.position.clone(),
                    depth: 5,
                    lower_bound: -20000,
                    upper_bound: 20000,
                }))
                .unwrap();
            self.is_thinking = true;
            self.is_pondering = ponder;
            self.infinite = infinite;
        }
    }

    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: do interactive deepening instead.
        if !self.is_pondering {
            if let Ok(_) = self.results.try_recv() {
                self.stop();
            }
        }

        self.replies.pop()
    }
}

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::from_u32(0);
            let mut v = Vec::with_capacity(128);
            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.move16();
        }
        if m != 0 {
            let move_type = ((m & 0b1100000000000000) >> 14) as MoveType;
            let orig_square = ((m & 0b0011111100000000) >> 8) as Square;
            let dest_square = ((m & 0b0000000011111100) >> 2) as Square;
            let promoted_piece = match m & 0b11 {
                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 }),
            ("MultiPV".to_string(), OptionDescription::Spin { min: 1, max: 32, default: 1 }),
        ]
    }

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


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

use cell::{Cell};
use dungeon::{Dungeon, DungeonCells, SurroundingCells, Surrounding};

pub type EvaluationFn = Box<Fn(&Dungeon) -> f64 + 'static + Send + Sync + Copy>;

pub fn check_1x1_rooms(dungeon: &Dungeon) -> f64 {
    let mut hits = 0;
    for i in 0..dungeon.width {
        for j in 0..dungeon.height {
            let ref cell = dungeon.cells[i][j];
            if cell.is_empty() {
                let sc = SurroundingCells::new(dungeon, cell, Surrounding::Cardinal);
                // check to see if all walls surround us.
                let surrounded = sc.fold(true, |accum, c| accum && !c.is_empty());
                if surrounded {
                    hits += 1;
                }
            }
        }
    }
    hits as f64
}

pub fn has_entrance_exit(dungeon: &Dungeon) -> f64 {
    let mut has_entrance = false;
    let mut has_exit = false;
    for i in 0..dungeon.width {
        for j in 0..dungeon.height {
            let ref cell = dungeon.cells[i][j];
            if cell.has_attribute("entrance") {
                has_entrance = true;
            }
            if cell.has_attribute("exit") {
                has_exit = true;
            }
        }
    }
    match (has_entrance, has_exit) {
        (false, false) => 2.0,
        (false, true) => 1.0,
        (true, false) => 1.0,
        _ => 0.0,
    }
}

pub fn doors_are_useful(dungeon: &Dungeon) -> f64 {
    let mut hits = 0;
    for i in 0..dungeon.width {
        for j in 0..dungeon.height {
            let ref cell = dungeon.cells[i][j];
            // check to see if doors have exactly two walls abutting them
            if cell.has_attribute("door") {
                let mut count = 0;
                for sc in SurroundingCells::new(dungeon, cell, Surrounding::Cardinal) {
                    if sc.has_attribute("wall") {
                        count += 1;
                    }
                }
                if count != 2 {
                    hits += 1;
                }
            }
        }
    }
    hits as f64
}

pub type Coords = HashSet<(u32, u32)>;

// helper function.
fn is_accessible(cell: &Cell) -> bool {
    cell.has_attribute("floor") || cell.has_attribute("door")
}

// recursively traverse the dungeon (depth-first)
fn search(dungeon: &Dungeon, visited: &Coords, x: u32, y: u32) -> Coords {
    let ref cell = dungeon.cells[x as usize][y as usize];
    let mut local = visited.clone();
    assert!(!local.contains(&(x, y)));
    local.insert((x, y));
    for sc in SurroundingCells::new(dungeon, cell, Surrounding::AllDirections) {
        let new_x = sc.x;
        let new_y = sc.y;
        if is_accessible(&sc) && !local.contains(&(new_x, new_y)) {
            let results = search(&dungeon, &local, new_x, new_y);
            let search_results: Coords = local.union(&results).cloned().collect();
            local = search_results.clone();
        }
    }
    local.clone()
}

pub fn rooms_are_accessible(dungeon: &Dungeon) -> f64 {
    let mut hits = 0;
    let dc = DungeonCells::new(&dungeon);
    let mut visited: Coords = HashSet::new();
    let floors: Coords = dc.filter(|ref cell| is_accessible(cell)).map(|cell| (cell.x, cell.y)).collect();
    while {
        let remaining: Coords = floors.difference(&visited).cloned().collect();
        match remaining.len() {
            0 => false,
            _ => {
                let &(x, y) = remaining.iter().next().unwrap();
                let swept = search(&dungeon, &visited, x, y);
                visited = visited.union(&swept).cloned().collect();
                let total: Coords = visited.intersection(&floors).cloned().collect();
                total.len() != floors.len()
            }
        }
    }
    {
        hits += 1;
    }
    hits as f64
}
<commit_msg>Minor fix to door evaluation function.<commit_after>use std::collections::{HashSet};

use cell::{Cell};
use dungeon::{Dungeon, DungeonCells, SurroundingCells, Surrounding};

pub type EvaluationFn = Box<Fn(&Dungeon) -> f64 + 'static + Send + Sync + Copy>;

pub fn check_1x1_rooms(dungeon: &Dungeon) -> f64 {
    let mut hits = 0;
    for i in 0..dungeon.width {
        for j in 0..dungeon.height {
            let ref cell = dungeon.cells[i][j];
            if cell.is_empty() {
                let sc = SurroundingCells::new(dungeon, cell, Surrounding::Cardinal);
                // check to see if all walls surround us.
                let surrounded = sc.fold(true, |accum, c| accum && !c.is_empty());
                if surrounded {
                    hits += 1;
                }
            }
        }
    }
    hits as f64
}

pub fn has_entrance_exit(dungeon: &Dungeon) -> f64 {
    let mut has_entrance = false;
    let mut has_exit = false;
    for i in 0..dungeon.width {
        for j in 0..dungeon.height {
            let ref cell = dungeon.cells[i][j];
            if cell.has_attribute("entrance") {
                has_entrance = true;
            }
            if cell.has_attribute("exit") {
                has_exit = true;
            }
        }
    }
    match (has_entrance, has_exit) {
        (false, false) => 2.0,
        (false, true) => 1.0,
        (true, false) => 1.0,
        _ => 0.0,
    }
}

pub fn doors_are_useful(dungeon: &Dungeon) -> f64 {
    let mut hits = 0;
    for i in 0..dungeon.width {
        for j in 0..dungeon.height {
            let ref cell = dungeon.cells[i][j];
            // check to see if doors have exactly two walls or gaps abutting them
            if cell.has_attribute("door") {
                let mut count = 0;
                for sc in SurroundingCells::new(dungeon, cell, Surrounding::Cardinal) {
                    if sc.tile.is_none() || sc.has_attribute("wall") {
                        count += 1;
                    }
                }
                if count != 2 {
                    hits += 1;
                }
            }
        }
    }
    hits as f64
}

pub type Coords = HashSet<(u32, u32)>;

// helper function.
fn is_accessible(cell: &Cell) -> bool {
    cell.has_attribute("floor") || cell.has_attribute("door")
}

// recursively traverse the dungeon (depth-first)
fn search(dungeon: &Dungeon, visited: &Coords, x: u32, y: u32) -> Coords {
    let ref cell = dungeon.cells[x as usize][y as usize];
    let mut local = visited.clone();
    assert!(!local.contains(&(x, y)));
    local.insert((x, y));
    for sc in SurroundingCells::new(dungeon, cell, Surrounding::AllDirections) {
        let new_x = sc.x;
        let new_y = sc.y;
        if is_accessible(&sc) && !local.contains(&(new_x, new_y)) {
            let results = search(&dungeon, &local, new_x, new_y);
            let search_results: Coords = local.union(&results).cloned().collect();
            local = search_results.clone();
        }
    }
    local.clone()
}

pub fn rooms_are_accessible(dungeon: &Dungeon) -> f64 {
    let mut hits = 0;
    let dc = DungeonCells::new(&dungeon);
    let mut visited: Coords = HashSet::new();
    let floors: Coords = dc.filter(|ref cell| is_accessible(cell)).map(|cell| (cell.x, cell.y)).collect();
    while {
        let remaining: Coords = floors.difference(&visited).cloned().collect();
        match remaining.len() {
            0 => false,
            _ => {
                let &(x, y) = remaining.iter().next().unwrap();
                let swept = search(&dungeon, &visited, x, y);
                visited = visited.union(&swept).cloned().collect();
                let total: Coords = visited.intersection(&floors).cloned().collect();
                total.len() != floors.len()
            }
        }
    }
    {
        hits += 1;
    }
    hits as f64
}
<|endoftext|>


<commit_before>use std::cell::Cell;

pub trait Incrementor {
    fn increment(&self);
}

impl Incrementor for Cell<u8> {
    fn increment(&self) {
        self.set(self.get() + 1);
    }
}

impl Incrementor for Cell<u16> {
    fn increment(&self) {
        self.set(self.get() + 1);
    }
}
<commit_msg>Change incrementor to be more generic<commit_after>use std::cell::Cell;
use std::num::One;

pub trait Incrementor {
    fn increment(&self);
}

impl<T: Copy + Num> Incrementor for Cell<T> {
    fn increment(&self) {
        self.set(self.get() + One::one());
    }
}
<|endoftext|>


<commit_before>use std::fs::Metadata;
use std::path::{Path, PathBuf};

use log::warn;

pub fn walk_dir<P: AsRef<Path>>(root: P) -> std::io::Result<WalkDir> {
    let metadata = std::fs::symlink_metadata(&root)?;

    Ok(WalkDir {
        root: Some(WalkEntry {
            path: root.as_ref().to_path_buf(),
            metadata: metadata,
        }),
        pending: vec!(list_dir(&root)?),
    })
}

fn list_dir<P: AsRef<Path>>(path: P) -> std::io::Result<ListDir> {
    let iter = std::fs::read_dir(path)?;

    Ok(ListDir {
        iter: iter,
    })
}

pub struct WalkDir {
    // TODO: Add support for stopping at device boundaries.
    root: Option<WalkEntry>,
    pending: Vec<ListDir>,
}

pub struct WalkEntry {
    path: PathBuf,
    metadata: Metadata,
}

impl WalkDir {

    fn push(&mut self, entry: &WalkEntry) {
        match list_dir(&entry.path) {
            Ok(iter) => {
                self.pending.push(iter);
            },
            Err(error) => {
                warn!("failed to read '{}': {}", entry.path.display(), error);
            },
        }
    }

    fn pop(&mut self) -> Option<WalkEntry> {
        while let Some(iter) = self.pending.last_mut() {
            for entry in iter {
                return Some(entry);
            }

            self.pending.pop();
        }

        None
    }
}

impl std::iter::Iterator for WalkDir {

    type Item = WalkEntry;

    fn next(&mut self) -> Option<WalkEntry> {
        if self.root.is_some() {
            return self.root.take();
        }

        let entry = self.pop()?;

        if entry.metadata.is_dir() {
            self.push(&entry);
        }

        Some(entry)
    }
}

/// Iterator over the entries in a directory.
///
/// This iterator is very similar to the standard `ReadDir` iterator, except
/// that it is forgetful and only yields entries that did not cause any errors.
///
/// Unlike the `ReadDir` iterator entries, `ListDir` entries are guaranteed to
/// have valid metadata objects attached.
struct ListDir {
    iter: std::fs::ReadDir,
}

impl std::iter::Iterator for ListDir {

    type Item = WalkEntry;

    fn next(&mut self) -> Option<WalkEntry> {
        for entry in &mut self.iter {
            let entry = match entry {
                Ok(entry) => entry,
                Err(error) => {
                    warn!("directory iteration error: {}", error);
                    continue
                },
            };

            let path = entry.path();
            let metadata = match entry.metadata() {
                Ok(metadata) => metadata,
                Err(error) => {
                    warn!("failed to stat '{}': {}", path.display(), error);
                    continue
                },
            };

            return Some(WalkEntry {
                path: path,
                metadata: metadata,
            });
        }

        None
    }
}

#[cfg(test)]
mod tests {

    use std::fs::File;

    use super::*;

    // TODO: Add test case for non-existing directory.

    #[test]
    fn test_list_dir_empty() {
        let tempdir = tempfile::tempdir().unwrap();

        let mut iter = list_dir(&tempdir).unwrap();

        assert!(iter.next().is_none());
    }

    #[test]
    fn test_list_dir_with_files() {
        let tempdir = tempfile::tempdir().unwrap();
        File::create(tempdir.path().join("abc")).unwrap();
        File::create(tempdir.path().join("def")).unwrap();
        File::create(tempdir.path().join("ghi")).unwrap();

        let mut results = list_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path().join("abc"));
        assert!(results[0].metadata.is_file());

        assert_eq!(results[1].path, tempdir.path().join("def"));
        assert!(results[1].metadata.is_file());

        assert_eq!(results[2].path, tempdir.path().join("ghi"));
        assert!(results[2].metadata.is_file());
    }

    #[test]
    fn test_list_dir_with_dirs() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("abc")).unwrap();
        std::fs::create_dir(tempdir.path().join("def")).unwrap();

        let mut results = list_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);

        assert_eq!(results[0].path, tempdir.path().join("abc"));
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("def"));
        assert!(results[1].metadata.is_dir());
    }

    #[cfg(target_family = "unix")]
    #[test]
    fn test_list_dir_with_links() {
        let tempdir = tempfile::tempdir().unwrap();
        let source = tempdir.path().join("abc");
        let target = tempdir.path().join("def");

        File::create(&source).unwrap();
        std::os::unix::fs::symlink(&source, &target).unwrap();

        let mut results = list_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);

        assert_eq!(results[0].path, source);
        assert!(results[0].metadata.file_type().is_file());

        assert_eq!(results[1].path, target);
        assert!(results[1].metadata.file_type().is_symlink());
    }

    #[test]
    fn test_walk_dir_empty() {
        let tempdir = tempfile::tempdir().unwrap();

        let mut iter = walk_dir(&tempdir).unwrap();

        let entry = iter.next().unwrap();
        assert_eq!(entry.path, tempdir.path());
        assert!(entry.metadata.is_dir());

        assert!(iter.next().is_none());
    }

    #[test]
    fn test_walk_dir_with_flat_files() {
        let tempdir = tempfile::tempdir().unwrap();
        File::create(tempdir.path().join("abc")).unwrap();
        File::create(tempdir.path().join("def")).unwrap();
        File::create(tempdir.path().join("ghi")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 4);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("abc"));
        assert!(results[1].metadata.is_file());

        assert_eq!(results[2].path, tempdir.path().join("def"));
        assert!(results[2].metadata.is_file());

        assert_eq!(results[3].path, tempdir.path().join("ghi"));
        assert!(results[3].metadata.is_file());
    }

    #[test]
    fn test_walk_dir_with_flat_dirs() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("abc")).unwrap();
        std::fs::create_dir(tempdir.path().join("def")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("abc"));
        assert!(results[1].metadata.is_dir());

        assert_eq!(results[2].path, tempdir.path().join("def"));
        assert!(results[2].metadata.is_dir());
    }

    #[test]
    fn test_walk_dir_with_nested_dirs() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("abc")).unwrap();
        std::fs::create_dir(tempdir.path().join("abc").join("def")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("abc"));
        assert!(results[1].metadata.is_dir());

        assert_eq!(results[2].path, tempdir.path().join("abc").join("def"));
        assert!(results[2].metadata.is_dir());
    }

    #[test]
    fn test_walk_dir_with_deeply_nested_dirs() {
        let tempdir = tempfile::tempdir().unwrap();

        let mut dir = tempdir.path().to_path_buf();
        for _ in 0..512 {
            dir.push("foo");
            std::fs::create_dir(&dir).unwrap();
        }

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 513);

        for entry in &results[1..] {
            assert!(entry.path.starts_with(&tempdir));
            assert!(entry.path.ends_with("foo"));
            assert!(entry.metadata.is_dir());
        }
    }

    #[test]
    fn test_walk_dir_with_files_inside_dir() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("foo")).unwrap();
        File::create(tempdir.path().join("foo").join("abc")).unwrap();
        File::create(tempdir.path().join("foo").join("def")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 4);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("foo"));
        assert!(results[1].metadata.is_dir());

        assert_eq!(results[2].path, tempdir.path().join("foo").join("abc"));
        assert!(results[2].metadata.is_file());

        assert_eq!(results[3].path, tempdir.path().join("foo").join("def"));
        assert!(results[3].metadata.is_file());
    }

    #[cfg(target_family = "unix")]
    #[test]
    fn test_walk_dir_with_dir_symlinks() {
        let tempdir = tempfile::tempdir().unwrap();
        let dir = tempdir.path().join("abc");
        let file = dir.join("def");
        let symlink = tempdir.path().join("ghi");

        std::fs::create_dir(&dir).unwrap();
        File::create(&file).unwrap();
        std::os::unix::fs::symlink(&dir, &symlink).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 4);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.file_type().is_dir());

        assert_eq!(results[1].path, dir);
        assert!(results[1].metadata.file_type().is_dir());

        assert_eq!(results[2].path, file);
        assert!(results[2].metadata.file_type().is_file());

        assert_eq!(results[3].path, symlink);
        assert!(results[3].metadata.file_type().is_symlink());
    }

    #[cfg(target_family = "unix")]
    #[test]
    fn test_walk_dir_with_circular_symlinks() {
        let tempdir = tempfile::tempdir().unwrap();
        let dir = tempdir.path().join("foo");
        let symlink = tempdir.path().join("foo").join("bar");

        std::fs::create_dir(&dir).unwrap();
        std::os::unix::fs::symlink(&dir, &symlink).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.file_type().is_dir());

        assert_eq!(results[1].path, dir);
        assert!(results[1].metadata.file_type().is_dir());

        assert_eq!(results[2].path, symlink);
        assert!(results[2].metadata.file_type().is_symlink());
    }

    #[test]
    fn test_walk_dir_with_unicode_names() {
        let tempdir = tempfile::tempdir().unwrap();
        File::create(tempdir.path().join("zażółć gęślą jaźń")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);
        assert_eq!(results[0].path, tempdir.path());
        assert_eq!(results[1].path, tempdir.path().join("zażółć gęślą jaźń"));
    }

    #[test]
    fn test_walk_dir_metadata_size() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::write(tempdir.path().join("foo"), b"123456789").unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);
        assert_eq!(results[1].metadata.len(), 9);
    }
}
<commit_msg>Avoid extra allocation in case of walk error.<commit_after>use std::fs::Metadata;
use std::path::{Path, PathBuf};

use log::warn;

pub fn walk_dir<P: AsRef<Path>>(root: P) -> std::io::Result<WalkDir> {
    let metadata = std::fs::symlink_metadata(&root)?;
    let pending = vec!(list_dir(&root)?);

    Ok(WalkDir {
        root: Some(WalkEntry {
            path: root.as_ref().to_path_buf(),
            metadata: metadata,
        }),
        pending: pending,
    })
}

fn list_dir<P: AsRef<Path>>(path: P) -> std::io::Result<ListDir> {
    let iter = std::fs::read_dir(path)?;

    Ok(ListDir {
        iter: iter,
    })
}

pub struct WalkDir {
    // TODO: Add support for stopping at device boundaries.
    root: Option<WalkEntry>,
    pending: Vec<ListDir>,
}

pub struct WalkEntry {
    path: PathBuf,
    metadata: Metadata,
}

impl WalkDir {

    fn push(&mut self, entry: &WalkEntry) {
        match list_dir(&entry.path) {
            Ok(iter) => {
                self.pending.push(iter);
            },
            Err(error) => {
                warn!("failed to read '{}': {}", entry.path.display(), error);
            },
        }
    }

    fn pop(&mut self) -> Option<WalkEntry> {
        while let Some(iter) = self.pending.last_mut() {
            for entry in iter {
                return Some(entry);
            }

            self.pending.pop();
        }

        None
    }
}

impl std::iter::Iterator for WalkDir {

    type Item = WalkEntry;

    fn next(&mut self) -> Option<WalkEntry> {
        if self.root.is_some() {
            return self.root.take();
        }

        let entry = self.pop()?;

        if entry.metadata.is_dir() {
            self.push(&entry);
        }

        Some(entry)
    }
}

/// Iterator over the entries in a directory.
///
/// This iterator is very similar to the standard `ReadDir` iterator, except
/// that it is forgetful and only yields entries that did not cause any errors.
///
/// Unlike the `ReadDir` iterator entries, `ListDir` entries are guaranteed to
/// have valid metadata objects attached.
struct ListDir {
    iter: std::fs::ReadDir,
}

impl std::iter::Iterator for ListDir {

    type Item = WalkEntry;

    fn next(&mut self) -> Option<WalkEntry> {
        for entry in &mut self.iter {
            let entry = match entry {
                Ok(entry) => entry,
                Err(error) => {
                    warn!("directory iteration error: {}", error);
                    continue
                },
            };

            let path = entry.path();
            let metadata = match entry.metadata() {
                Ok(metadata) => metadata,
                Err(error) => {
                    warn!("failed to stat '{}': {}", path.display(), error);
                    continue
                },
            };

            return Some(WalkEntry {
                path: path,
                metadata: metadata,
            });
        }

        None
    }
}

#[cfg(test)]
mod tests {

    use std::fs::File;

    use super::*;

    // TODO: Add test case for non-existing directory.

    #[test]
    fn test_list_dir_empty() {
        let tempdir = tempfile::tempdir().unwrap();

        let mut iter = list_dir(&tempdir).unwrap();

        assert!(iter.next().is_none());
    }

    #[test]
    fn test_list_dir_with_files() {
        let tempdir = tempfile::tempdir().unwrap();
        File::create(tempdir.path().join("abc")).unwrap();
        File::create(tempdir.path().join("def")).unwrap();
        File::create(tempdir.path().join("ghi")).unwrap();

        let mut results = list_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path().join("abc"));
        assert!(results[0].metadata.is_file());

        assert_eq!(results[1].path, tempdir.path().join("def"));
        assert!(results[1].metadata.is_file());

        assert_eq!(results[2].path, tempdir.path().join("ghi"));
        assert!(results[2].metadata.is_file());
    }

    #[test]
    fn test_list_dir_with_dirs() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("abc")).unwrap();
        std::fs::create_dir(tempdir.path().join("def")).unwrap();

        let mut results = list_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);

        assert_eq!(results[0].path, tempdir.path().join("abc"));
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("def"));
        assert!(results[1].metadata.is_dir());
    }

    #[cfg(target_family = "unix")]
    #[test]
    fn test_list_dir_with_links() {
        let tempdir = tempfile::tempdir().unwrap();
        let source = tempdir.path().join("abc");
        let target = tempdir.path().join("def");

        File::create(&source).unwrap();
        std::os::unix::fs::symlink(&source, &target).unwrap();

        let mut results = list_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);

        assert_eq!(results[0].path, source);
        assert!(results[0].metadata.file_type().is_file());

        assert_eq!(results[1].path, target);
        assert!(results[1].metadata.file_type().is_symlink());
    }

    #[test]
    fn test_walk_dir_empty() {
        let tempdir = tempfile::tempdir().unwrap();

        let mut iter = walk_dir(&tempdir).unwrap();

        let entry = iter.next().unwrap();
        assert_eq!(entry.path, tempdir.path());
        assert!(entry.metadata.is_dir());

        assert!(iter.next().is_none());
    }

    #[test]
    fn test_walk_dir_with_flat_files() {
        let tempdir = tempfile::tempdir().unwrap();
        File::create(tempdir.path().join("abc")).unwrap();
        File::create(tempdir.path().join("def")).unwrap();
        File::create(tempdir.path().join("ghi")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 4);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("abc"));
        assert!(results[1].metadata.is_file());

        assert_eq!(results[2].path, tempdir.path().join("def"));
        assert!(results[2].metadata.is_file());

        assert_eq!(results[3].path, tempdir.path().join("ghi"));
        assert!(results[3].metadata.is_file());
    }

    #[test]
    fn test_walk_dir_with_flat_dirs() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("abc")).unwrap();
        std::fs::create_dir(tempdir.path().join("def")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("abc"));
        assert!(results[1].metadata.is_dir());

        assert_eq!(results[2].path, tempdir.path().join("def"));
        assert!(results[2].metadata.is_dir());
    }

    #[test]
    fn test_walk_dir_with_nested_dirs() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("abc")).unwrap();
        std::fs::create_dir(tempdir.path().join("abc").join("def")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("abc"));
        assert!(results[1].metadata.is_dir());

        assert_eq!(results[2].path, tempdir.path().join("abc").join("def"));
        assert!(results[2].metadata.is_dir());
    }

    #[test]
    fn test_walk_dir_with_deeply_nested_dirs() {
        let tempdir = tempfile::tempdir().unwrap();

        let mut dir = tempdir.path().to_path_buf();
        for _ in 0..512 {
            dir.push("foo");
            std::fs::create_dir(&dir).unwrap();
        }

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 513);

        for entry in &results[1..] {
            assert!(entry.path.starts_with(&tempdir));
            assert!(entry.path.ends_with("foo"));
            assert!(entry.metadata.is_dir());
        }
    }

    #[test]
    fn test_walk_dir_with_files_inside_dir() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::create_dir(tempdir.path().join("foo")).unwrap();
        File::create(tempdir.path().join("foo").join("abc")).unwrap();
        File::create(tempdir.path().join("foo").join("def")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 4);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.is_dir());

        assert_eq!(results[1].path, tempdir.path().join("foo"));
        assert!(results[1].metadata.is_dir());

        assert_eq!(results[2].path, tempdir.path().join("foo").join("abc"));
        assert!(results[2].metadata.is_file());

        assert_eq!(results[3].path, tempdir.path().join("foo").join("def"));
        assert!(results[3].metadata.is_file());
    }

    #[cfg(target_family = "unix")]
    #[test]
    fn test_walk_dir_with_dir_symlinks() {
        let tempdir = tempfile::tempdir().unwrap();
        let dir = tempdir.path().join("abc");
        let file = dir.join("def");
        let symlink = tempdir.path().join("ghi");

        std::fs::create_dir(&dir).unwrap();
        File::create(&file).unwrap();
        std::os::unix::fs::symlink(&dir, &symlink).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 4);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.file_type().is_dir());

        assert_eq!(results[1].path, dir);
        assert!(results[1].metadata.file_type().is_dir());

        assert_eq!(results[2].path, file);
        assert!(results[2].metadata.file_type().is_file());

        assert_eq!(results[3].path, symlink);
        assert!(results[3].metadata.file_type().is_symlink());
    }

    #[cfg(target_family = "unix")]
    #[test]
    fn test_walk_dir_with_circular_symlinks() {
        let tempdir = tempfile::tempdir().unwrap();
        let dir = tempdir.path().join("foo");
        let symlink = tempdir.path().join("foo").join("bar");

        std::fs::create_dir(&dir).unwrap();
        std::os::unix::fs::symlink(&dir, &symlink).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 3);

        assert_eq!(results[0].path, tempdir.path());
        assert!(results[0].metadata.file_type().is_dir());

        assert_eq!(results[1].path, dir);
        assert!(results[1].metadata.file_type().is_dir());

        assert_eq!(results[2].path, symlink);
        assert!(results[2].metadata.file_type().is_symlink());
    }

    #[test]
    fn test_walk_dir_with_unicode_names() {
        let tempdir = tempfile::tempdir().unwrap();
        File::create(tempdir.path().join("zażółć gęślą jaźń")).unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);
        assert_eq!(results[0].path, tempdir.path());
        assert_eq!(results[1].path, tempdir.path().join("zażółć gęślą jaźń"));
    }

    #[test]
    fn test_walk_dir_metadata_size() {
        let tempdir = tempfile::tempdir().unwrap();
        std::fs::write(tempdir.path().join("foo"), b"123456789").unwrap();

        let mut results = walk_dir(&tempdir).unwrap().collect::<Vec<_>>();
        results.sort_by_key(|entry| entry.path.clone());

        assert_eq!(results.len(), 2);
        assert_eq!(results[1].metadata.len(), 9);
    }
}
<|endoftext|>