<commit_before>extern crate pwx;
extern crate docopt;
extern crate rustc_serialize;
extern crate uuid;
extern crate chrono;

use pwx::{Pwx, Field, Value};
use std::io::{Write, stderr};
use std::process::exit;
use std::path::PathBuf;
use docopt::Docopt;
use uuid::Uuid;
use pwx::util::{fuzzy_eq, from_time_t, abspath, get_password_from_user};
use std::str::from_utf8;
use chrono::Local;
use chrono::duration::Duration;

// Get pkg version at compile time
const VERSION: &'static str = env!("CARGO_PKG_VERSION");

#[derive(RustcDecodable, Debug)]
struct Args {
    arg_file: String,
    arg_fieldname: String,
    arg_uuid: String,
    arg_filter: Vec<String>,
    flag_url: String,
    flag_group: String,
    flag_password_age: u32,
    flag_username: String,
    flag_title: String,
    flag_quiet: bool,
    cmd_list: bool,
    cmd_get: bool,
    cmd_info: bool,
    flag_version: bool,
    flag_pass_interactive: bool,
    flag_no_pinentry: bool,
}

/// Get password
/// 1. If --pass-interactive read from console
/// 2. If PWX_PASSWORD is set use it
/// 3. Otherwise read from console
///
/// This function may panic on encoding issues
fn get_password(args: &Args, description: &str) -> Option<String> {
    let var = std::env::var("PWX_PASSWORD");
    if args.flag_pass_interactive || !var.is_ok() {
        get_password_from_user(description, args.flag_no_pinentry)
    } else {
        Some(var.unwrap())
    }
}

/// A filter to match multiple keywords
///
/// The default state is for the filter to be unmatched.
/// Push content with 'push()', if it matches the keywords
/// the filter will match.
struct KeywordFilter<'a> {
    m_filter: Vec<bool>,
    args: &'a Args,
}

impl<'a> KeywordFilter<'a> {
    fn new(args: &Args) -> KeywordFilter {
        let mut f = KeywordFilter {
            m_filter: Vec::with_capacity(args.arg_filter.len()),
            args: args,
        };
        f.m_filter.resize(args.arg_filter.len(), false);
        f
    }

    /// Attempt to match this filter against a field
    fn push(&mut self, field: &Field) {
        let utf8 = match *field {
            Field::Group(ref v) => from_utf8(v.as_ref()),
            Field::Title(ref v) => from_utf8(v.as_ref()),
            Field::Username(ref v) => from_utf8(v.as_ref()),
            Field::Notes(ref v) => from_utf8(v.as_ref()),
            Field::Password(_) => return,
            Field::Url(ref v) => from_utf8(v.as_ref()),
            _ => return,
        };

        if let Ok(s) = utf8 {
            for (idx, word) in self.args
                                   .arg_filter
                                   .iter()
                                   .enumerate() {
                self.m_filter[idx] = self.m_filter[idx] || fuzzy_eq(&word, s);
            }
        }
    }

    /// Returns true if all the keywords matched
    fn matched(&self) -> bool {
        self.m_filter.iter().all(|v| *v)
    }
}

fn real_main() -> i32 {

    let args: Args = Docopt::new(include_str!(concat!(env!("CARGO_MANIFEST_DIR"),
                                                      "/doc/pwx.docopt")))
                         .and_then(|d| d.decode())
                         .unwrap_or_else(|e| e.exit());

    if args.flag_version {
        println!("pwx {}", VERSION);
        return 0;
    }

    // The password safe is one of
    // 1. Command line [<file>]
    // 2. PWX_DATABASE env var
    // 3. ~/.pwsafe/psafe.psafe3
    let env_db = std::env::var("PWX_DATABASE").unwrap_or(String::new());

    let mut path = if !args.arg_file.is_empty() {
        PathBuf::from(&args.arg_file)
    } else if !env_db.is_empty() {
        PathBuf::from(&env_db)
    } else {
        std::env::home_dir()
            .expect("Cannot find your HOME path")
            .join(".pwsafe")
            .join("pwsafe.psafe3")
    };

    path = match abspath(&path) {
        Ok(abs) => abs,
        Err(_) => path,
    };

    // FIXME: PathBuf::exists is still Unstable, but we need to verify
    // if the file exists before prompting for password
    match std::fs::File::open(&path) {
        Err(err) => {
            let _ = writeln!(stderr(), "{}: {}", path.to_string_lossy(), err);
            return -1;
        }
        _ => (),
    }

    path = match abspath(&path) {
        Ok(abs) => abs,
        Err(_) => path,
    };

    let description = format!("Opening {}", path.to_string_lossy());
    let mut p = match Pwx::open(&path,
                                get_password(&args, &description)
                                    .expect("Unable to get user password")
                                    .as_bytes()) {
        Err(f) => {
            let _ = writeln!(stderr(), "Error: {} {}", f, path.to_string_lossy());
            return -1;
        }
        Ok(p) => p,
    };

    if !p.is_authentic() {
        return -1;
    }

    if args.cmd_list {
        let min_pw_age = Duration::days(args.flag_password_age as i64);

        for record in p.iter().unwrap() {
            let mut uuid = String::new();
            let mut title = String::new();
            let mut username = String::new();

            // Field filters
            let mut f_username = args.flag_username.is_empty();
            let mut f_title = args.flag_title.is_empty();
            let mut f_url = args.flag_url.is_empty();
            let mut f_group = args.flag_group.is_empty();

            // The password-age filter can use either ctime or ptime
            let mut creation_time = None;
            let mut password_age = None;

            // Keyword filters
            let mut f_keywords = KeywordFilter::new(&args);

            for field in record {
                f_keywords.push(&field);
                match field {
                    Field::Uuid(_) => uuid = format!("{}", field),
                    Field::Title(_) => {
                        title = format!("{}", field);
                        f_title = f_title || fuzzy_eq(&args.flag_title, &title);
                    }
                    Field::Username(_) => {
                        username = format!("{}", field);
                        f_username = f_username || fuzzy_eq(&args.flag_username, &username);
                    }
                    Field::Url(_) => {
                        let url = format!("{}", field);
                        f_url = f_url || fuzzy_eq(&args.flag_url, &url);
                    }
                    Field::Group(_) => {
                        let group = format!("{}", field);
                        f_group = f_group || fuzzy_eq(&args.flag_group, &group);
                    }
                    Field::CreationTime(val) => {
                        let ts = from_time_t(val.as_ref()).unwrap();
                        let diff = Local::now().naive_local() - ts;
                        creation_time = Some(diff);
                    }
                    Field::PasswordModificationTime(val) => {
                        let ts = from_time_t(val.as_ref()).unwrap();
                        let diff = Local::now().naive_local() - ts;
                        password_age = Some(diff);
                    }
                    _ => (),
                }
            }

            if !f_username || !f_title || !f_url || !f_group {
                // Skip, record filter did not match
                continue;
            }

            match (creation_time, password_age) {
                // No fields were found treat as a match
                (None, None) => (),
                // Password was never modified
                (Some(diff), None) if diff < min_pw_age => continue,
                // Password modification time
                (_, Some(diff)) if diff < min_pw_age => continue,
                _ => (),
            }

            if !f_keywords.matched() {
                // Skip, generic keyword filter did not match
                continue;
            }

            println!("{} {}[{}]", uuid, title, username);
        }
    } else if args.cmd_info {

        let info = p.info().unwrap();
        println!("{} {} {}@{}", info.uuid, info.mtime, info.user, info.host);
    } else if args.cmd_get {
        let get_uuid = Field::Uuid(Value::from(Uuid::parse_str(&args.arg_uuid)
                                                   .expect("Invalid UUID")
                                                   .as_bytes()
                                                   .to_vec()));

        for record in p.iter().unwrap() {
            // Find record by UUID
            let mut found = false;
            for field in &record {
                if *field == get_uuid {
                    found = true;
                    break;
                }
            }

            if !found {
                continue;
            }

            // Get field value
            for field in &record {
                if let Some(name) = field.name() {
                    if name == args.arg_fieldname {
                        println!("{}", field);
                        return 0;
                    }
                }
            }
            let _ = writeln!(stderr(), "Unknown field: {}", args.arg_fieldname);
        }

        let _ = writeln!(stderr(), "Unknown record: {}", args.arg_uuid);
        return -1;
    }

    return 0;
}

// We want to make sure we exit only after all the destructors
// are called, wait for real_main() to be done before actually
// calling exit.
fn main() {
    let exit_code = real_main();
    exit(exit_code);
}
<commit_msg>Use PathBuf::exists()<commit_after>extern crate pwx;
extern crate docopt;
extern crate rustc_serialize;
extern crate uuid;
extern crate chrono;

use pwx::{Pwx, Field, Value};
use std::io::{Write, stderr};
use std::process::exit;
use std::path::PathBuf;
use docopt::Docopt;
use uuid::Uuid;
use pwx::util::{fuzzy_eq, from_time_t, abspath, get_password_from_user};
use std::str::from_utf8;
use chrono::Local;
use chrono::duration::Duration;

// Get pkg version at compile time
const VERSION: &'static str = env!("CARGO_PKG_VERSION");

#[derive(RustcDecodable, Debug)]
struct Args {
    arg_file: String,
    arg_fieldname: String,
    arg_uuid: String,
    arg_filter: Vec<String>,
    flag_url: String,
    flag_group: String,
    flag_password_age: u32,
    flag_username: String,
    flag_title: String,
    flag_quiet: bool,
    cmd_list: bool,
    cmd_get: bool,
    cmd_info: bool,
    flag_version: bool,
    flag_pass_interactive: bool,
    flag_no_pinentry: bool,
}

/// Get password
/// 1. If --pass-interactive read from console
/// 2. If PWX_PASSWORD is set use it
/// 3. Otherwise read from console
///
/// This function may panic on encoding issues
fn get_password(args: &Args, description: &str) -> Option<String> {
    let var = std::env::var("PWX_PASSWORD");
    if args.flag_pass_interactive || !var.is_ok() {
        get_password_from_user(description, args.flag_no_pinentry)
    } else {
        Some(var.unwrap())
    }
}

