<commit_before>// Copyright (c) 2016 Nikita Pekin and the smexybot contributors
// See the README.md file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Provides functionality for the `tag` command.

extern crate uuid;

use std::collections::HashMap;
use std::fs::{self, File};
use std::io::{ErrorKind, Read, Write};
use std::sync::Mutex;

use chrono::{DateTime, UTC};
use self::uuid::Uuid;
use serde_json;
use serenity::client::{rest, Context};
use serenity::model::{GuildId, Message, UserId};
use serenity::utils::builder::CreateEmbed;

use ::CONFIG;
use util::{check_msg, merge, timestamp_to_string};

lazy_static! {
    // TODO: ensure this doesn't get loaded lazily (compared to the bot).
    static ref TAGS: Tags = Tags {
        // TODO: get the config name from the global config?
        config: Mutex::new(Config::new("tags.json")),
    };
}

#[cfg(feature = "nightly")]
include!("tag.in.rs");

#[cfg(feature = "with-syntex")]
include!(concat!(env!("OUT_DIR"), "/tag.rs"));

impl Tag {
    // TODO: perhaps take references here?
    fn new(
        name: String,
        content: String,
        owner_id: u64,
        uses: Option<u32>,
        location: Option<String>,
        created_at: Option<DateTime<UTC>>,
    ) -> Self {
        Tag {
            name: name,
            content: content,
            owner_id: owner_id,
            uses: uses.unwrap_or(0),
            location: location,
            created_at: created_at.unwrap_or_else(UTC::now),
        }
    }

    fn as_embed(&self, embed: CreateEmbed) -> CreateEmbed {
        embed
            .title(&self.name)
            .field(|f| f.name("Owner").value(&format!("<@!{}>", self.owner_id)))
            .field(|f| f.name("Uses").value(&self.uses.to_string()))
            .author(|a| {
                let owner_id = UserId(self.owner_id);
                let (name, avatar_url) = match owner_id.find() {
                    Some(user) => (user.name.clone(), user.avatar_url()),
                    None => match rest::get_user(owner_id.0) {
                        Ok(user) => (user.name.clone(), user.avatar_url()),
                        // TODO: handle this failure case better.
                        Err(_) => return a,
                    },
                };
                let mut a = a.name(&name);
                if let Some(avatar_url) = avatar_url {
                    a = a.icon_url(&avatar_url);
                }
                a
            })
            .timestamp(timestamp_to_string(&self.created_at))
            .footer(|f| f
                .text(
                    if self.is_generic() {
                        "Generic"
                    } else {
                        "Server-specific"
                    }
                )
            )
    }

    fn is_generic(&self) -> bool {
        self.location.is_none()
    }
}

#[derive(Debug)]
struct Config {
    name: String,
    tags: HashMap<String, HashMap<String, Tag>>,
}

impl Config {
    fn new(name: &str) -> Self {
        let mut config = Config {
            name: name.to_owned(),
            tags: HashMap::new(),
        };

        // TODO: do this async.
        config.load();

        config
    }

    fn get(&self, key: &str) -> Option<&HashMap<String, Tag>> {
        self.tags.get(key)
    }

    fn insert(&mut self, key: String, value: HashMap<String, Tag>) {
        self.tags.insert(key, value);
        self.save();
    }

    fn load(&mut self) {
        let mut file = match File::open(&self.name) {
            Ok(file) => file,
            // If no file is present, assume this is a fresh config.
            Err(ref err) if err.kind() == ErrorKind::NotFound => return,
            Err(_) => panic!("Failed to open file: {}", self.name),
        };
        let mut tags = String::new();
        file.read_to_string(&mut tags)
            .expect(&format!("Failed to read from file: {}", self.name));
        self.tags = serde_json::from_str(&tags)
            .expect("Failed to deserialize Config");
        debug!("Loaded config from: {}", self.name);
    }

    // TODO: do this async
    fn save(&self) {
        let temp = format!("{}-{}.tmp", Uuid::new_v4(), self.name);
        let mut file = File::create(&temp)
            .expect(&format!("Failed to create file: {}", temp));
        file.write_all(
            serde_json::to_string(&self.tags)
                .expect("Failed to serialize Config")
                .as_bytes(),
        ).expect(&format!("Failed to write to file: {}", temp));

        // Atomically copy the new config.
        fs::rename(temp, &self.name)
            .expect("Failed to write new Config");
        trace!("Saved config to: {}", self.name);
    }
}

#[derive(Debug)]
struct Tags {
    config: Mutex<Config>,
}

impl Tags {
    fn get_possible_tags(&self, guild: Option<GuildId>)
        -> HashMap<String, Tag>
    {
        // TODO: handle this unwrap.
        let config = self.config.lock().expect("Failed to lock Config");
        // TODO: properly handle a missing hashmap case.
        let generic = config.get("generic")
            .cloned()
            .unwrap_or_else(HashMap::new);

        match guild {
            None => generic,
            Some(guild) => merge(
                generic,
                config.get(&guild.to_string())
                    .cloned()
                    .unwrap_or_else(HashMap::new),
            ),
        }
    }

    // TODO: return a list of suggestions.
    fn get_tag(&self, guild: Option<GuildId>, name: String)
        -> Result<Tag, String>
    {
        self.get_possible_tags(guild)
            .get(&name)
            .cloned()
            .ok_or_else(|| "Tag not found".to_owned())
    }

    fn put_tag(&self, guild: Option<GuildId>, name: String, tag: Tag) {
        // Load the actual tag so we can modify it.
        // TODO: handle this better because this can cause an issue
        // if the tag was deleted since being retrieved.
        let mut config = TAGS.config.lock()
            .expect("Failed to lock Config");
        {
            let database = config
                .tags
                .get_mut(&get_database_location(guild))
                .unwrap();
            database.insert(name, tag);
        }
        config.save();
    }

    fn delete_tag(&self, guild: Option<GuildId>, name: &str) {
        // TODO: handle this better because this can cause an issue
        // if the tag was deleted since being retrieved.
        let mut config = TAGS.config.lock()
            .expect("Failed to lock Config");
        {
            let database = config
                .tags
                .get_mut(&get_database_location(guild))
                .unwrap();
            database.remove(name);
        }
        config.save();
    }
}

// TODO: sort tags by server
// TODO: add edit/delete functionality
pub fn handler(context: &Context, message: &Message, args: Vec<String>)
    -> Result<(), String>
{
    let mut args = args.into_iter();

    match args.next().as_ref().map(String::as_ref) {
        Some("create") => create(context, message, args.collect()),
        Some("info") => info(context, message, args.collect()),
        Some("list") => list(context, message, args.collect()),
        Some("edit") => edit(context, message, args.collect()),
        Some("delete") => delete(context, message, args.collect()),
        Some(name) => {
            // TODO: ensure we get the GuildId even if it's not in the cache.
            let guild_id = message.guild_id();

            let lookup = name.to_lowercase();
            match TAGS.get_tag(guild_id, lookup.clone()) {
                Ok(tag) => {
                    let mut tag = tag.clone();
                    tag.uses += 1;
                    TAGS.put_tag(guild_id, lookup, tag.clone());
                    check_msg(context.say(&tag.content));

                    Ok(())
                },
                Err(err) => Err(err),
            }
        },
        _ => Err("Either specify a tag name or use one of the available commands.".to_owned()),
    }
}

pub fn create(context: &Context, message: &Message, args: Vec<String>)
    -> Result<(), String>
{
    let mut args = args.into_iter();

    let name = match args.next() {
        Some(name) => name,
        None => return Err("Please specify a name for the tag.".to_owned()),
    };

    let content = args.collect::<Vec<String>>();
    let content = if content.is_empty() {
        return Err("Please specify some content for the tag.".to_owned());
    } else {
        content.join(" ")
    };

    // TODO: clean the tag content here.

    let name = name.trim().to_lowercase().to_owned();
    verify_tag_name(&name)?;

    let location = get_database_location(message.guild_id());
    let mut config = TAGS.config.lock().expect("Failed to lock Config");
    let mut database = config.get(&location)
        .cloned()
        .unwrap_or_else(HashMap::new);
    if database.contains_key(&name) {
        return Err("Tag already exists.".to_owned());
    }

    database.insert(name.clone(), Tag::new(
            name.clone(),
            content,
            message.author.id.0,
            None,
            Some(location.clone()),
            None,
    ));
    // TODO: perhaps do this more efficiently by modifying the DB in place?
    // TODO: or maybe use `put_tag`?
    config.insert(location, database);
    check_msg(context.say(&format!("Tag \"{}\" successfully created.", name)));

    Ok(())
}

pub fn info(context: &Context, message: &Message, args: Vec<String>)
    -> Result<(), String>
{
    let mut args = args.into_iter();

    let name = match args.next() {
        Some(name) => name,
        None => return Err("Please specify a name for the tag to get info on.".to_owned()),
    };

    let name = name.trim().to_lowercase().to_owned();
    // TODO: ensure we get the GuildId even if it's not in the cache.
    let guild_id = message.guild_id();
    let tag = TAGS.get_tag(guild_id, name)?;

    check_msg(context.send_message(
        message.channel_id,
        |m| m.embed(|e| tag.as_embed(e)),
    ));

    Ok(())
}

pub fn list(context: &Context, message: &Message, _args: Vec<String>)
    -> Result<(), String>
{
    // TODO: ensure we get the GuildId even if it's not in the cache.
    let guild_id = message.guild_id();
    let mut tags = TAGS.get_possible_tags(guild_id);
    let mut tags = tags
        .drain()
        .map(|(k, _)| k)
        .collect::<Vec<String>>();
    tags.sort();

    let response = if tags.is_empty() {
        "No tags available.".to_owned()
    } else {
        format!("Available tags: {}", tags.join(", "))
    };
    check_msg(context.say(&response));

    Ok(())
}

pub fn edit(context: &Context, message: &Message, args: Vec<String>)
    -> Result<(), String>
{
    let mut args = args.into_iter();

    let name = match args.next() {
        Some(name) => name,
        None => return Err("Please specify a tag to edit.".to_owned()),
    };

    let name = name.trim().to_lowercase().to_owned();

    let guild_id = message.guild_id();
    let mut tag = match TAGS.get_tag(guild_id, name.clone()) {
        Ok(tag) => tag,
        Err(err) => return Err(err),
    };

    if !owner_check(message, &tag) {
        return Err("You do not have permission to do that.".to_owned());
    }

    let content = args.collect::<Vec<String>>();
    let content = if content.is_empty() {
        return Err("Please specify some content for the tag.".to_owned());
    } else {
        content.join(" ")
    };

    // TODO: clean the tag content here.

    tag.content = content;
    TAGS.put_tag(guild_id, name.clone(), tag);

    check_msg(context.say(&format!("Tag \"{}\" successfully updated.", name)));

    Ok(())
}

