inline_fn/
lib.rs

1#![doc = include_str!("../README.md")]
2mod parser;
3
4use std::path::PathBuf;
5
6use anyhow::Context;
7use proc_macro::TokenStream;
8use syn::{parse::Parse, parse_macro_input};
9
10struct ItemFn {
11    fn_name: syn::Ident,
12    _comma: Option<syn::Token![,]>,
13    relative_path: Option<syn::LitStr>,
14}
15
16/// Inline a function from a file.
17/// Expands to the function body fetched from the specified file.
18/// # Usage
19/// * Function from a specific file relative from the workspace root (this is best for performance reasons)
20///   `inline_fn!(fn_name, "crates/app/src/file.rs")`
21/// * Function from the first file with the matching function name (this will cause a recursive scan in the workspace):
22///   `inline_fn!(fn_name)`
23///
24/// # Examples
25///
26/// ```ignore
27/// fn add_two(a: usize, b: usize) -> usize {
28///    a + b
29/// }
30///
31/// fn add_three(a: usize, b: usize, c: usize) -> usize {
32///    inline_fn::inline_fn!(add_two) + c // no need to repeat the code
33/// }
34/// ```
35#[proc_macro]
36pub fn inline_fn(input: TokenStream) -> TokenStream {
37    let input: ItemFn = parse_macro_input!(input);
38    let relative_path = input
39        .relative_path
40        .map(|p| PathBuf::from(p.token().to_string().trim_matches('"')));
41
42    let fn_name = input.fn_name.to_string();
43    let workspace_root = std::env::current_dir().expect("Failed to get current directory");
44
45    find_function(fn_name, workspace_root, relative_path)
46        .parse()
47        .expect("parse")
48}
49
50fn find_function(
51    fn_name: String,
52    workspace_root: PathBuf,
53    relative_path: Option<PathBuf>,
54) -> String {
55    if let Some(source_file) = relative_path {
56        match find_in_file(fn_name.clone(), workspace_root.join(source_file.clone())) {
57            Ok(Some(c)) => return c.join("\n"),
58            Ok(None) => panic!(
59                "Function `{}` not found, location searched: `{}`",
60                fn_name,
61                source_file.display()
62            ),
63            Err(e) => panic!("{}, location searched: {}", e, source_file.display()),
64        }
65    }
66
67    // search all codes
68    walkdir::WalkDir::new(&workspace_root)
69        .into_iter()
70        .filter_map(|e| e.ok())
71        .filter(|e| e.file_type().is_file())
72        .filter(|e| e.path().extension().map(|e| e == "rs").unwrap_or(false))
73        .map(|e| find_in_file(fn_name.clone(), e.path().to_path_buf()))
74        .filter_map(|e| e.ok())
75        .filter(|e| e.is_some())
76        .map(|e| {
77            e.with_context(|| {
78                format!(
79                    "Failed to find function `{}` in folder: {}",
80                    fn_name,
81                    workspace_root.display()
82                )
83            })
84            .unwrap()
85        })
86        .flatten()
87        .collect::<Vec<String>>()
88        .join("\n")
89}
90
91fn find_in_file(fn_name: String, file: PathBuf) -> anyhow::Result<Option<Vec<String>>> {
92    if file.extension().map(|e| e != "rs").unwrap_or(false) {
93        return Err(anyhow::anyhow!("File is not a rust file"));
94    }
95
96    let sc = std::fs::read_to_string(file.clone())
97        .with_context(|| format!("Failed to read file XX: {}", file.display()))?;
98
99    Ok(parser::extract_function(&sc, &fn_name).map(|c| {
100        c.split('\n')
101            .map(|s| s.to_string())
102            .collect::<Vec<String>>()
103    }))
104}
105
106impl Parse for ItemFn {
107    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
108        Ok(Self {
109            fn_name: input.parse()?,
110            _comma: input.parse()?,
111            relative_path: input.parse()?,
112        })
113    }
114}