/// A filter to match multiple keywords
///
/// The default state is for the filter to be unmatched.
/// Push content with 'push()', if it matches the keywords
/// the filter will match.
struct KeywordFilter<'a> {
    m_filter: Vec<bool>,
    args: &'a Args,
}

impl<'a> KeywordFilter<'a> {
    fn new(args: &Args) -> KeywordFilter {
        let mut f = KeywordFilter {
            m_filter: Vec::with_capacity(args.arg_filter.len()),
            args: args,
        };
        f.m_filter.resize(args.arg_filter.len(), false);
        f
    }

    /// Attempt to match this filter against a field
    fn push(&mut self, field: &Field) {
        let utf8 = match *field {
            Field::Group(ref v) => from_utf8(v.as_ref()),
            Field::Title(ref v) => from_utf8(v.as_ref()),
            Field::Username(ref v) => from_utf8(v.as_ref()),
            Field::Notes(ref v) => from_utf8(v.as_ref()),
            Field::Password(_) => return,
            Field::Url(ref v) => from_utf8(v.as_ref()),
            _ => return,
        };

        if let Ok(s) = utf8 {
            for (idx, word) in self.args
                                   .arg_filter
                                   .iter()
                                   .enumerate() {
                self.m_filter[idx] = self.m_filter[idx] || fuzzy_eq(&word, s);
            }
        }
    }

    /// Returns true if all the keywords matched
    fn matched(&self) -> bool {
        self.m_filter.iter().all(|v| *v)
    }
}

fn real_main() -> i32 {

    let args: Args = Docopt::new(include_str!(concat!(env!("CARGO_MANIFEST_DIR"),
                                                      "/doc/pwx.docopt")))
                         .and_then(|d| d.decode())
                         .unwrap_or_else(|e| e.exit());

    if args.flag_version {
        println!("pwx {}", VERSION);
        return 0;
    }

    // The password safe is one of
    // 1. Command line [<file>]
    // 2. PWX_DATABASE env var
    // 3. ~/.pwsafe/psafe.psafe3
    let env_db = std::env::var("PWX_DATABASE").unwrap_or(String::new());

    let mut path = if !args.arg_file.is_empty() {
        PathBuf::from(&args.arg_file)
    } else if !env_db.is_empty() {
        PathBuf::from(&env_db)
    } else {
        std::env::home_dir()
            .expect("Cannot find your HOME path")
            .join(".pwsafe")
            .join("pwsafe.psafe3")
    };

    path = match abspath(&path) {
        Ok(abs) => abs,
        Err(_) => path,
    };

    if !path.exists() {
        let _ = writeln!(stderr(), "File does not exist: {}", path.to_string_lossy());
        return -1;
    }

    path = match abspath(&path) {
        Ok(abs) => abs,
        Err(_) => path,
    };

    let description = format!("Opening {}", path.to_string_lossy());
    let mut p = match Pwx::open(&path,
                                get_password(&args, &description)
                                    .expect("Unable to get user password")
                                    .as_bytes()) {
        Err(f) => {
            let _ = writeln!(stderr(), "Error: {} {}", f, path.to_string_lossy());
            return -1;
        }
        Ok(p) => p,
    };

    if !p.is_authentic() {
        return -1;
    }

    if args.cmd_list {
        let min_pw_age = Duration::days(args.flag_password_age as i64);

        for record in p.iter().unwrap() {
            let mut uuid = String::new();
            let mut title = String::new();
            let mut username = String::new();

            // Field filters
            let mut f_username = args.flag_username.is_empty();
            let mut f_title = args.flag_title.is_empty();
            let mut f_url = args.flag_url.is_empty();
            let mut f_group = args.flag_group.is_empty();

            // The password-age filter can use either ctime or ptime
            let mut creation_time = None;
            let mut password_age = None;

            // Keyword filters
            let mut f_keywords = KeywordFilter::new(&args);

            for field in record {
                f_keywords.push(&field);
                match field {
                    Field::Uuid(_) => uuid = format!("{}", field),
                    Field::Title(_) => {
                        title = format!("{}", field);
                        f_title = f_title || fuzzy_eq(&args.flag_title, &title);
                    }
                    Field::Username(_) => {
                        username = format!("{}", field);
                        f_username = f_username || fuzzy_eq(&args.flag_username, &username);
                    }
                    Field::Url(_) => {
                        let url = format!("{}", field);
                        f_url = f_url || fuzzy_eq(&args.flag_url, &url);
                    }
                    Field::Group(_) => {
                        let group = format!("{}", field);
                        f_group = f_group || fuzzy_eq(&args.flag_group, &group);
                    }
                    Field::CreationTime(val) => {
                        let ts = from_time_t(val.as_ref()).unwrap();
                        let diff = Local::now().naive_local() - ts;
                        creation_time = Some(diff);
                    }
                    Field::PasswordModificationTime(val) => {
                        let ts = from_time_t(val.as_ref()).unwrap();
                        let diff = Local::now().naive_local() - ts;
                        password_age = Some(diff);
                    }
                    _ => (),
                }
            }

            if !f_username || !f_title || !f_url || !f_group {
                // Skip, record filter did not match
                continue;
            }

            match (creation_time, password_age) {
                // No fields were found treat as a match
                (None, None) => (),
                // Password was never modified
                (Some(diff), None) if diff < min_pw_age => continue,
                // Password modification time
                (_, Some(diff)) if diff < min_pw_age => continue,
                _ => (),
            }

            if !f_keywords.matched() {
                // Skip, generic keyword filter did not match
                continue;
            }

            println!("{} {}[{}]", uuid, title, username);
        }
    } else if args.cmd_info {

        let info = p.info().unwrap();
        println!("{} {} {}@{}", info.uuid, info.mtime, info.user, info.host);
    } else if args.cmd_get {
        let get_uuid = Field::Uuid(Value::from(Uuid::parse_str(&args.arg_uuid)
                                                   .expect("Invalid UUID")
                                                   .as_bytes()
                                                   .to_vec()));

        for record in p.iter().unwrap() {
            // Find record by UUID
            let mut found = false;
            for field in &record {
                if *field == get_uuid {
                    found = true;
                    break;
                }
            }

            if !found {
                continue;
            }

            // Get field value
            for field in &record {
                if let Some(name) = field.name() {
                    if name == args.arg_fieldname {
                        println!("{}", field);
                        return 0;
                    }
                }
            }
            let _ = writeln!(stderr(), "Unknown field: {}", args.arg_fieldname);
        }

        let _ = writeln!(stderr(), "Unknown record: {}", args.arg_uuid);
        return -1;
    }

    return 0;
}

// We want to make sure we exit only after all the destructors
// are called, wait for real_main() to be done before actually
// calling exit.
fn main() {
    let exit_code = real_main();
    exit(exit_code);
}
<|endoftext|>


<commit_before>use hir::Semantics;
use ide_db::{
    base_db::{FileId, FilePosition, FileRange, SourceDatabase},
    RootDatabase, SymbolKind,
};
use syntax::TextRange;

use crate::{
    file_structure::file_structure,
    fn_references::find_all_methods,
    goto_implementation::goto_implementation,
    references::find_all_refs,
    runnables::{runnables, Runnable},
    NavigationTarget, RunnableKind,
};

// Feature: Annotations
//
// Provides user with annotations above items for looking up references or impl blocks
// and running/debugging binaries.
#[derive(Debug)]
pub struct Annotation {
    pub range: TextRange,
    pub kind: AnnotationKind,
}

#[derive(Debug)]
pub enum AnnotationKind {
    Runnable { debug: bool, runnable: Runnable },
    HasImpls { position: FilePosition, data: Option<Vec<NavigationTarget>> },
    HasReferences { position: FilePosition, data: Option<Vec<FileRange>> },
}

pub struct AnnotationConfig {
    pub binary_target: bool,
    pub annotate_runnables: bool,
    pub annotate_impls: bool,
    pub annotate_references: bool,
    pub annotate_method_references: bool,
    pub run: bool,
    pub debug: bool,
}

pub(crate) fn annotations(
    db: &RootDatabase,
    file_id: FileId,
    config: AnnotationConfig,
) -> Vec<Annotation> {
    let mut annotations = Vec::default();

    if config.annotate_runnables {
        for runnable in runnables(db, file_id) {
            if should_skip_runnable(&runnable.kind, config.binary_target) {
                continue;
            }

            let action = runnable.action();
            let range = runnable.nav.full_range;

            if action.debugee && config.debug {
                annotations.push(Annotation {
                    range,

                    // FIXME: This one allocates without reason if run is enabled, but debug is disabled
                    kind: AnnotationKind::Runnable { debug: true, runnable: runnable.clone() },
                });
            }

            if config.run {
                annotations.push(Annotation {
                    range,
                    kind: AnnotationKind::Runnable { debug: false, runnable },
                });
            }
        }
    }

    file_structure(&db.parse(file_id).tree())
        .into_iter()
        .filter(|node| {
            matches!(
                node.kind,
                SymbolKind::Trait
                    | SymbolKind::Struct
                    | SymbolKind::Enum
                    | SymbolKind::Union
                    | SymbolKind::Const
            )
        })
        .for_each(|node| {
            if config.annotate_impls && node.kind != SymbolKind::Const {
                annotations.push(Annotation {
                    range: node.node_range,
                    kind: AnnotationKind::HasImpls {
                        position: FilePosition { file_id, offset: node.navigation_range.start() },
                        data: None,
                    },
                });
            }

            if config.annotate_references {
                annotations.push(Annotation {
                    range: node.node_range,
                    kind: AnnotationKind::HasReferences {
                        position: FilePosition { file_id, offset: node.navigation_range.start() },
                        data: None,
                    },
                });
            }
        });

    if config.annotate_method_references {
        annotations.extend(find_all_methods(db, file_id).into_iter().map(|method| Annotation {
            range: method.range,
            kind: AnnotationKind::HasReferences {
                position: FilePosition { file_id, offset: method.range.start() },
                data: None,
            },
        }));
    }

    annotations
}