pub fn delete(context: &Context, message: &Message, args: Vec<String>)
    -> Result<(), String>
{
    let mut args = args.into_iter();

    let name = match args.next() {
        Some(name) => name,
        None => return Err("Please specify a tag to delete.".to_owned()),
    };

    let name = name.trim().to_lowercase().to_owned();

    let guild_id = message.guild_id();
    let tag = match TAGS.get_tag(guild_id, name.clone()) {
        Ok(tag) => tag,
        Err(err) => return Err(err),
    };

    if !owner_check(message, &tag) {
        return Err("You do not have permission to do that.".to_owned());
    }

    TAGS.delete_tag(guild_id, &name);

    check_msg(context.say(&format!("Tag \"{}\" successfully deleted.", name)));

    Ok(())
}

// Denies certain tag names from being used as keys.
fn verify_tag_name(name: &str) -> Result<(), String> {
    if name.contains("@everyone") || name.contains("@here") {
        return Err("Tag contains blocked words".to_owned());
    }

    if name.len() > 100 {
        return Err("Tag name limit is 100 characters".to_owned());
    }

    Ok(())
}

fn owner_check(message: &Message, tag: &Tag) -> bool {
    CONFIG.owners.contains(&message.author.id.0) ||
        message.author.id == tag.owner_id
}

fn get_database_location(guild: Option<GuildId>) -> String {
    // TODO: ensure we get the GuildId even if it's not in the cache.
    guild.map(|g| g.to_string())
        .unwrap_or_else(|| "generic".to_owned())
}
<commit_msg>impr: use fn pointers to simplify match arms<commit_after>// Copyright (c) 2016 Nikita Pekin and the smexybot contributors
// See the README.md file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Provides functionality for the `tag` command.

extern crate uuid;

use std::collections::HashMap;
use std::fs::{self, File};
use std::io::{ErrorKind, Read, Write};
use std::sync::Mutex;

use chrono::{DateTime, UTC};
use self::uuid::Uuid;
use serde_json;
use serenity::client::{rest, Context};
use serenity::model::{GuildId, Message, UserId};
use serenity::utils::builder::CreateEmbed;

use ::CONFIG;
use util::{check_msg, merge, timestamp_to_string};

lazy_static! {
    // TODO: ensure this doesn't get loaded lazily (compared to the bot).
    static ref TAGS: Tags = Tags {
        // TODO: get the config name from the global config?
        config: Mutex::new(Config::new("tags.json")),
    };
}

#[cfg(feature = "nightly")]
include!("tag.in.rs");

#[cfg(feature = "with-syntex")]
include!(concat!(env!("OUT_DIR"), "/tag.rs"));

impl Tag {
    // TODO: perhaps take references here?
    fn new(
        name: String,
        content: String,
        owner_id: u64,
        uses: Option<u32>,
        location: Option<String>,
        created_at: Option<DateTime<UTC>>,
    ) -> Self {
        Tag {
            name: name,
            content: content,
            owner_id: owner_id,
            uses: uses.unwrap_or(0),
            location: location,
            created_at: created_at.unwrap_or_else(UTC::now),
        }
    }

    fn as_embed(&self, embed: CreateEmbed) -> CreateEmbed {
        embed
            .title(&self.name)
            .field(|f| f.name("Owner").value(&format!("<@!{}>", self.owner_id)))
            .field(|f| f.name("Uses").value(&self.uses.to_string()))
            .author(|a| {
                let owner_id = UserId(self.owner_id);
                let (name, avatar_url) = match owner_id.find() {
                    Some(user) => (user.name.clone(), user.avatar_url()),
                    None => match rest::get_user(owner_id.0) {
                        Ok(user) => (user.name.clone(), user.avatar_url()),
                        // TODO: handle this failure case better.
                        Err(_) => return a,
                    },
                };
                let mut a = a.name(&name);
                if let Some(avatar_url) = avatar_url {
                    a = a.icon_url(&avatar_url);
                }
                a
            })
            .timestamp(timestamp_to_string(&self.created_at))
            .footer(|f| f
                .text(
                    if self.is_generic() {
                        "Generic"
                    } else {
                        "Server-specific"
                    }
                )
            )
    }

    fn is_generic(&self) -> bool {
        self.location.is_none()
    }
}

#[derive(Debug)]
struct Config {
    name: String,
    tags: HashMap<String, HashMap<String, Tag>>,
}

impl Config {
    fn new(name: &str) -> Self {
        let mut config = Config {
            name: name.to_owned(),
            tags: HashMap::new(),
        };

        // TODO: do this async.
        config.load();

        config
    }

    fn get(&self, key: &str) -> Option<&HashMap<String, Tag>> {
        self.tags.get(key)
    }

    fn insert(&mut self, key: String, value: HashMap<String, Tag>) {
        self.tags.insert(key, value);
        self.save();
    }

    fn load(&mut self) {
        let mut file = match File::open(&self.name) {
            Ok(file) => file,
            // If no file is present, assume this is a fresh config.
            Err(ref err) if err.kind() == ErrorKind::NotFound => return,
            Err(_) => panic!("Failed to open file: {}", self.name),
        };
        let mut tags = String::new();
        file.read_to_string(&mut tags)
            .expect(&format!("Failed to read from file: {}", self.name));
        self.tags = serde_json::from_str(&tags)
            .expect("Failed to deserialize Config");
        debug!("Loaded config from: {}", self.name);
    }

    // TODO: do this async
    fn save(&self) {
        let temp = format!("{}-{}.tmp", Uuid::new_v4(), self.name);
        let mut file = File::create(&temp)
            .expect(&format!("Failed to create file: {}", temp));
        file.write_all(
            serde_json::to_string(&self.tags)
                .expect("Failed to serialize Config")
                .as_bytes(),
        ).expect(&format!("Failed to write to file: {}", temp));

        // Atomically copy the new config.
        fs::rename(temp, &self.name)
            .expect("Failed to write new Config");
        trace!("Saved config to: {}", self.name);
    }
}

#[derive(Debug)]
struct Tags {
    config: Mutex<Config>,
}

impl Tags {
    fn get_possible_tags(&self, guild: Option<GuildId>)
        -> HashMap<String, Tag>
    {
        // TODO: handle this unwrap.
        let config = self.config.lock().expect("Failed to lock Config");
        // TODO: properly handle a missing hashmap case.
        let generic = config.get("generic")
            .cloned()
            .unwrap_or_else(HashMap::new);

        match guild {
            None => generic,
            Some(guild) => merge(
                generic,
                config.get(&guild.to_string())
                    .cloned()
                    .unwrap_or_else(HashMap::new),
            ),
        }
    }

    // TODO: return a list of suggestions.
    fn get_tag(&self, guild: Option<GuildId>, name: String)
        -> Result<Tag, String>
    {
        self.get_possible_tags(guild)
            .get(&name)
            .cloned()
            .ok_or_else(|| "Tag not found".to_owned())
    }

    fn put_tag(&self, guild: Option<GuildId>, name: String, tag: Tag) {
        // Load the actual tag so we can modify it.
        // TODO: handle this better because this can cause an issue
        // if the tag was deleted since being retrieved.
        let mut config = TAGS.config.lock()
            .expect("Failed to lock Config");
        {
            let database = config
                .tags
                .get_mut(&get_database_location(guild))
                .unwrap();
            database.insert(name, tag);
        }
        config.save();
    }

    fn delete_tag(&self, guild: Option<GuildId>, name: &str) {
        // TODO: handle this better because this can cause an issue
        // if the tag was deleted since being retrieved.
        let mut config = TAGS.config.lock()
            .expect("Failed to lock Config");
        {
            let database = config
                .tags
                .get_mut(&get_database_location(guild))
                .unwrap();
            database.remove(name);
        }
        config.save();
    }
}

// TODO: sort tags by server
// TODO: add edit/delete functionality
pub fn handler(context: &Context, message: &Message, args: Vec<String>)
    -> Result<(), String>
{
    let mut args = args.into_iter();

    let f = match args.next().as_ref().map(String::as_ref) {
        Some("create") => create,
        Some("info") => info,
        Some("list") => list,
        Some("edit") => edit,
        Some("delete") => delete,
        Some(name) => return {
            // TODO: ensure we get the GuildId even if it's not in the cache.
            let guild_id = message.guild_id();

            let lookup = name.to_lowercase();
            match TAGS.get_tag(guild_id, lookup.clone()) {
                Ok(tag) => {
                    let mut tag = tag.clone();
                    tag.uses += 1;
                    TAGS.put_tag(guild_id, lookup, tag.clone());
                    check_msg(context.say(&tag.content));

                    Ok(())
                },
                Err(err) => Err(err),
            }
        },
        _ => return Err("Either specify a tag name or use one of the available commands.".to_owned()),
    };

    f(context, message, args.collect())
}

pub fn create(context: &Context, message: &Message, args: Vec<String>)
    -> Result<(), String>
{
    let mut args = args.into_iter();

    let name = match args.next() {
        Some(name) => name,
        None => return Err("Please specify a name for the tag.".to_owned()),
    };

    let content = args.collect::<Vec<String>>();
    let content = if content.is_empty() {
        return Err("Please specify some content for the tag.".to_owned());
    } else {
        content.join(" ")
    };

    // TODO: clean the tag content here.

    let name = name.trim().to_lowercase().to_owned();
    verify_tag_name(&name)?;

    let location = get_database_location(message.guild_id());
    let mut config = TAGS.config.lock().expect("Failed to lock Config");
    let mut database = config.get(&location)
        .cloned()
        .unwrap_or_else(HashMap::new);
    if database.contains_key(&name) {
        return Err("Tag already exists.".to_owned());
    }

    database.insert(name.clone(), Tag::new(
            name.clone(),
            content,
            message.author.id.0,
            None,
            Some(location.clone()),
            None,
    ));
    // TODO: perhaps do this more efficiently by modifying the DB in place?
    // TODO: or maybe use `put_tag`?
    config.insert(location, database);
    check_msg(context.say(&format!("Tag \"{}\" successfully created.", name)));

    Ok(())
}

pub fn info(context: &Context, message: &Message, args: Vec<String>)
    -> Result<(), String>
{
    let mut args = args.into_iter();

    let name = match args.next() {
        Some(name) => name,
        None => return Err("Please specify a name for the tag to get info on.".to_owned()),
    };

    let name = name.trim().to_lowercase().to_owned();
    // TODO: ensure we get the GuildId even if it's not in the cache.
    let guild_id = message.guild_id();
    let tag = TAGS.get_tag(guild_id, name)?;

    check_msg(context.send_message(
        message.channel_id,
        |m| m.embed(|e| tag.as_embed(e)),
    ));

    Ok(())
}

pub fn list(context: &Context, message: &Message, _args: Vec<String>)
    -> Result<(), String>
{
    // TODO: ensure we get the GuildId even if it's not in the cache.
    let guild_id = message.guild_id();
    let mut tags = TAGS.get_possible_tags(guild_id);
    let mut tags = tags
        .drain()
        .map(|(k, _)| k)
        .collect::<Vec<String>>();
    tags.sort();

    let response = if tags.is_empty() {
        "No tags available.".to_owned()
    } else {
        format!("Available tags: {}", tags.join(", "))
    };
    check_msg(context.say(&response));

    Ok(())
}

