Skip to content

Insert Code to Leetcode Generated Code

When I wrote leetcode in vscode with the plugin, I named generated leetcode file with format as[0-9]+(e.g, as123.rs) and put them under src. However, I have encountered some painpoints:

  1. The generated file isn't used by main.rs so the rust analyzer will ignore it, so I need to add mod as123; in main.rs every time.
  2. The generated code doesn't contain struct Soltuion definition so I need to add it every time.
  3. When I want to debug locally, I need to write the test function from scratch. It doesn't cost much but it's verbose.

These painpoints are trivial, but they're quite verbose. By using build script build.rs, I'm succeeded solving my painpoints!

The build.rs will be executed by rust-analyzer and generate a generated_mods.rs file, which contains all the mod for lc generated as*.rs mod. It ensures that the lc generated files are imported by main.rs. Then, we parse the file and insert struct declaration and test function. The input of the function is empty and could be filled later.

Finally, we call println!("cargo:rerun-if-changed=src"); to let cargo re-run the build.rs if the src changes.

include!("generated_mods.rs");
fn main(){}
[build-dependencies]
syn = { version = "2.0", features = ["full"] }
build.rs
use std::env;
use std::fs;
use std::fs::File;
use std::io::Read;
use std::io::Write;
use std::path::Path;
use std::path::PathBuf;
use syn::{File as SynFile, Item, ItemFn, ItemImpl, ItemStruct};

fn main() {
    let src_dir = Path::new(&env::var("CARGO_MANIFEST_DIR").unwrap()).join("src");

    let mut mod_declarations = String::new();
    for entry in fs::read_dir(&src_dir).unwrap() {
        let entry = entry.unwrap();
        let path = entry.path();
        if let Some(file_name) = path.file_name() {
            if let Some(file_name_str) = file_name.to_str() {
                if file_name_str.starts_with("as") && file_name_str.ends_with(".rs") {
                    let mod_name = file_name_str.trim_end_matches(".rs");
                    mod_declarations.push_str(&format!("pub mod {};\n", mod_name));

                    insert_code(path);
                }
            }
        }
    }

    let out_file = src_dir.join("generated_mods.rs");
    fs::write(&out_file, mod_declarations).unwrap();

    // Tell Cargo to re-run this script if the `src` directory changes
    println!("cargo:rerun-if-changed=src");
}

fn insert_code(path: PathBuf) {
    let mut content = String::new();
    File::open(&path)
        .unwrap()
        .read_to_string(&mut content)
        .unwrap();

    let syntax_tree: SynFile = syn::parse_file(&content).unwrap();
    let mut has_solution_struct = false;
    let mut function_names = Vec::new();

    for item in &syntax_tree.items {
        match item {
            Item::Struct(ItemStruct { ident, .. }) if ident == "Solution" => {
                has_solution_struct = true;
            }
            Item::Impl(ItemImpl { self_ty, items, .. }) => {
                if let syn::Type::Path(type_path) = &**self_ty {
                    if let Some(ident) = type_path.path.get_ident() {
                        if ident == "Solution" {
                            for impl_item in items {
                                if let syn::ImplItem::Fn(method) = impl_item {
                                    function_names.push(method.sig.ident.to_string());
                                }
                            }
                        }
                    }
                }
            }
            _ => {}
        }
    }

    if has_solution_struct {
        return;
    }
    // Generate new content
    let mut new_content = String::new();

    new_content.push_str("struct Solution;\n\n");

    // Append the original content of the file
    new_content.push_str(&content);

    if !function_names.is_empty() {
        new_content.push_str("\n#[cfg(test)]\nmod tests {\n");
        new_content.push_str("    use super::*;\n\n");
        for func_name in &function_names {
            new_content.push_str(&format!("    #[test]\n    fn test_{}() {{\n", func_name));
            new_content.push_str(&format!(
                "        let result = Solution::{}();\n",
                func_name
            ));
            new_content.push_str("        println!(\"Result: {:?}\", result);\n");
            new_content.push_str("    }\n");
        }
        new_content.push_str("}\n");
    }

    // Write the modified content back to the file
    let mut file = File::create(&path).unwrap();
    file.write_all(new_content.as_bytes()).unwrap();
}