pub(crate) fn resolve_annotation(db: &RootDatabase, mut annotation: Annotation) -> Annotation {
    match annotation.kind {
        AnnotationKind::HasImpls { position, ref mut data } => {
            *data = goto_implementation(db, position).map(|range| range.info);
        }
        AnnotationKind::HasReferences { position, ref mut data } => {
            *data = find_all_refs(&Semantics::new(db), position, None).map(|result| {
                result
                    .references
                    .into_iter()
                    .map(|(file_id, access)| {
                        access.into_iter().map(move |(range, _)| FileRange { file_id, range })
                    })
                    .flatten()
                    .collect()
            });
        }
        _ => {}
    };

    annotation
}

fn should_skip_runnable(kind: &RunnableKind, binary_target: bool) -> bool {
    match kind {
        RunnableKind::Bin => !binary_target,
        _ => false,
    }
}

#[cfg(test)]
mod tests {
    use ide_db::base_db::{FileId, FileRange};
    use syntax::{TextRange, TextSize};

    use crate::{fixture, Annotation, AnnotationConfig, AnnotationKind, RunnableKind};

    fn get_annotations(
        ra_fixture: &str,
        annotation_config: AnnotationConfig,
    ) -> (FileId, Vec<Annotation>) {
        let (analysis, file_id) = fixture::file(ra_fixture);

        let annotations: Vec<Annotation> = analysis
            .annotations(file_id, annotation_config)
            .unwrap()
            .into_iter()
            .map(move |annotation| analysis.resolve_annotation(annotation).unwrap())
            .collect();

        if annotations.len() == 0 {
            panic!("unresolved annotations")
        }

        (file_id, annotations)
    }

    macro_rules! check_annotation {
        ( $ra_fixture:expr, $config:expr, $item_positions:expr, $pattern:pat, $checker:expr ) => {
            let (file_id, annotations) = get_annotations($ra_fixture, $config);

            annotations.into_iter().for_each(|annotation| {
                assert!($item_positions.contains(&annotation.range));

                match annotation.kind {
                    $pattern => $checker(file_id),
                    _ => panic!("Unexpected annotation kind"),
                }
            });
        };
    }

    #[test]
    fn const_annotations() {
        check_annotation!(
            r#"
const DEMO: i32 = 123;

fn main() {
    let hello = DEMO;
}
            "#,
            AnnotationConfig {
                binary_target: false,
                annotate_runnables: false,
                annotate_impls: false,
                annotate_references: true,
                annotate_method_references: false,
                run: false,
                debug: false,
            },
            &[TextRange::new(TextSize::from(0), TextSize::from(22))],
            AnnotationKind::HasReferences { data: Some(ranges), .. },
            |file_id| assert_eq!(
                *ranges.first().unwrap(),
                FileRange {
                    file_id,
                    range: TextRange::new(TextSize::from(52), TextSize::from(56))
                }
            )
        );
    }

    #[test]
    fn unused_const_annotations() {
        check_annotation!(
            r#"
const DEMO: i32 = 123;

fn main() {}
            "#,
            AnnotationConfig {
                binary_target: false,
                annotate_runnables: false,
                annotate_impls: false,
                annotate_references: true,
                annotate_method_references: false,
                run: false,
                debug: false,
            },
            &[TextRange::new(TextSize::from(0), TextSize::from(22))],
            AnnotationKind::HasReferences { data: Some(ranges), .. },
            |_| assert_eq!(ranges.len(), 0)
        );
    }

    #[test]
    fn struct_references_annotations() {
        check_annotation!(
            r#"
struct Test;

fn main() {
    let test = Test;
}
            "#,
            AnnotationConfig {
                binary_target: false,
                annotate_runnables: false,
                annotate_impls: false,
                annotate_references: true,
                annotate_method_references: false,
                run: false,
                debug: false,
            },
            &[TextRange::new(TextSize::from(0), TextSize::from(12))],
            AnnotationKind::HasReferences { data: Some(ranges), .. },
            |file_id| assert_eq!(
                *ranges.first().unwrap(),
                FileRange {
                    file_id,
                    range: TextRange::new(TextSize::from(41), TextSize::from(45))
                }
            )
        );
    }

    #[test]
    fn struct_and_trait_impls_annotations() {
        check_annotation!(
            r#"
struct Test;

trait MyCoolTrait {}

impl MyCoolTrait for Test {}

fn main() {
    let test = Test;
}
            "#,
            AnnotationConfig {
                binary_target: false,
                annotate_runnables: false,
                annotate_impls: true,
                annotate_references: false,
                annotate_method_references: false,
                run: false,
                debug: false,
            },
            &[
                TextRange::new(TextSize::from(0), TextSize::from(12)),
                TextRange::new(TextSize::from(14), TextSize::from(34))
            ],
            AnnotationKind::HasImpls { data: Some(ranges), .. },
            |_| assert_eq!(
                ranges.first().unwrap().full_range,
                TextRange::new(TextSize::from(36), TextSize::from(64))
            )
        );
    }

    #[test]
    fn run_annotation() {
        check_annotation!(
            r#"
fn main() {}
            "#,
            AnnotationConfig {
                binary_target: true,
                annotate_runnables: true,
                annotate_impls: false,
                annotate_references: false,
                annotate_method_references: false,
                run: true,
                debug: false,
            },
            &[TextRange::new(TextSize::from(0), TextSize::from(12))],
            AnnotationKind::Runnable { debug: false, runnable },
            |_| {
                assert!(matches!(runnable.kind, RunnableKind::Bin));
                assert!(runnable.action().run_title.contains("Run"));
            }
        );
    }

    #[test]
    fn debug_annotation() {
        check_annotation!(
            r#"
fn main() {}
            "#,
            AnnotationConfig {
                binary_target: true,
                annotate_runnables: true,
                annotate_impls: false,
                annotate_references: false,
                annotate_method_references: false,
                run: false,
                debug: true,
            },
            &[TextRange::new(TextSize::from(0), TextSize::from(12))],
            AnnotationKind::Runnable { debug: true, runnable },
            |_| {
                assert!(matches!(runnable.kind, RunnableKind::Bin));
                assert!(runnable.action().debugee);
            }
        );
    }

    #[test]
    fn method_annotations() {
        // We actually want to skip `fn main` annotation, as it has no references in it
        // but just ignoring empty reference slices would lead to false-positive if something
        // goes wrong in annotation resolving mechanism. By tracking if we iterated before finding
        // an empty slice we can track if everything is settled.
        let mut iterated_once = false;

        check_annotation!(
            r#"
struct Test;

impl Test {
    fn self_by_ref(&self) {}
}

fn main() {
    Test.self_by_ref();
}
            "#,
            AnnotationConfig {
                binary_target: false,
                annotate_runnables: false,
                annotate_impls: false,
                annotate_references: false,
                annotate_method_references: true,
                run: false,
                debug: false,
            },
            &[
                TextRange::new(TextSize::from(33), TextSize::from(44)),
                TextRange::new(TextSize::from(61), TextSize::from(65))
            ],
            AnnotationKind::HasReferences { data: Some(ranges), .. },
            |file_id| {
                match ranges.as_slice() {
                    [first, ..] => {
                        assert_eq!(
                            *first,
                            FileRange {
                                file_id,
                                range: TextRange::new(TextSize::from(79), TextSize::from(90))
                            }
                        );

                        iterated_once = true;
                    }
                    [] if iterated_once => {}
                    [] => panic!("One reference was expected but not found"),
                }
            }
        );
    }
}
<commit_msg>Make annotations tests similar to those in runnables<commit_after>use hir::Semantics;
use ide_db::{
    base_db::{FileId, FilePosition, FileRange, SourceDatabase},
    RootDatabase, SymbolKind,
};
use syntax::TextRange;

use crate::{
    file_structure::file_structure,
    fn_references::find_all_methods,
    goto_implementation::goto_implementation,
    references::find_all_refs,
    runnables::{runnables, Runnable},
    NavigationTarget, RunnableKind,
};

// Feature: Annotations
//
// Provides user with annotations above items for looking up references or impl blocks
// and running/debugging binaries.
#[derive(Debug)]
pub struct Annotation {
    pub range: TextRange,
    pub kind: AnnotationKind,
}

#[derive(Debug)]
pub enum AnnotationKind {
    Runnable { debug: bool, runnable: Runnable },
    HasImpls { position: FilePosition, data: Option<Vec<NavigationTarget>> },
    HasReferences { position: FilePosition, data: Option<Vec<FileRange>> },
}

pub struct AnnotationConfig {
    pub binary_target: bool,
    pub annotate_runnables: bool,
    pub annotate_impls: bool,
    pub annotate_references: bool,
    pub annotate_method_references: bool,
    pub run: bool,
    pub debug: bool,
}

pub(crate) fn annotations(
    db: &RootDatabase,
    file_id: FileId,
    config: AnnotationConfig,
) -> Vec<Annotation> {
    let mut annotations = Vec::default();

    if config.annotate_runnables {
        for runnable in runnables(db, file_id) {
            if should_skip_runnable(&runnable.kind, config.binary_target) {
                continue;
            }

            let action = runnable.action();
            let range = runnable.nav.full_range;

            if action.debugee && config.debug {
                annotations.push(Annotation {
                    range,

                    // FIXME: This one allocates without reason if run is enabled, but debug is disabled
                    kind: AnnotationKind::Runnable { debug: true, runnable: runnable.clone() },
                });
            }

            if config.run {
                annotations.push(Annotation {
                    range,
                    kind: AnnotationKind::Runnable { debug: false, runnable },
                });
            }
        }
    }

    file_structure(&db.parse(file_id).tree())
        .into_iter()
        .filter(|node| {
            matches!(
                node.kind,
                SymbolKind::Trait
                    | SymbolKind::Struct
                    | SymbolKind::Enum
                    | SymbolKind::Union
                    | SymbolKind::Const
            )
        })
        .for_each(|node| {
            if config.annotate_impls && node.kind != SymbolKind::Const {
                annotations.push(Annotation {
                    range: node.node_range,
                    kind: AnnotationKind::HasImpls {
                        position: FilePosition { file_id, offset: node.navigation_range.start() },
                        data: None,
                    },
                });
            }

            if config.annotate_references {
                annotations.push(Annotation {
                    range: node.node_range,
                    kind: AnnotationKind::HasReferences {
                        position: FilePosition { file_id, offset: node.navigation_range.start() },
                        data: None,
                    },
                });
            }
        });

    if config.annotate_method_references {
        annotations.extend(find_all_methods(db, file_id).into_iter().map(|method| Annotation {
            range: method.range,
            kind: AnnotationKind::HasReferences {
                position: FilePosition { file_id, offset: method.range.start() },
                data: None,
            },
        }));
    }

    annotations
}

