<commit_before>#![crate_name = "locale"]
#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![feature(core, collections, libc, old_io, old_path, os, std_misc)]

//! Localisation is hard.
//!
//! Getting your program to work well in multiple languages is a world fraught with edge-cases,
//! minor grammatical errors, and most importantly, subtle things that don't map over well that you
//! have absolutely no idea are different in other cultures.
//!
//! Many people are aware of the simpler ones, such as whether to use decimal points or decimal
//! commas, or that the names of the months are different in other languages. But there are also
//! different ways to format dates and times, or variations on what day the week begins. It's
//! perfectly possible to write your program unaware of how these things have to be changed at all,
//! and that's why it's so hard.

use std::old_io::fs::PathExtensions;
use std::old_io::{IoResult, File, BufferedReader};
use std::env::var;
use std::num::{Int, Float};
use std::fmt::Display;

/// Trait defining how to obtain various components of a locale.
///
/// Use implementation of this trait to construct parts of the `Locale` object.
///
/// There may be various methods for obtaining locale data. The lowest common denominator is
/// standard C library. It is however quite limited and some systems (notably Android) don't
/// actually contain the corresponding data. Many systems also provide additional configurability
/// for the locale setting (Windows, KDE, etc.) that are only accessible via that system's specific
/// interface. So this trait exists to allow combining the methods for obtaining the data.
///
/// The implementations for individual locale categories are returned boxed, because they may need
/// to be polymorphic _and_ in options to allow combining partial implementations. Creating locale
/// data is not a performance critical operation, so dynamic polymrphism is used for sake of
/// simplicity.
///
/// All methods default to simply returning None, again so partial implementations that delegate to
/// another factory are possible. See `CompositeLocaleFactory`.
pub trait LocaleFactory {
    /// Get implementation of the Numeric locale category.
    fn get_numeric(&mut self) -> Option<Box<Numeric>> { None }

    /// Get implementation of the Time locale category.
    fn get_time(&mut self) -> Option<Box<Time>> { None }
}

/// Auxiliary class for creating composing partial implementations of locale factories.
// FIXME: Create (doc) test when there actually is another implementation to substitute.
#[derive(Debug, Clone)]
pub struct CompositeLocaleFactory<First: LocaleFactory, Second: LocaleFactory> {
    first: First,
    second: Second,
}

impl<F: LocaleFactory, S: LocaleFactory> CompositeLocaleFactory<F, S> {
    pub fn new(first: F, second: S) -> Self {
        CompositeLocaleFactory::<F, S> {
            first: first, second: second
        }
    }
}

impl<F: LocaleFactory, S: LocaleFactory> LocaleFactory for CompositeLocaleFactory<F, S> {
    // XXX: Make a macro for this
    fn get_numeric(&mut self) -> Option<Box<Numeric>> {
        if let Some(v) = self.first.get_numeric() {
            Some(v)
        } else {
            self.second.get_numeric()
        }
    }

    fn get_time(&mut self) -> Option<Box<Time>> {
        if let Some(v) = self.first.get_time() {
            Some(v)
        } else {
            self.second.get_time()
        }
    }
}

/// Factory of invariant locales.
///
/// Invariant locale, called "C" or "POSIX" by standard C library locale functions, is default
/// locale definitions for when no information about desired locale is available or localization is
/// turned off.
#[derive(Debug, Clone, Default)]
pub struct InvariantLocaleFactory;

impl InvariantLocaleFactory {
    /// Constructs invariant locale factory.
    ///
    /// The signature is just so that it matches the other locale factories so the classes can be
    /// substituted depending on target operating system and the code using them does not have to
    /// care.
    #[allow(unused_variables)]
    pub fn new(locale: &str) -> Result<Self, i32> {
        Ok(InvariantLocaleFactory)
    }
}

impl LocaleFactory for InvariantLocaleFactory {
    // NOTE: Yep, it's empty. This just returns nothing and the Locale constructor will take care
    // of the actual defaults.
}

#[cfg(target_os = "linux")]
pub mod linux;

#[cfg(target_os = "linux")]
pub use linux::LibCLocaleFactory as SystemLocaleFactory;

#[cfg(not(target_os = "linux"))]
pub use InvariantLocaleFactory as SystemLocaleFactory;

/// Return LocaleFactory appropriate for default user locale, as far as it can be determined.
///
/// The returned locale factory provides locale facets implemented using standard localization
/// functionality of the underlying operating system and configured for user's default locale.
pub fn user_locale_factory() -> SystemLocaleFactory {
    // FIXME: Error handling? Constructing locale with "" should never fail as far as I can tell.
    SystemLocaleFactory::new("").unwrap()
}

// ---- locale facets ----

/// The directory inside which locale files are found.
///
/// For example, the set of Korean files will be in
/// `/usr/share/locale/ko`.
static LOCALE_DIR: &'static str = "/usr/share/locale";

#[derive(Debug, Clone)]
enum LocaleType {
    Numeric, Time,
}

fn find_locale_path(locale_type: LocaleType) -> Option<Path> {
    let file_name = match locale_type {
        LocaleType::Numeric => "LC_NUMERIC",
        LocaleType::Time    => "LC_TIME",
    };

    let locale_dir = Path::new(LOCALE_DIR);

    if let Ok(specific_path) = var(file_name) {
        let path = locale_dir.join(Path::new(specific_path)).join(Path::new(file_name));

        if path.exists() {
            return Some(path);
        }
    }

    if let Ok(all_path) = var("LC_ALL") {
        let path = locale_dir.join(Path::new(all_path)).join(Path::new(file_name));

        if path.exists() {
            return Some(path);
        }
    }

    if let Ok(lang) = var("LANG") {
        let path = locale_dir.join(Path::new(lang)).join(Path::new(file_name));

        if path.exists() {
            return Some(path);
        }
    }

    None
}

// ---- numeric stuff ----

/// Information on how to format numbers.
#[derive(Debug, Clone)]
pub struct Numeric {
    /// The punctuation that separates the decimal part of a non-integer number. Usually a decimal
    /// point or a decimal comma.
    pub decimal_sep: String,

    /// The punctuation that separates groups of digits in long numbers.
    pub thousands_sep: String,
}

impl Numeric {
    pub fn load_user_locale() -> IoResult<Numeric> {
        let path = find_locale_path(LocaleType::Numeric);

        if let Some(path) = path {
            let mut file = BufferedReader::new(File::open(&path));
            let lines: Vec<String> = file.lines().map(|x| x.unwrap()).collect();

            Ok(Numeric {
                decimal_sep: lines[0].trim().to_string(),
                thousands_sep: lines[1].trim().to_string(),
            })
        }
        else {
            return Ok(Numeric::english());
        }
    }

    pub fn english() -> Numeric {
        Numeric::new(".", ",")
    }

    pub fn new(decimal_sep: &str, thousands_sep: &str) -> Numeric {
        Numeric {
            decimal_sep: decimal_sep.to_string(),
            thousands_sep: thousands_sep.to_string(),
        }
    }

    pub fn format_int<I: Int + Display>(&self, input: I) -> String {
        let s = input.to_string();
        let mut buf = String::new();

        for (i, c) in s.chars().enumerate() {
            buf.push(c);
            if (s.len() - i - 1) % 3 == 0 && i != s.len() - 1 {
                buf.push_str(&self.thousands_sep[..]);
            }
        }

        buf
    }

    pub fn format_float<F: Float>(&self, input: F, decimal_places: usize) -> String {
        use std::num::strconv;
        strconv::float_to_str_common(input, 10, false, strconv::SignFormat::SignNone, strconv::SignificantDigits::DigExact(decimal_places), strconv::ExponentFormat::ExpNone, false).0.replace(".", self.decimal_sep.as_slice())
    }
}

// ---- time stuff ---

#[derive(Debug, Clone)]
pub struct Time {
    month_names: Vec<String>,
    long_month_names: Vec<String>,
    day_names: Vec<String>,
    long_day_names: Vec<String>,
}

impl Time {
    pub fn load_user_locale() -> IoResult<Time> {
        let path = find_locale_path(LocaleType::Time);

        if let Some(path) = path {
            let mut file = BufferedReader::new(File::open(&path));
            let mut iter = file.lines().map(|x| x.unwrap().trim().to_string());

            let month_names      = iter.by_ref().take(12).collect();
            let long_month_names = iter.by_ref().take(12).collect();
            let day_names        = iter.by_ref().take(7).collect();
            let long_day_names   = iter.by_ref().take(7).collect();

            Ok(Time {
                month_names:      month_names,
                long_month_names: long_month_names,
                day_names:        day_names,
                long_day_names:   long_day_names,
            })
        }
        else {
            return Ok(Time::english());
        }
    }

    pub fn english() -> Time {
        Time {
            month_names: vec![
                "Jan".to_string(),  "Feb".to_string(),  "Mar".to_string(),
                "Apr".to_string(),  "May".to_string(),  "Jun".to_string(),
                "Jul".to_string(),  "Aug".to_string(),  "Sep".to_string(),
                "Oct".to_string(),  "Nov".to_string(),  "Dec".to_string(),
            ],
            long_month_names: vec![
                "January".to_string(),    "February".to_string(),
                "March".to_string(),      "April".to_string(),
                "May".to_string(),        "June".to_string(),
                "July".to_string(),       "August".to_string(),
                "September".to_string(),  "October".to_string(),
                "November".to_string(),   "December".to_string(),
            ],
            day_names: vec![
                "Sun".to_string(),
                "Mon".to_string(),  "Tue".to_string(),  "Wed".to_string(),
                "Thu".to_string(),  "Fri".to_string(),  "Sat".to_string(),
            ],
            long_day_names: vec![
                "Sunday".to_string(),
                "Monday".to_string(),    "Tuesday".to_string(),  "Wednesday".to_string(),
                "Thursday".to_string(),  "Friday".to_string(),   "Saturday".to_string(),
            ],
        }
    }

    pub fn long_month_name(&self, months_from_january: usize) -> String {
        self.long_month_names[months_from_january].clone()
    }

    pub fn short_month_name(&self, months_from_january: usize) -> String {
        self.month_names[months_from_january].clone()
    }

    pub fn long_day_name(&self, days_from_sunday: usize) -> String {
        self.day_names[days_from_sunday].clone()
    }

    pub fn short_day_name(&self, days_from_sunday: usize) -> String {
        self.day_names[days_from_sunday].clone()
    }

}

// ---- tests ----

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

    #[test]
    fn thousands_separator() {
        let numeric_options = Numeric::new("/", "=");
        assert_eq!("1=234=567".to_string(), numeric_options.format_int(1234567))
    }

    #[test]
    fn thousands_separator_2() {
        let numeric_options = Numeric::new("/", "=");
        assert_eq!("123=456".to_string(), numeric_options.format_int(123456))
    }

    #[test]
    fn thousands_separator_3() {
        let numeric_options = Numeric::new("/", "=");
        assert_eq!("12=345=678".to_string(), numeric_options.format_int(12345678))
    }
}
<commit_msg>Upgrade to new IO and Path libraries<commit_after>#![crate_name = "locale"]
#![crate_type = "rlib"]
#![crate_type = "dylib"]
#![feature(core, path_ext, std_misc)]

