<commit_before>use pueue_lib::log::clean_log_handles;
use pueue_lib::network::message::*;
use pueue_lib::state::SharedState;
use pueue_lib::task::{TaskResult, TaskStatus};

use super::*;
use crate::ok_or_return_failure_message;
use crate::state_helper::{is_task_removable, save_state};

fn construct_success_clean_message(message: CleanMessage) -> String {
    let successfull_only_fix = if message.successful_only {
        " successfully"
    } else {
        ""
    };

    format!(
        "All{} finished tasks have been removed",
        successfull_only_fix
    )
}

/// Invoked when calling `pueue clean`.
/// Remove all failed or done tasks from the state.
pub fn clean(message: CleanMessage, state: &SharedState) -> Message {
    let mut state = state.lock().unwrap();
    ok_or_return_failure_message!(save_state(&state));

    let (matching, _) = state.filter_tasks(|task| matches!(task.status, TaskStatus::Done(_)), None);

    for task_id in &matching {
        // Ensure the task is removable, i.e. there are no dependant tasks.
        if !is_task_removable(&state, task_id, &[]) {
            continue;
        }
        // Check if we should ignore this task, if only successful tasks should be removed.
        if message.successful_only {
            if let Some(task) = state.tasks.get(task_id) {
                if !matches!(task.status, TaskStatus::Done(TaskResult::Success)) {
                    continue;
                }
            }
        }
        let _ = state.tasks.remove(task_id).unwrap();
        clean_log_handles(*task_id, &state.settings.shared.pueue_directory());
    }

    ok_or_return_failure_message!(save_state(&state));

    create_success_message(construct_success_clean_message(message))
}

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

    use pretty_assertions::assert_eq;
    use tempfile::TempDir;

    fn get_message(successful_only: bool) -> CleanMessage {
        CleanMessage { successful_only }
    }

    fn get_clean_test_state() -> (SharedState, TempDir) {
        let (state, tempdir) = get_state();

        {
            let mut state = state.lock().unwrap();
            let task = get_stub_task("0", TaskStatus::Done(TaskResult::Success));
            state.add_task(task);

            let task = get_stub_task("1", TaskStatus::Done(TaskResult::Failed(1)));
            state.add_task(task);

            let task = get_stub_task(
                "2",
                TaskStatus::Done(TaskResult::FailedToSpawn("error".to_string())),
            );
            state.add_task(task);

            let task = get_stub_task("3", TaskStatus::Done(TaskResult::Killed));
            state.add_task(task);

            let task = get_stub_task("4", TaskStatus::Done(TaskResult::Errored));
            state.add_task(task);

            let task = get_stub_task("5", TaskStatus::Done(TaskResult::DependencyFailed));
            state.add_task(task);
        }

        (state, tempdir)
    }

    #[test]
    fn clean_normal() {
        let (state, _tempdir) = get_stub_state();

        // Only task 1 will be removed, since it's the only TaskStatus with `Done`.
        let message = clean(get_message(false), &state);

        // Return message is correct
        assert!(matches!(message, Message::Success(_)));
        if let Message::Success(text) = message {
            assert_eq!(text, "All finished tasks have been removed");
        };

        let state = state.lock().unwrap();
        assert_eq!(state.tasks.len(), 4);
    }

    #[test]
    fn clean_normal_for_all_results() {
        let (state, _tempdir) = get_clean_test_state();

        // All finished tasks should removed when calling default `clean`.
        let message = clean(get_message(false), &state);

        // Return message is correct
        assert!(matches!(message, Message::Success(_)));
        if let Message::Success(text) = message {
            assert_eq!(text, "All finished tasks have been removed");
        };

        let state = state.lock().unwrap();
        assert!(state.tasks.is_empty());
    }

    #[test]
    fn clean_successful_only() {
        let (state, _tempdir) = get_clean_test_state();

        // Only successfully finished tasks should get removed when
        // calling `clean` with the `successful_only` flag.
        let message = clean(get_message(true), &state);

        // Return message is correct
        assert!(matches!(message, Message::Success(_)));
        if let Message::Success(text) = message {
            assert_eq!(text, "All successfully finished tasks have been removed");
        };

        // Assert that only the first entry has been deleted (TaskResult::Success)
        let state = state.lock().unwrap();
        assert_eq!(state.tasks.len(), 5);
        assert!(state.tasks.get(&0).is_none());
    }
}
<commit_msg>daemon, message_handler: add a new helper for add stub tasks into state<commit_after>use pueue_lib::log::clean_log_handles;
use pueue_lib::network::message::*;
use pueue_lib::state::SharedState;
use pueue_lib::task::{TaskResult, TaskStatus};

use super::*;
use crate::ok_or_return_failure_message;
use crate::state_helper::{is_task_removable, save_state};

fn construct_success_clean_message(message: CleanMessage) -> String {
    let successfull_only_fix = if message.successful_only {
        " successfully"
    } else {
        ""
    };

    format!(
        "All{} finished tasks have been removed",
        successfull_only_fix
    )
}

/// Invoked when calling `pueue clean`.
/// Remove all failed or done tasks from the state.
pub fn clean(message: CleanMessage, state: &SharedState) -> Message {
    let mut state = state.lock().unwrap();
    ok_or_return_failure_message!(save_state(&state));

    let (matching, _) = state.filter_tasks(|task| matches!(task.status, TaskStatus::Done(_)), None);

    for task_id in &matching {
        // Ensure the task is removable, i.e. there are no dependant tasks.
        if !is_task_removable(&state, task_id, &[]) {
            continue;
        }
        // Check if we should ignore this task, if only successful tasks should be removed.
        if message.successful_only {
            if let Some(task) = state.tasks.get(task_id) {
                if !matches!(task.status, TaskStatus::Done(TaskResult::Success)) {
                    continue;
                }
            }
        }
        let _ = state.tasks.remove(task_id).unwrap();
        clean_log_handles(*task_id, &state.settings.shared.pueue_directory());
    }

    ok_or_return_failure_message!(save_state(&state));

    create_success_message(construct_success_clean_message(message))
}

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

    use pretty_assertions::assert_eq;
    use tempfile::TempDir;

    fn get_message(successful_only: bool) -> CleanMessage {
        CleanMessage { successful_only }
    }

    trait TaskAddable {
        fn add_stub_task(&mut self, id: &str, group: &str, task_result: TaskResult);
    }

    impl TaskAddable for State {
        fn add_stub_task(&mut self, id: &str, group: &str, task_result: TaskResult) {
            let task = get_stub_task_in_group(id, group, TaskStatus::Done(task_result));
            self.add_task(task);
        }
    }

    fn get_clean_test_state() -> (SharedState, TempDir) {
        let (state, tempdir) = get_state();

        {
            let mut state = state.lock().unwrap();

            state.add_stub_task("0", group, TaskResult::Success);
            state.add_stub_task("1", group, TaskResult::Failed(1));
            state.add_stub_task("2", group, TaskResult::FailedToSpawn("error".to_string()));
            state.add_stub_task("3", group, TaskResult::Killed);
            state.add_stub_task("4", group, TaskResult::Errored);
            state.add_stub_task("5", group, TaskResult::DependencyFailed);
        }

        (state, tempdir)
    }

    #[test]
    fn clean_normal() {
        let (state, _tempdir) = get_stub_state();

        // Only task 1 will be removed, since it's the only TaskStatus with `Done`.
        let message = clean(get_message(false), &state);

        // Return message is correct
        assert!(matches!(message, Message::Success(_)));
        if let Message::Success(text) = message {
            assert_eq!(text, "All finished tasks have been removed");
        };

        let state = state.lock().unwrap();
        assert_eq!(state.tasks.len(), 4);
    }

    #[test]
    fn clean_normal_for_all_results() {
        let (state, _tempdir) = get_clean_test_state();

        // All finished tasks should removed when calling default `clean`.
        let message = clean(get_message(false), &state);

        // Return message is correct
        assert!(matches!(message, Message::Success(_)));
        if let Message::Success(text) = message {
            assert_eq!(text, "All finished tasks have been removed");
        };

        let state = state.lock().unwrap();
        assert!(state.tasks.is_empty());
    }

    #[test]
    fn clean_successful_only() {
        let (state, _tempdir) = get_clean_test_state();

        // Only successfully finished tasks should get removed when
        // calling `clean` with the `successful_only` flag.
        let message = clean(get_message(true), &state);

        // Return message is correct
        assert!(matches!(message, Message::Success(_)));
        if let Message::Success(text) = message {
            assert_eq!(text, "All successfully finished tasks have been removed");
        };

        // Assert that only the first entry has been deleted (TaskResult::Success)
        let state = state.lock().unwrap();
        assert_eq!(state.tasks.len(), 5);
        assert!(state.tasks.get(&0).is_none());
    }
}
<|endoftext|>


<commit_before>use regex::Regex;
use reqwest;
use serde_json;

use web::data::{SolrTermSummary, SolrReferenceSummary};
use web::config::Config;

pub struct Search {
    solr_url: String,
    close_synonym_boost: f32,
    distant_synonym_boost: f32,
}

#[derive(Deserialize, Debug)]
struct SolrTermResponse {
    pub docs: Vec<SolrTermSummary>,
}

#[derive(Deserialize, Debug)]
struct SolrTermResponseContainer {
    pub response: SolrTermResponse,
}

#[derive(Deserialize, Debug)]
struct SolrRefResponse {
    pub docs: Vec<SolrReferenceSummary>,
}

#[derive(Deserialize, Debug)]
struct SolrRefResponseContainer {
    pub response: SolrRefResponse,
}

impl Search {
    pub fn new(config: &Config) -> Search {
        Search {
            solr_url: config.server.solr_url.clone(),
            close_synonym_boost: config.server.close_synonym_boost,
            distant_synonym_boost: config.server.distant_synonym_boost,
        }
    }

    fn get_query_part(&self, words: &[String]) -> String {
        let mut ret = String::new();

        let words_length = words.len();

        for (i, word) in words.iter().enumerate() {
            if i == words_length - 1 {
                ret += &format!("{} {}~0.8 {}*", word, word, word);
            } else {
                ret += &format!("{} {}~0.8 ", word, word);
            }
        }

        ret
    }

