eval_macro_internal/

lib.rs

#![allow(clippy::panic)]
#![allow(clippy::expect_used)]

use proc_macro2::Delimiter;
use proc_macro2::LineColumn;
use proc_macro2::TokenStream;
use proc_macro2::TokenTree;
use quote::quote;
use std::collections::hash_map::DefaultHasher;
use std::fs::File;
use std::fs;
use std::hash::Hash;
use std::hash::Hasher;
use std::io::Write;
use std::path::Path;
use std::path::PathBuf;
use std::process::Command;

// =================
// === Constants ===
// =================

/// Set to 'true' to enable debug prints.
const DEBUG: bool = false;

const NAME: &str = "eval_macro";

/// Rust keywords for special handling. This is not needed for this macro to work, it is only used
/// to make `IntelliJ` / `RustRover` work correctly, as their `TokenStream` spans are incorrect.
const KEYWORDS: &[&str] = &[
    "as", "async", "await", "break", "const", "continue", "crate", "dyn", "else", "enum",
    "extern", "false", "fn", "for", "if", "impl", "in", "let", "loop", "match", "mod", "move",
    "mut", "pub", "ref", "return", "self", "Self", "static", "struct", "super", "trait", "true",
    "type", "unsafe", "use", "where", "while", "abstract", "become", "box", "do", "final", "macro",
    "override", "priv", "typeof", "unsized", "virtual", "yield", "try",
];