pub(crate) fn resolve_annotation(db: &RootDatabase, mut annotation: Annotation) -> Annotation {
    match annotation.kind {
        AnnotationKind::HasImpls { position, ref mut data } => {
            *data = goto_implementation(db, position).map(|range| range.info);
        }
        AnnotationKind::HasReferences { position, ref mut data } => {
            *data = find_all_refs(&Semantics::new(db), position, None).map(|result| {
                result
                    .references
                    .into_iter()
                    .map(|(file_id, access)| {
                        access.into_iter().map(move |(range, _)| FileRange { file_id, range })
                    })
                    .flatten()
                    .collect()
            });
        }
        _ => {}
    };

    annotation
}

fn should_skip_runnable(kind: &RunnableKind, binary_target: bool) -> bool {
    match kind {
        RunnableKind::Bin => !binary_target,
        _ => false,
    }
}

#[cfg(test)]
mod tests {
    use expect_test::{expect, Expect};

    use crate::{fixture, Annotation, AnnotationConfig};

    fn check(ra_fixture: &str, expect: Expect) {
        let (analysis, file_id) = fixture::file(ra_fixture);

        let annotations: Vec<Annotation> = analysis
            .annotations(
                file_id,
                AnnotationConfig {
                    binary_target: true,
                    annotate_runnables: true,
                    annotate_impls: true,
                    annotate_references: true,
                    annotate_method_references: true,
                    run: true,
                    debug: true,
                },
            )
            .unwrap()
            .into_iter()
            .map(|annotation| analysis.resolve_annotation(annotation).unwrap())
            .collect();

        expect.assert_debug_eq(&annotations);
    }

