format_env 1.0.1

use proc_macro2::{Literal, TokenStream, TokenTree};
use std::iter::{once, FromIterator};

#[proc_macro]
pub fn format_env(ts: proc_macro::TokenStream) -> proc_macro::TokenStream {
    format_env_impl(ts.into()).into()
}

fn format_env_impl(ts: TokenStream) -> TokenStream {
    let mut tsi = ts.into_iter();
    let tt = tsi
        .next()
        .expect("format_env needs a string literal argument");
    if let Some(ntt) = tsi.next() {
        let ntts = ntt.to_string();
        let first_char = ntts.get(..1).unwrap_or("input");
        panic!("Unexpected '{}' while expanding format_env", first_char);
    }
    let tts = tt.to_string();
    let (open, mut s, close) = match tt {
        TokenTree::Literal(_) if tts.starts_with('"') => {
            (&tts[..1], &tts[1..tts.len() - 1], &tts[tts.len() - 1..])
        }
        TokenTree::Literal(_) if tts.starts_with("r#") => extract_raw_string(&tts),
        _ => panic!("format_env expects a string literal"),
    };
    let mut subbed = String::from(open);
    while let Some(i) = s.find('$') {
        subbed.push_str(&s[0..i]);
        s = &s[i..];
        if let Some(c) = s.chars().nth(1) {
            match c {
                '(' => {
                    let end = s.find(')');
                    let var_name = end
                        .and_then(|end| s.get(2..end))
                        .expect("Reached end of file while parsing environment variable");
                    let var_val = std::env::var(var_name).unwrap_or_else(|e| {
                        panic!(
                            "Failed to substitute environment variable \"{}\": {}",
                            var_name, e
                        )
                    });
                    subbed.push_str(&var_val);
                    s = &s[end.unwrap() + 1..];
                }
                '$' => {
                    subbed.push('$');
                    s = &s[2..];
                }
                _ => {
                    subbed.push('$');
                    s = &s[1..];
                }
            }
        }
    }
    subbed.push_str(s);
    subbed.push_str(close);
    subbed.parse().unwrap()
}

fn extract_raw_string(s: &str) -> (&str, &str, &str) {
    let mut chars = s.chars();
    assert_eq!(chars.next().unwrap(), 'r');
    let mut i = 1;
    while let Some('#') = chars.next() {
        i += 1
    }
    (&s[..i + 1], &s[i + 1..s.len() - i], &s[s.len() - i..])
}

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

    #[test]
    fn test_substitution() {
        let tokens_in: TokenStream = r#""$(CARGO_PKG_NAME)""#.parse().unwrap();
        let tokens_out: TokenStream = r#""format_env""#.parse().unwrap();
        assert_eq!(format_env_impl(tokens_in).to_string(), tokens_out.to_string());
        let tokens_in: TokenStream = r#""[$(CARGO_PKG_NAME)], [$(CARGO_PKG_LICENSE)]""#.parse().unwrap();
        let tokens_out: TokenStream = r#""[format_env], [MIT]""#.parse().unwrap();
        assert_eq!(format_env_impl(tokens_in).to_string(), tokens_out.to_string());
    }

    #[test]
    fn test_no_substitution() {
        let tokens: TokenStream = r#""lorem ipsum""#.parse().unwrap();
        assert_eq!(format_env_impl(tokens.clone()).to_string(), tokens.to_string());
    }

    #[test]
    fn test_escape() {
        let tokens_in: TokenStream = r#""$$(CARGO_PKG_NAME)""#.parse().unwrap();
        let tokens_out: TokenStream = r#""$(CARGO_PKG_NAME)""#.parse().unwrap();
        assert_eq!(format_env_impl(tokens_in).to_string(), tokens_out.to_string());
    }

    #[test]
    fn test_raw_string() {
        let tokens_in: TokenStream = r##"r#"$(CARGO_PKG_NAME)"#"##.parse().unwrap();
        let tokens_out: TokenStream = r##"r#"format_env"#"##.parse().unwrap();
        assert_eq!(format_env_impl(tokens_in).to_string(), tokens_out.to_string());
        let tokens_in: TokenStream = r#####"r####"$(CARGO_PKG_NAME)"####"#####.parse().unwrap();
        let tokens_out: TokenStream = r#####"r####"format_env"####"#####.parse().unwrap();
        assert_eq!(format_env_impl(tokens_in).to_string(), tokens_out.to_string());
        let tokens_in: TokenStream = r##"r#"
            $(CARGO_PKG_NAME)
        "#"##
            .parse()
            .unwrap();
        let tokens_out: TokenStream = r##"r#"
            format_env
        "#"##
            .parse()
            .unwrap();
        assert_eq!(format_env_impl(tokens_in).to_string(), tokens_out.to_string());
    }

    #[test]
    #[should_panic]
    fn test_reject_zero_tokens() {
        let tokens: TokenStream = "".parse().unwrap();
        format_env_impl(tokens);
    }

    #[test]
    #[should_panic]
    fn test_reject_two_tokens() {
        let tokens: TokenStream = r#""[$(CARGO_PKG_NAME)]" "[$(CARGO_PKG_LICENSE)]""#.parse().unwrap();
        format_env_impl(tokens);
    }

    #[test]
    #[should_panic]
    fn test_reject_unmatched() {
        let tokens: TokenStream = r#""$(CARGO_PKG_NAME""#.parse().unwrap();
        format_env_impl(tokens);
    }

    #[test]
    #[should_panic]
    fn test_reject_group() {
        let tokens: TokenStream = r#"{"abc" "def"}"#.parse().unwrap();
        format_env_impl(tokens);
    }

    #[test]
    #[should_panic]
    fn test_reject_ident() {
        let tokens: TokenStream = "foobar".parse().unwrap();
        format_env_impl(tokens);
    }

    #[test]
    #[should_panic]
    fn test_reject_punct() {
        let tokens: TokenStream = "+".parse().unwrap();
        format_env_impl(tokens);
    }

    #[test]
    #[should_panic]
    fn test_reject_number() {
        let tokens: TokenStream = "3.14159".parse().unwrap();
        format_env_impl(tokens);
    }
}