    pub fn term_complete(&self, cv_name: &str, q: &str)
                         -> Result<Vec<SolrTermSummary>, String>
    {
        let mut terms_url =
            self.solr_url.to_owned() + "/terms/select?wt=json&q=";

        let termid_re_string = r"^(?P<prefix>[\w_]+):(?P<accession>\d+)$";
        let termid_re = Regex::new(termid_re_string).unwrap();

        let parent_re_string = r"^\[".to_owned() + termid_re_string + r"\]$";
        let parent_re = Regex::new(&parent_re_string).unwrap();

        let maybe_captures = parent_re.captures(cv_name);

        if let Some(captures) = maybe_captures {
            let prefix = captures.name("prefix").unwrap().as_str();
            let accession = captures.name("accession").unwrap().as_str();
            terms_url += &format!("(interesting_parents:{}\\:{} OR id:{}\\:{})",
                                  prefix, accession, prefix, accession);
        } else {
            terms_url += "cv_name:";
            terms_url += cv_name;
        }

        if let Some(captures) = termid_re.captures(q) {
            let prefix = captures.name("prefix").unwrap().as_str();
            let accession = captures.name("accession").unwrap().as_str();
            terms_url += " AND id:";
            terms_url += prefix;
            terms_url += r"\:";
            terms_url += accession;
        } else {
            let substring = |s: &str, len: usize| s.chars().take(len).collect::<String>();
            let lower_q = substring(&q.to_lowercase(), 200);

            terms_url += " AND (name:(";

            let clean_words: Vec<String> =
                Regex::new(r"([\w-]+)").unwrap().captures_iter(&lower_q)
                .map(|cap| cap.get(1).unwrap().as_str().to_owned()).collect();

            if clean_words.is_empty() {
                return Ok(vec![]);
            }

            let clean_words_length = clean_words.len();

            for (i, word) in clean_words.iter().enumerate() {
                if i == clean_words_length - 1 {
                    terms_url += &format!("{} {}~0.8 {}*", word, word, word);
                } else {
                    terms_url += &format!("{} {}~0.8 ", word, word);
                }
            }

            let query_part = self.get_query_part(&clean_words);

            terms_url += &format!(") OR close_synonym_words:({})^{} OR distant_synonym_words:({})^{})",
                                  query_part, self.close_synonym_boost,
                                  query_part, self.distant_synonym_boost);
        }
        print!("Solr URL: {:?}\n", terms_url);

        match reqwest::get(&terms_url) {
            Ok(res) => {
                println!("Status: {}", res.status());

                if res.status().is_success() {
                    match serde_json::from_reader(res) {
                        Ok(container) => {
                            let solr_response_container: SolrTermResponseContainer = container;
                            Ok(solr_response_container.response.docs)
                        },
                        Err(err) => {
                            Err(format!("Error parsing response from Solr: {:?}", err))
                        }
                    }
                } else {
                    if let Some(reason) = res.status().canonical_reason() {
                        Err(format!("HTTP request to Solr failed: {} - {}", res.status(), reason))
                    } else {
                        Err(format!("HTTP request to Solr failed with status code: {}",
                                    res.status()))
                    }
                }
            },
            Err(err) => {
                Err(format!("Error from Reqwest: {:?}", err))
            }
        }
    }

    pub fn ref_complete(&self, q: &str) -> Result<Vec<SolrReferenceSummary>, String> {
        let mut refs_url =
            self.solr_url.to_owned() + "/refs/select?wt=json&q=";

        let id_re_string = r"^(?:(?P<prefix>[\w_]+):)?(?P<rest>\d\d\d\d\d+)$";
        let id_re = Regex::new(id_re_string).unwrap();

        if let Some(captures) = id_re.captures(q) {
            let mut prefix =
                if let Some(prefix_capture) = captures.name("prefix") {
                    prefix_capture.as_str()
                } else {
                    "PMID"
                };
            let accession = captures.name("rest").unwrap().as_str();
            refs_url += "id:";
            refs_url += prefix;
            refs_url += r"\:";
            refs_url += accession;
        } else {
            let substring = |s: &str, len: usize| s.chars().take(len).collect::<String>();
            let lower_q = substring(&q.to_lowercase(), 200);

            let clean_words: Vec<String> =
                Regex::new(r"([\w-]+)").unwrap().captures_iter(&lower_q)
                .map(|cap| cap.get(1).unwrap().as_str().to_owned()).collect();

            if clean_words.is_empty() {
                return Ok(vec![]);
            }

            let clean_words_length = clean_words.len();
            let mut clean_words_for_url = String::new();

            for (i, word) in clean_words.iter().enumerate() {
                if i == clean_words_length - 1 {
                    clean_words_for_url += &format!("{} {}~0.8 {}*", word, word, word);
                } else {
                    clean_words_for_url += &format!("{} {}~0.8 ", word, word);
                }
            }

            let url_parts =
                vec!["title", "citation", "authors", "authors_abbrev"]
                .iter()
                .map(|field_name| format!("{}:({})", field_name, clean_words_for_url))
                .collect::<Vec<String>>();

            refs_url += &url_parts.join(" OR ");

            for word in clean_words {
                if word.len() == 4 && (word.starts_with("19") || word.starts_with("20")) {
                    if let Ok(num) = word.parse::<u32>() {
                        refs_url += &format!(" OR publication_year:{}^20", num);
                    }
                }
            }
        }

        print!("Solr refs URL: {:?}\n", refs_url);

        match reqwest::get(&refs_url) {
            Ok(res) => {
                println!("Status: {}", res.status());

                if res.status().is_success() {
                    match serde_json::from_reader(res) {
                        Ok(container) => {
                            let solr_response_container: SolrRefResponseContainer = container;
                            Ok(solr_response_container.response.docs)
                        },
                        Err(err) => {
                            Err(format!("Error parsing response from Solr: {:?}", err))
                        }
                    }
                } else {
                    if let Some(reason) = res.status().canonical_reason() {
                        Err(format!("HTTP request to Solr failed: {} - {}", res.status(), reason))
                    } else {
                        Err(format!("HTTP request to Solr failed with status code: {}",
                                    res.status()))
                    }
                }
            },
            Err(err) => {
                Err(format!("Error from Reqwest: {:?}", err))
            }
        }
    }

}
<commit_msg>Don't do fuzzy searching for terms or refs<commit_after>use regex::Regex;
use reqwest;
use serde_json;

use web::data::{SolrTermSummary, SolrReferenceSummary};
use web::config::Config;

pub struct Search {
    solr_url: String,
    close_synonym_boost: f32,
    distant_synonym_boost: f32,
}

#[derive(Deserialize, Debug)]
struct SolrTermResponse {
    pub docs: Vec<SolrTermSummary>,
}

#[derive(Deserialize, Debug)]
struct SolrTermResponseContainer {
    pub response: SolrTermResponse,
}

#[derive(Deserialize, Debug)]
struct SolrRefResponse {
    pub docs: Vec<SolrReferenceSummary>,
}

#[derive(Deserialize, Debug)]
struct SolrRefResponseContainer {
    pub response: SolrRefResponse,
}

impl Search {
    pub fn new(config: &Config) -> Search {
        Search {
            solr_url: config.server.solr_url.clone(),
            close_synonym_boost: config.server.close_synonym_boost,
            distant_synonym_boost: config.server.distant_synonym_boost,
        }
    }

    fn get_query_part(&self, words: &[String]) -> String {
        let mut ret = String::new();

        let words_length = words.len();

        for (i, word) in words.iter().enumerate() {
            if i == words_length - 1 {
                ret += &format!("{} {}~0.8 {}*", word, word, word);
            } else {
                ret += &format!("{} {}~0.8 ", word, word);
            }
        }

        ret
    }

    pub fn term_complete(&self, cv_name: &str, q: &str)
                         -> Result<Vec<SolrTermSummary>, String>
    {
        let mut terms_url =
            self.solr_url.to_owned() + "/terms/select?wt=json&q=";

        let termid_re_string = r"^(?P<prefix>[\w_]+):(?P<accession>\d+)$";
        let termid_re = Regex::new(termid_re_string).unwrap();

        let parent_re_string = r"^\[".to_owned() + termid_re_string + r"\]$";
        let parent_re = Regex::new(&parent_re_string).unwrap();

        let maybe_captures = parent_re.captures(cv_name);

        if let Some(captures) = maybe_captures {
            let prefix = captures.name("prefix").unwrap().as_str();
            let accession = captures.name("accession").unwrap().as_str();
            terms_url += &format!("(interesting_parents:{}\\:{} OR id:{}\\:{})",
                                  prefix, accession, prefix, accession);
        } else {
            terms_url += "cv_name:";
            terms_url += cv_name;
        }

        if let Some(captures) = termid_re.captures(q) {
            let prefix = captures.name("prefix").unwrap().as_str();
            let accession = captures.name("accession").unwrap().as_str();
            terms_url += " AND id:";
            terms_url += prefix;
            terms_url += r"\:";
            terms_url += accession;
        } else {
            let substring = |s: &str, len: usize| s.chars().take(len).collect::<String>();
            let lower_q = substring(&q.to_lowercase(), 200);

            terms_url += " AND (name:(";

            let clean_words: Vec<String> =
                Regex::new(r"([\w\d\-]+)").unwrap().captures_iter(&lower_q)
                .map(|cap| cap.get(1).unwrap().as_str().to_owned()).collect();

            if clean_words.is_empty() {
                return Ok(vec![]);
            }

            let clean_words_length = clean_words.len();

            for (i, word) in clean_words.iter().enumerate() {
                if i == clean_words_length - 1 {
                    terms_url += &format!("{} {}*", word, word);
                } else {
                    terms_url += &format!("{} ", word);
                }
            }

            let query_part = self.get_query_part(&clean_words);

            terms_url += &format!(") OR close_synonym_words:({})^{} OR distant_synonym_words:({})^{})",
                                  query_part, self.close_synonym_boost,
                                  query_part, self.distant_synonym_boost);
        }
        print!("Solr URL: {:?}\n", terms_url);

        match reqwest::get(&terms_url) {
            Ok(res) => {
                println!("Status: {}", res.status());

                if res.status().is_success() {
                    match serde_json::from_reader(res) {
                        Ok(container) => {
                            let solr_response_container: SolrTermResponseContainer = container;
                            Ok(solr_response_container.response.docs)
                        },
                        Err(err) => {
                            Err(format!("Error parsing response from Solr: {:?}", err))
                        }
                    }
                } else {
                    if let Some(reason) = res.status().canonical_reason() {
                        Err(format!("HTTP request to Solr failed: {} - {}", res.status(), reason))
                    } else {
                        Err(format!("HTTP request to Solr failed with status code: {}",
                                    res.status()))
                    }
                }
            },
            Err(err) => {
                Err(format!("Error from Reqwest: {:?}", err))
            }
        }
    }