//! Localisation is hard.
//!
//! Getting your program to work well in multiple languages is a world fraught with edge-cases,
//! minor grammatical errors, and most importantly, subtle things that don't map over well that you
//! have absolutely no idea are different in other cultures.
//!
//! Many people are aware of the simpler ones, such as whether to use decimal points or decimal
//! commas, or that the names of the months are different in other languages. But there are also
//! different ways to format dates and times, or variations on what day the week begins. It's
//! perfectly possible to write your program unaware of how these things have to be changed at all,
//! and that's why it's so hard.

use std::path::{Path, PathBuf};
use std::fs::{PathExt, File};
use std::io::{BufRead, Result, BufReader};
use std::env::var;
use std::num::{Int, Float};
use std::fmt::Display;

/// Trait defining how to obtain various components of a locale.
///
/// Use implementation of this trait to construct parts of the `Locale` object.
///
/// There may be various methods for obtaining locale data. The lowest common denominator is
/// standard C library. It is however quite limited and some systems (notably Android) don't
/// actually contain the corresponding data. Many systems also provide additional configurability
/// for the locale setting (Windows, KDE, etc.) that are only accessible via that system's specific
/// interface. So this trait exists to allow combining the methods for obtaining the data.
///
/// The implementations for individual locale categories are returned boxed, because they may need
/// to be polymorphic _and_ in options to allow combining partial implementations. Creating locale
/// data is not a performance critical operation, so dynamic polymrphism is used for sake of
/// simplicity.
///
/// All methods default to simply returning None, again so partial implementations that delegate to
/// another factory are possible. See `CompositeLocaleFactory`.
pub trait LocaleFactory {
    /// Get implementation of the Numeric locale category.
    fn get_numeric(&mut self) -> Option<Box<Numeric>> { None }

    /// Get implementation of the Time locale category.
    fn get_time(&mut self) -> Option<Box<Time>> { None }
}

/// Auxiliary class for creating composing partial implementations of locale factories.
// FIXME: Create (doc) test when there actually is another implementation to substitute.
#[derive(Debug, Clone)]
pub struct CompositeLocaleFactory<First: LocaleFactory, Second: LocaleFactory> {
    first: First,
    second: Second,
}

impl<F: LocaleFactory, S: LocaleFactory> CompositeLocaleFactory<F, S> {
    pub fn new(first: F, second: S) -> Self {
        CompositeLocaleFactory::<F, S> {
            first: first, second: second
        }
    }
}

impl<F: LocaleFactory, S: LocaleFactory> LocaleFactory for CompositeLocaleFactory<F, S> {
    // XXX: Make a macro for this
    fn get_numeric(&mut self) -> Option<Box<Numeric>> {
        if let Some(v) = self.first.get_numeric() {
            Some(v)
        } else {
            self.second.get_numeric()
        }
    }

    fn get_time(&mut self) -> Option<Box<Time>> {
        if let Some(v) = self.first.get_time() {
            Some(v)
        } else {
            self.second.get_time()
        }
    }
}

/// Factory of invariant locales.
///
/// Invariant locale, called "C" or "POSIX" by standard C library locale functions, is default
/// locale definitions for when no information about desired locale is available or localization is
/// turned off.
#[derive(Debug, Clone, Default)]
pub struct InvariantLocaleFactory;

impl InvariantLocaleFactory {
    /// Constructs invariant locale factory.
    ///
    /// The signature is just so that it matches the other locale factories so the classes can be
    /// substituted depending on target operating system and the code using them does not have to
    /// care.
    #[allow(unused_variables)]
    pub fn new(locale: &str) -> std::result::Result<Self, i32> {
        Ok(InvariantLocaleFactory)
    }
}

impl LocaleFactory for InvariantLocaleFactory {
    // NOTE: Yep, it's empty. This just returns nothing and the Locale constructor will take care
    // of the actual defaults.
}

#[cfg(target_os = "linux")]
pub mod linux;

#[cfg(target_os = "linux")]
pub use linux::LibCLocaleFactory as SystemLocaleFactory;

#[cfg(not(target_os = "linux"))]
pub use InvariantLocaleFactory as SystemLocaleFactory;

/// Return LocaleFactory appropriate for default user locale, as far as it can be determined.
///
/// The returned locale factory provides locale facets implemented using standard localization
/// functionality of the underlying operating system and configured for user's default locale.
pub fn user_locale_factory() -> SystemLocaleFactory {
    // FIXME: Error handling? Constructing locale with "" should never fail as far as I can tell.
    SystemLocaleFactory::new("").unwrap()
}

// ---- locale facets ----

/// The directory inside which locale files are found.
///
/// For example, the set of Korean files will be in
/// `/usr/share/locale/ko`.
static LOCALE_DIR: &'static str = "/usr/share/locale";

#[derive(Debug, Clone)]
enum LocaleType {
    Numeric, Time,
}

fn find_locale_path(locale_type: LocaleType) -> Option<PathBuf> {
    let file_name = match locale_type {
        LocaleType::Numeric => "LC_NUMERIC",
        LocaleType::Time    => "LC_TIME",
    };

    let locale_dir = Path::new(LOCALE_DIR);

    if let Ok(specific_path) = var(file_name) {
        let path = locale_dir.join(Path::new(&specific_path)).join(Path::new(file_name));

        if path.exists() {
            return Some(path);
        }
    }

    if let Ok(all_path) = var("LC_ALL") {
        let path = locale_dir.join(Path::new(&all_path)).join(Path::new(file_name));

        if path.exists() {
            return Some(path);
        }
    }

    if let Ok(lang) = var("LANG") {
        let path = locale_dir.join(Path::new(&lang)).join(Path::new(file_name));

        if path.exists() {
            return Some(path);
        }
    }

    None
}

// ---- numeric stuff ----

/// Information on how to format numbers.
#[derive(Debug, Clone)]
pub struct Numeric {
    /// The punctuation that separates the decimal part of a non-integer number. Usually a decimal
    /// point or a decimal comma.
    pub decimal_sep: String,

    /// The punctuation that separates groups of digits in long numbers.
    pub thousands_sep: String,
}

impl Numeric {
    pub fn load_user_locale() -> Result<Numeric> {
        let path = find_locale_path(LocaleType::Numeric);

        if let Some(path) = path {
            let file = BufReader::new(try!(File::open(&path)));
            let lines: Vec<String> = file.lines().map(|x| x.unwrap()).collect();

            Ok(Numeric {
                decimal_sep: lines[0].trim().to_string(),
                thousands_sep: lines[1].trim().to_string(),
            })
        }
        else {
            return Ok(Numeric::english());
        }
    }

    pub fn english() -> Numeric {
        Numeric::new(".", ",")
    }

    pub fn new(decimal_sep: &str, thousands_sep: &str) -> Numeric {
        Numeric {
            decimal_sep: decimal_sep.to_string(),
            thousands_sep: thousands_sep.to_string(),
        }
    }

    pub fn format_int<I: Int + Display>(&self, input: I) -> String {
        let s = input.to_string();
        let mut buf = String::new();

        for (i, c) in s.chars().enumerate() {
            buf.push(c);
            if (s.len() - i - 1) % 3 == 0 && i != s.len() - 1 {
                buf.push_str(&self.thousands_sep[..]);
            }
        }

        buf
    }

    pub fn format_float<F: Float>(&self, input: F, decimal_places: usize) -> String {
        use std::num::strconv;
        strconv::float_to_str_common(input, 10, false, strconv::SignFormat::SignNone, strconv::SignificantDigits::DigExact(decimal_places), strconv::ExponentFormat::ExpNone, false).0.replace(".", self.decimal_sep.as_slice())
    }
}

// ---- time stuff ---

#[derive(Debug, Clone)]
pub struct Time {
    month_names: Vec<String>,
    long_month_names: Vec<String>,
    day_names: Vec<String>,
    long_day_names: Vec<String>,
}

impl Time {
    pub fn load_user_locale() -> Result<Time> {
        let path = find_locale_path(LocaleType::Time);

        if let Some(path) = path {
            let file = BufReader::new(try!(File::open(&path)));
            let mut iter = file.lines().map(|x| x.unwrap().trim().to_string());

            let month_names      = iter.by_ref().take(12).collect();
            let long_month_names = iter.by_ref().take(12).collect();
            let day_names        = iter.by_ref().take(7).collect();
            let long_day_names   = iter.by_ref().take(7).collect();

            Ok(Time {
                month_names:      month_names,
                long_month_names: long_month_names,
                day_names:        day_names,
                long_day_names:   long_day_names,
            })
        }
        else {
            return Ok(Time::english());
        }
    }

    pub fn english() -> Time {
        Time {
            month_names: vec![
                "Jan".to_string(),  "Feb".to_string(),  "Mar".to_string(),
                "Apr".to_string(),  "May".to_string(),  "Jun".to_string(),
                "Jul".to_string(),  "Aug".to_string(),  "Sep".to_string(),
                "Oct".to_string(),  "Nov".to_string(),  "Dec".to_string(),
            ],
            long_month_names: vec![
                "January".to_string(),    "February".to_string(),
                "March".to_string(),      "April".to_string(),
                "May".to_string(),        "June".to_string(),
                "July".to_string(),       "August".to_string(),
                "September".to_string(),  "October".to_string(),
                "November".to_string(),   "December".to_string(),
            ],
            day_names: vec![
                "Sun".to_string(),
                "Mon".to_string(),  "Tue".to_string(),  "Wed".to_string(),
                "Thu".to_string(),  "Fri".to_string(),  "Sat".to_string(),
            ],
            long_day_names: vec![
                "Sunday".to_string(),
                "Monday".to_string(),    "Tuesday".to_string(),  "Wednesday".to_string(),
                "Thursday".to_string(),  "Friday".to_string(),   "Saturday".to_string(),
            ],
        }
    }

    pub fn long_month_name(&self, months_from_january: usize) -> String {
        self.long_month_names[months_from_january].clone()
    }

    pub fn short_month_name(&self, months_from_january: usize) -> String {
        self.month_names[months_from_january].clone()
    }

    pub fn long_day_name(&self, days_from_sunday: usize) -> String {
        self.day_names[days_from_sunday].clone()
    }

    pub fn short_day_name(&self, days_from_sunday: usize) -> String {
        self.day_names[days_from_sunday].clone()
    }

}

// ---- tests ----

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

    #[test]
    fn thousands_separator() {
        let numeric_options = Numeric::new("/", "=");
        assert_eq!("1=234=567".to_string(), numeric_options.format_int(1234567))
    }

    #[test]
    fn thousands_separator_2() {
        let numeric_options = Numeric::new("/", "=");
        assert_eq!("123=456".to_string(), numeric_options.format_int(123456))
    }

    #[test]
    fn thousands_separator_3() {
        let numeric_options = Numeric::new("/", "=");
        assert_eq!("12=345=678".to_string(), numeric_options.format_int(12345678))
    }
}
<|endoftext|>