pub fn edit(context: &Context, message: &Message, args: Vec<String>)
    -> Result<(), String>
{
    let mut args = args.into_iter();

    let name = match args.next() {
        Some(name) => name,
        None => return Err("Please specify a tag to edit.".to_owned()),
    };

    let name = name.trim().to_lowercase().to_owned();

    let guild_id = message.guild_id();
    let mut tag = match TAGS.get_tag(guild_id, name.clone()) {
        Ok(tag) => tag,
        Err(err) => return Err(err),
    };

    if !owner_check(message, &tag) {
        return Err("You do not have permission to do that.".to_owned());
    }

    let content = args.collect::<Vec<String>>();
    let content = if content.is_empty() {
        return Err("Please specify some content for the tag.".to_owned());
    } else {
        content.join(" ")
    };

    // TODO: clean the tag content here.

    tag.content = content;
    TAGS.put_tag(guild_id, name.clone(), tag);

    check_msg(context.say(&format!("Tag \"{}\" successfully updated.", name)));

    Ok(())
}

pub fn delete(context: &Context, message: &Message, args: Vec<String>)
    -> Result<(), String>
{
    let mut args = args.into_iter();

    let name = match args.next() {
        Some(name) => name,
        None => return Err("Please specify a tag to delete.".to_owned()),
    };

    let name = name.trim().to_lowercase().to_owned();

    let guild_id = message.guild_id();
    let tag = match TAGS.get_tag(guild_id, name.clone()) {
        Ok(tag) => tag,
        Err(err) => return Err(err),
    };

    if !owner_check(message, &tag) {
        return Err("You do not have permission to do that.".to_owned());
    }

    TAGS.delete_tag(guild_id, &name);

    check_msg(context.say(&format!("Tag \"{}\" successfully deleted.", name)));

    Ok(())
}

// Denies certain tag names from being used as keys.
fn verify_tag_name(name: &str) -> Result<(), String> {
    if name.contains("@everyone") || name.contains("@here") {
        return Err("Tag contains blocked words".to_owned());
    }

    if name.len() > 100 {
        return Err("Tag name limit is 100 characters".to_owned());
    }

    Ok(())
}

fn owner_check(message: &Message, tag: &Tag) -> bool {
    CONFIG.owners.contains(&message.author.id.0) ||
        message.author.id == tag.owner_id
}

fn get_database_location(guild: Option<GuildId>) -> String {
    // TODO: ensure we get the GuildId even if it's not in the cache.
    guild.map(|g| g.to_string())
        .unwrap_or_else(|| "generic".to_owned())
}
<|endoftext|>


<commit_before>// Std
use std::io::Write;

// Internal
use app::parser::Parser;
use args::{ArgSettings, OptBuilder};
use shell::Shell;

macro_rules! w {
    ($buf:expr, $to_w:expr) => {
        match $buf.write_all($to_w) {
            Ok(..) => (),
            Err(..) => panic!("Failed to write to file completions file"),
        }
    };
}

pub struct ComplGen<'a, 'b>
    where 'a: 'b
{
    p: &'b Parser<'a, 'b>,
}

impl<'a, 'b> ComplGen<'a, 'b> {
    pub fn new(p: &'b Parser<'a, 'b>) -> Self {
        ComplGen { p: p }
    }

    pub fn generate<W: Write>(&self, for_shell: Shell, buf: &mut W) {
        match for_shell {
            Shell::Bash => self.gen_bash(buf),
            Shell::Fish => self.gen_fish(buf),
        }
    }

    fn gen_bash<W: Write>(&self, buf: &mut W) {
        w!(buf,
           format!("_{name}() {{
    local i cur prev opts cmds
    COMPREPLY=()
    cur=\"${{COMP_WORDS[COMP_CWORD]}}\"
    prev=\"${{COMP_WORDS[COMP_CWORD-1]}}\"
    cmd=\"\"
    opts=\"\"

    for i in ${{COMP_WORDS[@]}}
    do
        case \"${{i}}\" in
            {name})
                cmd=\"{name}\"
                ;;
            {subcmds}
            *)
                ;;
        esac
    done

    case \"${{cmd}}\" in
        {name})
            opts=\"{name_opts}\"
            if [[ ${{cur}} == -* || ${{COMP_CWORD}} -eq 1 ]] ; then
                COMPREPLY=( $(compgen -W \"${{opts}}\" -- ${{cur}}) )
                return 0
            fi
            case \"${{prev}}\" in
                {name_opts_details}
                *)
                    COMPREPLY=()
                    ;;
            esac
            COMPREPLY=( $(compgen -W \"${{opts}}\" -- ${{cur}}) )
            return 0
            ;;
        {subcmd_details}
    esac
}}

complete -F _{name} {name}
",
                   name = self.p.meta.bin_name.as_ref().unwrap(),
                   name_opts = self.all_options_for_path(self.p.meta.bin_name.as_ref().unwrap()),
                   name_opts_details =
                       self.option_details_for_path(self.p.meta.bin_name.as_ref().unwrap()),
                   subcmds = self.all_subcommands(),
                   subcmd_details = self.subcommand_details())
               .as_bytes());
    }

    fn all_subcommands(&self) -> String {
        let mut subcmds = String::new();
        let scs = get_all_subcommands(self.p);

        for sc in &scs {
            subcmds = format!("{}
            {name})
                cmd+=\"_{name}\"
                ;;",
                              subcmds,
                              name = sc.replace("-", "_"));
        }

        subcmds
    }

    fn subcommand_details(&self) -> String {
        let mut subcmd_dets = String::new();
        let mut scs = get_all_subcommand_paths(self.p, true);
        scs.sort();
        scs.dedup();

        for sc in &scs {
            subcmd_dets = format!("{}
        {subcmd})
            opts=\"{sc_opts}\"
            if [[ ${{cur}} == -* || ${{COMP_CWORD}} -eq {level} ]] ; then
                COMPREPLY=( $(compgen -W \"${{opts}}\" -- ${{cur}}) )
                return 0
            fi
            case \"${{prev}}\" in
                {opts_details}
                *)
                    COMPREPLY=()
                    ;;
            esac
            COMPREPLY=( $(compgen -W \"${{opts}}\" -- ${{cur}}) )
            return 0
            ;;",
                                  subcmd_dets,
                                  subcmd = sc.replace("-", "_"),
                                  sc_opts = self.all_options_for_path(&*sc),
                                  level = sc.split("_").map(|_| 1).fold(0, |acc, n| acc + n),
                                  opts_details = self.option_details_for_path(&*sc));
        }

        subcmd_dets
    }

    fn all_options_for_path(&self, path: &str) -> String {
        let mut p = self.p;
        for sc in path.split('_').skip(1) {
            debugln!("iter;sc={}", sc);
            p = &p.subcommands
                .iter()
                .find(|s| {
                    s.p.meta.name == sc ||
                    (s.p.meta.aliases.is_some() &&
                     s.p
                        .meta
                        .aliases
                        .as_ref()
                        .unwrap()
                        .iter()
                        .any(|&(n, _)| n == sc))
                })
                .unwrap()
                .p;
        }
        let mut opts = p.short_list.iter().fold(String::new(), |acc, s| format!("{} -{}", acc, s));
        opts = format!("{} {}",
                       opts,
                       p.long_list
                           .iter()
                           .fold(String::new(), |acc, l| format!("{} --{}", acc, l)));
        opts = format!("{} {}",
                       opts,
                       p.positionals
                           .values()
                           .fold(String::new(), |acc, p| format!("{} {}", acc, p)));
        opts = format!("{} {}",
                       opts,
                       p.subcommands
                           .iter()
                           .fold(String::new(), |acc, s| format!("{} {}", acc, s.p.meta.name)));
        for sc in &p.subcommands {
            if let Some(ref aliases) = sc.p.meta.aliases {
                opts = format!("{} {}",
                               opts,
                               aliases.iter()
                                   .map(|&(n, _)| n)
                                   .fold(String::new(), |acc, a| format!("{} {}", acc, a)));
            }
        }
        opts
    }

    fn option_details_for_path(&self, path: &str) -> String {
        let mut p = self.p;
        for sc in path.split('_').skip(1) {
            debugln!("iter;sc={}", sc);
            p = &p.subcommands
                .iter()
                .find(|s| {
                    s.p.meta.name == sc ||
                    (s.p.meta.aliases.is_some() &&
                     s.p
                        .meta
                        .aliases
                        .as_ref()
                        .unwrap()
                        .iter()
                        .any(|&(n, _)| n == sc))
                })
                .unwrap()
                .p;
        }
        let mut opts = String::new();
        for o in &p.opts {
            if let Some(l) = o.long {
                opts = format!("{}
                --{})
                    COMPREPLY=({})
                    return 0
                    ;;",
                               opts,
                               l,
                               vals_for(o));
            }
            if let Some(s) = o.short {
                opts = format!("{}
                    -{})
                    COMPREPLY=({})
                    return 0
                    ;;",
                               opts,
                               s,
                               vals_for(o));
            }
        }
        opts
    }

    fn gen_fish<W: Write>(&self, buf: &mut W) {
        let command = self.p.meta.bin_name.as_ref().unwrap();
        let subcommands: Vec<_> = get_all_subcommands(self.p);
        let has_subcommands = subcommands.len() > 1;

        // function to detect subcommand
        let detect_subcommand_function = if has_subcommands {
            format!(
r#"function __fish_{}_no_subcommand --description "Test if there isn't given a subcommand"
    for i in (commandline -opc)
        if contains -- $i {}
            return 1
        end
    end
    return 0
end
"#, command, subcommands.join(" "))
        } else {
            "".to_string()
        };

        let mut buffer = detect_subcommand_function;
        gen_fish_inner(command, self, vec![], &mut buffer, has_subcommands);
        w!(buf, buffer.as_bytes());
    }
}

pub fn get_all_subcommands(p: &Parser) -> Vec<String> {
    let mut subcmds = vec![];
    if !p.has_subcommands() {
        let mut ret = vec![p.meta.name.clone()];
        if let Some(ref aliases) = p.meta.aliases {
            for &(n, _) in aliases {
                ret.push(n.to_owned());
            }
        }
        return ret;
    }
    for sc in &p.subcommands {
        if let Some(ref aliases) = sc.p.meta.aliases {
            for &(n, _) in aliases {
                subcmds.push(n.to_owned());
            }
        }
        subcmds.push(sc.p.meta.name.clone());
    }
    for sc_v in p.subcommands.iter().map(|s| get_all_subcommands(&s.p)) {
        subcmds.extend(sc_v);
    }
    subcmds.sort();
    subcmds.dedup();
    subcmds
}

pub fn get_all_subcommand_paths(p: &Parser, first: bool) -> Vec<String> {
    let mut subcmds = vec![];
    if !p.has_subcommands() {
        if !first {
            let name = &*p.meta.name;
            let path = p.meta.bin_name.as_ref().unwrap().clone().replace(" ", "_");
            let mut ret = vec![path.clone()];
            if let Some(ref aliases) = p.meta.aliases {
                for &(n, _) in aliases {
                    ret.push(path.replace(name, n));
                }
            }
            return ret;
        }
        return vec![];
    }
    for sc in &p.subcommands {
        let name = &*sc.p.meta.name;
        let path = sc.p.meta.bin_name.as_ref().unwrap().clone().replace(" ", "_");
        subcmds.push(path.clone());
        if let Some(ref aliases) = sc.p.meta.aliases {
            for &(n, _) in aliases {
                subcmds.push(path.replace(name, n));
            }
        }
    }
    for sc_v in p.subcommands.iter().map(|s| get_all_subcommand_paths(&s.p, false)) {
        subcmds.extend(sc_v);
    }
    subcmds
}