    pub fn ref_complete(&self, q: &str) -> Result<Vec<SolrReferenceSummary>, String> {
        let mut refs_url =
            self.solr_url.to_owned() + "/refs/select?wt=json&q=";

        let id_re_string = r"^(?:(?P<prefix>[\w_]+):\w*)?(?P<rest>\d\d\d\d\d+)$";
        let id_re = Regex::new(id_re_string).unwrap();

        if let Some(captures) = id_re.captures(q) {
            let mut prefix =
                if let Some(prefix_capture) = captures.name("prefix") {
                    prefix_capture.as_str()
                } else {
                    "PMID"
                };
            let accession = captures.name("rest").unwrap().as_str();
            refs_url += "id:";
            refs_url += prefix;
            refs_url += r"\:";
            refs_url += accession;
        } else {
            let substring = |s: &str, len: usize| s.chars().take(len).collect::<String>();
            let lower_q = substring(&q.to_lowercase(), 200);

            let clean_words: Vec<String> =
                Regex::new(r"([\w\d\-]+)").unwrap().captures_iter(&lower_q)
                .map(|cap| cap.get(1).unwrap().as_str().to_owned()).collect();

            if clean_words.is_empty() {
                return Ok(vec![]);
            }

            let clean_words_length = clean_words.len();
            let mut clean_words_for_url = String::new();

            for (i, word) in clean_words.iter().enumerate() {
                if i == clean_words_length - 1 {
                    clean_words_for_url += &format!("{} {}*", word, word);
                } else {
                    clean_words_for_url += &format!("{} ", word);
                }
            }

            let url_parts =
                vec!["title", "citation", "authors", "authors_abbrev"]
                .iter()
                .map(|field_name| format!("{}:({})", field_name, clean_words_for_url))
                .collect::<Vec<String>>();

            refs_url += &url_parts.join(" OR ");

            for word in clean_words {
                if word.len() == 4 && (word.starts_with("19") || word.starts_with("20")) {
                    if let Ok(num) = word.parse::<u32>() {
                        refs_url += &format!(" OR publication_year:{}^20", num);
                    }
                }
            }
        }

        print!("Solr refs URL: {:?}\n", refs_url);

        match reqwest::get(&refs_url) {
            Ok(res) => {
                println!("Status: {}", res.status());

                if res.status().is_success() {
                    match serde_json::from_reader(res) {
                        Ok(container) => {
                            let solr_response_container: SolrRefResponseContainer = container;
                            Ok(solr_response_container.response.docs)
                        },
                        Err(err) => {
                            Err(format!("Error parsing response from Solr: {:?}", err))
                        }
                    }
                } else {
                    if let Some(reason) = res.status().canonical_reason() {
                        Err(format!("HTTP request to Solr failed: {} - {}", res.status(), reason))
                    } else {
                        Err(format!("HTTP request to Solr failed with status code: {}",
                                    res.status()))
                    }
                }
            },
            Err(err) => {
                Err(format!("Error from Reqwest: {:?}", err))
            }
        }
    }

}
<|endoftext|>


<commit_before>#![recursion_limit = "1000"]

extern crate proc_macro;
extern crate syn;
#[macro_use]
extern crate quote;

use proc_macro::TokenStream;
use syn::Ident;
use quote::Tokens;

#[proc_macro_derive(IntoHeap)]
pub fn derive_into_heap(input: TokenStream) -> TokenStream {
    let source = input.to_string();
    let ast = syn::parse_derive_input(&source).unwrap();
    let expanded = impl_into_heap(&ast);

    // Uncomment to debug the generated code...
    // println!("{:?}", expanded);

    expanded.parse().unwrap()
}

fn impl_into_heap(ast: &syn::DeriveInput) -> Tokens {
    match ast.body {
        syn::Body::Struct(ref data) => impl_into_heap_for_struct(ast, data),
        syn::Body::Enum(ref variants) => impl_into_heap_for_enum(ast, variants),
    }
}

fn impl_into_heap_for_struct(ast: &syn::DeriveInput, data: &syn::VariantData) -> Tokens {
    let name = &ast.ident;
    let name_str: &str = name.as_ref();
    let storage_type_name: Ident = Ident::from(name_str.to_string() + "Storage");
    let vis = &ast.vis;
    let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();
    let heap_lifetime = &ast.generics
        .lifetimes
        .first()
        .expect("lifetime parameter required")
        .lifetime;

    match *data {
        syn::VariantData::Struct(ref fields) => {
            let field_vis: &Vec<_> = &fields.iter().map(|f| &f.vis).collect();
            let field_names: &Vec<_> = &fields.iter().map(|f| &f.ident).collect();
            let field_types: &Vec<_> = &fields.iter().map(|f| &f.ty).collect();
            let field_storage_types: Vec<_> = fields
                .iter()
                .map(|f| {
                    let field_ty = &f.ty;
                    quote! {
                        <#field_ty as ::cell_gc::traits::IntoHeapBase>::In
                    }
                })
                .collect();

            // 1. The in-heap representation of the struct.
            let storage_struct = quote! {
                #vis struct #storage_type_name #impl_generics #where_clause {
                    #( #field_vis #field_names: #field_storage_types ),*
                }
            };

            // 2. IntoHeap implementation.
            // Body of the trace() method.
            let trace_fields: Vec<Tokens> = fields
                .iter()
                .map(|f| {
                    let name = &f.ident;
                    let ty = &f.ty;
                    quote! {
                    <#ty as ::cell_gc::traits::IntoHeapBase>
                        ::trace(&storage.#name, tracer);
                }
                })
                .collect();

            // Oddly you can't use the same identifier more than once in the
            // same loop. So create an alias.
            let field_names_1 = field_names;

            let into_heap = quote! {
                impl #impl_generics ::cell_gc::traits::IntoHeapBase
                    for #name #ty_generics
                    #where_clause
                {
                    type In = #storage_type_name #ty_generics;

                    fn into_heap(self) -> Self::In {
                        #storage_type_name {
                            #(
                                #field_names:
                                    ::cell_gc::traits::IntoHeapBase::into_heap(
                                        self.#field_names_1)
                            ),*
                        }
                    }

                    unsafe fn trace<R>(storage: &Self::In, tracer: &mut R)
                        where R: ::cell_gc::traits::Tracer
                    {
                        #( #trace_fields )*

                        // Quiet unused variable warnings when `$(...)*` expands
                        // to nothing.
                        let _ = tracer;
                    }

                    unsafe fn from_heap(storage: &Self::In) -> Self {
                        #name {
                            #(
                                #field_names:
                                    ::cell_gc::traits::IntoHeapBase::from_heap(
                                        &storage.#field_names_1)
                            ),*
                        }
                    }
                }

                unsafe impl #impl_generics ::cell_gc::traits::IntoHeap<#heap_lifetime>
                    for #name #ty_generics
                    #where_clause
                {}
            };

            // 3. IntoHeapAllocation implementation.
            let ref_type_name: Ident = Ident::from(name_str.to_string() + "Ref");
            let into_heap_allocation = quote! {
                impl #impl_generics ::cell_gc::traits::IntoHeapAllocation<#heap_lifetime>
                    for #name #ty_generics
                    #where_clause
                {
                    type Ref = #ref_type_name #ty_generics;

                    fn wrap_gc_ref(gc_ref: ::cell_gc::GcRef<#heap_lifetime, #name #ty_generics>)
                        -> Self::Ref
                    {
                        #ref_type_name(gc_ref)
                    }

                    fn into_gc_ref(wrapped_ref: Self::Ref)
                        -> ::cell_gc::GcRef<#heap_lifetime, #name #ty_generics>
                    {
                        wrapped_ref.0
                    }
                }
            };

            // 4. #ref_type_name: A safe reference to the struct
            let ref_type = quote! {
                #[derive(Clone, Debug, PartialEq, Eq)]
                #vis struct #ref_type_name #impl_generics
                    (::cell_gc::GcRef<#heap_lifetime, #name #ty_generics>)
                    #where_clause;
            };

            // 5. The ref type also gets an IntoHeap impl...
            let ref_type_into_heap = quote! {
                impl #impl_generics ::cell_gc::traits::IntoHeapBase
                    for #ref_type_name #ty_generics
                    #where_clause
                {
                    type In = <::cell_gc::GcRef<#heap_lifetime, #name #ty_generics>
                               as ::cell_gc::traits::IntoHeapBase>::In;

                    fn into_heap(self) -> Self::In {
                        self.0.into_heap()
                    }

                    unsafe fn from_heap(storage: &Self::In) -> Self {
                        #ref_type_name(::cell_gc::GcRef::<#heap_lifetime, #name #ty_generics>::new(*storage))
                    }

                    unsafe fn trace<R>(storage: &Self::In, tracer: &mut R)
                        where R: ::cell_gc::traits::Tracer
                    {
                        <::cell_gc::GcRef<#heap_lifetime, #name #ty_generics>
                            as ::cell_gc::traits::IntoHeapBase>::trace(storage, tracer)
                    }
                }

                unsafe impl #impl_generics ::cell_gc::traits::IntoHeap<#heap_lifetime>
                    for #ref_type_name #ty_generics
                    #where_clause
                {}
            };

            // 6. Getters and setters.
            let field_setter_names: Vec<_> = fields
                .iter()
                .map(|f| {
                    let field_str: &str = f.ident.as_ref().unwrap().as_ref();
                    Ident::from(format!("set_{}", field_str))
                })
                .collect();
            let accessors = quote! {
                impl #impl_generics #ref_type_name #ty_generics #where_clause {
                    #(
                        #[allow(dead_code)]
                        #field_vis fn #field_names(&self) -> #field_types {
                            let ptr = self.0.as_ptr();
                            unsafe {
                                ::cell_gc::traits::IntoHeapBase::from_heap(
                                    &(*ptr).#field_names_1)
                            }
                        }
                    )*

                    #(
                        #[allow(dead_code)]
                        #field_vis fn #field_setter_names(&self, v: #field_types) {
                            let ptr = self.0.as_mut_ptr();
                            let u = ::cell_gc::traits::IntoHeapBase::into_heap(v);
                            unsafe {
                                (*ptr).#field_names = u;
                            }
                        }
                    )*

                    ///// Get all fields at once.
                    //pub fn get(&self) -> #name {
                    //    ::cell_gc::traits::IntoHeapBase::from_heap(self.0.ptr())
                    //}

                    #[allow(dead_code)]
                    pub fn as_mut_ptr(&self) -> *mut #storage_type_name #ty_generics {
                        self.0.as_mut_ptr()
                    }
                }
            };

            quote! {
                #storage_struct
                #into_heap
                #into_heap_allocation
                #ref_type
                #ref_type_into_heap
                #accessors
            }
        }
        syn::VariantData::Tuple(ref _fields) => {
            panic!("#[derive(IntoHeap)] does not support tuple structs");
        }
        syn::VariantData::Unit => {
            panic!("#[derive(IntoHeap)] does not support unit structs");
        }
    }
}