<commit_before>use std::error::Error;
use std::io;
use std::str;
use std::io::Write;

use itertools::Itertools;
use rust_htslib::bcf;
use rust_htslib::bcf::record::Numeric;

pub struct Writer {
    inner: io::BufWriter<io::Stdout>,
    field_count: usize
}

impl Writer {
    fn new(inner: io::BufWriter<io::Stdout>) -> Self {
        Writer {
            inner: inner,
            field_count: 0
        }
    }

    fn write_integer(&mut self, value: i32) -> Result<(), Box<Error>> {
        self.write_field(format!("{}", value).as_bytes())
    }

    fn write_float(&mut self, value: f32) -> Result<(), Box<Error>> {
        self.write_field(format!("{}", value).as_bytes())
    }

    fn write_flag(&mut self, value: bool) -> Result<(), Box<Error>> {
        self.write_field(format!("{}", value).as_bytes())
    }

    fn write_field(&mut self, value: &[u8]) -> Result<(), Box<Error>> {
        if self.field_count > 0 {
            try!(self.inner.write(b"\t"));
        }
        try!(self.inner.write(value));
        self.field_count += 1;
        Ok(())
    }

    fn newline(&mut self) -> Result<(), Box<Error>> {
        try!(self.inner.write(b"\n"));
        self.field_count = 0;
        Ok(())
    }
}

const HEADER_COMMON: &'static [u8] = b"VARIANT";

pub fn to_txt(
    info_tags: &[&str],
    format_tags: &[&str],
    show_genotypes: bool
) -> Result<(), Box<Error>> {
    let mut reader = bcf::Reader::from_stdin()?;
    let mut writer = Writer::new(io::BufWriter::new(io::stdout()));

    let common_n = 5 + info_tags.len();
    try!(writer.write_field(HEADER_COMMON));
    for _ in 1..common_n {
        try!(writer.write_field(HEADER_COMMON));
    }
    let show_samples = show_genotypes || !format_tags.is_empty();
    if show_samples {
        for sample in reader.header().samples() {
            try!(writer.write_field(sample));
            for _ in 1..format_tags.len() + show_genotypes as usize {
                try!(writer.write_field(sample));
            }
        }
    }
    try!(writer.newline());
    try!(writer.write_field(b"CHROM"));
    try!(writer.write_field(b"POS"));
    try!(writer.write_field(b"REF"));
    try!(writer.write_field(b"ALT"));
    try!(writer.write_field(b"QUAL"));
    for name in info_tags {
        try!(writer.write_field(name.as_bytes()));
    }
    if show_samples {
        for _ in 0..reader.header().sample_count() {
            if show_genotypes {
                try!(writer.write_field(b"GT"));
            }
            for name in format_tags {
                try!(writer.write_field(name.as_bytes()));
            }
        }
    }
    try!(writer.newline());

    let mut record = bcf::Record::new();
    loop {
        if let Err(e) = reader.read(&mut record) {
            if e.is_eof() {
                break;
            } else {
                return Err(Box::new(e));
            }
        }

        let alleles = record.alleles().into_iter().map(|a| a.to_owned()).collect_vec();
        for (i, allele) in alleles[1..].iter().enumerate() {
            try!(writer.write_field(reader.header().rid2name(record.rid().unwrap())));
            try!(writer.write_integer(record.pos() as i32 + 1));
            try!(writer.write_field(&alleles[0]));
            try!(writer.write_field(allele));
            match record.qual() {
                q if q.is_missing() => try!(writer.write_field(b"")),
                q => try!(writer.write_float(q))
            }

            for name in info_tags {
                let _name = name.as_bytes();
                if let Ok((tag_type, tag_length)) = reader.header().info_type(_name) {
                    let get_idx = || {
                        match tag_length {
                            bcf::header::TagLength::Fixed => {
                                Ok(0)
                            },
                            bcf::header::TagLength::AltAlleles => {
                                Ok(i)
                            },
                            bcf::header::TagLength::Alleles => {
                                Ok(i + 1)
                            },
                            bcf::header::TagLength::Variable => {
                                Ok(0)
                            },
                            _ => Err(Box::new(ParseError::UnsupportedTagLength))
                        }
                    };

                    match tag_type {
                        bcf::header::TagType::Flag => {
                            try!(writer.write_flag(try!(record.info(_name).flag())));
                        },
                        bcf::header::TagType::Integer => {
                            let i = try!(get_idx());
                            if let Some(values) = try!(record.info(_name).integer()) {
                                try!(writer.write_integer(values[i]));
                            } else {
                                try!(writer.write_field(b""));
                            }
                        },
                        bcf::header::TagType::Float => {
                            let i = try!(get_idx());
                            if let Some(values) = try!(record.info(_name).float()) {
                                try!(writer.write_float(values[i]));
                            } else {
                                try!(writer.write_field(b""));
                            }
                        },
                        bcf::header::TagType::String => {
                            let i = try!(get_idx());
                            if let Some(values) = try!(record.info(_name).string()) {
                                try!(writer.write_field(values[i]));
                            } else {
                                try!(writer.write_field(b""));
                            }
                        }
                    }
                } else {
                    // tag undefined, write NA
                    try!(writer.write_field(b""));
                }
            }

            let genotypes = if show_genotypes {
                let genotypes = try!(record.genotypes());
                Some(
                    (0..reader.header().sample_count() as usize).map(|s| {
                        format!("{}", genotypes.get(s))
                    }).collect_vec()
                )
            } else {
                None
            };

            for s in 0..reader.header().sample_count() as usize {
                if let Some(ref genotypes) = genotypes {
                    try!(writer.write_field(genotypes[s].as_bytes()));
                }
                for name in format_tags {
                    let _name = name.as_bytes();
                    if let Ok((tag_type, tag_length)) = reader.header().format_type(_name) {
                        let i = match tag_length {
                            bcf::header::TagLength::Fixed => {
                                0
                            },
                            bcf::header::TagLength::AltAlleles => {
                                i
                            },
                            bcf::header::TagLength::Alleles => {
                                i + 1
                            },
                            _ => return Err(Box::new(ParseError::UnsupportedTagLength))
                        };

                        match tag_type {
                            bcf::header::TagType::Flag => {
                                panic!("there is no flag type for format");
                            },
                            bcf::header::TagType::Integer => {
                                try!(writer.write_field(format!("{}", try!(record.format(_name).integer())[s][i]).as_bytes()));
                            },
                            bcf::header::TagType::Float => {
                                try!(writer.write_field(format!("{}", try!(record.format(_name).float())[s][i]).as_bytes()));
                            },
                            bcf::header::TagType::String => {
                                try!(writer.write_field(try!(record.format(_name).string())[s]));
                            }
                        }
                    } else {
                        // tag undefined, write NA
                        try!(writer.write_field(b""));
                    }
                }
            }
            try!(writer.newline());
        }
    }

    Ok(())
}

quick_error! {
    #[derive(Debug)]
    pub enum ParseError {
        UnsupportedTagLength {
            description("currently, only R, A, and 1 are supported multiplicities of tags")
        }

    }
}
<commit_msg>Handle missing values.<commit_after>use std::error::Error;
use std::io;
use std::str;
use std::io::Write;

use itertools::Itertools;
use rust_htslib::bcf;
use rust_htslib::bcf::record::Numeric;

pub struct Writer {
    inner: io::BufWriter<io::Stdout>,
    field_count: usize
}

impl Writer {
    fn new(inner: io::BufWriter<io::Stdout>) -> Self {
        Writer {
            inner: inner,
            field_count: 0
        }
    }

    fn write_integer(&mut self, value: i32) -> Result<(), Box<Error>> {
        let fmt = if value.is_missing() {
            "".to_owned()
        } else {
            format!("{}", value)
        };
        self.write_field(fmt.as_bytes())
    }

    fn write_float(&mut self, value: f32) -> Result<(), Box<Error>> {
        let fmt = if value.is_missing() {
            "".to_owned()
        } else {
            format!("{}", value)
        };
        self.write_field(fmt.as_bytes())
    }

    fn write_flag(&mut self, value: bool) -> Result<(), Box<Error>> {
        self.write_field(format!("{}", value).as_bytes())
    }

    fn write_field(&mut self, value: &[u8]) -> Result<(), Box<Error>> {
        if self.field_count > 0 {
            try!(self.inner.write(b"\t"));
        }
        try!(self.inner.write(value));
        self.field_count += 1;
        Ok(())
    }

    fn newline(&mut self) -> Result<(), Box<Error>> {
        try!(self.inner.write(b"\n"));
        self.field_count = 0;
        Ok(())
    }
}

const HEADER_COMMON: &'static [u8] = b"VARIANT";