fn vals_for(o: &OptBuilder) -> String {
    use args::AnyArg;
    let mut ret = String::new();
    let mut needs_quotes = true;
    if let Some(vals) = o.possible_vals() {
        needs_quotes = false;
        ret = format!("$(compgen -W \"{}\" -- ${{cur}})", vals.join(" "));
    } else if let Some(vec) = o.val_names() {
        let mut it = vec.iter().peekable();
        while let Some((_, val)) = it.next() {
            ret = format!("{}<{}>{}",
                          ret,
                          val,
                          if it.peek().is_some() { " " } else { "" });
        }
        let num = vec.len();
        if o.is_set(ArgSettings::Multiple) && num == 1 {
            ret = format!("{}...", ret);
        }
    } else if let Some(num) = o.num_vals() {
        let mut it = (0..num).peekable();
        while let Some(_) = it.next() {
            ret = format!("{}<{}>{}",
                          ret,
                          o.name(),
                          if it.peek().is_some() { " " } else { "" });
        }
        if o.is_set(ArgSettings::Multiple) && num == 1 {
            ret = format!("{}...", ret);
        }
    } else {
        ret = format!("<{}>", o.name());
        if o.is_set(ArgSettings::Multiple) {
            ret = format!("{}...", ret);
        }
    }
    if needs_quotes {
        ret = format!("\"{}\"", ret);
    }
    ret
}

fn gen_fish_inner(root_command: &str,
                  comp_gen: &ComplGen,
                  parent_cmds: Vec<&String>,
                  buffer: &mut String,
                  has_no_subcommand_fn: bool) {
    // example :
    //
    // complete
    //      -c {command}
    //      -d "{description}"
    //      -s {short}
    //      -l {long}
    //      -a "{possible_arguments}"
    //      -r # if require parameter
    //      -f # don't use file completion
    //      -n "__fish_seen_subcommand_from install" # complete for subcommand "install"

    let command = &comp_gen.p.meta.name;
    let subcommands: Vec<_> = get_all_subcommands(comp_gen.p);

    for option in &comp_gen.p.opts {
        let mut template = format!("complete -c {}", root_command);
        if !parent_cmds.is_empty() {
            template.push_str(format!(" -n '__fish_seen_subcommand_from {}'", command).as_str());
        } else if has_no_subcommand_fn {
            template.push_str(format!(" -n '__fish_{}_no_subcommand'",
                                      comp_gen.p.meta.bin_name.as_ref().unwrap())
                .as_str());
        }
        if let Some(data) = option.short {
            template.push_str(format!(" -s {}", data).as_str());
        }
        if let Some(data) = option.long {
            template.push_str(format!(" -l {}", data).as_str());
        }
        if let Some(data) = option.help {
            template.push_str(format!(" -d '{}'", data).as_str());
        }
        if let Some(ref data) = option.possible_vals {
            template.push_str(format!(" -r -f -a '{}'", data.join(" ")).as_str());
        }
        buffer.push_str(template.as_str());
        buffer.push_str("\n");
    }

    for flag in &comp_gen.p.flags {
        let mut template = format!("complete -c {}", root_command);
        if !parent_cmds.is_empty() {
            template.push_str(format!(" -n '__fish_seen_subcommand_from {}'", command).as_str());
        } else if has_no_subcommand_fn {
            template.push_str(format!(" -n '__fish_{}_no_subcommand'",
                                      comp_gen.p.meta.bin_name.as_ref().unwrap())
                .as_str());
        }
        if let Some(data) = flag.short {
            template.push_str(format!(" -s {}", data).as_str());
        }
        if let Some(data) = flag.long {
            template.push_str(format!(" -l {}", data).as_str());
        }
        if let Some(data) = flag.help {
            template.push_str(format!(" -d '{}'", data).as_str());
        }
        buffer.push_str(template.as_str());
        buffer.push_str("\n");
    }

    if subcommands.len() > 1 {
        for subcommand in subcommands {
            let mut template = format!("complete -c {}", root_command);
            if !parent_cmds.is_empty() {
                template.push_str(format!(" -n '__fish_seen_subcommand_from {}'",
                                          subcommand)
                    .as_str());
            } else if has_no_subcommand_fn {
                template.push_str(format!(" -n '__fish_{}_no_subcommand'",
                                          comp_gen.p.meta.bin_name.as_ref().unwrap())
                    .as_str());
            }
            template.push_str(" -f");
            template.push_str(format!(" -a '{}'", subcommand).as_str());
            buffer.push_str(template.as_str());
            buffer.push_str("\n");
        }
    }

    // generate options of subcommands
    for subcommand in &comp_gen.p.subcommands {
        let sub_comp_gen = ComplGen::new(&subcommand.p);
        let mut sub_parent_cmds = parent_cmds.clone();
        sub_parent_cmds.push(command);
        gen_fish_inner(root_command,
                       &sub_comp_gen,
                       sub_parent_cmds,
                       buffer,
                       has_no_subcommand_fn);
    }
}
<commit_msg>Auto merge of #687 - wdv4758h:issue_685, r=kbknapp<commit_after>// Std
use std::io::Write;

// Internal
use app::parser::Parser;
use args::{ArgSettings, OptBuilder};
use shell::Shell;

macro_rules! w {
    ($buf:expr, $to_w:expr) => {
        match $buf.write_all($to_w) {
            Ok(..) => (),
            Err(..) => panic!("Failed to write to file completions file"),
        }
    };
}

pub struct ComplGen<'a, 'b>
    where 'a: 'b
{
    p: &'b Parser<'a, 'b>,
}

impl<'a, 'b> ComplGen<'a, 'b> {
    pub fn new(p: &'b Parser<'a, 'b>) -> Self {
        ComplGen { p: p }
    }

    pub fn generate<W: Write>(&self, for_shell: Shell, buf: &mut W) {
        match for_shell {
            Shell::Bash => self.gen_bash(buf),
            Shell::Fish => self.gen_fish(buf),
        }
    }

    fn gen_bash<W: Write>(&self, buf: &mut W) {
        w!(buf,
           format!("_{name}() {{
    local i cur prev opts cmds
    COMPREPLY=()
    cur=\"${{COMP_WORDS[COMP_CWORD]}}\"
    prev=\"${{COMP_WORDS[COMP_CWORD-1]}}\"
    cmd=\"\"
    opts=\"\"

    for i in ${{COMP_WORDS[@]}}
    do
        case \"${{i}}\" in
            {name})
                cmd=\"{name}\"
                ;;
            {subcmds}
            *)
                ;;
        esac
    done

    case \"${{cmd}}\" in
        {name})
            opts=\"{name_opts}\"
            if [[ ${{cur}} == -* || ${{COMP_CWORD}} -eq 1 ]] ; then
                COMPREPLY=( $(compgen -W \"${{opts}}\" -- ${{cur}}) )
                return 0
            fi
            case \"${{prev}}\" in
                {name_opts_details}
                *)
                    COMPREPLY=()
                    ;;
            esac
            COMPREPLY=( $(compgen -W \"${{opts}}\" -- ${{cur}}) )
            return 0
            ;;
        {subcmd_details}
    esac
}}

complete -F _{name} {name}
",
                   name = self.p.meta.bin_name.as_ref().unwrap(),
                   name_opts = self.all_options_for_path(self.p.meta.bin_name.as_ref().unwrap()),
                   name_opts_details =
                       self.option_details_for_path(self.p.meta.bin_name.as_ref().unwrap()),
                   subcmds = self.all_subcommands(),
                   subcmd_details = self.subcommand_details())
               .as_bytes());
    }

    fn all_subcommands(&self) -> String {
        let mut subcmds = String::new();
        let scs = get_all_subcommands(self.p);

        for sc in &scs {
            subcmds = format!("{}
            {name})
                cmd+=\"_{name}\"
                ;;",
                              subcmds,
                              name = sc.replace("-", "_"));
        }

        subcmds
    }

    fn subcommand_details(&self) -> String {
        let mut subcmd_dets = String::new();
        let mut scs = get_all_subcommand_paths(self.p, true);
        scs.sort();
        scs.dedup();

        for sc in &scs {
            subcmd_dets = format!("{}
        {subcmd})
            opts=\"{sc_opts}\"
            if [[ ${{cur}} == -* || ${{COMP_CWORD}} -eq {level} ]] ; then
                COMPREPLY=( $(compgen -W \"${{opts}}\" -- ${{cur}}) )
                return 0
            fi
            case \"${{prev}}\" in
                {opts_details}
                *)
                    COMPREPLY=()
                    ;;
            esac
            COMPREPLY=( $(compgen -W \"${{opts}}\" -- ${{cur}}) )
            return 0
            ;;",
                                  subcmd_dets,
                                  subcmd = sc.replace("-", "_"),
                                  sc_opts = self.all_options_for_path(&*sc),
                                  level = sc.split("_").map(|_| 1).fold(0, |acc, n| acc + n),
                                  opts_details = self.option_details_for_path(&*sc));
        }

        subcmd_dets
    }

    fn all_options_for_path(&self, path: &str) -> String {
        let mut p = self.p;
        for sc in path.split('_').skip(1) {
            debugln!("iter;sc={}", sc);
            p = &p.subcommands
                .iter()
                .find(|s| {
                    s.p.meta.name == sc ||
                    (s.p.meta.aliases.is_some() &&
                     s.p
                        .meta
                        .aliases
                        .as_ref()
                        .unwrap()
                        .iter()
                        .any(|&(n, _)| n == sc))
                })
                .unwrap()
                .p;
        }
        let mut opts = p.short_list.iter().fold(String::new(), |acc, s| format!("{} -{}", acc, s));
        opts = format!("{} {}",
                       opts,
                       p.long_list
                           .iter()
                           .fold(String::new(), |acc, l| format!("{} --{}", acc, l)));
        opts = format!("{} {}",
                       opts,
                       p.positionals
                           .values()
                           .fold(String::new(), |acc, p| format!("{} {}", acc, p)));
        opts = format!("{} {}",
                       opts,
                       p.subcommands
                           .iter()
                           .fold(String::new(), |acc, s| format!("{} {}", acc, s.p.meta.name)));
        for sc in &p.subcommands {
            if let Some(ref aliases) = sc.p.meta.aliases {
                opts = format!("{} {}",
                               opts,
                               aliases.iter()
                                   .map(|&(n, _)| n)
                                   .fold(String::new(), |acc, a| format!("{} {}", acc, a)));
            }
        }
        opts
    }

    fn option_details_for_path(&self, path: &str) -> String {
        let mut p = self.p;
        for sc in path.split('_').skip(1) {
            debugln!("iter;sc={}", sc);
            p = &p.subcommands
                .iter()
                .find(|s| {
                    s.p.meta.name == sc ||
                    (s.p.meta.aliases.is_some() &&
                     s.p
                        .meta
                        .aliases
                        .as_ref()
                        .unwrap()
                        .iter()
                        .any(|&(n, _)| n == sc))
                })
                .unwrap()
                .p;
        }
        let mut opts = String::new();
        for o in &p.opts {
            if let Some(l) = o.long {
                opts = format!("{}
                --{})
                    COMPREPLY=({})
                    return 0
                    ;;",
                               opts,
                               l,
                               vals_for(o));
            }
            if let Some(s) = o.short {
                opts = format!("{}
                    -{})
                    COMPREPLY=({})
                    return 0
                    ;;",
                               opts,
                               s,
                               vals_for(o));
            }
        }
        opts
    }

    fn gen_fish<W: Write>(&self, buf: &mut W) {
        let command = self.p.meta.bin_name.as_ref().unwrap();

        // function to detect subcommand
        let detect_subcommand_function =
r#"function __fish_using_command
    set cmd (commandline -opc)
    if [ (count $cmd) -eq (count $argv) ]
        for i in (seq (count $argv))
            if [ $cmd[$i] != $argv[$i] ]
                return 1
            end
        end
        return 0
    end
    return 1
end

"#.to_string();

        let mut buffer = detect_subcommand_function;
        gen_fish_inner(command, self, &command.to_string(), &mut buffer);
        w!(buf, buffer.as_bytes());
    }
}