fn impl_into_heap_for_enum(ast: &syn::DeriveInput, variants: &[syn::Variant]) -> Tokens {
    let attrs = &ast.attrs;
    let name = &ast.ident;
    let name_str: &str = name.as_ref();
    let storage_type_name: Ident = Ident::from(name_str.to_string() + "Storage");
    let vis = &ast.vis;
    let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();
    let heap_lifetime = &ast.generics
        .lifetimes
        .first()
        .expect("lifetime parameter required")
        .lifetime;

    let variant_storage = variants.iter().map(|v| {
        let attrs = &v.attrs;
        let ident = &v.ident;
        if v.discriminant.is_some() {
            panic!("#[derive(IntoHeap)] does not support enum variants with discriminants");
        }
        match v.data {
            syn::VariantData::Struct(ref fields) => {
                let field_decls = fields.iter().map(|f| {
                    let field_attrs = &f.attrs;
                    let field_name = &f.ident;
                    let field_ty = &f.ty;
                    quote! {
                        #( #field_attrs )*
                        #field_name:
                            <#field_ty as ::cell_gc::traits::IntoHeapBase>::In
                    }
                });
                quote! {
                    #( #attrs )*
                    #ident { #( #field_decls ),* }
                }
            }
            syn::VariantData::Tuple(ref fields) => {
                let field_decls = fields.iter().map(|f| {
                    let field_ty = &f.ty;
                    let field_attrs = &f.attrs;
                    quote! {
                        #( #field_attrs )*
                        <#field_ty as ::cell_gc::traits::IntoHeapBase>::In
                    }
                });
                quote! {
                    #( #attrs )*
                    #ident( #( #field_decls ),* )
                }
            }
            syn::VariantData::Unit => {
                quote! { #( #attrs )* #ident }
            }
        }
    });
    let storage_enum = quote! {
        #( #attrs )*
        #vis enum #storage_type_name #impl_generics
            #where_clause
        {
            #( #variant_storage ),*
        }
    };

    let into_heap_arms = variants.iter().map(|v| {
        let ident = &v.ident;
        match v.data {
            syn::VariantData::Struct(ref fields) => {
                let field_names: &Vec<_> = &fields.iter().map(|f| &f.ident).collect();
                let field_names_1 = field_names;
                quote! {
                    #name::#ident { #(#field_names),* } => {
                        #storage_type_name::#ident {
                            #(
                                #field_names:
                                    ::cell_gc::traits::IntoHeapBase::into_heap(#field_names_1)
                            ),*
                        }
                    }
                }
            }
            syn::VariantData::Tuple(ref fields) => {
                let bindings: &Vec<Ident> = &(0..fields.len())
                    .map(|n| Ident::from(format!("x{}", n)))
                    .collect();
                quote! {
                    #name::#ident( #(#bindings),* ) => {
                        #storage_type_name::#ident(
                            #(
                                ::cell_gc::traits::IntoHeapBase::into_heap(#bindings)
                            ),*
                        )
                    }
                }
            }
            syn::VariantData::Unit => {
                quote! {
                    #name::#ident => #storage_type_name::#ident
                }
            }
        }
    });

    let from_heap_arms = variants.iter().map(|v| {
        let ident = &v.ident;
        match v.data {
            syn::VariantData::Struct(ref fields) => {
                let field_names: &Vec<_> = &fields.iter().map(|f| &f.ident).collect();
                let field_names_1 = field_names;
                let field_types = fields.iter().map(|f| &f.ty);
                quote! {
                    #storage_type_name::#ident { #(ref #field_names),* } => {
                        #name::#ident {
                            #(
                                #field_names:
                                    <#field_types as ::cell_gc::traits::IntoHeapBase>
                                    ::from_heap(#field_names_1)
                            ),*
                        }
                    }
                }
            }
            syn::VariantData::Tuple(ref fields) => {
                let field_types = fields.iter().map(|f| &f.ty);
                let bindings: &Vec<Ident> = &(0..fields.len())
                    .map(|n| Ident::from(format!("x{}", n)))
                    .collect();
                quote! {
                    #storage_type_name::#ident( #(ref #bindings),* ) => {
                        #name::#ident(
                            #(
                                <#field_types as ::cell_gc::traits::IntoHeapBase>
                                    ::from_heap(#bindings)
                            ),*
                        )
                    }
                }
            }
            syn::VariantData::Unit => {
                quote! {
                    #storage_type_name::#ident => #name::#ident
                }
            }
        }
    });

    let trace_arms = variants.iter().map(|v| {
        let ident = &v.ident;
        match v.data {
            syn::VariantData::Struct(ref fields) => {
                let field_names: &Vec<_> = &fields.iter().map(|f| &f.ident).collect();
                let field_types = fields.iter().map(|f| &f.ty);
                quote! {
                    #storage_type_name::#ident { #(ref #field_names),* } => {
                        #(
                            <#field_types as ::cell_gc::traits::IntoHeapBase>
                                ::trace(#field_names, tracer);
                        )*
                    }
                }
            }
            syn::VariantData::Tuple(ref fields) => {
                let field_types = fields.iter().map(|f| &f.ty);
                let bindings: &Vec<Ident> = &(0..fields.len())
                    .map(|n| Ident::from(format!("x{}", n)))
                    .collect();
                quote! {
                    #storage_type_name::#ident( #(ref #bindings),* ) => {
                        #(
                            <#field_types as ::cell_gc::traits::IntoHeapBase>
                                ::trace(#bindings, tracer);
                        )*
                    }
                }
            }
            syn::VariantData::Unit => {
                quote! { #storage_type_name::#ident => () }
            }
        }
    });

    let into_heap = quote! {
        impl #impl_generics
            ::cell_gc::traits::IntoHeapBase
            for #name #ty_generics
            #where_clause
        {
            type In = #storage_type_name #ty_generics;

            fn into_heap(self) -> Self::In {
                match self {
                    #( #into_heap_arms ),*
                }
            }

            unsafe fn from_heap(storage: &Self::In) -> Self {
                match *storage {
                    #( #from_heap_arms ),*
                }
            }

            unsafe fn trace<R>(storage: &Self::In, tracer: &mut R)
                where R: ::cell_gc::traits::Tracer
            {
                match *storage {
                    #( #trace_arms ),*
                }

                // Quiet unused variable warnings when `$(...)*` expands to
                // nothing.
                let _ = tracer;
            }
        }

        unsafe impl #impl_generics
            ::cell_gc::traits::IntoHeap<#heap_lifetime>
            for #name #ty_generics
            #where_clause
        {}
    };

    quote! {
        #storage_enum
        #into_heap
    }
}
<commit_msg>Fix indentation.<commit_after>#![recursion_limit = "1000"]

extern crate proc_macro;
extern crate syn;
#[macro_use]
extern crate quote;

use proc_macro::TokenStream;
use syn::Ident;
use quote::Tokens;

#[proc_macro_derive(IntoHeap)]
pub fn derive_into_heap(input: TokenStream) -> TokenStream {
    let source = input.to_string();
    let ast = syn::parse_derive_input(&source).unwrap();
    let expanded = impl_into_heap(&ast);

    // Uncomment to debug the generated code...
    // println!("{:?}", expanded);

    expanded.parse().unwrap()
}

fn impl_into_heap(ast: &syn::DeriveInput) -> Tokens {
    match ast.body {
        syn::Body::Struct(ref data) => impl_into_heap_for_struct(ast, data),
        syn::Body::Enum(ref variants) => impl_into_heap_for_enum(ast, variants),
    }
}