    #[test]
    fn const_annotations() {
        check(
            r#"
const DEMO: i32 = 123;

const UNUSED: i32 = 123;

fn main() {
    let hello = DEMO;
}
            "#,
            expect![[r#"
                [
                    Annotation {
                        range: 50..85,
                        kind: Runnable {
                            debug: true,
                            runnable: Runnable {
                                nav: NavigationTarget {
                                    file_id: FileId(
                                        0,
                                    ),
                                    full_range: 50..85,
                                    focus_range: 53..57,
                                    name: "main",
                                    kind: Function,
                                },
                                kind: Bin,
                                cfg: None,
                            },
                        },
                    },
                    Annotation {
                        range: 50..85,
                        kind: Runnable {
                            debug: false,
                            runnable: Runnable {
                                nav: NavigationTarget {
                                    file_id: FileId(
                                        0,
                                    ),
                                    full_range: 50..85,
                                    focus_range: 53..57,
                                    name: "main",
                                    kind: Function,
                                },
                                kind: Bin,
                                cfg: None,
                            },
                        },
                    },
                    Annotation {
                        range: 0..22,
                        kind: HasReferences {
                            position: FilePosition {
                                file_id: FileId(
                                    0,
                                ),
                                offset: 6,
                            },
                            data: Some(
                                [
                                    FileRange {
                                        file_id: FileId(
                                            0,
                                        ),
                                        range: 78..82,
                                    },
                                ],
                            ),
                        },
                    },
                    Annotation {
                        range: 24..48,
                        kind: HasReferences {
                            position: FilePosition {
                                file_id: FileId(
                                    0,
                                ),
                                offset: 30,
                            },
                            data: Some(
                                [],
                            ),
                        },
                    },
                    Annotation {
                        range: 53..57,
                        kind: HasReferences {
                            position: FilePosition {
                                file_id: FileId(
                                    0,
                                ),
                                offset: 53,
                            },
                            data: Some(
                                [],
                            ),
                        },
                    },
                ]
            "#]],
        );
    }

    #[test]
    fn struct_references_annotations() {
        check(
            r#"
struct Test;

fn main() {
    let test = Test;
}
            "#,
            expect![[r#"
                [
                    Annotation {
                        range: 14..48,
                        kind: Runnable {
                            debug: true,
                            runnable: Runnable {
                                nav: NavigationTarget {
                                    file_id: FileId(
                                        0,
                                    ),
                                    full_range: 14..48,
                                    focus_range: 17..21,
                                    name: "main",
                                    kind: Function,
                                },
                                kind: Bin,
                                cfg: None,
                            },
                        },
                    },
                    Annotation {
                        range: 14..48,
                        kind: Runnable {
                            debug: false,
                            runnable: Runnable {
                                nav: NavigationTarget {
                                    file_id: FileId(
                                        0,
                                    ),
                                    full_range: 14..48,
                                    focus_range: 17..21,
                                    name: "main",
                                    kind: Function,
                                },
                                kind: Bin,
                                cfg: None,
                            },
                        },
                    },
                    Annotation {
                        range: 0..12,
                        kind: HasImpls {
                            position: FilePosition {
                                file_id: FileId(
                                    0,
                                ),
                                offset: 7,
                            },
                            data: Some(
                                [],
                            ),
                        },
                    },
                    Annotation {
                        range: 0..12,
                        kind: HasReferences {
                            position: FilePosition {
                                file_id: FileId(
                                    0,
                                ),
                                offset: 7,
                            },
                            data: Some(
                                [
                                    FileRange {
                                        file_id: FileId(
                                            0,
                                        ),
                                        range: 41..45,
                                    },
                                ],
                            ),
                        },
                    },
                    Annotation {
                        range: 17..21,
                        kind: HasReferences {
                            position: FilePosition {
                                file_id: FileId(
                                    0,
                                ),
                                offset: 17,
                            },
                            data: Some(
                                [],
                            ),
                        },
                    },
                ]
            "#]],
        );
    }

    #[test]
    fn struct_and_trait_impls_annotations() {
        check(
            r#"
struct Test;

trait MyCoolTrait {}

impl MyCoolTrait for Test {}

fn main() {
    let test = Test;
}
            "#,
            expect![[r#"
                [
                    Annotation {
                        range: 66..100,
                        kind: Runnable {
                            debug: true,
                            runnable: Runnable {
                                nav: NavigationTarget {
                                    file_id: FileId(
                                        0,
                                    ),
                                    full_range: 66..100,
                                    focus_range: 69..73,
                                    name: "main",
                                    kind: Function,
                                },
                                kind: Bin,
                                cfg: None,
                            },
                        },
                    },
                    Annotation {
                        range: 66..100,
                        kind: Runnable {
                            debug: false,
                            runnable: Runnable {
                                nav: NavigationTarget {
                                    file_id: FileId(
                                        0,
                                    ),
                                    full_range: 66..100,
                                    focus_range: 69..73,
                                    name: "main",
                                    kind: Function,
                                },
                                kind: Bin,
                                cfg: None,
                            },
                        },
                    },
                    Annotation {
                        range: 0..12,
                        kind: HasImpls {
                            position: FilePosition {
                                file_id: FileId(
                                    0,
                                ),
                                offset: 7,
                            },
                            data: Some(
                                [
                                    NavigationTarget {
                                        file_id: FileId(
                                            0,
                                        ),
                                        full_range: 36..64,
                                        focus_range: 57..61,
                                        name: "impl",
                                        kind: Impl,
                                    },
                                ],
                            ),
                        },
                    },
                    Annotation {
                        range: 0..12,
                        kind: HasReferences {
                            position: FilePosition {
                                file_id: FileId(
                                    0,
                                ),
                                offset: 7,
                            },
                            data: Some(
                                [
                                    FileRange {
                                        file_id: FileId(
                                            0,
                                        ),
                                        range: 57..61,
                                    },
                                    FileRange {
                                        file_id: FileId(
                                            0,
                                        ),
                                        range: 93..97,
                                    },
                                ],
                            ),
                        },
                    },
                    Annotation {
                        range: 14..34,
                        kind: HasImpls {
                            position: FilePosition {
                                file_id: FileId(
                                    0,
                                ),
                                offset: 20,
                            },
                            data: Some(
                                [
                                    NavigationTarget {
                                        file_id: FileId(
                                            0,
                                        ),
                                        full_range: 36..64,
                                        focus_range: 57..61,
                                        name: "impl",
                                        kind: Impl,
                                    },
                                ],
                            ),
                        },
                    },
                    Annotation {
                        range: 14..34,
                        kind: HasReferences {
                            position: FilePosition {
                                file_id: FileId(
                                    0,
                                ),
                                offset: 20,
                            },
                            data: Some(
                                [
                                    FileRange {
                                        file_id: FileId(
                                            0,
                                        ),
                                        range: 41..52,
                                    },
                                ],
                            ),
                        },
                    },
                    Annotation {
                        range: 69..73,
                        kind: HasReferences {
                            position: FilePosition {
                                file_id: FileId(
                                    0,
                                ),
                                offset: 69,
                            },
                            data: Some(
                                [],
                            ),
                        },
                    },
                ]
            "#]],
        );
    }

    #[test]
    fn runnable_annotation() {
        check(
            r#"
fn main() {}
            "#,
            expect![[r#"
                [
                    Annotation {
                        range: 0..12,
                        kind: Runnable {
                            debug: true,
                            runnable: Runnable {
                                nav: NavigationTarget {
                                    file_id: FileId(
                                        0,
                                    ),
                                    full_range: 0..12,
                                    focus_range: 3..7,
                                    name: "main",
                                    kind: Function,
                                },
                                kind: Bin,
                                cfg: None,
                            },
                        },
                    },
                    Annotation {
                        range: 0..12,
                        kind: Runnable {
                            debug: false,
                            runnable: Runnable {
                                nav: NavigationTarget {
                                    file_id: FileId(
                                        0,
                                    ),
                                    full_range: 0..12,
                                    focus_range: 3..7,
                                    name: "main",
                                    kind: Function,
                                },
                                kind: Bin,
                                cfg: None,
                            },
                        },
                    },
                    Annotation {
                        range: 3..7,
                        kind: HasReferences {
                            position: FilePosition {
                                file_id: FileId(
                                    0,
                                ),
                                offset: 3,
                            },
                            data: Some(
                                [],
                            ),
                        },
                    },
                ]
            "#]],
        );
    }

    #[test]
    fn method_annotations() {
        check(
            r#"
struct Test;

impl Test {
    fn self_by_ref(&self) {}
}

fn main() {
    Test.self_by_ref();
}
            "#,
            expect![[r#"
                [
                    Annotation {
                        range: 58..95,
                        kind: Runnable {
                            debug: true,
                            runnable: Runnable {
                                nav: NavigationTarget {
                                    file_id: FileId(
                                        0,
                                    ),
                                    full_range: 58..95,
                                    focus_range: 61..65,
                                    name: "main",
                                    kind: Function,
                                },
                                kind: Bin,
                                cfg: None,
                            },
                        },
                    },
                    Annotation {
                        range: 58..95,
                        kind: Runnable {
                            debug: false,
                            runnable: Runnable {
                                nav: NavigationTarget {
                                    file_id: FileId(
                                        0,
                                    ),
                                    full_range: 58..95,
                                    focus_range: 61..65,
                                    name: "main",
                                    kind: Function,
                                },
                                kind: Bin,
                                cfg: None,
                            },
                        },
                    },
                    Annotation {
                        range: 0..12,
                        kind: HasImpls {
                            position: FilePosition {
                                file_id: FileId(
                                    0,
                                ),
                                offset: 7,
                            },
                            data: Some(
                                [
                                    NavigationTarget {
                                        file_id: FileId(
                                            0,
                                        ),
                                        full_range: 14..56,
                                        focus_range: 19..23,
                                        name: "impl",
                                        kind: Impl,
                                    },
                                ],
                            ),
                        },
                    },
                    Annotation {
                        range: 0..12,
                        kind: HasReferences {
                            position: FilePosition {
                                file_id: FileId(
                                    0,
                                ),
                                offset: 7,
                            },
                            data: Some(
                                [
                                    FileRange {
                                        file_id: FileId(
                                            0,
                                        ),
                                        range: 19..23,
                                    },
                                    FileRange {
                                        file_id: FileId(
                                            0,
                                        ),
                                        range: 74..78,
                                    },
                                ],
                            ),
                        },
                    },
                    Annotation {
                        range: 33..44,
                        kind: HasReferences {
                            position: FilePosition {
                                file_id: FileId(
                                    0,
                                ),
                                offset: 33,
                            },
                            data: Some(
                                [
                                    FileRange {
                                        file_id: FileId(
                                            0,
                                        ),
                                        range: 79..90,
                                    },
                                ],
                            ),
                        },
                    },
                    Annotation {
                        range: 61..65,
                        kind: HasReferences {
                            position: FilePosition {
                                file_id: FileId(
                                    0,
                                ),
                                offset: 61,
                            },
                            data: Some(
                                [],
                            ),
                        },
                    },
                ]
            "#]],
        );
    }

    #[test]
    fn test_annotations() {
        check(
            r#"
fn main() {}

mod tests {
    #[test]
    fn my_cool_test() {}
}
            "#,
            expect![[r#"
                [
                    Annotation {
                        range: 0..12,
                        kind: Runnable {
                            debug: true,
                            runnable: Runnable {
                                nav: NavigationTarget {
                                    file_id: FileId(
                                        0,
                                    ),
                                    full_range: 0..12,
                                    focus_range: 3..7,
                                    name: "main",
                                    kind: Function,
                                },
                                kind: Bin,
                                cfg: None,
                            },
                        },
                    },
                    Annotation {
                        range: 0..12,
                        kind: Runnable {
                            debug: false,
                            runnable: Runnable {
                                nav: NavigationTarget {
                                    file_id: FileId(
                                        0,
                                    ),
                                    full_range: 0..12,
                                    focus_range: 3..7,
                                    name: "main",
                                    kind: Function,
                                },
                                kind: Bin,
                                cfg: None,
                            },
                        },
                    },
                    Annotation {
                        range: 14..64,
                        kind: Runnable {
                            debug: true,
                            runnable: Runnable {
                                nav: NavigationTarget {
                                    file_id: FileId(
                                        0,
                                    ),
                                    full_range: 14..64,
                                    focus_range: 18..23,
                                    name: "tests",
                                    kind: Module,
                                },
                                kind: TestMod {
                                    path: "tests",
                                },
                                cfg: None,
                            },
                        },
                    },
                    Annotation {
                        range: 14..64,
                        kind: Runnable {
                            debug: false,
                            runnable: Runnable {
                                nav: NavigationTarget {
                                    file_id: FileId(
                                        0,
                                    ),
                                    full_range: 14..64,
                                    focus_range: 18..23,
                                    name: "tests",
                                    kind: Module,
                                },
                                kind: TestMod {
                                    path: "tests",
                                },
                                cfg: None,
                            },
                        },
                    },
                    Annotation {
                        range: 30..62,
                        kind: Runnable {
                            debug: true,
                            runnable: Runnable {
                                nav: NavigationTarget {
                                    file_id: FileId(
                                        0,
                                    ),
                                    full_range: 30..62,
                                    focus_range: 45..57,
                                    name: "my_cool_test",
                                    kind: Function,
                                },
                                kind: Test {
                                    test_id: Path(
                                        "tests::my_cool_test",
                                    ),
                                    attr: TestAttr {
                                        ignore: false,
                                    },
                                },
                                cfg: None,
                            },
                        },
                    },
                    Annotation {
                        range: 30..62,
                        kind: Runnable {
                            debug: false,
                            runnable: Runnable {
                                nav: NavigationTarget {
                                    file_id: FileId(
                                        0,
                                    ),
                                    full_range: 30..62,
                                    focus_range: 45..57,
                                    name: "my_cool_test",
                                    kind: Function,
                                },
                                kind: Test {
                                    test_id: Path(
                                        "tests::my_cool_test",
                                    ),
                                    attr: TestAttr {
                                        ignore: false,
                                    },
                                },
                                cfg: None,
                            },
                        },
                    },
                    Annotation {
                        range: 3..7,
                        kind: HasReferences {
                            position: FilePosition {
                                file_id: FileId(
                                    0,
                                ),
                                offset: 3,
                            },
                            data: Some(
                                [],
                            ),
                        },
                    },
                ]
            "#]],
        );
    }
}
<|endoftext|>


<commit_before>use error_chain::ChainedError;
use ffmpeg;
use std::ptr;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::{Mutex, RwLock, Once, ONCE_INIT};
use std::thread;

use encode;
use errors::*;

lazy_static! {
    static ref CAPTURING: RwLock<bool> = RwLock::new(false);

    static ref ENCODER: Mutex<Option<encode::Encoder>> = Mutex::new(None);

    /// Receives frames to write pixels to.
    static ref FRAME_RECEIVER: Mutex<Option<Receiver<Buffer>>> = Mutex::new(None);

    /// Sends frames to encode.
    static ref FRAME_SENDER: Mutex<Option<Sender<Buffer>>> = Mutex::new(None);
}

pub struct Buffer {
    pub data: Vec<u8>,
    width: u32,
    height: u32,
}

fn capture_thread(frame_sender: Sender<Buffer>, frame_receiver: Receiver<Buffer>) {
    let mut frame = ffmpeg::frame::Video::empty();
    let mut buf = Buffer {
        data: Vec::new(),
        width: 0,
        height: 0,
    };

    frame_sender.send(buf).unwrap();

    loop {
        buf = frame_receiver.recv().unwrap();

        // Make sure frame is of correct size.
        if buf.width != frame.width() || buf.height != frame.height() {
            frame = ffmpeg::frame::Video::new(ffmpeg::format::Pixel::RGB24, buf.width, buf.height);
        }

        // Copy the pixel data into the frame.
        {
            let linesize = unsafe { ((*frame.as_ptr()).linesize[0]) as u32 };
            let mut data = frame.data_mut(0);

            for y in 0..buf.height {
                unsafe {
                    ptr::copy_nonoverlapping(buf.data.as_ptr().offset((y * buf.width * 3) as isize),
                                             data.as_mut_ptr().offset(((buf.height - y - 1) * linesize) as isize),
                                             (buf.width * 3) as usize);
                }
            }
        }

        // We're done with buf, now it can receive the next pack of pixels.
        // TODO: figure out threading, currently this causes one extra frame to be left over,
        // so the encoder is dropped in the end and then reinitialized for one frame.
        // frame_sender.send(buf).unwrap();

        // Let's encode the frame we just received.
        let mut encoder = ENCODER.lock().unwrap();

        // If the encoder wasn't initialized or if the frame size changed, initialize it.
        if encoder.as_ref().map_or(true, |enc| {
                enc.width() != frame.width() || enc.height() != frame.height()
            }) {
            *encoder = start_encoder((frame.width(), frame.height()));
        }

        // If we couldn't initialize the encoder just carry on with the loop.
        if encoder.is_none() {
            continue;
        }

        // Encode the frame.
        if let Err(ref e) = encoder.as_mut().unwrap().encode(&frame)
            .chain_err(|| "could not encode a frame") {
            println!("{}", e.display());
            // TODO: while we were getting here the main thread potentially
            // got into the CAPTURING if the second time already.
            *CAPTURING.write().unwrap() = false;
        };

        // If the capture was stopped, drop the encoder.
        if !*CAPTURING.read().unwrap() {
            *encoder = None;
        }

        // TODO: see above.
        frame_sender.send(buf).unwrap()
    }
}

fn start_encoder((width, height): (u32, u32)) -> Option<encode::Encoder> {
    match encode::Encoder::start("/home/yalter/test.mp4",
                                 (width, height),
                                 (1, 60).into()) {
        Ok(enc) => Some(enc),
        Err(ref e) => {
            println!("could not create the encoder: {}", e.display());
            // TODO: while we were getting here the main thread potentially
            // got into the CAPTURING if the second time already. This means
            // that our error message will be printed twice.
            *CAPTURING.write().unwrap() = false;
            None
        }
    }
}

pub fn initialize() {
    static INIT: Once = ONCE_INIT;
    INIT.call_once(|| {
        let (tx, rx) = channel::<Buffer>();
        let (tx2, rx2) = channel::<Buffer>();

        *FRAME_RECEIVER.lock().unwrap() = Some(rx);
        *FRAME_SENDER.lock().unwrap() = Some(tx2);

        thread::spawn(move || capture_thread(tx, rx2));
    });
}

pub fn get_buffer((width, height): (u32, u32)) -> Buffer {
    let mut buf = FRAME_RECEIVER.lock().unwrap().as_ref().unwrap().recv().unwrap();

    if buf.width != width || buf.height != height {
        println!("Changing resolution from {:?} to {:?}.",
                 (buf.width, buf.height),
                 (width, height));

        buf.data.resize((width * height * 3) as usize, 0);
        buf.width = width;
        buf.height = height;
    }

    buf
}

pub fn capture(buf: Buffer) {
    FRAME_SENDER.lock().unwrap().as_ref().unwrap().send(buf).unwrap();
}

pub fn is_capturing() -> bool {
    *CAPTURING.read().unwrap()
}

command!(cap_start, |_engine| {
    *CAPTURING.write().unwrap() = true;
});

command!(cap_stop, |_engine| {
    *CAPTURING.write().unwrap() = false;
});
<commit_msg>Fix threading errors<commit_after>use error_chain::ChainedError;
use ffmpeg;
use std::ptr;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::{Mutex, RwLock, Once, ONCE_INIT};
use std::thread;

use encode;
use errors::*;

lazy_static! {
    static ref CAPTURING: RwLock<bool> = RwLock::new(false);

    static ref ENCODER: Mutex<Option<encode::Encoder>> = Mutex::new(None);

    /// Receives buffers to write pixels to.
    static ref BUF_RECEIVER: Mutex<Option<Receiver<Buffer>>> = Mutex::new(None);

    /// Sends events and frames to encode to the capture thread.
    static ref SEND_TO_CAPTURE_THREAD: Mutex<Option<Sender<CaptureThreadEvent>>> = Mutex::new(None);
}

enum CaptureThreadEvent {
    CaptureStart,
    CaptureStop,
    Frame(Buffer),
}

pub struct Buffer {
    pub data: Vec<u8>,
    width: u32,
    height: u32,
}

fn capture_thread(buf_sender: Sender<Buffer>, event_receiver: Receiver<CaptureThreadEvent>) {
    // Send the buffer to the game thread right away.
    buf_sender.send(Buffer {
        data: Vec::new(),
        width: 0,
        height: 0,
    }).unwrap();

    // This is our frame which will only be reallocated on resolution changes.
    let mut frame = ffmpeg::frame::Video::empty();

    // This is set to true on encoding error or cap_stop and reset on cap_start.
    // When this is true, ignore any received frames.
    let mut drop_frames = true;

    // Event loop for the capture thread.
    loop {
        match event_receiver.recv().unwrap() {
            CaptureThreadEvent::CaptureStart => {
                drop_frames = false;
            },

            CaptureThreadEvent::CaptureStop => {
                *ENCODER.lock().unwrap() = None;
                drop_frames = true;
            },

            CaptureThreadEvent::Frame(buffer) => {
                if drop_frames {
                    continue;
                }

                if let Err(ref e) = encode(buffer, &buf_sender, &mut frame) {
                    *CAPTURING.write().unwrap() = false;
                    drop_frames = true;
                    println!("Encoding error: {}", e.display());
                }
            },
        }
    }
}

fn start_encoder((width, height): (u32, u32)) -> Result<encode::Encoder> {
    encode::Encoder::start("/home/yalter/test.mp4", (width, height), (1, 60).into())
}

fn encode(buf: Buffer, buf_sender: &Sender<Buffer>, frame: &mut ffmpeg::frame::Video) -> Result<()> {
    // Make sure frame is of correct size.
    if buf.width != frame.width() || buf.height != frame.height() {
        *frame = ffmpeg::frame::Video::new(ffmpeg::format::Pixel::RGB24, buf.width, buf.height);
    }

    // Copy the pixel data into the frame.
    {
        let linesize = unsafe { ((*frame.as_ptr()).linesize[0]) as u32 };
        let mut data = frame.data_mut(0);

        for y in 0..buf.height {
            unsafe {
                ptr::copy_nonoverlapping(buf.data.as_ptr().offset((y * buf.width * 3) as isize),
                data.as_mut_ptr().offset(((buf.height - y - 1) * linesize) as isize),
                (buf.width * 3) as usize);
            }
        }
    }

    // We're done with buf, now it can receive the next pack of pixels.
    buf_sender.send(buf).unwrap();

    // Let's encode the frame we just received.
    let mut encoder = ENCODER.lock().unwrap();

    // If the encoder wasn't initialized or if the frame size changed, initialize it.
    if encoder.as_ref().map_or(true, |enc| {
        enc.width() != frame.width() || enc.height() != frame.height()
    }) {
        *encoder = Some(start_encoder((frame.width(), frame.height()))
            .chain_err(|| "could not start the video encoder")?);
    }

    // Encode the frame.
    encoder.as_mut().unwrap().encode(&frame)
        .chain_err(|| "could not encode the frame")?;

    Ok(())
}

pub fn initialize() {
    static INIT: Once = ONCE_INIT;
    INIT.call_once(|| {
        let (tx, rx) = channel::<Buffer>();
        let (tx2, rx2) = channel::<CaptureThreadEvent>();

        *BUF_RECEIVER.lock().unwrap() = Some(rx);
        *SEND_TO_CAPTURE_THREAD.lock().unwrap() = Some(tx2);

        thread::spawn(move || capture_thread(tx, rx2));
    });
}

pub fn get_buffer((width, height): (u32, u32)) -> Buffer {
    let mut buf = BUF_RECEIVER.lock().unwrap().as_ref().unwrap().recv().unwrap();

    if buf.width != width || buf.height != height {
        println!("Changing resolution from {:?} to {:?}.",
                 (buf.width, buf.height),
                 (width, height));

        buf.data.resize((width * height * 3) as usize, 0);
        buf.width = width;
        buf.height = height;
    }

    buf
}

pub fn capture(buf: Buffer) {
    SEND_TO_CAPTURE_THREAD.lock().unwrap()
        .as_ref().unwrap()
        .send(CaptureThreadEvent::Frame(buf)).unwrap();
}

pub fn is_capturing() -> bool {
    *CAPTURING.read().unwrap()
}

command!(cap_start, |_engine| {
    *CAPTURING.write().unwrap() = true;

    SEND_TO_CAPTURE_THREAD.lock().unwrap()
        .as_ref().unwrap()
        .send(CaptureThreadEvent::CaptureStart).unwrap();
});

command!(cap_stop, |_engine| {
    *CAPTURING.write().unwrap() = false;

    SEND_TO_CAPTURE_THREAD.lock().unwrap()
        .as_ref().unwrap()
        .send(CaptureThreadEvent::CaptureStop).unwrap();
});
<|endoftext|>


<commit_before>#![crate_name = "cat"]
#![allow(unstable)]

#![feature(box_syntax, unsafe_destructor)]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Jordi Boggiano <j.boggiano@seld.be>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

/* last synced with: cat (GNU coreutils) 8.13 */

extern crate getopts;

use std::io::{print, File};
use std::io::stdio::{stdout_raw, stdin_raw, stderr};
use std::io::{IoResult};
use std::ptr::{copy_nonoverlapping_memory};

pub fn uumain(args: Vec<String>) -> isize {
    let program = args[0].as_slice();
    let opts = [
        getopts::optflag("A", "show-all", "equivalent to -vET"),
        getopts::optflag("b", "number-nonblank",
                         "number nonempty output lines, overrides -n"),
        getopts::optflag("e", "", "equivalent to -vE"),
        getopts::optflag("E", "show-ends", "display $ at end of each line"),
        getopts::optflag("n", "number", "number all output lines"),
        getopts::optflag("s", "squeeze-blank", "suppress repeated empty output lines"),
        getopts::optflag("t", "", "equivalent to -vT"),
        getopts::optflag("T", "show-tabs", "display TAB characters as ^I"),
        getopts::optflag("v", "show-nonprinting",
                         "use ^ and M- notation, except for LF (\\n) and TAB (\\t)"),
        getopts::optflag("h", "help", "display this help and exit"),
        getopts::optflag("V", "version", "output version information and exit"),
    ];
    let matches = match getopts::getopts(args.tail(), &opts) {
        Ok(m) => m,
        Err(f) => panic!("Invalid options\n{}", f)
    };
    if matches.opt_present("help") {
        println!("cat 1.0.0");
        println!("");
        println!("Usage:");
        println!("  {0} [OPTION]... [FILE]...", program);
        println!("");
        print(getopts::usage("Concatenate FILE(s), or standard input, to \
                             standard output.", &opts).as_slice());
        println!("");
        println!("With no FILE, or when FILE is -, read standard input.");
        return 0;
    }
    if matches.opt_present("version") {
        println!("cat 1.0.0");
        return 0;
    }

    let mut number_mode = NumberingMode::NumberNone;
    if matches.opt_present("n") {
        number_mode = NumberingMode::NumberAll;
    }
    if matches.opt_present("b") {
        number_mode = NumberingMode::NumberNonEmpty;
    }
    let show_nonprint = matches.opts_present(&["A".to_string(), "e".to_string(),
                                              "t".to_string(), "v".to_string()]);
    let show_ends = matches.opts_present(&["E".to_string(), "A".to_string(),
                                          "e".to_string()]);
    let show_tabs = matches.opts_present(&["A".to_string(), "T".to_string(),
                                          "t".to_string()]);
    let squeeze_blank = matches.opt_present("s");
    let mut files = matches.free;
    if files.is_empty() {
        files.push("-".to_string());
    }

    exec(files, number_mode, show_nonprint, show_ends, show_tabs, squeeze_blank);

    0
}

#[derive(Eq, PartialEq)]
enum NumberingMode {
    NumberNone,
    NumberNonEmpty,
    NumberAll,
}

fn write_lines(files: Vec<String>, number: NumberingMode, squeeze_blank: bool,
               show_ends: bool) {

    let mut line_counter: usize = 1;

    for (mut reader, interactive) in files.iter().filter_map(|p| open(&p[])) {

        let mut in_buf  = [0; 1024 * 31];
        let mut out_buf = [0; 1024 * 64];
        let mut writer = UnsafeWriter::new(out_buf.as_mut_slice(), stdout_raw());
        let mut at_line_start = true;
        loop {
            let n = match reader.read(&mut in_buf) {
                Ok(n) if n != 0 => n,
                _ => break,
            };
            let in_buf = &in_buf[..n];
            let mut buf_pos = range(0, n);
            loop {
                writer.possibly_flush();
                let pos = match buf_pos.next() {
                    Some(p) => p,
                    None => break,
                };
                if in_buf[pos] == '\n' as u8 {
                    if !at_line_start || !squeeze_blank {
                        if at_line_start && number == NumberingMode::NumberAll {
                            (write!(&mut writer, "{0:6}\t", line_counter)).unwrap();
                            line_counter += 1;
                        }
                        if show_ends {
                            writer.write_u8('$' as u8).unwrap();
                        }
                        writer.write_u8('\n' as u8).unwrap();
                        if interactive {
                            writer.flush().unwrap();
                        }
                    }
                    at_line_start = true;
                    continue;
                }
                if at_line_start && number != NumberingMode::NumberNone {
                    (write!(&mut writer, "{0:6}\t", line_counter)).unwrap();
                    line_counter += 1;
                }
                match in_buf[pos..].iter().position(|c| *c == '\n' as u8) {
                    Some(p) => {
                        writer.write(&in_buf[pos..pos + p]).unwrap();
                        if show_ends {
                            writer.write_u8('$' as u8).unwrap();
                        }
                        writer.write_u8('\n' as u8).unwrap();
                        if interactive {
                            writer.flush().unwrap();
                        }
                        buf_pos = range(pos + p + 1, n);
                        at_line_start = true;
                    },
                    None => {
                        writer.write(&in_buf[pos..]).unwrap();
                        at_line_start = false;
                        break;
                    }
                };
            }
        }
    }
}

fn write_bytes(files: Vec<String>, number: NumberingMode, squeeze_blank: bool,
               show_ends: bool, show_nonprint: bool, show_tabs: bool) {

    let mut line_counter: usize = 1;

    for (mut reader, interactive) in files.iter().filter_map(|p| open(&p[])) {

        // Flush all 1024 iterations.
        let mut flush_counter = range(0us, 1024);

        let mut in_buf  = [0; 1024 * 32];
        let mut out_buf = [0; 1024 * 64];
        let mut writer = UnsafeWriter::new(out_buf.as_mut_slice(), stdout_raw());
        let mut at_line_start = true;
        loop {
            let n = match reader.read(&mut in_buf) {
                Ok(n) if n != 0 => n,
                _ => break,
            };
            for &byte in in_buf[..n].iter() {
                if flush_counter.next().is_none() {
                    writer.possibly_flush();
                    flush_counter = range(0us, 1024);
                }
                if byte == '\n' as u8 {
                    if !at_line_start || !squeeze_blank {
                        if at_line_start && number == NumberingMode::NumberAll {
                            (write!(&mut writer, "{0:6}\t", line_counter)).unwrap();
                            line_counter += 1;
                        }
                        if show_ends {
                            writer.write_u8('$' as u8).unwrap();
                        }
                        writer.write_u8('\n' as u8).unwrap();
                        if interactive {
                            writer.flush().unwrap();
                        }
                    }
                    at_line_start = true;
                    continue;
                }
                if at_line_start && number != NumberingMode::NumberNone {
                    (write!(&mut writer, "{0:6}\t", line_counter)).unwrap();
                    line_counter += 1;
                    at_line_start = false;
                }
                // This code is slow because of the many branches. cat in glibc avoids
                // this by having the whole loop inside show_nonprint.
                if byte == '\t' as u8 {
                    if show_tabs {
                        writer.write_str("^I")
                    } else {
                        writer.write_u8(byte)
                    }
                } else if show_nonprint {
                    let byte = match byte {
                        128 ... 255 => {
                            writer.write_str("M-").unwrap();
                            byte - 128
                        },
                        _ => byte,
                    };
                    match byte {
                        0 ... 31 => writer.write(&['^' as u8, byte + 64]),
                        127      => writer.write(&['^' as u8, byte - 64]),
                        _        => writer.write_u8(byte),
                    }
                } else {
                    writer.write_u8(byte)
                }.unwrap();
            }
        }
    }
}

fn write_fast(files: Vec<String>) {
    let mut writer = stdout_raw();
    let mut in_buf = [0; 1024 * 64];

    for (mut reader, _) in files.iter().filter_map(|p| open(&p[])) {
        loop {
            match reader.read(&mut in_buf) {
                Ok(n) if n != 0 => {
                    // This interface is completely broken.
                    writer.write(&in_buf[..n]).unwrap();
                },
                _ => break
            }
        }
    }
}

fn exec(files: Vec<String>, number: NumberingMode, show_nonprint: bool,
        show_ends: bool, show_tabs: bool, squeeze_blank: bool) {

    if show_nonprint || show_tabs {
        write_bytes(files, number, squeeze_blank, show_ends, show_nonprint, show_tabs);
    } else if number != NumberingMode::NumberNone || squeeze_blank || show_ends {
        write_lines(files, number, squeeze_blank, show_ends);
    } else {
        write_fast(files);
    }
}

fn open(path: &str) -> Option<(Box<Reader>, bool)> {
    if path == "-" {
        let stdin = stdin_raw();
        let interactive = stdin.isatty();
        return Some((box stdin as Box<Reader>, interactive));
    }

    match File::open(&std::path::Path::new(path)) {
        Ok(f) => Some((box f as Box<Reader>, false)),
        Err(e) => {
            (writeln!(&mut stderr(), "cat: {0}: {1}", path, e.to_string())).unwrap();
            None
        },
    }
}

struct UnsafeWriter<'a, W> {
    inner: W,
    buf: &'a mut [u8],
    pos: usize,
    threshold: usize,
}

impl<'a, W: Writer> UnsafeWriter<'a, W> {
    fn new(buf: &'a mut [u8], inner: W) -> UnsafeWriter<'a, W> {
        let threshold = buf.len()/2;
        UnsafeWriter {
            inner: inner,
            buf: buf,
            pos: 0,
            threshold: threshold,
        }
    }

    fn flush_buf(&mut self) -> IoResult<()> {
        if self.pos != 0 {
            let ret = self.inner.write(&self.buf[..self.pos]);
            self.pos = 0;
            ret
        } else {
            Ok(())
        }
    }

    fn possibly_flush(&mut self) {
        if self.pos > self.threshold {
            self.inner.write(&self.buf[..self.pos]).unwrap();
            self.pos = 0;
        }
    }
}

#[inline(never)]
fn fail() -> ! {
    panic!("assertion failed");
}

impl<'a, W: Writer> Writer for UnsafeWriter<'a, W> {
    fn write(&mut self, buf: &[u8]) -> IoResult<()> {
        let dst = &mut self.buf[self.pos..];
        if buf.len() > dst.len() {
            fail();
        }
        unsafe {
            copy_nonoverlapping_memory(dst.as_mut_ptr(), buf.as_ptr(), buf.len())
        }
        self.pos += buf.len();
        Ok(())
    }

    fn flush(&mut self) -> IoResult<()> {
        self.flush_buf().and_then(|()| self.inner.flush())
    }
}

#[unsafe_destructor]
impl<'a, W: Writer> Drop for UnsafeWriter<'a, W> {
    fn drop(&mut self) {
        let _ = self.flush_buf();
    }
}

/* vim: set ai ts=4 sw=4 sts=4 et : */
<commit_msg>Cleanup cat: use while let instead of loop then match.<commit_after>#![crate_name = "cat"]
#![allow(unstable)]

#![feature(box_syntax, unsafe_destructor)]

/*
 * This file is part of the uutils coreutils package.
 *
 * (c) Jordi Boggiano <j.boggiano@seld.be>
 *
 * For the full copyright and license information, please view the LICENSE
 * file that was distributed with this source code.
 */

/* last synced with: cat (GNU coreutils) 8.13 */

extern crate getopts;

use std::io::{print, File};
use std::io::stdio::{stdout_raw, stdin_raw, stderr};
use std::io::{IoResult};
use std::ptr::{copy_nonoverlapping_memory};

pub fn uumain(args: Vec<String>) -> isize {
    let program = args[0].as_slice();
    let opts = [
        getopts::optflag("A", "show-all", "equivalent to -vET"),
        getopts::optflag("b", "number-nonblank",
                         "number nonempty output lines, overrides -n"),
        getopts::optflag("e", "", "equivalent to -vE"),
        getopts::optflag("E", "show-ends", "display $ at end of each line"),
        getopts::optflag("n", "number", "number all output lines"),
        getopts::optflag("s", "squeeze-blank", "suppress repeated empty output lines"),
        getopts::optflag("t", "", "equivalent to -vT"),
        getopts::optflag("T", "show-tabs", "display TAB characters as ^I"),
        getopts::optflag("v", "show-nonprinting",
                         "use ^ and M- notation, except for LF (\\n) and TAB (\\t)"),
        getopts::optflag("h", "help", "display this help and exit"),
        getopts::optflag("V", "version", "output version information and exit"),
    ];
    let matches = match getopts::getopts(args.tail(), &opts) {
        Ok(m) => m,
        Err(f) => panic!("Invalid options\n{}", f)
    };
    if matches.opt_present("help") {
        println!("cat 1.0.0");
        println!("");
        println!("Usage:");
        println!("  {0} [OPTION]... [FILE]...", program);
        println!("");
        print(getopts::usage("Concatenate FILE(s), or standard input, to \
                             standard output.", &opts).as_slice());
        println!("");
        println!("With no FILE, or when FILE is -, read standard input.");
        return 0;
    }
    if matches.opt_present("version") {
        println!("cat 1.0.0");
        return 0;
    }

    let mut number_mode = NumberingMode::NumberNone;
    if matches.opt_present("n") {
        number_mode = NumberingMode::NumberAll;
    }
    if matches.opt_present("b") {
        number_mode = NumberingMode::NumberNonEmpty;
    }
    let show_nonprint = matches.opts_present(&["A".to_string(), "e".to_string(),
                                              "t".to_string(), "v".to_string()]);
    let show_ends = matches.opts_present(&["E".to_string(), "A".to_string(),
                                          "e".to_string()]);
    let show_tabs = matches.opts_present(&["A".to_string(), "T".to_string(),
                                          "t".to_string()]);
    let squeeze_blank = matches.opt_present("s");
    let mut files = matches.free;
    if files.is_empty() {
        files.push("-".to_string());
    }

    exec(files, number_mode, show_nonprint, show_ends, show_tabs, squeeze_blank);

    0
}

#[derive(Eq, PartialEq)]
enum NumberingMode {
    NumberNone,
    NumberNonEmpty,
    NumberAll,
}

fn write_lines(files: Vec<String>, number: NumberingMode, squeeze_blank: bool,
               show_ends: bool) {

    let mut line_counter: usize = 1;

    for (mut reader, interactive) in files.iter().filter_map(|p| open(&p[])) {

        let mut in_buf  = [0; 1024 * 31];
        let mut out_buf = [0; 1024 * 64];
        let mut writer = UnsafeWriter::new(out_buf.as_mut_slice(), stdout_raw());
        let mut at_line_start = true;
        while let Ok(n) = reader.read(&mut in_buf) {
            if n == 0 { break }

            let in_buf = &in_buf[..n];
            let mut buf_pos = range(0, n);
            loop {
                writer.possibly_flush();
                let pos = match buf_pos.next() {
                    Some(p) => p,
                    None => break,
                };
                if in_buf[pos] == '\n' as u8 {
                    if !at_line_start || !squeeze_blank {
                        if at_line_start && number == NumberingMode::NumberAll {
                            (write!(&mut writer, "{0:6}\t", line_counter)).unwrap();
                            line_counter += 1;
                        }
                        if show_ends {
                            writer.write_u8('$' as u8).unwrap();
                        }
                        writer.write_u8('\n' as u8).unwrap();
                        if interactive {
                            writer.flush().unwrap();
                        }
                    }
                    at_line_start = true;
                    continue;
                }
                if at_line_start && number != NumberingMode::NumberNone {
                    (write!(&mut writer, "{0:6}\t", line_counter)).unwrap();
                    line_counter += 1;
                }
                match in_buf[pos..].iter().position(|c| *c == '\n' as u8) {
                    Some(p) => {
                        writer.write(&in_buf[pos..pos + p]).unwrap();
                        if show_ends {
                            writer.write_u8('$' as u8).unwrap();
                        }
                        writer.write_u8('\n' as u8).unwrap();
                        if interactive {
                            writer.flush().unwrap();
                        }
                        buf_pos = range(pos + p + 1, n);
                        at_line_start = true;
                    },
                    None => {
                        writer.write(&in_buf[pos..]).unwrap();
                        at_line_start = false;
                        break;
                    }
                };
            }
        }
    }
}

fn write_bytes(files: Vec<String>, number: NumberingMode, squeeze_blank: bool,
               show_ends: bool, show_nonprint: bool, show_tabs: bool) {

    let mut line_counter: usize = 1;

    for (mut reader, interactive) in files.iter().filter_map(|p| open(&p[])) {

        // Flush all 1024 iterations.
        let mut flush_counter = range(0us, 1024);

        let mut in_buf  = [0; 1024 * 32];
        let mut out_buf = [0; 1024 * 64];
        let mut writer = UnsafeWriter::new(out_buf.as_mut_slice(), stdout_raw());
        let mut at_line_start = true;
        while let Ok(n) = reader.read(&mut in_buf) {
            if n == 0 { break }

            for &byte in in_buf[..n].iter() {
                if flush_counter.next().is_none() {
                    writer.possibly_flush();
                    flush_counter = range(0us, 1024);
                }
                if byte == '\n' as u8 {
                    if !at_line_start || !squeeze_blank {
                        if at_line_start && number == NumberingMode::NumberAll {
                            (write!(&mut writer, "{0:6}\t", line_counter)).unwrap();
                            line_counter += 1;
                        }
                        if show_ends {
                            writer.write_u8('$' as u8).unwrap();
                        }
                        writer.write_u8('\n' as u8).unwrap();
                        if interactive {
                            writer.flush().unwrap();
                        }
                    }
                    at_line_start = true;
                    continue;
                }
                if at_line_start && number != NumberingMode::NumberNone {
                    (write!(&mut writer, "{0:6}\t", line_counter)).unwrap();
                    line_counter += 1;
                    at_line_start = false;
                }
                // This code is slow because of the many branches. cat in glibc avoids
                // this by having the whole loop inside show_nonprint.
                if byte == '\t' as u8 {
                    if show_tabs {
                        writer.write_str("^I")
                    } else {
                        writer.write_u8(byte)
                    }
                } else if show_nonprint {
                    let byte = match byte {
                        128 ... 255 => {
                            writer.write_str("M-").unwrap();
                            byte - 128
                        },
                        _ => byte,
                    };
                    match byte {
                        0 ... 31 => writer.write(&['^' as u8, byte + 64]),
                        127      => writer.write(&['^' as u8, byte - 64]),
                        _        => writer.write_u8(byte),
                    }
                } else {
                    writer.write_u8(byte)
                }.unwrap();
            }
        }
    }
}

fn write_fast(files: Vec<String>) {
    let mut writer = stdout_raw();
    let mut in_buf = [0; 1024 * 64];

    for (mut reader, _) in files.iter().filter_map(|p| open(&p[])) {
        while let Ok(n) = reader.read(&mut in_buf) {
            if n == 0 { break }
            // This interface is completely broken.
            writer.write(&in_buf[..n]).unwrap();
        }
    }
}

fn exec(files: Vec<String>, number: NumberingMode, show_nonprint: bool,
        show_ends: bool, show_tabs: bool, squeeze_blank: bool) {

    if show_nonprint || show_tabs {
        write_bytes(files, number, squeeze_blank, show_ends, show_nonprint, show_tabs);
    } else if number != NumberingMode::NumberNone || squeeze_blank || show_ends {
        write_lines(files, number, squeeze_blank, show_ends);
    } else {
        write_fast(files);
    }
}

fn open(path: &str) -> Option<(Box<Reader>, bool)> {
    if path == "-" {
        let stdin = stdin_raw();
        let interactive = stdin.isatty();
        return Some((box stdin as Box<Reader>, interactive));
    }

    match File::open(&std::path::Path::new(path)) {
        Ok(f) => Some((box f as Box<Reader>, false)),
        Err(e) => {
            (writeln!(&mut stderr(), "cat: {0}: {1}", path, e.to_string())).unwrap();
            None
        },
    }
}

struct UnsafeWriter<'a, W> {
    inner: W,
    buf: &'a mut [u8],
    pos: usize,
    threshold: usize,
}

impl<'a, W: Writer> UnsafeWriter<'a, W> {
    fn new(buf: &'a mut [u8], inner: W) -> UnsafeWriter<'a, W> {
        let threshold = buf.len()/2;
        UnsafeWriter {
            inner: inner,
            buf: buf,
            pos: 0,
            threshold: threshold,
        }
    }

    fn flush_buf(&mut self) -> IoResult<()> {
        if self.pos != 0 {
            let ret = self.inner.write(&self.buf[..self.pos]);
            self.pos = 0;
            ret
        } else {
            Ok(())
        }
    }

    fn possibly_flush(&mut self) {
        if self.pos > self.threshold {
            self.inner.write(&self.buf[..self.pos]).unwrap();
            self.pos = 0;
        }
    }
}

#[inline(never)]
fn fail() -> ! {
    panic!("assertion failed");
}

impl<'a, W: Writer> Writer for UnsafeWriter<'a, W> {
    fn write(&mut self, buf: &[u8]) -> IoResult<()> {
        let dst = &mut self.buf[self.pos..];
        if buf.len() > dst.len() {
            fail();
        }
        unsafe {
            copy_nonoverlapping_memory(dst.as_mut_ptr(), buf.as_ptr(), buf.len())
        }
        self.pos += buf.len();
        Ok(())
    }

    fn flush(&mut self) -> IoResult<()> {
        self.flush_buf().and_then(|()| self.inner.flush())
    }
}

#[unsafe_destructor]
impl<'a, W: Writer> Drop for UnsafeWriter<'a, W> {
    fn drop(&mut self) {
        let _ = self.flush_buf();
    }
}

/* vim: set ai ts=4 sw=4 sts=4 et : */
<|endoftext|>