pub fn get_all_subcommands(p: &Parser) -> Vec<String> {
    let mut subcmds = vec![];
    if !p.has_subcommands() {
        let mut ret = vec![p.meta.name.clone()];
        if let Some(ref aliases) = p.meta.aliases {
            for &(n, _) in aliases {
                ret.push(n.to_owned());
            }
        }
        return ret;
    }
    for sc in &p.subcommands {
        if let Some(ref aliases) = sc.p.meta.aliases {
            for &(n, _) in aliases {
                subcmds.push(n.to_owned());
            }
        }
        subcmds.push(sc.p.meta.name.clone());
    }
    for sc_v in p.subcommands.iter().map(|s| get_all_subcommands(&s.p)) {
        subcmds.extend(sc_v);
    }
    subcmds.sort();
    subcmds.dedup();
    subcmds
}

pub fn get_all_subcommand_paths(p: &Parser, first: bool) -> Vec<String> {
    let mut subcmds = vec![];
    if !p.has_subcommands() {
        if !first {
            let name = &*p.meta.name;
            let path = p.meta.bin_name.as_ref().unwrap().clone().replace(" ", "_");
            let mut ret = vec![path.clone()];
            if let Some(ref aliases) = p.meta.aliases {
                for &(n, _) in aliases {
                    ret.push(path.replace(name, n));
                }
            }
            return ret;
        }
        return vec![];
    }
    for sc in &p.subcommands {
        let name = &*sc.p.meta.name;
        let path = sc.p.meta.bin_name.as_ref().unwrap().clone().replace(" ", "_");
        subcmds.push(path.clone());
        if let Some(ref aliases) = sc.p.meta.aliases {
            for &(n, _) in aliases {
                subcmds.push(path.replace(name, n));
            }
        }
    }
    for sc_v in p.subcommands.iter().map(|s| get_all_subcommand_paths(&s.p, false)) {
        subcmds.extend(sc_v);
    }
    subcmds
}

fn vals_for(o: &OptBuilder) -> String {
    use args::AnyArg;
    let mut ret = String::new();
    let mut needs_quotes = true;
    if let Some(vals) = o.possible_vals() {
        needs_quotes = false;
        ret = format!("$(compgen -W \"{}\" -- ${{cur}})", vals.join(" "));
    } else if let Some(vec) = o.val_names() {
        let mut it = vec.iter().peekable();
        while let Some((_, val)) = it.next() {
            ret = format!("{}<{}>{}",
                          ret,
                          val,
                          if it.peek().is_some() { " " } else { "" });
        }
        let num = vec.len();
        if o.is_set(ArgSettings::Multiple) && num == 1 {
            ret = format!("{}...", ret);
        }
    } else if let Some(num) = o.num_vals() {
        let mut it = (0..num).peekable();
        while let Some(_) = it.next() {
            ret = format!("{}<{}>{}",
                          ret,
                          o.name(),
                          if it.peek().is_some() { " " } else { "" });
        }
        if o.is_set(ArgSettings::Multiple) && num == 1 {
            ret = format!("{}...", ret);
        }
    } else {
        ret = format!("<{}>", o.name());
        if o.is_set(ArgSettings::Multiple) {
            ret = format!("{}...", ret);
        }
    }
    if needs_quotes {
        ret = format!("\"{}\"", ret);
    }
    ret
}

fn gen_fish_inner(root_command: &str,
                  comp_gen: &ComplGen,
                  parent_cmds: &String,
                  buffer: &mut String) {

    // example :
    //
    // complete
    //      -c {command}
    //      -d "{description}"
    //      -s {short}
    //      -l {long}
    //      -a "{possible_arguments}"
    //      -r # if require parameter
    //      -f # don't use file completion
    //      -n "__fish_using_command myprog subcmd1" # complete for command "myprog subcmd1"

    let basic_template = format!("complete -c {} -n '__fish_using_command {}'",
                                 root_command,
                                 parent_cmds);

    for option in &comp_gen.p.opts {
        let mut template = basic_template.clone();
        if let Some(data) = option.short {
            template.push_str(format!(" -s {}", data).as_str());
        }
        if let Some(data) = option.long {
            template.push_str(format!(" -l {}", data).as_str());
        }
        if let Some(data) = option.help {
            template.push_str(format!(" -d '{}'", data).as_str());
        }
        if let Some(ref data) = option.possible_vals {
            template.push_str(format!(" -r -f -a '{}'", data.join(" ")).as_str());
        }
        buffer.push_str(template.as_str());
        buffer.push_str("\n");
    }

    for flag in &comp_gen.p.flags {
        let mut template = basic_template.clone();
        if let Some(data) = flag.short {
            template.push_str(format!(" -s {}", data).as_str());
        }
        if let Some(data) = flag.long {
            template.push_str(format!(" -l {}", data).as_str());
        }
        if let Some(data) = flag.help {
            template.push_str(format!(" -d '{}'", data).as_str());
        }
        buffer.push_str(template.as_str());
        buffer.push_str("\n");
    }

    for subcommand in &comp_gen.p.subcommands {
        let mut template = basic_template.clone();
        template.push_str(" -f");
        template.push_str(format!(" -a '{}'", &subcommand.p.meta.name).as_str());
        buffer.push_str(template.as_str());
        buffer.push_str("\n");
    }

    // generate options of subcommands
    for subcommand in &comp_gen.p.subcommands {
        let sub_comp_gen = ComplGen::new(&subcommand.p);
        // make new "parent_cmds" for different subcommands
        let mut sub_parent_cmds = parent_cmds.clone();
        if !sub_parent_cmds.is_empty() {
            sub_parent_cmds.push_str(" ");
        }
        sub_parent_cmds.push_str(&subcommand.p.meta.name);
        gen_fish_inner(root_command,
                       &sub_comp_gen,
                       &mut sub_parent_cmds,
                       buffer);
    }
}
<|endoftext|>


<commit_before>use document;
use document::{YamlNode, YamlNodeData};
use ffi::{YamlErrorType, YamlScalarStyle};
use error::{YamlMark, YamlError, YamlErrorContext};

use std::f64;
use std::char;
use regex::Regex;

pub trait YamlConstructor<T, E> {
    fn construct_scalar(&self, scalar: document::YamlScalarData) -> Result<T, E>;
    fn construct_sequence(&self, sequence: document::YamlSequenceData) -> Result<T, E>;
    fn construct_mapping(&self, mapping: document::YamlMappingData) -> Result<T, E>;

    fn construct<'r>(&self, node: document::YamlNode<'r>) -> Result<T, E> {
        match node {
            YamlNode::YamlScalarNode(scalar) => self.construct_scalar(scalar),
            YamlNode::YamlSequenceNode(sequence) => self.construct_sequence(sequence),
            YamlNode::YamlMappingNode(mapping) => self.construct_mapping(mapping)
        }
    }
}

#[derive(PartialEq, Debug)]
pub enum YamlStandardData {
    YamlInteger(isize),
    YamlFloat(f64),
    YamlString(String),
    YamlNull,
    YamlBool(bool),
    YamlSequence(Vec<YamlStandardData>),
    YamlMapping(Vec<(YamlStandardData, YamlStandardData)>),
}

#[derive(Clone, Copy)]
pub struct YamlStandardConstructor;

fn standard_error(message: String, mark: &YamlMark) -> YamlError {
    let context = YamlErrorContext {
        byte_offset: mark.index,
        problem_mark: *mark,
        context: None,
        context_mark: *mark,
    };

    YamlError {
        kind: YamlErrorType::YAML_PARSER_ERROR,
        problem: Some(message),
        io_error: None,
        context: Some(context)
    }
}

fn take(iter: &mut Iterator<Item=char>, n: usize) -> String
{
    let mut s = String::new();
    for _ in 0..n {
        match iter.next() {
            Some(c) => s.push(c),
            None => break
        }
    }
    return s;
}

impl YamlStandardConstructor {
    pub fn new() -> YamlStandardConstructor {
        YamlStandardConstructor
    }

    fn parse_double_quoted(value: &str, mark: &YamlMark) -> Result<String, YamlError> {
        let mut buf = String::new();
        let mut it = value.chars();

        loop {
            match it.next() {
                None => return Ok(buf),
                Some('\\') => {
                    // escape sequences
                    match it.next() {
                        None => return Err(standard_error(
                                    "unexpected end of string after escape".to_string(),
                                    mark
                                )),
                        Some('0') => buf.push('\x00'),              // null
                        Some('a') => buf.push('\x07'),              // ASCII bell
                        Some('b') => buf.push('\x08'),              // backspace
                        Some('t') | Some('\t') => buf.push('\t'),   // horizontal tab
                        Some('n') => buf.push('\n'),                // linefeed
                        Some('v') => buf.push('\x0b'),              // vertical tab
                        Some('f') => buf.push('\x0c'),              // form feed
                        Some('r') => buf.push('\x0d'),              // carriage return
                        Some('e') => buf.push('\x1b'),              // ASCII escape
                        Some('N') => buf.push('\u{85}'),            // unicode next line
                        Some('_') => buf.push('\u{a0}'),            // unicode non-breaking space
                        Some('L') => buf.push('\u{2028}'),          // unicode line separator
                        Some('P') => buf.push('\u{2029}'),          // unicode paragraph separator
                        Some('x') => {
                            let code:String = take(&mut it, 2);
                            match parse_escape_sequence(&code[..], 2) {
                                Some(c) => buf.push(c),
                                None => return Err(standard_error(
                                            format!("invalid escape sequence {}", code),
                                            mark
                                        ))
                            }
                        },
                        Some('u') => {
                            let code:String = take(&mut it, 4);
                            match parse_escape_sequence(&code[..], 4) {
                                Some(c) => buf.push(c),
                                None => return Err(standard_error(
                                            format!("invalid escape sequence {}", code),
                                            mark
                                        ))
                            }
                        },
                        Some('U') => {
                            let code:String = take(&mut it, 8);
                            match parse_escape_sequence(&code[..], 8) {
                                Some(c) => buf.push(c),
                                None => return Err(standard_error(
                                            format!("invalid escape sequence {}", code),
                                            mark
                                        ))
                            }
                        },
                        Some(c) => buf.push(c)
                    }
                },
                Some(c) => buf.push(c)
            }
        }
    }
}