fn impl_into_heap_for_struct(ast: &syn::DeriveInput, data: &syn::VariantData) -> Tokens {
    let name = &ast.ident;
    let name_str: &str = name.as_ref();
    let storage_type_name: Ident = Ident::from(name_str.to_string() + "Storage");
    let vis = &ast.vis;
    let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();
    let heap_lifetime = &ast.generics
        .lifetimes
        .first()
        .expect("lifetime parameter required")
        .lifetime;

    match *data {
        syn::VariantData::Struct(ref fields) => {
            let field_vis: &Vec<_> = &fields.iter().map(|f| &f.vis).collect();
            let field_names: &Vec<_> = &fields.iter().map(|f| &f.ident).collect();
            let field_types: &Vec<_> = &fields.iter().map(|f| &f.ty).collect();
            let field_storage_types: Vec<_> = fields
                .iter()
                .map(|f| {
                    let field_ty = &f.ty;
                    quote! {
                        <#field_ty as ::cell_gc::traits::IntoHeapBase>::In
                    }
                })
                .collect();

            // 1. The in-heap representation of the struct.
            let storage_struct = quote! {
                #vis struct #storage_type_name #impl_generics #where_clause {
                    #( #field_vis #field_names: #field_storage_types ),*
                }
            };

            // 2. IntoHeap implementation.
            // Body of the trace() method.
            let trace_fields: Vec<Tokens> = fields
                .iter()
                .map(|f| {
                    let name = &f.ident;
                    let ty = &f.ty;
                    quote! {
                        <#ty as ::cell_gc::traits::IntoHeapBase>
                            ::trace(&storage.#name, tracer);
                    }
                })
                .collect();

            // Oddly you can't use the same identifier more than once in the
            // same loop. So create an alias.
            let field_names_1 = field_names;

            let into_heap = quote! {
                impl #impl_generics ::cell_gc::traits::IntoHeapBase
                    for #name #ty_generics
                    #where_clause
                {
                    type In = #storage_type_name #ty_generics;

                    fn into_heap(self) -> Self::In {
                        #storage_type_name {
                            #(
                                #field_names:
                                    ::cell_gc::traits::IntoHeapBase::into_heap(
                                        self.#field_names_1)
                            ),*
                        }
                    }

                    unsafe fn trace<R>(storage: &Self::In, tracer: &mut R)
                        where R: ::cell_gc::traits::Tracer
                    {
                        #( #trace_fields )*

                        // Quiet unused variable warnings when `$(...)*` expands
                        // to nothing.
                        let _ = tracer;
                    }

                    unsafe fn from_heap(storage: &Self::In) -> Self {
                        #name {
                            #(
                                #field_names:
                                    ::cell_gc::traits::IntoHeapBase::from_heap(
                                        &storage.#field_names_1)
                            ),*
                        }
                    }
                }

                unsafe impl #impl_generics ::cell_gc::traits::IntoHeap<#heap_lifetime>
                    for #name #ty_generics
                    #where_clause
                {}
            };

            // 3. IntoHeapAllocation implementation.
            let ref_type_name: Ident = Ident::from(name_str.to_string() + "Ref");
            let into_heap_allocation = quote! {
                impl #impl_generics ::cell_gc::traits::IntoHeapAllocation<#heap_lifetime>
                    for #name #ty_generics
                    #where_clause
                {
                    type Ref = #ref_type_name #ty_generics;

                    fn wrap_gc_ref(gc_ref: ::cell_gc::GcRef<#heap_lifetime, #name #ty_generics>)
                        -> Self::Ref
                    {
                        #ref_type_name(gc_ref)
                    }

                    fn into_gc_ref(wrapped_ref: Self::Ref)
                        -> ::cell_gc::GcRef<#heap_lifetime, #name #ty_generics>
                    {
                        wrapped_ref.0
                    }
                }
            };

            // 4. #ref_type_name: A safe reference to the struct
            let ref_type = quote! {
                #[derive(Clone, Debug, PartialEq, Eq)]
                #vis struct #ref_type_name #impl_generics
                    (::cell_gc::GcRef<#heap_lifetime, #name #ty_generics>)
                    #where_clause;
            };

            // 5. The ref type also gets an IntoHeap impl...
            let ref_type_into_heap = quote! {
                impl #impl_generics ::cell_gc::traits::IntoHeapBase
                    for #ref_type_name #ty_generics
                    #where_clause
                {
                    type In = <::cell_gc::GcRef<#heap_lifetime, #name #ty_generics>
                               as ::cell_gc::traits::IntoHeapBase>::In;

                    fn into_heap(self) -> Self::In {
                        self.0.into_heap()
                    }

                    unsafe fn from_heap(storage: &Self::In) -> Self {
                        #ref_type_name(::cell_gc::GcRef::<#heap_lifetime, #name #ty_generics>::new(*storage))
                    }

                    unsafe fn trace<R>(storage: &Self::In, tracer: &mut R)
                        where R: ::cell_gc::traits::Tracer
                    {
                        <::cell_gc::GcRef<#heap_lifetime, #name #ty_generics>
                            as ::cell_gc::traits::IntoHeapBase>::trace(storage, tracer)
                    }
                }

                unsafe impl #impl_generics ::cell_gc::traits::IntoHeap<#heap_lifetime>
                    for #ref_type_name #ty_generics
                    #where_clause
                {}
            };

            // 6. Getters and setters.
            let field_setter_names: Vec<_> = fields
                .iter()
                .map(|f| {
                    let field_str: &str = f.ident.as_ref().unwrap().as_ref();
                    Ident::from(format!("set_{}", field_str))
                })
                .collect();
            let accessors = quote! {
                impl #impl_generics #ref_type_name #ty_generics #where_clause {
                    #(
                        #[allow(dead_code)]
                        #field_vis fn #field_names(&self) -> #field_types {
                            let ptr = self.0.as_ptr();
                            unsafe {
                                ::cell_gc::traits::IntoHeapBase::from_heap(
                                    &(*ptr).#field_names_1)
                            }
                        }
                    )*

                    #(
                        #[allow(dead_code)]
                        #field_vis fn #field_setter_names(&self, v: #field_types) {
                            let ptr = self.0.as_mut_ptr();
                            let u = ::cell_gc::traits::IntoHeapBase::into_heap(v);
                            unsafe {
                                (*ptr).#field_names = u;
                            }
                        }
                    )*

                    ///// Get all fields at once.
                    //pub fn get(&self) -> #name {
                    //    ::cell_gc::traits::IntoHeapBase::from_heap(self.0.ptr())
                    //}

                    #[allow(dead_code)]
                    pub fn as_mut_ptr(&self) -> *mut #storage_type_name #ty_generics {
                        self.0.as_mut_ptr()
                    }
                }
            };

            quote! {
                #storage_struct
                #into_heap
                #into_heap_allocation
                #ref_type
                #ref_type_into_heap
                #accessors
            }
        }
        syn::VariantData::Tuple(ref _fields) => {
            panic!("#[derive(IntoHeap)] does not support tuple structs");
        }
        syn::VariantData::Unit => {
            panic!("#[derive(IntoHeap)] does not support unit structs");
        }
    }
}

fn impl_into_heap_for_enum(ast: &syn::DeriveInput, variants: &[syn::Variant]) -> Tokens {
    let attrs = &ast.attrs;
    let name = &ast.ident;
    let name_str: &str = name.as_ref();
    let storage_type_name: Ident = Ident::from(name_str.to_string() + "Storage");
    let vis = &ast.vis;
    let (impl_generics, ty_generics, where_clause) = ast.generics.split_for_impl();
    let heap_lifetime = &ast.generics
        .lifetimes
        .first()
        .expect("lifetime parameter required")
        .lifetime;

    let variant_storage = variants.iter().map(|v| {
        let attrs = &v.attrs;
        let ident = &v.ident;
        if v.discriminant.is_some() {
            panic!("#[derive(IntoHeap)] does not support enum variants with discriminants");
        }
        match v.data {
            syn::VariantData::Struct(ref fields) => {
                let field_decls = fields.iter().map(|f| {
                    let field_attrs = &f.attrs;
                    let field_name = &f.ident;
                    let field_ty = &f.ty;
                    quote! {
                        #( #field_attrs )*
                        #field_name:
                            <#field_ty as ::cell_gc::traits::IntoHeapBase>::In
                    }
                });
                quote! {
                    #( #attrs )*
                    #ident { #( #field_decls ),* }
                }
            }
            syn::VariantData::Tuple(ref fields) => {
                let field_decls = fields.iter().map(|f| {
                    let field_ty = &f.ty;
                    let field_attrs = &f.attrs;
                    quote! {
                        #( #field_attrs )*
                        <#field_ty as ::cell_gc::traits::IntoHeapBase>::In
                    }
                });
                quote! {
                    #( #attrs )*
                    #ident( #( #field_decls ),* )
                }
            }
            syn::VariantData::Unit => {
                quote! { #( #attrs )* #ident }
            }
        }
    });
    let storage_enum = quote! {
        #( #attrs )*
        #vis enum #storage_type_name #impl_generics
            #where_clause
        {
            #( #variant_storage ),*
        }
    };

    let into_heap_arms = variants.iter().map(|v| {
        let ident = &v.ident;
        match v.data {
            syn::VariantData::Struct(ref fields) => {
                let field_names: &Vec<_> = &fields.iter().map(|f| &f.ident).collect();
                let field_names_1 = field_names;
                quote! {
                    #name::#ident { #(#field_names),* } => {
                        #storage_type_name::#ident {
                            #(
                                #field_names:
                                    ::cell_gc::traits::IntoHeapBase::into_heap(#field_names_1)
                            ),*
                        }
                    }
                }
            }
            syn::VariantData::Tuple(ref fields) => {
                let bindings: &Vec<Ident> = &(0..fields.len())
                    .map(|n| Ident::from(format!("x{}", n)))
                    .collect();
                quote! {
                    #name::#ident( #(#bindings),* ) => {
                        #storage_type_name::#ident(
                            #(
                                ::cell_gc::traits::IntoHeapBase::into_heap(#bindings)
                            ),*
                        )
                    }
                }
            }
            syn::VariantData::Unit => {
                quote! {
                    #name::#ident => #storage_type_name::#ident
                }
            }
        }
    });

    let from_heap_arms = variants.iter().map(|v| {
        let ident = &v.ident;
        match v.data {
            syn::VariantData::Struct(ref fields) => {
                let field_names: &Vec<_> = &fields.iter().map(|f| &f.ident).collect();
                let field_names_1 = field_names;
                let field_types = fields.iter().map(|f| &f.ty);
                quote! {
                    #storage_type_name::#ident { #(ref #field_names),* } => {
                        #name::#ident {
                            #(
                                #field_names:
                                    <#field_types as ::cell_gc::traits::IntoHeapBase>
                                    ::from_heap(#field_names_1)
                            ),*
                        }
                    }
                }
            }
            syn::VariantData::Tuple(ref fields) => {
                let field_types = fields.iter().map(|f| &f.ty);
                let bindings: &Vec<Ident> = &(0..fields.len())
                    .map(|n| Ident::from(format!("x{}", n)))
                    .collect();
                quote! {
                    #storage_type_name::#ident( #(ref #bindings),* ) => {
                        #name::#ident(
                            #(
                                <#field_types as ::cell_gc::traits::IntoHeapBase>
                                    ::from_heap(#bindings)
                            ),*
                        )
                    }
                }
            }
            syn::VariantData::Unit => {
                quote! {
                    #storage_type_name::#ident => #name::#ident
                }
            }
        }
    });

    let trace_arms = variants.iter().map(|v| {
        let ident = &v.ident;
        match v.data {
            syn::VariantData::Struct(ref fields) => {
                let field_names: &Vec<_> = &fields.iter().map(|f| &f.ident).collect();
                let field_types = fields.iter().map(|f| &f.ty);
                quote! {
                    #storage_type_name::#ident { #(ref #field_names),* } => {
                        #(
                            <#field_types as ::cell_gc::traits::IntoHeapBase>
                                ::trace(#field_names, tracer);
                        )*
                    }
                }
            }
            syn::VariantData::Tuple(ref fields) => {
                let field_types = fields.iter().map(|f| &f.ty);
                let bindings: &Vec<Ident> = &(0..fields.len())
                    .map(|n| Ident::from(format!("x{}", n)))
                    .collect();
                quote! {
                    #storage_type_name::#ident( #(ref #bindings),* ) => {
                        #(
                            <#field_types as ::cell_gc::traits::IntoHeapBase>
                                ::trace(#bindings, tracer);
                        )*
                    }
                }
            }
            syn::VariantData::Unit => {
                quote! { #storage_type_name::#ident => () }
            }
        }
    });

    let into_heap = quote! {
        impl #impl_generics
            ::cell_gc::traits::IntoHeapBase
            for #name #ty_generics
            #where_clause
        {
            type In = #storage_type_name #ty_generics;

            fn into_heap(self) -> Self::In {
                match self {
                    #( #into_heap_arms ),*
                }
            }

            unsafe fn from_heap(storage: &Self::In) -> Self {
                match *storage {
                    #( #from_heap_arms ),*
                }
            }

            unsafe fn trace<R>(storage: &Self::In, tracer: &mut R)
                where R: ::cell_gc::traits::Tracer
            {
                match *storage {
                    #( #trace_arms ),*
                }

                // Quiet unused variable warnings when `$(...)*` expands to
                // nothing.
                let _ = tracer;
            }
        }

        unsafe impl #impl_generics
            ::cell_gc::traits::IntoHeap<#heap_lifetime>
            for #name #ty_generics
            #where_clause
        {}
    };

    quote! {
        #storage_enum
        #into_heap
    }
}
<|endoftext|>