pub fn to_txt(
    info_tags: &[&str],
    format_tags: &[&str],
    show_genotypes: bool
) -> Result<(), Box<Error>> {
    let mut reader = bcf::Reader::from_stdin()?;
    let mut writer = Writer::new(io::BufWriter::new(io::stdout()));

    let common_n = 5 + info_tags.len();
    try!(writer.write_field(HEADER_COMMON));
    for _ in 1..common_n {
        try!(writer.write_field(HEADER_COMMON));
    }
    let show_samples = show_genotypes || !format_tags.is_empty();
    if show_samples {
        for sample in reader.header().samples() {
            try!(writer.write_field(sample));
            for _ in 1..format_tags.len() + show_genotypes as usize {
                try!(writer.write_field(sample));
            }
        }
    }
    try!(writer.newline());
    try!(writer.write_field(b"CHROM"));
    try!(writer.write_field(b"POS"));
    try!(writer.write_field(b"REF"));
    try!(writer.write_field(b"ALT"));
    try!(writer.write_field(b"QUAL"));
    for name in info_tags {
        try!(writer.write_field(name.as_bytes()));
    }
    if show_samples {
        for _ in 0..reader.header().sample_count() {
            if show_genotypes {
                try!(writer.write_field(b"GT"));
            }
            for name in format_tags {
                try!(writer.write_field(name.as_bytes()));
            }
        }
    }
    try!(writer.newline());

    let mut record = bcf::Record::new();
    loop {
        if let Err(e) = reader.read(&mut record) {
            if e.is_eof() {
                break;
            } else {
                return Err(Box::new(e));
            }
        }

        let alleles = record.alleles().into_iter().map(|a| a.to_owned()).collect_vec();
        for (i, allele) in alleles[1..].iter().enumerate() {
            try!(writer.write_field(reader.header().rid2name(record.rid().unwrap())));
            try!(writer.write_integer(record.pos() as i32 + 1));
            try!(writer.write_field(&alleles[0]));
            try!(writer.write_field(allele));
            match record.qual() {
                q if q.is_missing() => try!(writer.write_field(b"")),
                q => try!(writer.write_float(q))
            }

            for name in info_tags {
                let _name = name.as_bytes();
                if let Ok((tag_type, tag_length)) = reader.header().info_type(_name) {
                    let get_idx = || {
                        match tag_length {
                            bcf::header::TagLength::Fixed => {
                                Ok(0)
                            },
                            bcf::header::TagLength::AltAlleles => {
                                Ok(i)
                            },
                            bcf::header::TagLength::Alleles => {
                                Ok(i + 1)
                            },
                            bcf::header::TagLength::Variable => {
                                Ok(0)
                            },
                            _ => Err(Box::new(ParseError::UnsupportedTagLength))
                        }
                    };

                    match tag_type {
                        bcf::header::TagType::Flag => {
                            try!(writer.write_flag(try!(record.info(_name).flag())));
                        },
                        bcf::header::TagType::Integer => {
                            let i = try!(get_idx());
                            if let Some(values) = try!(record.info(_name).integer()) {
                                try!(writer.write_integer(values[i]));
                            } else {
                                try!(writer.write_field(b""));
                            }
                        },
                        bcf::header::TagType::Float => {
                            let i = try!(get_idx());
                            if let Some(values) = try!(record.info(_name).float()) {
                                try!(writer.write_float(values[i]));
                            } else {
                                try!(writer.write_field(b""));
                            }
                        },
                        bcf::header::TagType::String => {
                            let i = try!(get_idx());
                            if let Some(values) = try!(record.info(_name).string()) {
                                try!(writer.write_field(values[i]));
                            } else {
                                try!(writer.write_field(b""));
                            }
                        }
                    }
                } else {
                    // tag undefined, write NA
                    try!(writer.write_field(b""));
                }
            }

            let genotypes = if show_genotypes {
                let genotypes = try!(record.genotypes());
                Some(
                    (0..reader.header().sample_count() as usize).map(|s| {
                        format!("{}", genotypes.get(s))
                    }).collect_vec()
                )
            } else {
                None
            };

            for s in 0..reader.header().sample_count() as usize {
                if let Some(ref genotypes) = genotypes {
                    try!(writer.write_field(genotypes[s].as_bytes()));
                }
                for name in format_tags {
                    let _name = name.as_bytes();
                    if let Ok((tag_type, tag_length)) = reader.header().format_type(_name) {
                        let i = match tag_length {
                            bcf::header::TagLength::Fixed => {
                                0
                            },
                            bcf::header::TagLength::AltAlleles => {
                                i
                            },
                            bcf::header::TagLength::Alleles => {
                                i + 1
                            },
                            _ => return Err(Box::new(ParseError::UnsupportedTagLength))
                        };

                        match tag_type {
                            bcf::header::TagType::Flag => {
                                panic!("there is no flag type for format");
                            },
                            bcf::header::TagType::Integer => {
                                try!(writer.write_field(format!("{}", try!(record.format(_name).integer())[s][i]).as_bytes()));
                            },
                            bcf::header::TagType::Float => {
                                try!(writer.write_field(format!("{}", try!(record.format(_name).float())[s][i]).as_bytes()));
                            },
                            bcf::header::TagType::String => {
                                try!(writer.write_field(try!(record.format(_name).string())[s]));
                            }
                        }
                    } else {
                        // tag undefined, write NA
                        try!(writer.write_field(b""));
                    }
                }
            }
            try!(writer.newline());
        }
    }

    Ok(())
}

quick_error! {
    #[derive(Debug)]
    pub enum ParseError {
        UnsupportedTagLength {
            description("currently, only R, A, and 1 are supported multiplicities of tags")
        }

    }
}
<|endoftext|>


<commit_before>use std::f64::{MAX as F64_MAX};
use std::usize::{MAX as USIZE_MAX};
use std::ops::{Index, IndexMut};
use cgmath::prelude::*;
use event::CellEvent;
use super::Point;

const HEIGHT: usize = 6;

#[derive(Copy, Clone, Default)]
struct Spherical {
    theta: f64,
    phi: f64,
}

#[derive(Copy, Clone, PartialEq)]
pub struct ArcId(u32);

impl ArcId {
    pub const NONE: Self = ArcId(::std::u32::MAX);
}

impl Default for ArcId {
    fn default() -> Self {
        ArcId::NONE
    }
}

#[derive(Copy, Clone, Default)]
pub struct Link {
    pub prev: ArcId,
    pub next: ArcId,
}

pub struct Arc {
    pub cell_index: usize,
    pub focus: Point,
    pub vertex: Point,
    pub theta: f64,
    cached_intersection: Spherical,
    links: [Link; HEIGHT],
    heap_index: usize,
}

impl Arc {
    #[inline]
    pub fn neighbors(&self) -> Link {
        self.links[0]
    }

    #[inline]
    fn next(&self) -> ArcId {
        self.links[0].next
    }

    fn height(&self) -> usize {
        (0 .. HEIGHT).filter(|&level| self.links[level].next == ArcId::NONE).next().unwrap_or(HEIGHT)
    }
}

#[derive(Default)]
pub struct BeachLine {
    arcs: Vec<Arc>,
    free: Vec<ArcId>,
    head: ArcId,
    len: usize,
    levels: [usize; HEIGHT],
    heap: Vec<ArcId>,
}

impl BeachLine {
    pub fn insert(&mut self, cell_index: usize, event: &CellEvent) -> ArcId {
        let arc = self.create_arc(cell_index, event.point);
        if self.len > 1 {
            let mut current = self.head;
            let mut next = ArcId::NONE;
            let mut height = HEIGHT;
            let mut links = [ArcId::NONE; HEIGHT];
            let sin_theta = event.theta.sin();
            while height > 0 {
                next = self[current].links[height - 1].next;
                if self.intersect(current, event, sin_theta) < self.intersect(next, event, sin_theta) {
                    current = next;
                } else {
                    height -= 1;
                    links[height] = current;
                }
            }
            let twin_index = self[next].cell_index;
            let twin_focus = self[next].focus;
            let twin = self.create_arc(twin_index, twin_focus);
            self.add_links(twin, &mut links);
            self.add_links(arc, &mut links);
        } else {
            if self.len == 0 {
                self.head = arc;
            }
            let head = self.head;
            self.add_links(arc, &mut [head; HEIGHT]);
        }
        arc
    }

    pub fn remove(&mut self, arc: ArcId) {
        let head = self.head;
        if arc == head {
            let mut new_head = self[self.head].links[HEIGHT - 1].next;
            if new_head == self.head {
                new_head = self[self.head].next();
                // promote new_head to HEIGHT
                let height = self[new_head].height();
                self.levels[height - 1] -= 1;
                self.levels[HEIGHT - 1] += 1;
                for level in height .. HEIGHT {
                    let next = self[self.head].links[level].next;
                    self[next].links[level].prev = new_head;
                    self[new_head].links[level] = Link { prev: head, next };
                    self[head].links[level].next = new_head;
                }
            }
            self.head = new_head;
        }
        let height = self[arc].height();
        for level in 0 .. height {
            let link = self[arc].links[level];
            self[link.next].links[level].prev = link.prev;
            self[link.prev].links[level].next = link.next;
        }
        self.len -= 1;
        self.levels[height - 1] -= 1;
        self.free.push(arc);
    }

    pub fn heap_update(&mut self, arc: ArcId) {
        let mut index = self[arc].heap_index;
        let theta = self[arc].theta;
        if index == USIZE_MAX {
            index = self.heap.len();
            self.heap.push(arc);
        }
        while index > 0 {
            let parent_index = (index - 1) / 2;
            let parent = self.heap[parent_index];
            if theta < self[parent].theta {
                self.heap[index] = parent;
                self[parent].heap_index = index;
            } else {
                break;
            }
            index = parent_index;
        }
        self.heap[index] = arc;
        self[arc].heap_index = index;
    }

    pub fn has_vertices(&self) -> bool {
        !self.heap.is_empty()
    }

    pub fn heap_pop(&mut self) -> ArcId {
        let first = self.heap[0];
        let last = self.heap.pop().unwrap();
        let theta = self[last].theta;
        let mut index = 0;
        while 2 * index + 1 < self.heap.len() {
            // left child by default
            let mut child_index = 2 * index + 1;
            let mut child = self.heap[child_index];
            // check right child
            let right_index = child_index + 1;
            if right_index < self.heap.len() {
                let right_child = self.heap[right_index];
                if self[right_child].theta < self[child].theta {
                    child_index = right_index;
                    child = right_child;
                }
            }
            if theta < self[child].theta {
                break;
            }
            self.heap[index] = child;
            self[child].heap_index = index;
            index = child_index;
        }
        if self.heap.len() > 0 {
            self.heap[index] = last;
            self[last].heap_index = index;
        }
        first
    }

    pub fn clear(&mut self) {
        self.head = ArcId::NONE;
        self.len = 0;
        self.levels = [0; HEIGHT];
    }

    pub fn top_theta(&self) -> f64 {
        if let Some(&first) = self.heap.first() {
            self[first].theta
        } else {
            F64_MAX
        }
    }

    fn create_arc(&mut self, cell_index: usize, focus: Point) -> ArcId {
        let data = Arc {
            cell_index,
            focus,
            vertex: Point::zero(),
            theta: F64_MAX,
            cached_intersection: Spherical::default(),
            links: Default::default(),
            heap_index: USIZE_MAX,
        };
        if let Some(arc) = self.free.pop() {
            self[arc] = data;
            arc
        } else {
            self.arcs.push(data);
            ArcId(self.arcs.len() as u32 - 1)
        }
    }

    fn intersect(&mut self, arc: ArcId, event: &CellEvent, sin_theta: f64) -> f64 {
        let cached = self[arc].cached_intersection;
        let theta = event.theta;
        if cached.theta >= theta {
            return cached.phi;
        }
        let point0 = self[arc].focus;
        let point1 = self[self[arc].next()].focus;
        let phi = event.intersect(&point0, &point1, sin_theta);
        self[arc].cached_intersection = Spherical { theta, phi };
        phi
    }

    fn add_links(&mut self, arc: ArcId, links: &mut [ArcId; HEIGHT]) {
        let height = self.insertion_height();
        for level in 0 .. height {
            let prev = links[level];
            let mut next = self[prev].links[level].next;
            if next == ArcId::NONE {
                next = prev;
            }
            self[arc].links[level] = Link { prev, next };
            self[next].links[level].prev = arc;
            self[prev].links[level].next = arc;
            links[level] = arc;
        }
        self.len += 1;
        self.levels[height - 1] += 1;
    }

    fn insertion_height(&self) -> usize {
        if self.len == 0 {
            HEIGHT
        } else {
            1 + (0 .. HEIGHT).min_by_key(|&level| self.levels[level] * (1 << level)).unwrap()
        }
    }
}

impl Index<ArcId> for BeachLine {
    type Output = Arc;

    fn index(&self, arc: ArcId) -> &Self::Output {
        &self.arcs[arc.0 as usize]
    }
}

impl IndexMut<ArcId> for BeachLine {
    fn index_mut(&mut self, arc: ArcId) -> &mut Self::Output {
        &mut self.arcs[arc.0 as usize]
    }
}
<commit_msg>Smaller arc id<commit_after>use std::f64::{MAX as F64_MAX};
use std::usize::{MAX as USIZE_MAX};
use std::ops::{Index, IndexMut};
use cgmath::prelude::*;
use event::CellEvent;
use super::Point;

const HEIGHT: usize = 6;