fn parse_escape_sequence(rep: &str, expected_len: usize) -> Option<char> {
    match u32::from_str_radix(rep, 16) {
        Ok(code) if rep.len() == expected_len => char::from_u32(code),
        _ => None
    }
}

fn parse_int(sign: &str, data: &str, radix: u32) -> isize {
    let sign_flag = if sign == "-" {
            -1
        } else {
            1
        };

    let filtered:String = data.chars().filter(|&c| c != '_').collect();
    let unsigned:isize = isize::from_str_radix(&filtered[..], radix).unwrap();
    return unsigned * sign_flag;
}

fn parse_float(sign: &str, data: &str) -> f64 {
    let unsigned:f64 = data.parse().unwrap();
    if sign == "-" {
        return -unsigned;
    } else {
        return unsigned;
    }
}

static DEC_INT:Regex = regex!(r"^[-+]?(0|[1-9][0-9_]*)$");
static OCT_INT:Regex = regex!(r"^([-+]?)0o?([0-7_]+)$");
static HEX_INT:Regex = regex!(r"^([-+]?)0x([0-9a-fA-F_]+)$");
static BIN_INT:Regex = regex!(r"^([-+]?)0b([0-1_]+)$");
static FLOAT_PATTERN:Regex = regex!(r"^([-+]?)(\.[0-9]+|[0-9]+(\.[0-9]*)?([eE][-+]?[0-9]+)?)$");
static POS_INF:Regex = regex!(r"^[+]?(\.inf|\.Inf|\.INF)$");
static NEG_INF:Regex = regex!(r"^-(\.inf|\.Inf|\.INF)$");
static NAN_PATTERN:Regex = regex!(r"^(\.nan|\.NaN|\.NAN)$");
static NULL_PATTERN:Regex = regex!(r"^(null|Null|NULL|~)$");
static TRUE_PATTERN:Regex = regex!(r"^(true|True|TRUE|yes|Yes|YES)$");
static FALSE_PATTERN:Regex = regex!(r"^(false|False|FALSE|no|No|NO)$");

impl YamlConstructor<YamlStandardData, YamlError> for YamlStandardConstructor {
    fn construct_scalar(&self, scalar: document::YamlScalarData) -> Result<YamlStandardData, YamlError> {
        let value = scalar.get_value();
        let mark = scalar.start_mark();

        match scalar.style() {
            YamlScalarStyle::YamlPlainScalarStyle => {
                match BIN_INT.captures(&value[..]) {
                    Some(caps) => return Ok(YamlStandardData::YamlInteger(parse_int(
                                caps.at(1).unwrap(), caps.at(2).unwrap(), 2))),
                    None => ()
                };
                match OCT_INT.captures(&value[..]) {
                    Some(caps) => return Ok(YamlStandardData::YamlInteger(parse_int(
                                caps.at(1).unwrap(), caps.at(2).unwrap(), 8))),
                    None => ()
                };
                match HEX_INT.captures(&value[..]) {
                    Some(caps) => return Ok(YamlStandardData::YamlInteger(parse_int(
                                caps.at(1).unwrap(), caps.at(2).unwrap(), 16))),
                    None => ()
                };

                if DEC_INT.is_match(&value[..]) {
                    return Ok(YamlStandardData::YamlInteger(parse_int("", &value[..], 10)));
                }

                match FLOAT_PATTERN.captures(&value[..]) {
                    Some(caps) => return Ok(YamlStandardData::YamlFloat(parse_float(
                                caps.at(1).unwrap(), caps.at(2).unwrap()))),
                    None => ()
                };

                if POS_INF.is_match(&value[..]) {
                    Ok(YamlStandardData::YamlFloat(f64::INFINITY))
                } else if NEG_INF.is_match(&value[..]) {
                    Ok(YamlStandardData::YamlFloat(f64::NEG_INFINITY))
                } else if NAN_PATTERN.is_match(&value[..]) {
                    Ok(YamlStandardData::YamlFloat(f64::NAN))
                } else if NULL_PATTERN.is_match(&value[..]) {
                    Ok(YamlStandardData::YamlNull)
                } else if TRUE_PATTERN.is_match(&value[..]) {
                    Ok(YamlStandardData::YamlBool(true))
                } else if FALSE_PATTERN.is_match(&value[..]) {
                    Ok(YamlStandardData::YamlBool(false))
                } else {
                    Ok(YamlStandardData::YamlString(value))
                }
            },
            YamlScalarStyle::YamlDoubleQuotedScalarStyle => {
                YamlStandardConstructor::parse_double_quoted(&value[..], &mark).map(YamlStandardData::YamlString)
            },
            _ => {
                Ok(YamlStandardData::YamlString(value))
            }
        }
    }

    fn construct_sequence(&self, sequence: document::YamlSequenceData) -> Result<YamlStandardData, YamlError> {
        let res:Result<Vec<YamlStandardData>, YamlError> = sequence.values().map(|node| { self.construct(node) }).collect();
        res.map(|list| YamlStandardData::YamlSequence(list))
    }

    fn construct_mapping(&self, mapping: document::YamlMappingData) -> Result<YamlStandardData, YamlError> {
        let pairs = mapping.pairs().map(|(key_node, value_node)| {
            match self.construct(key_node) {
                Ok(key) => match self.construct(value_node) {
                    Ok(value) => Ok((key, value)),
                    Err(e) => Err(e)
                },
                Err(e) => Err(e)
            }
        });
        let res:Result<Vec<(YamlStandardData, YamlStandardData)>, YamlError> = pairs.collect();
        res.map(YamlStandardData::YamlMapping)
    }
}

#[cfg(test)]
mod test {
    use super::YamlStandardData::*;
    use parser::{YamlParser, YamlByteParser};
    use std::f64;
    use ffi::YamlEncoding::YamlUtf8Encoding;
    use constructor::{YamlConstructor, YamlStandardConstructor};

    #[test]
    fn test_standard_constructor() {
        let data = "[1, 2, 3]";
        let parser = YamlByteParser::init(data.as_bytes(), YamlUtf8Encoding);

        match parser.load().next() {
            Some(Ok(doc)) => {
                let ctor = YamlStandardConstructor::new();
                assert_eq!(Ok(YamlSequence(vec![YamlInteger(1), YamlInteger(2), YamlInteger(3)])), ctor.construct(doc.root().unwrap()))
            },
            _ => panic!("unexpected result")
        }
    }

    #[test]
    fn test_integer_parser() {
        let data = "[0o10, 0x21, -30]";
        let parser = YamlByteParser::init(data.as_bytes(), YamlUtf8Encoding);

        match parser.load().next() {
            Some(Ok(doc)) => {
                let ctor = YamlStandardConstructor::new();
                assert_eq!(Ok(YamlSequence(vec![YamlInteger(0o10), YamlInteger(0x21), YamlInteger(-30)])), ctor.construct(doc.root().unwrap()))
            },
            _ => panic!("unexpected result")
        }
    }

    #[test]
    fn test_float_parser() {
        let data = "[0.3, -.4, 1e+2, -1.2e-3]";
        let parser = YamlByteParser::init(data.as_bytes(), YamlUtf8Encoding);

        match parser.load().next() {
            Some(Ok(doc)) => {
                let ctor = YamlStandardConstructor::new();
                let value = ctor.construct(doc.root().unwrap());
                match value {
                    Ok(YamlSequence(seq)) => {
                        match &seq[..] {
                            [YamlFloat(f1), YamlFloat(f2), YamlFloat(f3), YamlFloat(f4)] => {
                                assert!((f1 - 0.3).abs() < 1.0e-6);
                                assert!((f2 + 0.4) < 1.0e-6);
                                assert!((f3 - 1e+2) < 1.0e-6);
                                assert!((f4 + 1.2e-3) < 1.0e-6);
                            },
                            _ => panic!("unexpected sequence")
                        }
                    },
                    _ => panic!("unexpected result")
                }
            },
            _ => panic!("document parse failure")
        }
    }

    #[test]
    fn test_inf_parser() {
        let data = "[.inf, -.INF]";
        let parser = YamlByteParser::init(data.as_bytes(), YamlUtf8Encoding);

        match parser.load().next() {
            Some(Ok(doc)) => {
                let ctor = YamlStandardConstructor::new();
                assert_eq!(Ok(YamlSequence(vec![YamlFloat(f64::INFINITY), YamlFloat(f64::NEG_INFINITY)])), ctor.construct(doc.root().unwrap()))
            },
            _ => panic!("document parse failure")
        }
    }

    #[test]
    fn test_misc_parser() {
        let data = "[yes, False, ~]";
        let parser = YamlByteParser::init(data.as_bytes(), YamlUtf8Encoding);

        match parser.load().next() {
            Some(Ok(doc)) => {
                let ctor = YamlStandardConstructor::new();
                assert_eq!(Ok(YamlSequence(vec![YamlBool(true), YamlBool(false), YamlNull])), ctor.construct(doc.root().unwrap()))
            },
            _ => panic!("document parse failure")
        }
    }

    #[test]
    fn test_double_quoted_parser() {
        let data = r#""hello, \"world\"""#;
        let parser = YamlByteParser::init(data.as_bytes(), YamlUtf8Encoding);

        match parser.load().next() {
            Some(Ok(doc)) => {
                let ctor = YamlStandardConstructor::new();
                assert_eq!(Ok(YamlString("hello, \"world\"".to_string())), ctor.construct(doc.root().unwrap()))
            },
            _ => panic!("document parse failure")
        }
    }