<commit_before>use std::cmp::max;
use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter, Write};
use std::path::{Path, PathBuf};

use grep::matcher::{Captures, Matcher};
use grep::regex::{RegexMatcher, RegexMatcherBuilder};
use grep::searcher::{
    BinaryDetection, Searcher, SearcherBuilder, Sink, SinkContext, SinkContextKind, SinkMatch,
};
use ignore::Walk;
use structopt::StructOpt;
use text_io::read;

#[derive(Debug, StructOpt)]
#[structopt(name = "fnr")]
/// Look for things, optionally replace them.
struct Opts {
    /// Run search in case insensitive mode
    #[structopt(short, long)]
    insensitive: bool,

    /// Print out replacements without actually performing them
    #[structopt(short, long)]
    dry_run: bool,

    #[structopt(short = "A", long, default_value = "0")]
    after: usize,
    #[structopt(short = "B", long, default_value = "0")]
    before: usize,
    #[structopt(short = "C", long, default_value = "0")]
    context: usize,

    #[structopt(short, long)]
    prompt: bool,

    /// What to search for.
    #[structopt(name = "FIND")]
    find: String,

    /// What to replace it with.
    #[structopt(name = "REPLACE")]
    replace: String,

    /// Locations to search. Current directory if not specified.
    #[structopt(name = "PATH", parse(from_os_str))]
    paths: Vec<PathBuf>,
}

#[derive(Debug, Clone)]
struct Line(u64, String);

#[derive(Debug)]
struct SearchMatch {
    line: Line,
    context_pre: Vec<Line>,
    context_post: Vec<Line>,
}

impl SearchMatch {
    fn as_change(&self, line: String) -> Change {
        return Change {
            line_number: self.line.0,
            new_line: line,
        };
    }

    fn display_change(&self, replacement_line: &str) {
        for line in &self.context_pre {
            print!("    {}: {}", line.0, line.1);
        }

        // TODO: Multiple matches on same line
        // TODO: Disable colors when not atty
        // TODO: Align line number with padding.
        print!("\x1B[31m-   {}: {}\x1B[0m", self.line.0, self.line.1);
        print!("\x1B[32m+   {}: {}\x1B[0m", self.line.0, replacement_line);

        for line in &self.context_post {
            print!("    {}: {}", line.0, line.1);
        }
    }
}

#[derive(Debug, Copy, Clone, PartialEq)]
enum MatchState {
    Init,
    Before,
    Match,
    After,
}

#[derive(Debug)]
struct SearchMatchCollector {
    state: MatchState,

    cur_context_pre: Vec<Line>,
    cur_context_post: Vec<Line>,
    cur_match_line: Option<Line>,

    matches: Vec<SearchMatch>,
}

impl SearchMatchCollector {
    fn new() -> SearchMatchCollector {
        return SearchMatchCollector {
            state: MatchState::Init,
            cur_match_line: None,
            cur_context_pre: Vec::new(),
            cur_context_post: Vec::new(),
            matches: Vec::new(),
        };
    }

    fn maybe_emit(&mut self) {
        if let Some(line) = &self.cur_match_line {
            // TODO: clones are wasteful here.
            let m = SearchMatch {
                line: line.clone(),
                context_pre: self.cur_context_pre.clone(),
                context_post: self.cur_context_post.clone(),
            };

            self.matches.push(m);
            self.reset();
        }
    }

    fn reset(&mut self) {
        self.state = MatchState::Init;

        self.cur_match_line = None;
        self.cur_context_pre.clear();
        self.cur_context_post.clear();
    }

    #[inline]
    fn transition(&mut self, next: MatchState) {
        match (self.state, next) {
            // Beginning a new match or ending a previous one
            (MatchState::Match, MatchState::Before)       // Have before context, no after context
            | (MatchState::Match, MatchState::Match)      // No before context, no after context
            | (MatchState::After, MatchState::Before)     // Have before context, have after context
            | (MatchState::After, MatchState::Match) => { // Have after context, no before context
                self.maybe_emit();
            }

            (_prev, next) => {
                self.state = next;
            }
        }
    }

    fn collect(&mut self) -> &Vec<SearchMatch> {
        self.maybe_emit();
        return &self.matches;
    }
}

impl Drop for SearchMatchCollector {
    fn drop(&mut self) {
        self.maybe_emit();
    }
}

impl Sink for SearchMatchCollector {
    type Error = std::io::Error;

    fn matched(
        &mut self,
        _searcher: &Searcher,
        mat: &SinkMatch<'_>,
    ) -> Result<bool, std::io::Error> {
        self.transition(MatchState::Match);

        // TODO: handle errors
        let line = Line(
            mat.line_number().expect("grab line number"),
            String::from_utf8_lossy(mat.bytes()).to_string(),
        );

        self.cur_match_line = Some(line);

        Ok(true)
    }

    fn context(
        &mut self,
        _searcher: &Searcher,
        ctx: &SinkContext<'_>,
    ) -> Result<bool, std::io::Error> {
        // TODO: handle errors.
        // TODO: can we avoid the additional .to_string() by storing a COW instead?
        let line = Line(
            ctx.line_number().expect("grab line number"),
            String::from_utf8_lossy(ctx.bytes()).to_string(),
        );

        match ctx.kind() {
            &SinkContextKind::Before => {
                self.transition(MatchState::Before);
                self.cur_context_pre.push(line);
            }
            &SinkContextKind::After => {
                self.transition(MatchState::After);
                self.cur_context_post.push(line);
            }
            &SinkContextKind::Other => {}
        }

        Ok(true)
    }
}

trait Replacer {
    fn replace(&self, input: &str) -> Option<String>;
}

struct RegexReplacer<'a> {
    matcher: &'a RegexMatcher,
    template: &'a str,
}

impl<'a> Replacer for RegexReplacer<'a> {
    fn replace(&self, input: &str) -> Option<String> {
        let mut caps = self.matcher.new_captures().unwrap();
        let mut dst = vec![];

        self.matcher
            .replace_with_captures(input.as_bytes(), &mut caps, &mut dst, |caps, dst| {
                caps.interpolate(
                    |name| self.matcher.capture_index(name),
                    input.as_bytes(),
                    self.template.as_bytes(),
                    dst,
                );
                false
            })
            .unwrap();

        Some(String::from_utf8_lossy(&dst).to_string())
    }
}

#[derive(Debug)]
struct Change {
    line_number: u64,
    // TODO: use slice here
    new_line: String,
}

struct SearchMatchProcessor<'a> {
    replacer: &'a dyn Replacer,
    replacement_decider: ReplacementDecider,

    total_matches: usize,
    total_replacements: usize,
}

impl<'a> SearchMatchProcessor<'a> {
    fn new(
        replacer: &'a dyn Replacer,
        replacement_decider: ReplacementDecider,
    ) -> SearchMatchProcessor<'a> {
        return SearchMatchProcessor {
            replacer,
            replacement_decider,

            total_matches: 0,
            total_replacements: 0,
        };
    }

    fn handle_path(
        &mut self,
        path: &Path,
        matches: &Vec<SearchMatch>,
    ) -> Result<(), std::io::Error> {
        if matches.is_empty() {
            return Ok(());
        }
        println!(
            "--- {}: {} match{}",
            path.display(),
            matches.len(),
            if matches.len() == 1 { "" } else { "es" }
        );
        self.total_matches += matches.len();
        self.replacement_decider.reset_local_decision();

        let mut change_list = vec![];
        for m in matches {
            let replacement = if let Some(r) = self.replacer.replace(&m.line.1) {
                r
            } else {
                println!("TODO: replacer failed on line");
                continue;
            };

            let decision = self.replacement_decider.decide(m, &replacement);
            let change = match decision {
                ReplacementDecision::Accept => m.as_change(replacement),
                ReplacementDecision::Ignore => continue,
                ReplacementDecision::Edit => {
                    let mut line = read_input("Replace with [^D to skip] ")?;
                    if line == "" {
                        println!("... skipped ...");
                        continue;
                    }

                    line.push('\n');
                    m.display_change(&line);
                    println!("--");
                    m.as_change(line)
                }
            };

            change_list.push(change);
            self.total_replacements += 1;
        }

        // TODO: For --prompt, confirm before applying here
        if !change_list.is_empty() {
            self.apply_changes(path, &mut change_list)?;
        }

        Ok(())
    }

    fn apply_changes(&self, path: &Path, mut changes: &[Change]) -> Result<(), std::io::Error> {
        /*println!(
            "{}: {} change{}",
            path.display(),
            changes.len(),
            if changes.len() != 1 { "s" } else { "" }
        );*/
        let dst_path = path.with_extension("~");
        let src = File::open(path)?;
        let dst = File::create(&dst_path)?;

        let mut reader = BufReader::new(src);
        let mut writer = BufWriter::new(dst);

        let mut line_num = 0;
        loop {
            let mut line = String::new();
            let bytes_read = reader.read_line(&mut line)?;
            // EOF reached
            if bytes_read == 0 {
                break;
            }

            line_num += 1;

            if !changes.is_empty() && changes[0].line_number == line_num {
                writer.write(changes[0].new_line.as_bytes())?;
                changes = &changes[1..];
            } else {
                writer.write(line.as_bytes())?;
            }
        }

        std::fs::rename(dst_path, path)?;
        return Ok(());
    }

    fn finalize(&self) {
        println!(
            "All done. Replaced {} of {} matches",
            self.total_replacements, self.total_matches
        );
    }
}

#[derive(Debug, Copy, Clone)]
enum ReplacementDecision {
    Accept,
    Ignore,
    Edit,
}

fn read_input(prompt: &str) -> Result<String, std::io::Error> {
    print!("{}", prompt);
    std::io::stdout().flush()?;

    Ok(read!("{}\n"))
}

struct ReplacementDecider {
    should_prompt: bool,
    global_decision: Option<ReplacementDecision>,
    local_decision: Option<ReplacementDecision>,
}

impl ReplacementDecider {
    fn with_prompt() -> ReplacementDecider {
        return ReplacementDecider {
            should_prompt: true,
            global_decision: None,
            local_decision: None,
        };
    }

    fn constantly(decision: ReplacementDecision) -> ReplacementDecider {
        return ReplacementDecider {
            should_prompt: false,
            global_decision: Some(decision),
            local_decision: None,
        };
    }

    fn reset_local_decision(&mut self) {
        self.local_decision = None;
    }