#[derive(Copy, Clone, Default)]
struct Spherical {
    theta: f64,
    phi: f64,
}

#[derive(Copy, Clone, PartialEq)]
pub struct ArcId(u16);

impl ArcId {
    pub const NONE: Self = ArcId(::std::u16::MAX);
}

impl Default for ArcId {
    fn default() -> Self {
        ArcId::NONE
    }
}

#[derive(Copy, Clone, Default)]
pub struct Link {
    pub prev: ArcId,
    pub next: ArcId,
}

pub struct Arc {
    pub cell_index: usize,
    pub focus: Point,
    pub vertex: Point,
    pub theta: f64,
    cached_intersection: Spherical,
    links: [Link; HEIGHT],
    heap_index: usize,
}

impl Arc {
    #[inline]
    pub fn neighbors(&self) -> Link {
        self.links[0]
    }

    #[inline]
    fn next(&self) -> ArcId {
        self.links[0].next
    }

    fn height(&self) -> usize {
        (0 .. HEIGHT).filter(|&level| self.links[level].next == ArcId::NONE).next().unwrap_or(HEIGHT)
    }
}

#[derive(Default)]
pub struct BeachLine {
    arcs: Vec<Arc>,
    free: Vec<ArcId>,
    head: ArcId,
    len: usize,
    levels: [usize; HEIGHT],
    heap: Vec<ArcId>,
}

impl BeachLine {
    pub fn insert(&mut self, cell_index: usize, event: &CellEvent) -> ArcId {
        let arc = self.create_arc(cell_index, event.point);
        if self.len > 1 {
            let mut current = self.head;
            let mut next = ArcId::NONE;
            let mut height = HEIGHT;
            let mut links = [ArcId::NONE; HEIGHT];
            let sin_theta = event.theta.sin();
            while height > 0 {
                next = self[current].links[height - 1].next;
                if self.intersect(current, event, sin_theta) < self.intersect(next, event, sin_theta) {
                    current = next;
                } else {
                    height -= 1;
                    links[height] = current;
                }
            }
            let twin_index = self[next].cell_index;
            let twin_focus = self[next].focus;
            let twin = self.create_arc(twin_index, twin_focus);
            self.add_links(twin, &mut links);
            self.add_links(arc, &mut links);
        } else {
            if self.len == 0 {
                self.head = arc;
            }
            let head = self.head;
            self.add_links(arc, &mut [head; HEIGHT]);
        }
        arc
    }

    pub fn remove(&mut self, arc: ArcId) {
        let head = self.head;
        if arc == head {
            let mut new_head = self[self.head].links[HEIGHT - 1].next;
            if new_head == self.head {
                new_head = self[self.head].next();
                // promote new_head to HEIGHT
                let height = self[new_head].height();
                self.levels[height - 1] -= 1;
                self.levels[HEIGHT - 1] += 1;
                for level in height .. HEIGHT {
                    let next = self[self.head].links[level].next;
                    self[next].links[level].prev = new_head;
                    self[new_head].links[level] = Link { prev: head, next };
                    self[head].links[level].next = new_head;
                }
            }
            self.head = new_head;
        }
        let height = self[arc].height();
        for level in 0 .. height {
            let link = self[arc].links[level];
            self[link.next].links[level].prev = link.prev;
            self[link.prev].links[level].next = link.next;
        }
        self.len -= 1;
        self.levels[height - 1] -= 1;
        self.free.push(arc);
    }

    pub fn heap_update(&mut self, arc: ArcId) {
        let mut index = self[arc].heap_index;
        let theta = self[arc].theta;
        if index == USIZE_MAX {
            index = self.heap.len();
            self.heap.push(arc);
        }
        while index > 0 {
            let parent_index = (index - 1) / 2;
            let parent = self.heap[parent_index];
            if theta < self[parent].theta {
                self.heap[index] = parent;
                self[parent].heap_index = index;
            } else {
                break;
            }
            index = parent_index;
        }
        self.heap[index] = arc;
        self[arc].heap_index = index;
    }

    pub fn has_vertices(&self) -> bool {
        !self.heap.is_empty()
    }

    pub fn heap_pop(&mut self) -> ArcId {
        let first = self.heap[0];
        let last = self.heap.pop().unwrap();
        let theta = self[last].theta;
        let mut index = 0;
        while 2 * index + 1 < self.heap.len() {
            // left child by default
            let mut child_index = 2 * index + 1;
            let mut child = self.heap[child_index];
            // check right child
            let right_index = child_index + 1;
            if right_index < self.heap.len() {
                let right_child = self.heap[right_index];
                if self[right_child].theta < self[child].theta {
                    child_index = right_index;
                    child = right_child;
                }
            }
            if theta < self[child].theta {
                break;
            }
            self.heap[index] = child;
            self[child].heap_index = index;
            index = child_index;
        }
        if self.heap.len() > 0 {
            self.heap[index] = last;
            self[last].heap_index = index;
        }
        first
    }

    pub fn clear(&mut self) {
        self.head = ArcId::NONE;
        self.len = 0;
        self.levels = [0; HEIGHT];
    }

    pub fn top_theta(&self) -> f64 {
        if let Some(&first) = self.heap.first() {
            self[first].theta
        } else {
            F64_MAX
        }
    }

    fn create_arc(&mut self, cell_index: usize, focus: Point) -> ArcId {
        let data = Arc {
            cell_index,
            focus,
            vertex: Point::zero(),
            theta: F64_MAX,
            cached_intersection: Spherical::default(),
            links: Default::default(),
            heap_index: USIZE_MAX,
        };
        if let Some(arc) = self.free.pop() {
            self[arc] = data;
            arc
        } else {
            self.arcs.push(data);
            ArcId(self.arcs.len() as u16 - 1)
        }
    }

    fn intersect(&mut self, arc: ArcId, event: &CellEvent, sin_theta: f64) -> f64 {
        let cached = self[arc].cached_intersection;
        let theta = event.theta;
        if cached.theta >= theta {
            return cached.phi;
        }
        let point0 = self[arc].focus;
        let point1 = self[self[arc].next()].focus;
        let phi = event.intersect(&point0, &point1, sin_theta);
        self[arc].cached_intersection = Spherical { theta, phi };
        phi
    }

    fn add_links(&mut self, arc: ArcId, links: &mut [ArcId; HEIGHT]) {
        let height = self.insertion_height();
        for level in 0 .. height {
            let prev = links[level];
            let mut next = self[prev].links[level].next;
            if next == ArcId::NONE {
                next = prev;
            }
            self[arc].links[level] = Link { prev, next };
            self[next].links[level].prev = arc;
            self[prev].links[level].next = arc;
            links[level] = arc;
        }
        self.len += 1;
        self.levels[height - 1] += 1;
    }

    fn insertion_height(&self) -> usize {
        if self.len == 0 {
            HEIGHT
        } else {
            1 + (0 .. HEIGHT).min_by_key(|&level| self.levels[level] * (1 << level)).unwrap()
        }
    }
}

impl Index<ArcId> for BeachLine {
    type Output = Arc;

    fn index(&self, arc: ArcId) -> &Self::Output {
        &self.arcs[arc.0 as usize]
    }
}

impl IndexMut<ArcId> for BeachLine {
    fn index_mut(&mut self, arc: ArcId) -> &mut Self::Output {
        &mut self.arcs[arc.0 as usize]
    }
}
<|endoftext|>


<commit_before>extern crate mdbook;
#[macro_use]
extern crate clap;
extern crate log;
extern crate env_logger;
extern crate open;

// Dependencies for the Watch feature
#[cfg(feature = "watch")]
extern crate notify;
#[cfg(feature = "watch")]
extern crate time;
#[cfg(feature = "watch")]
extern crate crossbeam;

// Dependencies for the Serve feature
#[cfg(feature = "serve")]
extern crate iron;
#[cfg(feature = "serve")]
extern crate staticfile;
#[cfg(feature = "serve")]
extern crate ws;

use std::env;
use std::error::Error;
use std::ffi::OsStr;
use std::io::{self, Write};
use std::path::{Path, PathBuf};

use clap::{App, ArgMatches, SubCommand, AppSettings};

// Uses for the Watch feature
#[cfg(feature = "watch")]
use notify::Watcher;
#[cfg(feature = "watch")]
use std::time::Duration;
#[cfg(feature = "watch")]
use std::sync::mpsc::channel;

use mdbook::MDBook;

const NAME: &'static str = "mdbook";