/// Common functions. After this lib stabilizes, they should be imported from this library, not
/// injected.
const PRELUDE: &str = "
    macro_rules! write_ln {
        ($target:ident, $($ts:tt)*) => {
            $target.push_str(&format!( $($ts)* ));
            $target.push_str(\"\n\");
        };
    }

    fn sum_combinations(n: usize) -> Vec<Vec<usize>> {
        let mut result = Vec::new();

        fn generate(n: usize, current: Vec<usize>, result: &mut Vec<Vec<usize>>) {
            if n == 0 {
                if current.len() > 1 {
                    result.push(current);
                }
                return;
            }

            for i in 1..=n {
                let mut next = current.clone();
                next.push(i);
                generate(n - i, next, result);
            }
        }

        generate(n, vec![], &mut result);
        result
    }
";

// ============
// === Path ===
// ============

/// Output directory for projects generated by this macro.
fn get_output_dir() -> PathBuf {
    let home_dir = std::env::var("HOME")
        .expect("HOME environment variable not set — this is required to locate ~/.cargo.");

    let eval_macro_dir = PathBuf::from(home_dir)
        .join(".cargo")
        .join(NAME);

    if !eval_macro_dir.exists() {
        fs::create_dir_all(&eval_macro_dir)
            .expect("Failed to create ~/.cargo/eval_macro directory.");
    }

    eval_macro_dir
}

// ==========================
// === Project Management ===
// ==========================

#[derive(Debug)]
struct CargoConfig {
    edition: Option<String>,
    dependencies: Vec<String>
}

impl CargoConfig {
    fn new(dependencies: Vec<String>) -> Self {
        Self {
            edition: None,
            dependencies
        }
    }

    fn print(&self) -> String {
        let edition = self.edition.as_ref().map_or("2024", |t| t.as_str());
        let dependencies = self.dependencies.join("\n");
        format!("
            [package]
            name = \"eval_project\"
            version = \"1.0.0\"
            edition = \"{edition}\"

            [dependencies]
            {dependencies}
        ")
    }
}

fn project_name_from_input(input_str: &str) -> String {
    let mut hasher = DefaultHasher::new();
    input_str.hash(&mut hasher);
    format!("project_{:016x}", hasher.finish())
}

fn create_project_skeleton(project_dir: &Path, cargo_config: CargoConfig, main_content: &str) {
    let src_dir = project_dir.join("src");
    if !src_dir.exists() {
        fs::create_dir_all(&src_dir).expect("Failed to create src directory.");
    }

    let cargo_toml = project_dir.join("Cargo.toml");
    let cargo_toml_content = cargo_config.print();
    fs::write(&cargo_toml, cargo_toml_content).expect("Failed to write Cargo.toml.");

    let main_rs = src_dir.join("main.rs");
    let mut file = File::create(&main_rs).expect("Failed to create main.rs");
    file.write_all(main_content.as_bytes()).expect("Failed to write main.rs");
}

fn run_cargo_project(project_dir: &PathBuf) -> String {
    let output = Command::new("cargo")
        .arg("run")
        .current_dir(project_dir)
        .output()
        .expect("Failed to execute cargo run");

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        eprintln!("{stderr}");
        panic!("Cargo project failed to compile or run.");
    }

    String::from_utf8_lossy(&output.stdout).to_string()
}

// ===================================
// === Top-level Attribute Parsing ===
// ===================================

fn extract_dependencies(tokens: TokenStream) -> (TokenStream, Vec<String>) {
    let (tokens2, attributes) = extract_top_level_attributes(tokens);
    let mut dependencies = Vec::new();
    for attr in attributes {
        let pfx = "dependency(";
        if attr.starts_with(pfx) && attr.ends_with(")") {
            dependencies.push(attr[pfx.len()..attr.len()-1].to_string());
        } else {
            panic!("Invalid attribute: {attr}");
        }
    }
    (tokens2, dependencies)
}

/// Extract all top-level attributes (`#![...]`) from the input `TokenStream` and return the
/// remaining `TokenStream`.
fn extract_top_level_attributes(tokens: TokenStream) -> (TokenStream, Vec<String>) {
    let mut output = TokenStream::new();
    let mut attributes = Vec::new();
    let mut iter = tokens.into_iter().peekable();
    while let Some(_token) = iter.peek() {
        if let Some(dep) = try_parse_inner_attr(&mut iter) {
            attributes.push(dep);
        } else if let Some(token) = iter.next() {
            output.extend(Some(token));
        }
    }
    (output, attributes)
}

/// Try to parse `#![...]` as an inner attribute and return the parsed content if successful.
fn try_parse_inner_attr(iter: &mut std::iter::Peekable<impl Iterator<Item = TokenTree>>) -> Option<String> {
    // Check for '#'.
    let Some(TokenTree::Punct(pound)) = iter.peek() else { return None; };
    if pound.as_char() != '#' { return None; }
    iter.next();

    // Check for '!'.
    let Some(TokenTree::Punct(bang)) = iter.peek() else { return None; };
    if bang.as_char() != '!' { return None; }
    iter.next();

    // Check for [ ... ] group.
    let Some(TokenTree::Group(group)) = iter.peek() else { return None; };
    if group.delimiter() != Delimiter::Bracket { return None; }
    let content = group.stream().to_string();
    iter.next();

    Some(content)
}

// ====================
// === Output Macro ===
// ====================

/// Find and expand the `output!` macro in the input `TokenStream`. After this lib stabilizes, this
/// should be rewritten to standard macro and imported by the generated code.
fn expand_output_macro(input: TokenStream) -> TokenStream {
    let tokens: Vec<TokenTree> = input.into_iter().collect();
    let mut output = TokenStream::new();
    let mut i = 0;
    while i < tokens.len() {
        if let TokenTree::Ident(ref ident) = tokens[i] {
            if *ident == "output" && i + 1 < tokens.len() {
                if let TokenTree::Punct(ref excl) = tokens[i + 1] {
                    if excl.as_char() == '!' && i + 2 < tokens.len() {
                        if let TokenTree::Group(ref group) = tokens[i + 2] {
                            let inner_rewritten = expand_output_macro(group.stream());
                            let content_str = print(&inner_rewritten);
                            let lit = syn::LitStr::new(&content_str, proc_macro2::Span::call_site());
                            let new_tokens = quote! { write_ln!(output_buffer, #lit); };
                            output.extend(new_tokens);
                            i += 3;
                            continue;
                        }
                    }
                }
            }
        }
        match &tokens[i] {
            TokenTree::Group(group) => {
                let new_stream = expand_output_macro(group.stream());
                let new_group = TokenTree::Group(proc_macro2::Group::new(group.delimiter(), new_stream));
                output.extend(std::iter::once(new_group));
            }
            _ => {
                output.extend(std::iter::once(tokens[i].clone()));
            }
        }
        i += 1;
    }
    output
}

// =============
// === Print ===
// =============

#[derive(Debug)]
struct PrintOutput {
    output: String,
    start_token: Option<LineColumn>,
    end_token: Option<LineColumn>,
}

// Used to indicate that `{{` and `}}` was already collapsed to `{` and `}`. So, for example, the
// following transformations will be performed:
// `{ {{a}} }` -> `{ {{{a}}} }` -> `{ %%%{a}%%% }` -> `{{ %%%{a}%%% }}` -> `{{ {a} }}`
const SPACER: &str = "%%%";

/// Prints the token stream as a string ready to be used by the format macro. The following
/// transformations are performed:
/// - A trailing space is added after printing each token.
/// - If `prev_token.line == next_token.line` and `prev_token.end_column >= next_token.start_column`,
///   the prev token trailing space is removed. This basically preserves spaces from the input code,
///   which is needed to distinguish between such inputs as `MyName{x}`, `MyName {x}`, etc.
///   The `>=` is used because sometimes a few tokens can have the same start column. For example,
///   for lifetimes (`'t`), the apostrophe and the lifetime ident have the same start column.
/// - Every occurrence of `{{` and `}}` is replaced with `{` and `}` respectively. Also, every
///   occurrence of `{` and `}` is replaced with `{{` and `}}` respectively. This prepares the
///   string to be used in the format macro.
///
/// There is also a special transformation for `IntellIJ` / `RustRover`. Their spans are different from
/// rustc ones, so sometimes tokens are glued together. This is why we discover Rust keywords and add
/// additional spacing around. This is probably not covering all the bugs. If there will be a bug
/// report, this is the place to look at.
fn print(tokens: &TokenStream) -> String {
    print_internal(tokens).output.replace("{%%%", "{ %%%").replace("%%%}", "%%% }").replace(SPACER, "")
}

fn print_internal(tokens: &TokenStream) -> PrintOutput {
    let token_vec: Vec<TokenTree> = tokens.clone().into_iter().collect();
    let mut output = String::new();
    let mut first_token_start = None;
    let mut prev_token_end: Option<LineColumn> = None;
    for (i, token) in token_vec.iter().enumerate() {
        let mut token_start = token.span().start();
        let mut token_end = token.span().end();
        let mut is_keyword = false;
        let token_str = match token {
            TokenTree::Group(g) => {
                let content = print_internal(&g.stream());
                let mut content_str = content.output;
                content_str.pop();
                let (open, close) = match g.delimiter() {
                    Delimiter::Brace =>{
                        if content_str.starts_with('{') && content_str.ends_with('}') {
                            // We already replaced the internal `{` and `}` with `{{` and `}}`.
                            content_str.pop();
                            content_str.remove(0);
                            (SPACER, SPACER)
                        } else {
                            ("{{", "}}")
                        }
                    },
                    Delimiter::Parenthesis => ("(", ")"),
                    Delimiter::Bracket => ("[", "]"),
                    _ => ("", ""),
                };

                if let Some(content_first_token_start) = content.start_token {
                    token_start.line = content_first_token_start.line;
                    if content_first_token_start.column > 0 {
                        token_start.column = content_first_token_start.column - 1;
                    }
                }
                if let Some(content_end) = content.end_token {
                    token_end.line = content_end.line;
                    token_end.column = content_end.column + 1;
                }
                format!("{open}{content_str}{close}")
            }
            TokenTree::Ident(ident) => {
                let str = ident.to_string();
                is_keyword = KEYWORDS.contains(&str.as_str());
                str
            },
            TokenTree::Literal(lit) => lit.to_string(),
            TokenTree::Punct(punct) => punct.as_char().to_string(),
        };
        if DEBUG {
            println!("{i}: [{token_start:?}-{token_end:?}] [{prev_token_end:?}]: {token}");
        }
        if let Some(prev_token_end) = prev_token_end {
            if prev_token_end.line == token_start.line && prev_token_end.column >= token_start.column {
                output.pop();
            }
        }

        // Pushing a space before and after keywords is for IntelliJ only. Their token spans are invalid.
        if is_keyword { output.push(' '); }
        output.push_str(&token_str);
        output.push(' ');
        if is_keyword { output.push(' '); }

        first_token_start.get_or_insert(token_start);
        prev_token_end = Some(token_end);
    }
    PrintOutput {
        output,
        start_token: first_token_start,
        end_token: prev_token_end,
    }
}

// ==================
// === Eval Macro ===
// ==================

#[proc_macro]
pub fn eval(input_raw: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let (input, dependencies) = extract_dependencies(input_raw.into());
    let cargo_config = CargoConfig::new(dependencies);

    let input_str = expand_output_macro(input).to_string();
    if DEBUG { println!("REWRITTEN INPUT: {input_str}"); }

    let out_dir = get_output_dir();
    let project_name = project_name_from_input(&input_str);
    let project_dir = out_dir.join(&project_name);

    if !project_dir.exists() {
        fs::create_dir_all(&project_dir)
            .expect("Failed to create project directory.");
    }

    let main_content = format!(
        "{PRELUDE}
        fn main() {{
            let mut output_buffer = String::new();
            {input_str}
            println!(\"{{output_buffer}}\");
        }}",
    );

    create_project_skeleton(&project_dir, cargo_config, &main_content);
    let output = run_cargo_project(&project_dir);
    fs::remove_dir_all(&project_dir).ok();
    let out: TokenStream = output.parse().expect("Failed to parse generated code.");
    if DEBUG {
        println!("OUT: {out}");
    }
    out.into()
}