    fn decide(&mut self, match_: &SearchMatch, replacement: &str) -> ReplacementDecision {
        if let Some(decision) = self.global_decision {
            return decision;
        } else if let Some(decision) = self.local_decision {
            return decision;
        }

        if !self.should_prompt {
            panic!("invalid state: no decision, but should not prompt");
        }

        // TODO: Should we display this always?
        match_.display_change(replacement);
        return self.prompt_for_decision();
    }

    fn prompt_for_decision(&mut self) -> ReplacementDecision {
        loop {
            let line = read_input("Stage this replacement [y,n,q,a,e,d,?] ").unwrap();

            return match line.as_str() {
                "y" => ReplacementDecision::Accept,
                "n" => ReplacementDecision::Ignore,
                "q" => {
                    self.global_decision = Some(ReplacementDecision::Ignore);
                    ReplacementDecision::Ignore
                }
                "a" => {
                    self.local_decision = Some(ReplacementDecision::Ignore);
                    ReplacementDecision::Accept
                }
                "d" => {
                    self.local_decision = Some(ReplacementDecision::Ignore);
                    ReplacementDecision::Accept
                }
                "e" => {
                    self.local_decision = Some(ReplacementDecision::Edit);
                    ReplacementDecision::Edit
                }

                "?" | _ => {
                    println!(
                        "\x1B[31m
Y - replace this line
n - do not replace this line
q - quit; do not replace this line or any remaining ones
a - replace this line and all remaining ones in this file
d - do not replace this line nor any remaining ones in this file
e - edit this replacement
? - show help
\x1B[0m"
                    );
                    continue;
                }
            };
        }
    }
}

fn main() {
    let opts: Opts = Opts::from_args();
    println!("Parsed opts: {:?}", opts);

    let matcher = RegexMatcherBuilder::new()
        .case_insensitive(opts.insensitive)
        .build(&opts.find)
        .expect("bad pattern");

    let mut searcher = SearcherBuilder::new()
        .binary_detection(BinaryDetection::quit(0x00))
        .line_number(true)
        .before_context(max(opts.context, opts.before))
        .after_context(max(opts.context, opts.after))
        .build();

    // TODO: Confirm that template does not reference more capture groups than exist.
    let replacer = RegexReplacer {
        matcher: &matcher,
        template: &opts.replace,
    };

    if opts.dry_run && opts.prompt {
        println!("WARN: --prompt does not make sense with --dry-run, skipping");
    }

    let replacement_decider = if opts.dry_run {
        ReplacementDecider::constantly(ReplacementDecision::Ignore)
    } else if opts.prompt {
        ReplacementDecider::with_prompt()
    } else {
        ReplacementDecider::constantly(ReplacementDecision::Accept)
    };

    let mut proc = SearchMatchProcessor::new(&replacer, replacement_decider);

    let paths = if opts.paths.is_empty() {
        vec![Path::new("./").to_owned()]
    } else {
        opts.paths
    };

    for path in paths {
        // TODO: Support files as well as directories here
        for dir_entry in Walk::new(path) {
            let dir_entry = dir_entry.unwrap();
            let path = dir_entry.path();
            if !path.is_file() {
                continue;
            }

            let mut sink = SearchMatchCollector::new();

            searcher
                .search_path(&matcher, dir_entry.path(), &mut sink)
                .expect("search failed");

            let matches = sink.collect();
            proc.handle_path(path, matches).expect("handle path");
        }
    }

    if opts.dry_run {
        println!("--dry-run enabled, no changes applied.");
    }

    proc.finalize();
}
<commit_msg>Replace all instances on line<commit_after>use std::cmp::max;
use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter, Write};
use std::path::{Path, PathBuf};

use grep::matcher::{Captures, Matcher};
use grep::regex::{RegexMatcher, RegexMatcherBuilder};
use grep::searcher::{
    BinaryDetection, Searcher, SearcherBuilder, Sink, SinkContext, SinkContextKind, SinkMatch,
};
use ignore::Walk;
use structopt::StructOpt;
use text_io::read;

#[derive(Debug, StructOpt)]
#[structopt(name = "fnr")]
/// Look for things, optionally replace them.
struct Opts {
    /// Run search in case insensitive mode
    #[structopt(short, long)]
    insensitive: bool,

    /// Print out replacements without actually performing them
    #[structopt(short, long)]
    dry_run: bool,

    #[structopt(short = "A", long, default_value = "0")]
    after: usize,
    #[structopt(short = "B", long, default_value = "0")]
    before: usize,
    #[structopt(short = "C", long, default_value = "0")]
    context: usize,

    #[structopt(short, long)]
    prompt: bool,

    /// What to search for.
    #[structopt(name = "FIND")]
    find: String,

    /// What to replace it with.
    #[structopt(name = "REPLACE")]
    replace: String,

    /// Locations to search. Current directory if not specified.
    #[structopt(name = "PATH", parse(from_os_str))]
    paths: Vec<PathBuf>,
}

#[derive(Debug, Clone)]
struct Line(u64, String);

#[derive(Debug)]
struct SearchMatch {
    line: Line,
    context_pre: Vec<Line>,
    context_post: Vec<Line>,
}

impl SearchMatch {
    fn as_change(&self, line: String) -> Change {
        return Change {
            line_number: self.line.0,
            new_line: line,
        };
    }

    fn display_change(&self, replacement_line: &str) {
        for line in &self.context_pre {
            print!("    {}: {}", line.0, line.1);
        }

        // TODO: Multiple matches on same line
        // TODO: Disable colors when not atty
        // TODO: Align line number with padding.
        print!("\x1B[31m-   {}: {}\x1B[0m", self.line.0, self.line.1);
        print!("\x1B[32m+   {}: {}\x1B[0m", self.line.0, replacement_line);

        for line in &self.context_post {
            print!("    {}: {}", line.0, line.1);
        }
    }
}

#[derive(Debug, Copy, Clone, PartialEq)]
enum MatchState {
    Init,
    Before,
    Match,
    After,
}

#[derive(Debug)]
struct SearchMatchCollector {
    state: MatchState,

    cur_context_pre: Vec<Line>,
    cur_context_post: Vec<Line>,
    cur_match_line: Option<Line>,

    matches: Vec<SearchMatch>,
}

impl SearchMatchCollector {
    fn new() -> SearchMatchCollector {
        return SearchMatchCollector {
            state: MatchState::Init,
            cur_match_line: None,
            cur_context_pre: Vec::new(),
            cur_context_post: Vec::new(),
            matches: Vec::new(),
        };
    }

    fn maybe_emit(&mut self) {
        if let Some(line) = &self.cur_match_line {
            // TODO: clones are wasteful here.
            let m = SearchMatch {
                line: line.clone(),
                context_pre: self.cur_context_pre.clone(),
                context_post: self.cur_context_post.clone(),
            };

            self.matches.push(m);
            self.reset();
        }
    }

    fn reset(&mut self) {
        self.state = MatchState::Init;

        self.cur_match_line = None;
        self.cur_context_pre.clear();
        self.cur_context_post.clear();
    }

    #[inline]
    fn transition(&mut self, next: MatchState) {
        match (self.state, next) {
            // Beginning a new match or ending a previous one
            (MatchState::Match, MatchState::Before)       // Have before context, no after context
            | (MatchState::Match, MatchState::Match)      // No before context, no after context
            | (MatchState::After, MatchState::Before)     // Have before context, have after context
            | (MatchState::After, MatchState::Match) => { // Have after context, no before context
                self.maybe_emit();
            }

            (_prev, next) => {
                self.state = next;
            }
        }
    }

    fn collect(&mut self) -> &Vec<SearchMatch> {
        self.maybe_emit();
        return &self.matches;
    }
}

impl Drop for SearchMatchCollector {
    fn drop(&mut self) {
        self.maybe_emit();
    }
}

impl Sink for SearchMatchCollector {
    type Error = std::io::Error;

    fn matched(
        &mut self,
        _searcher: &Searcher,
        mat: &SinkMatch<'_>,
    ) -> Result<bool, std::io::Error> {
        self.transition(MatchState::Match);

        // TODO: handle errors
        let line = Line(
            mat.line_number().expect("grab line number"),
            String::from_utf8_lossy(mat.bytes()).to_string(),
        );

        self.cur_match_line = Some(line);

        Ok(true)
    }

    fn context(
        &mut self,
        _searcher: &Searcher,
        ctx: &SinkContext<'_>,
    ) -> Result<bool, std::io::Error> {
        // TODO: handle errors.
        // TODO: can we avoid the additional .to_string() by storing a COW instead?
        let line = Line(
            ctx.line_number().expect("grab line number"),
            String::from_utf8_lossy(ctx.bytes()).to_string(),
        );

        match ctx.kind() {
            &SinkContextKind::Before => {
                self.transition(MatchState::Before);
                self.cur_context_pre.push(line);
            }
            &SinkContextKind::After => {
                self.transition(MatchState::After);
                self.cur_context_post.push(line);
            }
            &SinkContextKind::Other => {}
        }

        Ok(true)
    }
}

trait Replacer {
    fn replace(&self, input: &str) -> Option<String>;
}

struct RegexReplacer<'a> {
    matcher: &'a RegexMatcher,
    template: &'a str,
}

impl<'a> Replacer for RegexReplacer<'a> {
    fn replace(&self, input: &str) -> Option<String> {
        let mut caps = self.matcher.new_captures().unwrap();
        let mut dst = vec![];

        self.matcher
            .replace_with_captures(input.as_bytes(), &mut caps, &mut dst, |caps, dst| {
                caps.interpolate(
                    |name| self.matcher.capture_index(name),
                    input.as_bytes(),
                    self.template.as_bytes(),
                    dst,
                );
                true
            })
            .unwrap();

        Some(String::from_utf8_lossy(&dst).to_string())
    }
}

#[derive(Debug)]
struct Change {
    line_number: u64,
    // TODO: use slice here
    new_line: String,
}

struct SearchMatchProcessor<'a> {
    replacer: &'a dyn Replacer,
    replacement_decider: ReplacementDecider,

    total_matches: usize,
    total_replacements: usize,
}