fn main() {
    env_logger::init().unwrap();

    // Create a list of valid arguments and sub-commands
    let matches = App::new(NAME)
                    .about("Create a book in form of a static website from markdown files")
                    .author("Mathieu David <mathieudavid@mathieudavid.org>")
                    // Get the version from our Cargo.toml using clap's crate_version!() macro
                    .version(&*format!("v{}", crate_version!()))
                    .setting(AppSettings::SubcommandRequired)
                    .after_help("For more information about a specific command, try `mdbook <command> --help`\nSource code for mdbook available at: https://github.com/azerupi/mdBook")
                    .subcommand(SubCommand::with_name("init")
                        .about("Create boilerplate structure and files in the directory")
                        // the {n} denotes a newline which will properly aligned in all help messages
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'")
                        .arg_from_usage("--theme 'Copies the default theme into your source folder'")
                        .arg_from_usage("--force 'skip confirmation prompts'"))
                    .subcommand(SubCommand::with_name("build")
                        .about("Build the book from the markdown files")
                        .arg_from_usage("-o, --open 'Open the compiled book in a web browser'")
                        .arg_from_usage("-d, --dest-dir=[dest-dir] 'The output directory for your book{n}(Defaults to ./book when omitted)'")
                        .arg_from_usage("--no-create 'Will not create non-existent files linked from SUMMARY.md'")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'"))
                    .subcommand(SubCommand::with_name("watch")
                        .about("Watch the files for changes")
                        .arg_from_usage("-o, --open 'Open the compiled book in a web browser'")
                        .arg_from_usage("-d, --dest-dir=[dest-dir] 'The output directory for your book{n}(Defaults to ./book when omitted)'")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'"))
                    .subcommand(SubCommand::with_name("serve")
                        .about("Serve the book at http://localhost:3000. Rebuild and reload on change.")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'")
                        .arg_from_usage("-d, --dest-dir=[dest-dir] 'The output directory for your book{n}(Defaults to ./book when omitted)'")
                        .arg_from_usage("-p, --port=[port] 'Use another port{n}(Defaults to 3000)'")
                        .arg_from_usage("-w, --websocket-port=[ws-port] 'Use another port for the websocket connection (livereload){n}(Defaults to 3001)'")
                        .arg_from_usage("-i, --interface=[interface] 'Interface to listen on{n}(Defaults to localhost)'")
                        .arg_from_usage("-a, --address=[address] 'Address that the browser can reach the websocket server from{n}(Defaults to the interface address)'")
                        .arg_from_usage("-o, --open 'Open the book server in a web browser'"))
                    .subcommand(SubCommand::with_name("test")
                        .about("Test that code samples compile"))
                    .get_matches();

    // Check which subcomamnd the user ran...
    let res = match matches.subcommand() {
        ("init", Some(sub_matches)) => init(sub_matches),
        ("build", Some(sub_matches)) => build(sub_matches),
        #[cfg(feature = "watch")]
        ("watch", Some(sub_matches)) => watch(sub_matches),
        #[cfg(feature = "serve")]
        ("serve", Some(sub_matches)) => serve(sub_matches),
        ("test", Some(sub_matches)) => test(sub_matches),
        (_, _) => unreachable!(),
    };

    if let Err(e) = res {
        writeln!(&mut io::stderr(), "An error occured:\n{}", e).ok();
        ::std::process::exit(101);
    }
}

// Simple function that user comfirmation
fn confirm() -> bool {
    io::stdout().flush().unwrap();
    let mut s = String::new();
    io::stdin().read_line(&mut s).ok();
    match &*s.trim() {
        "Y" | "y" | "yes" | "Yes" => true,
        _ => false,
    }
}

// Init command implementation
fn init(args: &ArgMatches) -> Result<(), Box<Error>> {

    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir);

    // Call the function that does the initialization
    book.init()?;

    // If flag `--theme` is present, copy theme to src
    if args.is_present("theme") {

        // Skip this if `--force` is present
        if !args.is_present("force") {
            // Print warning
            print!("\nCopying the default theme to {:?}", book.get_src());
            println!("could potentially overwrite files already present in that directory.");
            print!("\nAre you sure you want to continue? (y/n) ");

            // Read answer from user and exit if it's not 'yes'
            if !confirm() {
                println!("\nSkipping...\n");
                println!("All done, no errors...");
                ::std::process::exit(0);
            }
        }

        // Call the function that copies the theme
        book.copy_theme()?;
        println!("\nTheme copied.");

    }

    // Because of `src/book/mdbook.rs#L37-L39`, `dest` will always start with `root`
    let is_dest_inside_root = book.get_dest().starts_with(book.get_root());

    if !args.is_present("force") && is_dest_inside_root {
        println!("\nDo you want a .gitignore to be created? (y/n)");

        if confirm() {
            book.create_gitignore();
            println!("\n.gitignore created.");
        }
    }

    println!("\nAll done, no errors...");

    Ok(())
}

// Build command implementation
fn build(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let book = MDBook::new(&book_dir).read_config();

    let mut book = match args.value_of("dest-dir") {
        Some(dest_dir) => book.set_dest(Path::new(dest_dir)),
        None => book,
    };

    if args.is_present("no-create") {
        book.create_missing = false;
    }

    book.build()?;

    if args.is_present("open") {
        open(book.get_dest().join("index.html"));
    }

    Ok(())
}

// Watch command implementation
#[cfg(feature = "watch")]
fn watch(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let book = MDBook::new(&book_dir).read_config();

    let mut book = match args.value_of("dest-dir") {
        Some(dest_dir) => book.set_dest(Path::new(dest_dir)),
        None => book,
    };

    if args.is_present("open") {
        book.build()?;
        open(book.get_dest().join("index.html"));
    }

    trigger_on_change(&mut book, |path, book| {
        println!("File changed: {:?}\nBuilding book...\n", path);
        if let Err(e) = book.build() {
            println!("Error while building: {:?}", e);
        }
        println!("");
    });

    Ok(())
}

// Watch command implementation
#[cfg(feature = "serve")]
fn serve(args: &ArgMatches) -> Result<(), Box<Error>> {
    const RELOAD_COMMAND: &'static str = "reload";

    let book_dir = get_book_dir(args);
    let book = MDBook::new(&book_dir).read_config();

    let mut book = match args.value_of("dest-dir") {
        Some(dest_dir) => book.set_dest(Path::new(dest_dir)),
        None => book,
    };

    let port = args.value_of("port").unwrap_or("3000");
    let ws_port = args.value_of("ws-port").unwrap_or("3001");
    let interface = args.value_of("interface").unwrap_or("localhost");
    let public_address = args.value_of("address").unwrap_or(interface);
    let open_browser = args.is_present("open");

    let address = format!("{}:{}", interface, port);
    let ws_address = format!("{}:{}", interface, ws_port);

    book.set_livereload(format!(r#"
        <script type="text/javascript">
            var socket = new WebSocket("ws://{}:{}");
            socket.onmessage = function (event) {{
                if (event.data === "{}") {{
                    socket.close();
                    location.reload(true); // force reload from server (not from cache)
                }}
            }};

            window.onbeforeunload = function() {{
                socket.close();
            }}
        </script>
    "#,
                                public_address,
                                ws_port,
                                RELOAD_COMMAND));

    book.build()?;

    let staticfile = staticfile::Static::new(book.get_dest());
    let iron = iron::Iron::new(staticfile);
    let _iron = iron.http(&*address).unwrap();

    let ws_server = ws::WebSocket::new(|_| |_| Ok(())).unwrap();

    let broadcaster = ws_server.broadcaster();

    std::thread::spawn(move || { ws_server.listen(&*ws_address).unwrap(); });

    println!("\nServing on {}", address);

    if open_browser {
        open(format!("http://{}", address));
    }

    trigger_on_change(&mut book, move |path, book| {
        println!("File changed: {:?}\nBuilding book...\n", path);
        match book.build() {
            Err(e) => println!("Error while building: {:?}", e),
            _ => broadcaster.send(RELOAD_COMMAND).unwrap(),
        }
        println!("");
    });

    Ok(())
}

fn test(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config();

    book.test()?;

    Ok(())
}

fn get_book_dir(args: &ArgMatches) -> PathBuf {
    if let Some(dir) = args.value_of("dir") {
        // Check if path is relative from current dir, or absolute...
        let p = Path::new(dir);
        if p.is_relative() {
            env::current_dir().unwrap().join(dir)
        } else {
            p.to_path_buf()
        }
    } else {
        env::current_dir().unwrap()
    }
}

fn open<P: AsRef<OsStr>>(path: P) {
    if let Err(e) = open::that(path) {
        println!("Error opening web browser: {}", e);
    }
}

// Calls the closure when a book source file is changed. This is blocking!
#[cfg(feature = "watch")]
fn trigger_on_change<F>(book: &mut MDBook, closure: F) -> ()
    where F: Fn(&Path, &mut MDBook) -> ()
{
    use notify::RecursiveMode::*;
    use notify::DebouncedEvent::*;

    // Create a channel to receive the events.
    let (tx, rx) = channel();

    let mut watcher = match notify::watcher(tx, Duration::from_secs(1)) {
        Ok(w) => w,
        Err(e) => {
            println!("Error while trying to watch the files:\n\n\t{:?}", e);
            ::std::process::exit(0);
        },
    };

    // Add the source directory to the watcher
    if let Err(e) = watcher.watch(book.get_src(), Recursive) {
        println!("Error while watching {:?}:\n    {:?}", book.get_src(), e);
        ::std::process::exit(0);
    };

    // Add the book.{json,toml} file to the watcher if it exists, because it's not
    // located in the source directory
    if watcher
           .watch(book.get_root().join("book.json"), NonRecursive)
           .is_err() {
        // do nothing if book.json is not found
    }
    if watcher
           .watch(book.get_root().join("book.toml"), NonRecursive)
           .is_err() {
        // do nothing if book.toml is not found
    }

    println!("\nListening for changes...\n");

    loop {
        match rx.recv() {
            Ok(event) => {
                match event {
                    NoticeWrite(path) |
                    NoticeRemove(path) |
                    Create(path) |
                    Write(path) |
                    Remove(path) |
                    Rename(_, path) => {
                        closure(&path, book);
                    },
                    _ => {},
                }
            },
            Err(e) => {
                println!("An error occured: {:?}", e);
            },
        }
    }
}
<commit_msg>Fix websocket port option for serve<commit_after>extern crate mdbook;
#[macro_use]
extern crate clap;
extern crate log;
extern crate env_logger;
extern crate open;

// Dependencies for the Watch feature
#[cfg(feature = "watch")]
extern crate notify;
#[cfg(feature = "watch")]
extern crate time;
#[cfg(feature = "watch")]
extern crate crossbeam;

// Dependencies for the Serve feature
#[cfg(feature = "serve")]
extern crate iron;
#[cfg(feature = "serve")]
extern crate staticfile;
#[cfg(feature = "serve")]
extern crate ws;

use std::env;
use std::error::Error;
use std::ffi::OsStr;
use std::io::{self, Write};
use std::path::{Path, PathBuf};

use clap::{App, ArgMatches, SubCommand, AppSettings};

// Uses for the Watch feature
#[cfg(feature = "watch")]
use notify::Watcher;
#[cfg(feature = "watch")]
use std::time::Duration;
#[cfg(feature = "watch")]
use std::sync::mpsc::channel;

use mdbook::MDBook;

const NAME: &'static str = "mdbook";

fn main() {
    env_logger::init().unwrap();

    // Create a list of valid arguments and sub-commands
    let matches = App::new(NAME)
                    .about("Create a book in form of a static website from markdown files")
                    .author("Mathieu David <mathieudavid@mathieudavid.org>")
                    // Get the version from our Cargo.toml using clap's crate_version!() macro
                    .version(&*format!("v{}", crate_version!()))
                    .setting(AppSettings::SubcommandRequired)
                    .after_help("For more information about a specific command, try `mdbook <command> --help`\nSource code for mdbook available at: https://github.com/azerupi/mdBook")
                    .subcommand(SubCommand::with_name("init")
                        .about("Create boilerplate structure and files in the directory")
                        // the {n} denotes a newline which will properly aligned in all help messages
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'")
                        .arg_from_usage("--theme 'Copies the default theme into your source folder'")
                        .arg_from_usage("--force 'skip confirmation prompts'"))
                    .subcommand(SubCommand::with_name("build")
                        .about("Build the book from the markdown files")
                        .arg_from_usage("-o, --open 'Open the compiled book in a web browser'")
                        .arg_from_usage("-d, --dest-dir=[dest-dir] 'The output directory for your book{n}(Defaults to ./book when omitted)'")
                        .arg_from_usage("--no-create 'Will not create non-existent files linked from SUMMARY.md'")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'"))
                    .subcommand(SubCommand::with_name("watch")
                        .about("Watch the files for changes")
                        .arg_from_usage("-o, --open 'Open the compiled book in a web browser'")
                        .arg_from_usage("-d, --dest-dir=[dest-dir] 'The output directory for your book{n}(Defaults to ./book when omitted)'")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'"))
                    .subcommand(SubCommand::with_name("serve")
                        .about("Serve the book at http://localhost:3000. Rebuild and reload on change.")
                        .arg_from_usage("[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'")
                        .arg_from_usage("-d, --dest-dir=[dest-dir] 'The output directory for your book{n}(Defaults to ./book when omitted)'")
                        .arg_from_usage("-p, --port=[port] 'Use another port{n}(Defaults to 3000)'")
                        .arg_from_usage("-w, --ws-port=[ws-port] 'Use another port for the websocket connection (livereload){n}(Defaults to 3001)'")
                        .arg_from_usage("-i, --interface=[interface] 'Interface to listen on{n}(Defaults to localhost)'")
                        .arg_from_usage("-a, --address=[address] 'Address that the browser can reach the websocket server from{n}(Defaults to the interface address)'")
                        .arg_from_usage("-o, --open 'Open the book server in a web browser'"))
                    .subcommand(SubCommand::with_name("test")
                        .about("Test that code samples compile"))
                    .get_matches();

    // Check which subcomamnd the user ran...
    let res = match matches.subcommand() {
        ("init", Some(sub_matches)) => init(sub_matches),
        ("build", Some(sub_matches)) => build(sub_matches),
        #[cfg(feature = "watch")]
        ("watch", Some(sub_matches)) => watch(sub_matches),
        #[cfg(feature = "serve")]
        ("serve", Some(sub_matches)) => serve(sub_matches),
        ("test", Some(sub_matches)) => test(sub_matches),
        (_, _) => unreachable!(),
    };

    if let Err(e) = res {
        writeln!(&mut io::stderr(), "An error occured:\n{}", e).ok();
        ::std::process::exit(101);
    }
}

// Simple function that user comfirmation
fn confirm() -> bool {
    io::stdout().flush().unwrap();
    let mut s = String::new();
    io::stdin().read_line(&mut s).ok();
    match &*s.trim() {
        "Y" | "y" | "yes" | "Yes" => true,
        _ => false,
    }
}

// Init command implementation
fn init(args: &ArgMatches) -> Result<(), Box<Error>> {

    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir);

    // Call the function that does the initialization
    book.init()?;

    // If flag `--theme` is present, copy theme to src
    if args.is_present("theme") {

        // Skip this if `--force` is present
        if !args.is_present("force") {
            // Print warning
            print!("\nCopying the default theme to {:?}", book.get_src());
            println!("could potentially overwrite files already present in that directory.");
            print!("\nAre you sure you want to continue? (y/n) ");

            // Read answer from user and exit if it's not 'yes'
            if !confirm() {
                println!("\nSkipping...\n");
                println!("All done, no errors...");
                ::std::process::exit(0);
            }
        }

        // Call the function that copies the theme
        book.copy_theme()?;
        println!("\nTheme copied.");

    }

    // Because of `src/book/mdbook.rs#L37-L39`, `dest` will always start with `root`
    let is_dest_inside_root = book.get_dest().starts_with(book.get_root());

    if !args.is_present("force") && is_dest_inside_root {
        println!("\nDo you want a .gitignore to be created? (y/n)");

        if confirm() {
            book.create_gitignore();
            println!("\n.gitignore created.");
        }
    }

    println!("\nAll done, no errors...");

    Ok(())
}

// Build command implementation
fn build(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let book = MDBook::new(&book_dir).read_config();

    let mut book = match args.value_of("dest-dir") {
        Some(dest_dir) => book.set_dest(Path::new(dest_dir)),
        None => book,
    };

    if args.is_present("no-create") {
        book.create_missing = false;
    }

    book.build()?;

    if args.is_present("open") {
        open(book.get_dest().join("index.html"));
    }

    Ok(())
}

// Watch command implementation
#[cfg(feature = "watch")]
fn watch(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let book = MDBook::new(&book_dir).read_config();

    let mut book = match args.value_of("dest-dir") {
        Some(dest_dir) => book.set_dest(Path::new(dest_dir)),
        None => book,
    };

    if args.is_present("open") {
        book.build()?;
        open(book.get_dest().join("index.html"));
    }

    trigger_on_change(&mut book, |path, book| {
        println!("File changed: {:?}\nBuilding book...\n", path);
        if let Err(e) = book.build() {
            println!("Error while building: {:?}", e);
        }
        println!("");
    });

    Ok(())
}

// Watch command implementation
#[cfg(feature = "serve")]
fn serve(args: &ArgMatches) -> Result<(), Box<Error>> {
    const RELOAD_COMMAND: &'static str = "reload";

    let book_dir = get_book_dir(args);
    let book = MDBook::new(&book_dir).read_config();

    let mut book = match args.value_of("dest-dir") {
        Some(dest_dir) => book.set_dest(Path::new(dest_dir)),
        None => book,
    };

    let port = args.value_of("port").unwrap_or("3000");
    let ws_port = args.value_of("ws-port").unwrap_or("3001");
    let interface = args.value_of("interface").unwrap_or("localhost");
    let public_address = args.value_of("address").unwrap_or(interface);
    let open_browser = args.is_present("open");

    let address = format!("{}:{}", interface, port);
    let ws_address = format!("{}:{}", interface, ws_port);

    book.set_livereload(format!(r#"
        <script type="text/javascript">
            var socket = new WebSocket("ws://{}:{}");
            socket.onmessage = function (event) {{
                if (event.data === "{}") {{
                    socket.close();
                    location.reload(true); // force reload from server (not from cache)
                }}
            }};

            window.onbeforeunload = function() {{
                socket.close();
            }}
        </script>
    "#,
                                public_address,
                                ws_port,
                                RELOAD_COMMAND));

    book.build()?;

    let staticfile = staticfile::Static::new(book.get_dest());
    let iron = iron::Iron::new(staticfile);
    let _iron = iron.http(&*address).unwrap();

    let ws_server = ws::WebSocket::new(|_| |_| Ok(())).unwrap();

    let broadcaster = ws_server.broadcaster();

    std::thread::spawn(move || { ws_server.listen(&*ws_address).unwrap(); });

    println!("\nServing on {}", address);

    if open_browser {
        open(format!("http://{}", address));
    }

    trigger_on_change(&mut book, move |path, book| {
        println!("File changed: {:?}\nBuilding book...\n", path);
        match book.build() {
            Err(e) => println!("Error while building: {:?}", e),
            _ => broadcaster.send(RELOAD_COMMAND).unwrap(),
        }
        println!("");
    });

    Ok(())
}

fn test(args: &ArgMatches) -> Result<(), Box<Error>> {
    let book_dir = get_book_dir(args);
    let mut book = MDBook::new(&book_dir).read_config();

    book.test()?;

    Ok(())
}

fn get_book_dir(args: &ArgMatches) -> PathBuf {
    if let Some(dir) = args.value_of("dir") {
        // Check if path is relative from current dir, or absolute...
        let p = Path::new(dir);
        if p.is_relative() {
            env::current_dir().unwrap().join(dir)
        } else {
            p.to_path_buf()
        }
    } else {
        env::current_dir().unwrap()
    }
}

fn open<P: AsRef<OsStr>>(path: P) {
    if let Err(e) = open::that(path) {
        println!("Error opening web browser: {}", e);
    }
}

// Calls the closure when a book source file is changed. This is blocking!
#[cfg(feature = "watch")]
fn trigger_on_change<F>(book: &mut MDBook, closure: F) -> ()
    where F: Fn(&Path, &mut MDBook) -> ()
{
    use notify::RecursiveMode::*;
    use notify::DebouncedEvent::*;

    // Create a channel to receive the events.
    let (tx, rx) = channel();

    let mut watcher = match notify::watcher(tx, Duration::from_secs(1)) {
        Ok(w) => w,
        Err(e) => {
            println!("Error while trying to watch the files:\n\n\t{:?}", e);
            ::std::process::exit(0);
        },
    };

    // Add the source directory to the watcher
    if let Err(e) = watcher.watch(book.get_src(), Recursive) {
        println!("Error while watching {:?}:\n    {:?}", book.get_src(), e);
        ::std::process::exit(0);
    };

    // Add the book.{json,toml} file to the watcher if it exists, because it's not
    // located in the source directory
    if watcher
           .watch(book.get_root().join("book.json"), NonRecursive)
           .is_err() {
        // do nothing if book.json is not found
    }
    if watcher
           .watch(book.get_root().join("book.toml"), NonRecursive)
           .is_err() {
        // do nothing if book.toml is not found
    }

    println!("\nListening for changes...\n");

    loop {
        match rx.recv() {
            Ok(event) => {
                match event {
                    NoticeWrite(path) |
                    NoticeRemove(path) |
                    Create(path) |
                    Write(path) |
                    Remove(path) |
                    Rename(_, path) => {
                        closure(&path, book);
                    },
                    _ => {},
                }
            },
            Err(e) => {
                println!("An error occured: {:?}", e);
            },
        }
    }
}
<|endoftext|>


<commit_before>use console::{Console, ConsoleTextKind};
use memmap::{Mmap, Protection};
use pipeline_matcher::PathMatch;
use std::env;
use std::fs;
use std::io;
use std::io::{Write, Error};
use std::process;
use std::sync::mpsc::{Receiver, Sender};
use time;
use util::{catch, decode_error, PipelineInfo};
use tempfile::NamedTempFile;

// ---------------------------------------------------------------------------------------------------------------------
// PipelineReplacer
// ---------------------------------------------------------------------------------------------------------------------

pub trait PipelineReplacer {
    fn replace( &mut self, replacement: &[u8], rx: Receiver<PipelineInfo<PathMatch>>, tx: Sender<PipelineInfo<()>> );
}

// ---------------------------------------------------------------------------------------------------------------------
// SimplePipelineReplacer
// ---------------------------------------------------------------------------------------------------------------------

pub struct SimplePipelineReplacer {
    pub is_color      : bool,
    pub is_interactive: bool,
    pub print_file    : bool,
    pub print_column  : bool,
    pub infos         : Vec<String>,
    pub errors        : Vec<String>,
    console           : Console,
    all_replace       : bool,
    time_beg          : u64,
    time_end          : u64,
    time_bsy          : u64,
}

impl SimplePipelineReplacer {
    pub fn new() -> Self {
        SimplePipelineReplacer {
            is_color      : true,
            is_interactive: true,
            print_file    : true,
            print_column  : true,
            infos         : Vec::new(),
            errors        : Vec::new(),
            console       : Console::new(),
            all_replace   : false,
            time_beg      : 0,
            time_end      : 0,
            time_bsy      : 0,
        }
    }

    fn replace_match( &mut self, replacement: &[u8], pm: PathMatch ) {
        if pm.matches.is_empty() { return; }
        self.console.is_color = self.is_color;

        let result = catch::<_, (), Error> ( || {
            let tmpfile_dir = env::temp_dir();
            let mut tmpfile = try!( NamedTempFile::new_in( tmpfile_dir ) );

            {
                let mmap = try!( Mmap::open_path( &pm.path, Protection::Read ) );
                let src  = unsafe { mmap.as_slice() };

                let mut i       = 0;
                let mut pos     = 0;
                let mut column  = 0;
                let mut last_lf = 0;
                for m in &pm.matches {
                    try!( tmpfile.write_all( &src[i..m.beg] ) );

                    let mut do_replace = true;
                    if self.is_interactive & !self.all_replace {
                        if self.print_file {
                            self.console.write( ConsoleTextKind::Filename, pm.path.to_str().unwrap() );
                            self.console.write( ConsoleTextKind::Other, ": " );
                        }
                        if self.print_column {
                            while pos < m.beg {
                                if src[pos] == 0x0a {
                                    column += 1;
                                    last_lf = pos;
                                }
                                pos += 1;
                            }
                            self.console.write( ConsoleTextKind::Other, &format!( "{}:{}:", column + 1, m.beg - last_lf ) );
                        }

                        self.console.write_match_line( src, m );

                        loop {
                            self.console.write( ConsoleTextKind::Other, "Replace keyword? ( Yes[Y], No[N], All[A], Quit[Q] ):" );
                            let mut buf = String::new();
                            io::stdin().read_line( &mut buf ).unwrap();
                            match buf.trim().to_lowercase().as_ref() {
                                "y"    => { do_replace = true ; break },
                                "yes"  => { do_replace = true ; break },
                                "n"    => { do_replace = false; break },
                                "no"   => { do_replace = false; break },
                                "a"    => { self.all_replace = true ; break },
                                "all"  => { self.all_replace = true ; break },
                                "q"    => { let _ = tmpfile.close(); process::exit( 0 ) },
                                "quit" => { let _ = tmpfile.close(); process::exit( 0 ) },
                                _      => continue,
                            }
                        }
                    }

                    if do_replace {
                        try!( tmpfile.write_all( &replacement ) );
                    } else {
                        try!( tmpfile.write_all( &src[m.beg..m.end] ) );
                    }
                    i = m.end;
                }

                if i < src.len() {
                    try!( tmpfile.write_all( &src[i..src.len()] ) );
                }
                try!( tmpfile.flush() );
            }

            let metadata = try!( fs::metadata( &pm.path ) );

            try!( fs::rename( tmpfile.path(), &pm.path ) );
            try!( fs::set_permissions( &pm.path, metadata.permissions() ) );

            Ok(())
        } );
        match result {
            Ok ( _ ) => (),
            Err( e ) => self.console.write( ConsoleTextKind::Error, &format!( "Error: {} @ {:?}\n", decode_error( e.kind() ), pm.path ) ),
        }
    }
}

impl PipelineReplacer for SimplePipelineReplacer {
    fn replace( &mut self, replacement: &[u8], rx: Receiver<PipelineInfo<PathMatch>>, tx: Sender<PipelineInfo<()>> ) {
        loop {
            match rx.recv() {
                Ok( PipelineInfo::Ok( pm ) ) => {
                    let beg = time::precise_time_ns();

                    self.replace_match( replacement, pm );
                    let _ = tx.send( PipelineInfo::Ok( () ) );

                    let end = time::precise_time_ns();
                    self.time_bsy += end - beg;
                },

                Ok( PipelineInfo::Beg( x ) ) => {
                    self.infos  = Vec::new();
                    self.errors = Vec::new();

                    self.time_beg = time::precise_time_ns();
                    let _ = tx.send( PipelineInfo::Beg( x ) );
                },

                Ok( PipelineInfo::End( x ) ) => {
                    for i in &self.infos  { let _ = tx.send( PipelineInfo::Info( i.clone() ) ); }
                    for e in &self.errors { let _ = tx.send( PipelineInfo::Err ( e.clone() ) ); }

                    self.time_end = time::precise_time_ns();
                    let _ = tx.send( PipelineInfo::Time( self.time_bsy, self.time_end - self.time_beg ) );
                    let _ = tx.send( PipelineInfo::End( x ) );
                    break;
                },

                Ok ( PipelineInfo::Info( e      ) ) => { let _ = tx.send( PipelineInfo::Info( e      ) ); },
                Ok ( PipelineInfo::Err ( e      ) ) => { let _ = tx.send( PipelineInfo::Err ( e      ) ); },
                Ok ( PipelineInfo::Time( t0, t1 ) ) => { let _ = tx.send( PipelineInfo::Time( t0, t1 ) ); },
                Err( _                            ) => break,
            }
        }
    }
}

// ---------------------------------------------------------------------------------------------------------------------
// Test
// ---------------------------------------------------------------------------------------------------------------------

//#[cfg(test)]
//mod tests {
//    use super::*;
//    use matcher::QuickSearchMatcher;
//    use std::path::PathBuf;
//    use std::thread;
//    use std::sync::mpsc;
//    use util::PipelineInfo;
//
//    #[test]
//    fn test_simple_pipeline_printer() {
//    }
//}

<commit_msg>Files can't be renamed across filesystems<commit_after>use console::{Console, ConsoleTextKind};
use memmap::{Mmap, Protection};
use pipeline_matcher::PathMatch;
use std::fs;
use std::io;
use std::io::{Write, Error};
use std::process;
use std::sync::mpsc::{Receiver, Sender};
use time;
use util::{catch, decode_error, PipelineInfo};
use tempfile::NamedTempFile;

// ---------------------------------------------------------------------------------------------------------------------
// PipelineReplacer
// ---------------------------------------------------------------------------------------------------------------------

pub trait PipelineReplacer {
    fn replace( &mut self, replacement: &[u8], rx: Receiver<PipelineInfo<PathMatch>>, tx: Sender<PipelineInfo<()>> );
}

// ---------------------------------------------------------------------------------------------------------------------
// SimplePipelineReplacer
// ---------------------------------------------------------------------------------------------------------------------

pub struct SimplePipelineReplacer {
    pub is_color      : bool,
    pub is_interactive: bool,
    pub print_file    : bool,
    pub print_column  : bool,
    pub infos         : Vec<String>,
    pub errors        : Vec<String>,
    console           : Console,
    all_replace       : bool,
    time_beg          : u64,
    time_end          : u64,
    time_bsy          : u64,
}

impl SimplePipelineReplacer {
    pub fn new() -> Self {
        SimplePipelineReplacer {
            is_color      : true,
            is_interactive: true,
            print_file    : true,
            print_column  : true,
            infos         : Vec::new(),
            errors        : Vec::new(),
            console       : Console::new(),
            all_replace   : false,
            time_beg      : 0,
            time_end      : 0,
            time_bsy      : 0,
        }
    }

    fn replace_match( &mut self, replacement: &[u8], pm: PathMatch ) {
        if pm.matches.is_empty() { return; }
        self.console.is_color = self.is_color;

        let result = catch::<_, (), Error> ( || {
            let mut tmpfile = try!( NamedTempFile::new_in( pm.path.parent().unwrap_or( &pm.path ) )  );

            {
                let mmap = try!( Mmap::open_path( &pm.path, Protection::Read ) );
                let src  = unsafe { mmap.as_slice() };

                let mut i       = 0;
                let mut pos     = 0;
                let mut column  = 0;
                let mut last_lf = 0;
                for m in &pm.matches {
                    try!( tmpfile.write_all( &src[i..m.beg] ) );

                    let mut do_replace = true;
                    if self.is_interactive & !self.all_replace {
                        if self.print_file {
                            self.console.write( ConsoleTextKind::Filename, pm.path.to_str().unwrap() );
                            self.console.write( ConsoleTextKind::Other, ": " );
                        }
                        if self.print_column {
                            while pos < m.beg {
                                if src[pos] == 0x0a {
                                    column += 1;
                                    last_lf = pos;
                                }
                                pos += 1;
                            }
                            self.console.write( ConsoleTextKind::Other, &format!( "{}:{}:", column + 1, m.beg - last_lf ) );
                        }

                        self.console.write_match_line( src, m );

                        loop {
                            self.console.write( ConsoleTextKind::Other, "Replace keyword? ( Yes[Y], No[N], All[A], Quit[Q] ):" );
                            let mut buf = String::new();
                            io::stdin().read_line( &mut buf ).unwrap();
                            match buf.trim().to_lowercase().as_ref() {
                                "y"    => { do_replace = true ; break },
                                "yes"  => { do_replace = true ; break },
                                "n"    => { do_replace = false; break },
                                "no"   => { do_replace = false; break },
                                "a"    => { self.all_replace = true ; break },
                                "all"  => { self.all_replace = true ; break },
                                "q"    => { let _ = tmpfile.close(); process::exit( 0 ) },
                                "quit" => { let _ = tmpfile.close(); process::exit( 0 ) },
                                _      => continue,
                            }
                        }
                    }

                    if do_replace {
                        try!( tmpfile.write_all( &replacement ) );
                    } else {
                        try!( tmpfile.write_all( &src[m.beg..m.end] ) );
                    }
                    i = m.end;
                }

                if i < src.len() {
                    try!( tmpfile.write_all( &src[i..src.len()] ) );
                }
                try!( tmpfile.flush() );
            }

            let metadata = try!( fs::metadata( &pm.path ) );

            try!( fs::rename( tmpfile.path(), &pm.path ) );
            try!( fs::set_permissions( &pm.path, metadata.permissions() ) );

            Ok(())
        } );
        match result {
            Ok ( _ ) => (),
            Err( e ) => self.console.write( ConsoleTextKind::Error, &format!( "Error: {} @ {:?}\n", decode_error( e.kind() ), pm.path ) ),
        }
    }
}

impl PipelineReplacer for SimplePipelineReplacer {
    fn replace( &mut self, replacement: &[u8], rx: Receiver<PipelineInfo<PathMatch>>, tx: Sender<PipelineInfo<()>> ) {
        loop {
            match rx.recv() {
                Ok( PipelineInfo::Ok( pm ) ) => {
                    let beg = time::precise_time_ns();

                    self.replace_match( replacement, pm );
                    let _ = tx.send( PipelineInfo::Ok( () ) );

                    let end = time::precise_time_ns();
                    self.time_bsy += end - beg;
                },

                Ok( PipelineInfo::Beg( x ) ) => {
                    self.infos  = Vec::new();
                    self.errors = Vec::new();

                    self.time_beg = time::precise_time_ns();
                    let _ = tx.send( PipelineInfo::Beg( x ) );
                },

                Ok( PipelineInfo::End( x ) ) => {
                    for i in &self.infos  { let _ = tx.send( PipelineInfo::Info( i.clone() ) ); }
                    for e in &self.errors { let _ = tx.send( PipelineInfo::Err ( e.clone() ) ); }

                    self.time_end = time::precise_time_ns();
                    let _ = tx.send( PipelineInfo::Time( self.time_bsy, self.time_end - self.time_beg ) );
                    let _ = tx.send( PipelineInfo::End( x ) );
                    break;
                },

                Ok ( PipelineInfo::Info( e      ) ) => { let _ = tx.send( PipelineInfo::Info( e      ) ); },
                Ok ( PipelineInfo::Err ( e      ) ) => { let _ = tx.send( PipelineInfo::Err ( e      ) ); },
                Ok ( PipelineInfo::Time( t0, t1 ) ) => { let _ = tx.send( PipelineInfo::Time( t0, t1 ) ); },
                Err( _                            ) => break,
            }
        }
    }
}

// ---------------------------------------------------------------------------------------------------------------------
// Test
// ---------------------------------------------------------------------------------------------------------------------

//#[cfg(test)]
//mod tests {
//    use super::*;
//    use matcher::QuickSearchMatcher;
//    use std::path::PathBuf;
//    use std::thread;
//    use std::sync::mpsc;
//    use util::PipelineInfo;
//
//    #[test]
//    fn test_simple_pipeline_printer() {
//    }
//}

<|endoftext|>