    #[test]
    fn test_single_quoted_parser() {
        let data = r#"'here''s to "quotes"'"#;
        let parser = YamlByteParser::init(data.as_bytes(), YamlUtf8Encoding);

        match parser.load().next() {
            Some(Ok(doc)) => {
                let ctor = YamlStandardConstructor::new();
                assert_eq!(Ok(YamlString(r#"here's to "quotes""#.to_string())), ctor.construct(doc.root().unwrap()))
            },
            _ => panic!("document parse failure")
        }
    }

    #[test]
    fn test_underlined_integer() {
        let data = "[1_000, -2_000_000]";
        let parser = YamlByteParser::init(data.as_bytes(), YamlUtf8Encoding);

        match parser.load().next() {
            Some(Ok(doc)) => {
                let ctor = YamlStandardConstructor::new();
                assert_eq!(Ok(YamlSequence(vec![YamlInteger(1000), YamlInteger(-2000000)])), ctor.construct(doc.root().unwrap()))
            },
            _ => panic!("document parse failure")
        }
    }

    #[test]
    fn test_negative_radix() {
        let data = "[-0x30, -0700, -0b110]";
        let parser = YamlByteParser::init(data.as_bytes(), YamlUtf8Encoding);

        match parser.load().next() {
            Some(Ok(doc)) => {
                let ctor = YamlStandardConstructor::new();
                assert_eq!(Ok(YamlSequence(vec![YamlInteger(-48), YamlInteger(-448), YamlInteger(-6)])), ctor.construct(doc.root().unwrap()))
            },
            _ => panic!("document parse failure")
        }
    }
}
<commit_msg>Make YamlStandardData clonable<commit_after>use document;
use document::{YamlNode, YamlNodeData};
use ffi::{YamlErrorType, YamlScalarStyle};
use error::{YamlMark, YamlError, YamlErrorContext};

use std::f64;
use std::char;
use regex::Regex;

pub trait YamlConstructor<T, E> {
    fn construct_scalar(&self, scalar: document::YamlScalarData) -> Result<T, E>;
    fn construct_sequence(&self, sequence: document::YamlSequenceData) -> Result<T, E>;
    fn construct_mapping(&self, mapping: document::YamlMappingData) -> Result<T, E>;

    fn construct<'r>(&self, node: document::YamlNode<'r>) -> Result<T, E> {
        match node {
            YamlNode::YamlScalarNode(scalar) => self.construct_scalar(scalar),
            YamlNode::YamlSequenceNode(sequence) => self.construct_sequence(sequence),
            YamlNode::YamlMappingNode(mapping) => self.construct_mapping(mapping)
        }
    }
}

#[derive(PartialEq, Clone, Debug)]
pub enum YamlStandardData {
    YamlInteger(isize),
    YamlFloat(f64),
    YamlString(String),
    YamlNull,
    YamlBool(bool),
    YamlSequence(Vec<YamlStandardData>),
    YamlMapping(Vec<(YamlStandardData, YamlStandardData)>),
}

#[derive(Clone, Copy)]
pub struct YamlStandardConstructor;

fn standard_error(message: String, mark: &YamlMark) -> YamlError {
    let context = YamlErrorContext {
        byte_offset: mark.index,
        problem_mark: *mark,
        context: None,
        context_mark: *mark,
    };

    YamlError {
        kind: YamlErrorType::YAML_PARSER_ERROR,
        problem: Some(message),
        io_error: None,
        context: Some(context)
    }
}

fn take(iter: &mut Iterator<Item=char>, n: usize) -> String
{
    let mut s = String::new();
    for _ in 0..n {
        match iter.next() {
            Some(c) => s.push(c),
            None => break
        }
    }
    return s;
}

impl YamlStandardConstructor {
    pub fn new() -> YamlStandardConstructor {
        YamlStandardConstructor
    }

    fn parse_double_quoted(value: &str, mark: &YamlMark) -> Result<String, YamlError> {
        let mut buf = String::new();
        let mut it = value.chars();

        loop {
            match it.next() {
                None => return Ok(buf),
                Some('\\') => {
                    // escape sequences
                    match it.next() {
                        None => return Err(standard_error(
                                    "unexpected end of string after escape".to_string(),
                                    mark
                                )),
                        Some('0') => buf.push('\x00'),              // null
                        Some('a') => buf.push('\x07'),              // ASCII bell
                        Some('b') => buf.push('\x08'),              // backspace
                        Some('t') | Some('\t') => buf.push('\t'),   // horizontal tab
                        Some('n') => buf.push('\n'),                // linefeed
                        Some('v') => buf.push('\x0b'),              // vertical tab
                        Some('f') => buf.push('\x0c'),              // form feed
                        Some('r') => buf.push('\x0d'),              // carriage return
                        Some('e') => buf.push('\x1b'),              // ASCII escape
                        Some('N') => buf.push('\u{85}'),            // unicode next line
                        Some('_') => buf.push('\u{a0}'),            // unicode non-breaking space
                        Some('L') => buf.push('\u{2028}'),          // unicode line separator
                        Some('P') => buf.push('\u{2029}'),          // unicode paragraph separator
                        Some('x') => {
                            let code:String = take(&mut it, 2);
                            match parse_escape_sequence(&code[..], 2) {
                                Some(c) => buf.push(c),
                                None => return Err(standard_error(
                                            format!("invalid escape sequence {}", code),
                                            mark
                                        ))
                            }
                        },
                        Some('u') => {
                            let code:String = take(&mut it, 4);
                            match parse_escape_sequence(&code[..], 4) {
                                Some(c) => buf.push(c),
                                None => return Err(standard_error(
                                            format!("invalid escape sequence {}", code),
                                            mark
                                        ))
                            }
                        },
                        Some('U') => {
                            let code:String = take(&mut it, 8);
                            match parse_escape_sequence(&code[..], 8) {
                                Some(c) => buf.push(c),
                                None => return Err(standard_error(
                                            format!("invalid escape sequence {}", code),
                                            mark
                                        ))
                            }
                        },
                        Some(c) => buf.push(c)
                    }
                },
                Some(c) => buf.push(c)
            }
        }
    }
}

fn parse_escape_sequence(rep: &str, expected_len: usize) -> Option<char> {
    match u32::from_str_radix(rep, 16) {
        Ok(code) if rep.len() == expected_len => char::from_u32(code),
        _ => None
    }
}

fn parse_int(sign: &str, data: &str, radix: u32) -> isize {
    let sign_flag = if sign == "-" {
            -1
        } else {
            1
        };

    let filtered:String = data.chars().filter(|&c| c != '_').collect();
    let unsigned:isize = isize::from_str_radix(&filtered[..], radix).unwrap();
    return unsigned * sign_flag;
}

fn parse_float(sign: &str, data: &str) -> f64 {
    let unsigned:f64 = data.parse().unwrap();
    if sign == "-" {
        return -unsigned;
    } else {
        return unsigned;
    }
}

static DEC_INT:Regex = regex!(r"^[-+]?(0|[1-9][0-9_]*)$");
static OCT_INT:Regex = regex!(r"^([-+]?)0o?([0-7_]+)$");
static HEX_INT:Regex = regex!(r"^([-+]?)0x([0-9a-fA-F_]+)$");
static BIN_INT:Regex = regex!(r"^([-+]?)0b([0-1_]+)$");
static FLOAT_PATTERN:Regex = regex!(r"^([-+]?)(\.[0-9]+|[0-9]+(\.[0-9]*)?([eE][-+]?[0-9]+)?)$");
static POS_INF:Regex = regex!(r"^[+]?(\.inf|\.Inf|\.INF)$");
static NEG_INF:Regex = regex!(r"^-(\.inf|\.Inf|\.INF)$");
static NAN_PATTERN:Regex = regex!(r"^(\.nan|\.NaN|\.NAN)$");
static NULL_PATTERN:Regex = regex!(r"^(null|Null|NULL|~)$");
static TRUE_PATTERN:Regex = regex!(r"^(true|True|TRUE|yes|Yes|YES)$");
static FALSE_PATTERN:Regex = regex!(r"^(false|False|FALSE|no|No|NO)$");

impl YamlConstructor<YamlStandardData, YamlError> for YamlStandardConstructor {
    fn construct_scalar(&self, scalar: document::YamlScalarData) -> Result<YamlStandardData, YamlError> {
        let value = scalar.get_value();
        let mark = scalar.start_mark();

        match scalar.style() {
            YamlScalarStyle::YamlPlainScalarStyle => {
                match BIN_INT.captures(&value[..]) {
                    Some(caps) => return Ok(YamlStandardData::YamlInteger(parse_int(
                                caps.at(1).unwrap(), caps.at(2).unwrap(), 2))),
                    None => ()
                };
                match OCT_INT.captures(&value[..]) {
                    Some(caps) => return Ok(YamlStandardData::YamlInteger(parse_int(
                                caps.at(1).unwrap(), caps.at(2).unwrap(), 8))),
                    None => ()
                };
                match HEX_INT.captures(&value[..]) {
                    Some(caps) => return Ok(YamlStandardData::YamlInteger(parse_int(
                                caps.at(1).unwrap(), caps.at(2).unwrap(), 16))),
                    None => ()
                };

                if DEC_INT.is_match(&value[..]) {
                    return Ok(YamlStandardData::YamlInteger(parse_int("", &value[..], 10)));
                }

                match FLOAT_PATTERN.captures(&value[..]) {
                    Some(caps) => return Ok(YamlStandardData::YamlFloat(parse_float(
                                caps.at(1).unwrap(), caps.at(2).unwrap()))),
                    None => ()
                };

                if POS_INF.is_match(&value[..]) {
                    Ok(YamlStandardData::YamlFloat(f64::INFINITY))
                } else if NEG_INF.is_match(&value[..]) {
                    Ok(YamlStandardData::YamlFloat(f64::NEG_INFINITY))
                } else if NAN_PATTERN.is_match(&value[..]) {
                    Ok(YamlStandardData::YamlFloat(f64::NAN))
                } else if NULL_PATTERN.is_match(&value[..]) {
                    Ok(YamlStandardData::YamlNull)
                } else if TRUE_PATTERN.is_match(&value[..]) {
                    Ok(YamlStandardData::YamlBool(true))
                } else if FALSE_PATTERN.is_match(&value[..]) {
                    Ok(YamlStandardData::YamlBool(false))
                } else {
                    Ok(YamlStandardData::YamlString(value))
                }
            },
            YamlScalarStyle::YamlDoubleQuotedScalarStyle => {
                YamlStandardConstructor::parse_double_quoted(&value[..], &mark).map(YamlStandardData::YamlString)
            },
            _ => {
                Ok(YamlStandardData::YamlString(value))
            }
        }
    }

    fn construct_sequence(&self, sequence: document::YamlSequenceData) -> Result<YamlStandardData, YamlError> {
        let res:Result<Vec<YamlStandardData>, YamlError> = sequence.values().map(|node| { self.construct(node) }).collect();
        res.map(|list| YamlStandardData::YamlSequence(list))
    }

    fn construct_mapping(&self, mapping: document::YamlMappingData) -> Result<YamlStandardData, YamlError> {
        let pairs = mapping.pairs().map(|(key_node, value_node)| {
            match self.construct(key_node) {
                Ok(key) => match self.construct(value_node) {
                    Ok(value) => Ok((key, value)),
                    Err(e) => Err(e)
                },
                Err(e) => Err(e)
            }
        });
        let res:Result<Vec<(YamlStandardData, YamlStandardData)>, YamlError> = pairs.collect();
        res.map(YamlStandardData::YamlMapping)
    }
}

#[cfg(test)]
mod test {
    use super::YamlStandardData::*;
    use parser::{YamlParser, YamlByteParser};
    use std::f64;
    use ffi::YamlEncoding::YamlUtf8Encoding;
    use constructor::{YamlConstructor, YamlStandardConstructor};

    #[test]
    fn test_standard_constructor() {
        let data = "[1, 2, 3]";
        let parser = YamlByteParser::init(data.as_bytes(), YamlUtf8Encoding);

        match parser.load().next() {
            Some(Ok(doc)) => {
                let ctor = YamlStandardConstructor::new();
                assert_eq!(Ok(YamlSequence(vec![YamlInteger(1), YamlInteger(2), YamlInteger(3)])), ctor.construct(doc.root().unwrap()))
            },
            _ => panic!("unexpected result")
        }
    }

    #[test]
    fn test_integer_parser() {
        let data = "[0o10, 0x21, -30]";
        let parser = YamlByteParser::init(data.as_bytes(), YamlUtf8Encoding);

        match parser.load().next() {
            Some(Ok(doc)) => {
                let ctor = YamlStandardConstructor::new();
                assert_eq!(Ok(YamlSequence(vec![YamlInteger(0o10), YamlInteger(0x21), YamlInteger(-30)])), ctor.construct(doc.root().unwrap()))
            },
            _ => panic!("unexpected result")
        }
    }

    #[test]
    fn test_float_parser() {
        let data = "[0.3, -.4, 1e+2, -1.2e-3]";
        let parser = YamlByteParser::init(data.as_bytes(), YamlUtf8Encoding);

        match parser.load().next() {
            Some(Ok(doc)) => {
                let ctor = YamlStandardConstructor::new();
                let value = ctor.construct(doc.root().unwrap());
                match value {
                    Ok(YamlSequence(seq)) => {
                        match &seq[..] {
                            [YamlFloat(f1), YamlFloat(f2), YamlFloat(f3), YamlFloat(f4)] => {
                                assert!((f1 - 0.3).abs() < 1.0e-6);
                                assert!((f2 + 0.4) < 1.0e-6);
                                assert!((f3 - 1e+2) < 1.0e-6);
                                assert!((f4 + 1.2e-3) < 1.0e-6);
                            },
                            _ => panic!("unexpected sequence")
                        }
                    },
                    _ => panic!("unexpected result")
                }
            },
            _ => panic!("document parse failure")
        }
    }

    #[test]
    fn test_inf_parser() {
        let data = "[.inf, -.INF]";
        let parser = YamlByteParser::init(data.as_bytes(), YamlUtf8Encoding);

        match parser.load().next() {
            Some(Ok(doc)) => {
                let ctor = YamlStandardConstructor::new();
                assert_eq!(Ok(YamlSequence(vec![YamlFloat(f64::INFINITY), YamlFloat(f64::NEG_INFINITY)])), ctor.construct(doc.root().unwrap()))
            },
            _ => panic!("document parse failure")
        }
    }

    #[test]
    fn test_misc_parser() {
        let data = "[yes, False, ~]";
        let parser = YamlByteParser::init(data.as_bytes(), YamlUtf8Encoding);

        match parser.load().next() {
            Some(Ok(doc)) => {
                let ctor = YamlStandardConstructor::new();
                assert_eq!(Ok(YamlSequence(vec![YamlBool(true), YamlBool(false), YamlNull])), ctor.construct(doc.root().unwrap()))
            },
            _ => panic!("document parse failure")
        }
    }

    #[test]
    fn test_double_quoted_parser() {
        let data = r#""hello, \"world\"""#;
        let parser = YamlByteParser::init(data.as_bytes(), YamlUtf8Encoding);

        match parser.load().next() {
            Some(Ok(doc)) => {
                let ctor = YamlStandardConstructor::new();
                assert_eq!(Ok(YamlString("hello, \"world\"".to_string())), ctor.construct(doc.root().unwrap()))
            },
            _ => panic!("document parse failure")
        }
    }

    #[test]
    fn test_single_quoted_parser() {
        let data = r#"'here''s to "quotes"'"#;
        let parser = YamlByteParser::init(data.as_bytes(), YamlUtf8Encoding);

        match parser.load().next() {
            Some(Ok(doc)) => {
                let ctor = YamlStandardConstructor::new();
                assert_eq!(Ok(YamlString(r#"here's to "quotes""#.to_string())), ctor.construct(doc.root().unwrap()))
            },
            _ => panic!("document parse failure")
        }
    }

    #[test]
    fn test_underlined_integer() {
        let data = "[1_000, -2_000_000]";
        let parser = YamlByteParser::init(data.as_bytes(), YamlUtf8Encoding);

        match parser.load().next() {
            Some(Ok(doc)) => {
                let ctor = YamlStandardConstructor::new();
                assert_eq!(Ok(YamlSequence(vec![YamlInteger(1000), YamlInteger(-2000000)])), ctor.construct(doc.root().unwrap()))
            },
            _ => panic!("document parse failure")
        }
    }

    #[test]
    fn test_negative_radix() {
        let data = "[-0x30, -0700, -0b110]";
        let parser = YamlByteParser::init(data.as_bytes(), YamlUtf8Encoding);

        match parser.load().next() {
            Some(Ok(doc)) => {
                let ctor = YamlStandardConstructor::new();
                assert_eq!(Ok(YamlSequence(vec![YamlInteger(-48), YamlInteger(-448), YamlInteger(-6)])), ctor.construct(doc.root().unwrap()))
            },
            _ => panic!("document parse failure")
        }
    }
}
<|endoftext|>


<commit_before>use crate::atomicmin::AtomicMin;
use crate::error::PngError;
use crate::Deadline;
use crate::PngResult;
use indexmap::IndexSet;

#[cfg(feature = "zopfli")]
use std::num::NonZeroU8;

#[doc(hidden)]
pub mod miniz_stream;

#[cfg(feature = "libdeflater")]
mod deflater;
#[cfg(feature = "libdeflater")]
pub use deflater::deflate as libdeflater_deflate;

#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
pub mod cfzlib;

#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
pub mod cfzlib {
    pub fn is_supported() -> bool {
        return false;
    }
}

/// Decompress a data stream using the DEFLATE algorithm
pub fn inflate(data: &[u8]) -> PngResult<Vec<u8>> {
    miniz_oxide::inflate::decompress_to_vec_zlib(data)
        .map_err(|e| PngError::new(&format!("Error on decompress: {:?}", e)))
}

/// Compress a data stream using the DEFLATE algorithm
#[doc(hidden)]
pub fn deflate(
    data: &[u8],
    zc: u8,
    zs: u8,
    zw: u8,
    max_size: &AtomicMin,
    deadline: &Deadline,
) -> PngResult<Vec<u8>> {
    #[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
    {
        if cfzlib::is_supported() {
            return cfzlib::cfzlib_deflate(data, zc, zs, zw, max_size, deadline);
        }
    }

    miniz_stream::compress_to_vec_oxipng(data, zc, zw.into(), zs.into(), max_size, deadline)
}

#[cfg(feature = "zopfli")]
pub fn zopfli_deflate(data: &[u8], iterations: NonZeroU8) -> PngResult<Vec<u8>> {
    use std::cmp::max;

    let mut output = Vec::with_capacity(max(1024, data.len() / 20));
    let options = zopfli::Options {
        iteration_count: iterations,
        ..Default::default()
    };
    match zopfli::compress(&options, &zopfli::Format::Zlib, data, &mut output) {
        Ok(_) => (),
        Err(_) => return Err(PngError::new("Failed to compress in zopfli")),
    };
    output.shrink_to_fit();
    Ok(output)
}

#[derive(Clone, Debug, PartialEq, Eq)]
/// DEFLATE algorithms supported by oxipng
pub enum Deflaters {
    /// Use the Zlib/Miniz DEFLATE implementation
    Zlib {
        /// Which zlib compression levels to try on the file (1-9)
        ///
        /// Default: `9`
        compression: IndexSet<u8>,
        /// Which zlib compression strategies to try on the file (0-3)
        ///
        /// Default: `0-3`
        strategies: IndexSet<u8>,
        /// Window size to use when compressing the file, as `2^window` bytes.
        ///
        /// Doesn't affect compression but may affect speed and memory usage.
        /// 8-15 are valid values.
        ///
        /// Default: `15`
        window: u8,
    },
    #[cfg(feature = "zopfli")]
    /// Use the better but slower Zopfli implementation
    Zopfli {
        /// The number of compression iterations to do. 15 iterations are fine
        /// for small files, but bigger files will need to be compressed with
        /// less iterations, or else they will be too slow.
        iterations: NonZeroU8,
    },
    #[cfg(feature = "libdeflater")]
    /// Use libdeflater.
    Libdeflater,
}
<commit_msg>Avoid printing output bytes on decompress error<commit_after>use crate::atomicmin::AtomicMin;
use crate::error::PngError;
use crate::Deadline;
use crate::PngResult;
use indexmap::IndexSet;

#[cfg(feature = "zopfli")]
use std::num::NonZeroU8;

#[doc(hidden)]
pub mod miniz_stream;

#[cfg(feature = "libdeflater")]
mod deflater;
#[cfg(feature = "libdeflater")]
pub use deflater::deflate as libdeflater_deflate;

#[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
pub mod cfzlib;

#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
pub mod cfzlib {
    pub fn is_supported() -> bool {
        return false;
    }
}

/// Decompress a data stream using the DEFLATE algorithm
pub fn inflate(data: &[u8]) -> PngResult<Vec<u8>> {
    miniz_oxide::inflate::decompress_to_vec_zlib(data).map_err(|e| {
        PngError::new(&format!(
            "Error on decompress: {:?} (after {:?} decompressed bytes)",
            e.status,
            e.output.len()
        ))
    })
}

/// Compress a data stream using the DEFLATE algorithm
#[doc(hidden)]
pub fn deflate(
    data: &[u8],
    zc: u8,
    zs: u8,
    zw: u8,
    max_size: &AtomicMin,
    deadline: &Deadline,
) -> PngResult<Vec<u8>> {
    #[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
    {
        if cfzlib::is_supported() {
            return cfzlib::cfzlib_deflate(data, zc, zs, zw, max_size, deadline);
        }
    }

    miniz_stream::compress_to_vec_oxipng(data, zc, zw.into(), zs.into(), max_size, deadline)
}

#[cfg(feature = "zopfli")]
pub fn zopfli_deflate(data: &[u8], iterations: NonZeroU8) -> PngResult<Vec<u8>> {
    use std::cmp::max;

    let mut output = Vec::with_capacity(max(1024, data.len() / 20));
    let options = zopfli::Options {
        iteration_count: iterations,
        ..Default::default()
    };
    match zopfli::compress(&options, &zopfli::Format::Zlib, data, &mut output) {
        Ok(_) => (),
        Err(_) => return Err(PngError::new("Failed to compress in zopfli")),
    };
    output.shrink_to_fit();
    Ok(output)
}

#[derive(Clone, Debug, PartialEq, Eq)]
/// DEFLATE algorithms supported by oxipng
pub enum Deflaters {
    /// Use the Zlib/Miniz DEFLATE implementation
    Zlib {
        /// Which zlib compression levels to try on the file (1-9)
        ///
        /// Default: `9`
        compression: IndexSet<u8>,
        /// Which zlib compression strategies to try on the file (0-3)
        ///
        /// Default: `0-3`
        strategies: IndexSet<u8>,
        /// Window size to use when compressing the file, as `2^window` bytes.
        ///
        /// Doesn't affect compression but may affect speed and memory usage.
        /// 8-15 are valid values.
        ///
        /// Default: `15`
        window: u8,
    },
    #[cfg(feature = "zopfli")]
    /// Use the better but slower Zopfli implementation
    Zopfli {
        /// The number of compression iterations to do. 15 iterations are fine
        /// for small files, but bigger files will need to be compressed with
        /// less iterations, or else they will be too slow.
        iterations: NonZeroU8,
    },
    #[cfg(feature = "libdeflater")]
    /// Use libdeflater.
    Libdeflater,
}
<|endoftext|>