impl<'a> SearchMatchProcessor<'a> {
    fn new(
        replacer: &'a dyn Replacer,
        replacement_decider: ReplacementDecider,
    ) -> SearchMatchProcessor<'a> {
        return SearchMatchProcessor {
            replacer,
            replacement_decider,

            total_matches: 0,
            total_replacements: 0,
        };
    }

    fn handle_path(
        &mut self,
        path: &Path,
        matches: &Vec<SearchMatch>,
    ) -> Result<(), std::io::Error> {
        if matches.is_empty() {
            return Ok(());
        }
        println!(
            "--- {}: {} match{}",
            path.display(),
            matches.len(),
            if matches.len() == 1 { "" } else { "es" }
        );
        self.total_matches += matches.len();
        self.replacement_decider.reset_local_decision();

        let mut change_list = vec![];
        for m in matches {
            let replacement = if let Some(r) = self.replacer.replace(&m.line.1) {
                r
            } else {
                println!("TODO: replacer failed on line");
                continue;
            };

            let decision = self.replacement_decider.decide(m, &replacement);
            let change = match decision {
                ReplacementDecision::Accept => m.as_change(replacement),
                ReplacementDecision::Ignore => continue,
                ReplacementDecision::Edit => {
                    let mut line = read_input("Replace with [^D to skip] ")?;
                    if line == "" {
                        println!("... skipped ...");
                        continue;
                    }

                    line.push('\n');
                    m.display_change(&line);
                    println!("--");
                    m.as_change(line)
                }
            };

            change_list.push(change);
            self.total_replacements += 1;
        }

        // TODO: For --prompt, confirm before applying here
        if !change_list.is_empty() {
            self.apply_changes(path, &mut change_list)?;
        }

        Ok(())
    }

    fn apply_changes(&self, path: &Path, mut changes: &[Change]) -> Result<(), std::io::Error> {
        /*println!(
            "{}: {} change{}",
            path.display(),
            changes.len(),
            if changes.len() != 1 { "s" } else { "" }
        );*/
        let dst_path = path.with_extension("~");
        let src = File::open(path)?;
        let dst = File::create(&dst_path)?;

        let mut reader = BufReader::new(src);
        let mut writer = BufWriter::new(dst);

        let mut line_num = 0;
        loop {
            let mut line = String::new();
            let bytes_read = reader.read_line(&mut line)?;
            // EOF reached
            if bytes_read == 0 {
                break;
            }

            line_num += 1;

            if !changes.is_empty() && changes[0].line_number == line_num {
                writer.write(changes[0].new_line.as_bytes())?;
                changes = &changes[1..];
            } else {
                writer.write(line.as_bytes())?;
            }
        }

        std::fs::rename(dst_path, path)?;
        return Ok(());
    }

    fn finalize(&self) {
        println!(
            "All done. Replaced {} of {} matches",
            self.total_replacements, self.total_matches
        );
    }
}

#[derive(Debug, Copy, Clone)]
enum ReplacementDecision {
    Accept,
    Ignore,
    Edit,
}

fn read_input(prompt: &str) -> Result<String, std::io::Error> {
    print!("{}", prompt);
    std::io::stdout().flush()?;

    Ok(read!("{}\n"))
}

struct ReplacementDecider {
    should_prompt: bool,
    global_decision: Option<ReplacementDecision>,
    local_decision: Option<ReplacementDecision>,
}

impl ReplacementDecider {
    fn with_prompt() -> ReplacementDecider {
        return ReplacementDecider {
            should_prompt: true,
            global_decision: None,
            local_decision: None,
        };
    }

    fn constantly(decision: ReplacementDecision) -> ReplacementDecider {
        return ReplacementDecider {
            should_prompt: false,
            global_decision: Some(decision),
            local_decision: None,
        };
    }

    fn reset_local_decision(&mut self) {
        self.local_decision = None;
    }

    fn decide(&mut self, match_: &SearchMatch, replacement: &str) -> ReplacementDecision {
        if let Some(decision) = self.global_decision {
            return decision;
        } else if let Some(decision) = self.local_decision {
            return decision;
        }

        if !self.should_prompt {
            panic!("invalid state: no decision, but should not prompt");
        }

        // TODO: Should we display this always?
        match_.display_change(replacement);
        return self.prompt_for_decision();
    }

    fn prompt_for_decision(&mut self) -> ReplacementDecision {
        loop {
            let line = read_input("Stage this replacement [y,n,q,a,e,d,?] ").unwrap();

            return match line.as_str() {
                "y" => ReplacementDecision::Accept,
                "n" => ReplacementDecision::Ignore,
                "q" => {
                    self.global_decision = Some(ReplacementDecision::Ignore);
                    ReplacementDecision::Ignore
                }
                "a" => {
                    self.local_decision = Some(ReplacementDecision::Ignore);
                    ReplacementDecision::Accept
                }
                "d" => {
                    self.local_decision = Some(ReplacementDecision::Ignore);
                    ReplacementDecision::Accept
                }
                "e" => {
                    self.local_decision = Some(ReplacementDecision::Edit);
                    ReplacementDecision::Edit
                }

                "?" | _ => {
                    println!(
                        "\x1B[31m
Y - replace this line
n - do not replace this line
q - quit; do not replace this line or any remaining ones
a - replace this line and all remaining ones in this file
d - do not replace this line nor any remaining ones in this file
e - edit this replacement
? - show help
\x1B[0m"
                    );
                    continue;
                }
            };
        }
    }
}

fn main() {
    let opts: Opts = Opts::from_args();
    println!("Parsed opts: {:?}", opts);

    let matcher = RegexMatcherBuilder::new()
        .case_insensitive(opts.insensitive)
        .build(&opts.find)
        .expect("bad pattern");

    let mut searcher = SearcherBuilder::new()
        .binary_detection(BinaryDetection::quit(0x00))
        .line_number(true)
        .before_context(max(opts.context, opts.before))
        .after_context(max(opts.context, opts.after))
        .build();

    // TODO: Confirm that template does not reference more capture groups than exist.
    let replacer = RegexReplacer {
        matcher: &matcher,
        template: &opts.replace,
    };

    if opts.dry_run && opts.prompt {
        println!("WARN: --prompt does not make sense with --dry-run, skipping");
    }

    let replacement_decider = if opts.dry_run {
        ReplacementDecider::constantly(ReplacementDecision::Ignore)
    } else if opts.prompt {
        ReplacementDecider::with_prompt()
    } else {
        ReplacementDecider::constantly(ReplacementDecision::Accept)
    };

    let mut proc = SearchMatchProcessor::new(&replacer, replacement_decider);

    let paths = if opts.paths.is_empty() {
        vec![Path::new("./").to_owned()]
    } else {
        opts.paths
    };

    for path in paths {
        // TODO: Support files as well as directories here
        for dir_entry in Walk::new(path) {
            let dir_entry = dir_entry.unwrap();
            let path = dir_entry.path();
            if !path.is_file() {
                continue;
            }

            let mut sink = SearchMatchCollector::new();

            searcher
                .search_path(&matcher, dir_entry.path(), &mut sink)
                .expect("search failed");

            let matches = sink.collect();
            proc.handle_path(path, matches).expect("handle path");
        }
    }

    if opts.dry_run {
        println!("--dry-run enabled, no changes applied.");
    }

    proc.finalize();
}
<|endoftext|>


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

use std::vec;

/**
 * The Digest trait specifies an interface common to digest functions, such as SHA-1 and the SHA-2
 * family of digest functions.
 */
pub trait Digest {
    /**
     * Provide message data.
     *
     * # Arguments
     *
     * * input - A vector of message data
     */
    fn input(&mut self, input: &[u8]);

    /**
     * Retrieve the digest result. This method may be called multiple times.
     *
     * # Arguments
     *
     * * out - the vector to hold the result. Must be large enough to contain output_bits().
     */
    fn result(&mut self, out: &mut [u8]);

    /**
     * Reset the digest. This method must be called after result() and before supplying more
     * data.
     */
    fn reset(&mut self);

    /**
     * Get the output size in bits.
     */
    fn output_bits(&self) -> uint;

    /**
     * Get the output size in bytes.
     */
    fn output_bytes(&self) -> uint {
        return (self.output_bits() + 7) / 8;
    }

    /**
     * Get the block size in bytes.
     */
    fn block_size(&self) -> uint;

    /**
     * Convenience function that feeds a string into a digest.
     *
     * # Arguments
     *
     * * `input` The string to feed into the digest
     */
    fn input_str(&mut self, input: &str) {
        self.input(input.as_bytes());
    }

    /**
     * Convenience function that retrieves the result of a digest as a
     * ~str in hexadecimal format.
     */
    fn result_str(&mut self) -> ~str {
        let mut buf = vec::from_elem((self.output_bits()+7)/8, 0u8);
        self.result(buf);
        return to_hex(buf);
    }
}

fn to_hex(rr: &[u8]) -> ~str {
    let mut s = ~"";
    for b in rr.iter() {
        let hex = (*b as uint).to_str_radix(16u);
        if hex.len() == 1 {
            s.push_char('0');
        }
        s.push_str(hex);
    }
    return s;
}
<commit_msg>Explicitly import ToStrRadix now that its no longer in the prelude<commit_after>// Copyright 2012-2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::num::ToStrRadix;
use std::vec;

/**
 * The Digest trait specifies an interface common to digest functions, such as SHA-1 and the SHA-2
 * family of digest functions.
 */
pub trait Digest {
    /**
     * Provide message data.
     *
     * # Arguments
     *
     * * input - A vector of message data
     */
    fn input(&mut self, input: &[u8]);

    /**
     * Retrieve the digest result. This method may be called multiple times.
     *
     * # Arguments
     *
     * * out - the vector to hold the result. Must be large enough to contain output_bits().
     */
    fn result(&mut self, out: &mut [u8]);

    /**
     * Reset the digest. This method must be called after result() and before supplying more
     * data.
     */
    fn reset(&mut self);

    /**
     * Get the output size in bits.
     */
    fn output_bits(&self) -> uint;

    /**
     * Get the output size in bytes.
     */
    fn output_bytes(&self) -> uint {
        return (self.output_bits() + 7) / 8;
    }

    /**
     * Get the block size in bytes.
     */
    fn block_size(&self) -> uint;

    /**
     * Convenience function that feeds a string into a digest.
     *
     * # Arguments
     *
     * * `input` The string to feed into the digest
     */
    fn input_str(&mut self, input: &str) {
        self.input(input.as_bytes());
    }

    /**
     * Convenience function that retrieves the result of a digest as a
     * ~str in hexadecimal format.
     */
    fn result_str(&mut self) -> ~str {
        let mut buf = vec::from_elem((self.output_bits()+7)/8, 0u8);
        self.result(buf);
        return to_hex(buf);
    }
}

fn to_hex(rr: &[u8]) -> ~str {
    let mut s = ~"";
    for b in rr.iter() {
        let hex = (*b as uint).to_str_radix(16u);
        if hex.len() == 1 {
            s.push_char('0');
        }
        s.push_str(hex);
    }
    return s;
}
<|endoftext|>