no-block-pls 0.1.0

use quote::quote;
use syn::spanned::Spanned;
use syn::{
    Expr, ExprAsync, ExprAwait, ExprForLoop, ExprLoop, ExprMacro, ExprWhile, File, ImplItemFn,
    ItemFn, Macro, parse_quote,
    visit_mut::{self, VisitMut},
};

// Macros that internally await and should be instrumented
const ASYNC_MACROS: &[&str] = &["select", "select_biased", "join", "try_join", "race"];

fn orig_line<T: Spanned>(node: &T) -> u32 {
    node.span().start().line as u32
}

fn await_token_line(node: &ExprAwait) -> u32 {
    node.await_token.span().start().line as u32
}

pub struct AsyncInstrumenter {
    guard_depth: usize,
}

impl AsyncInstrumenter {
    pub fn new(_threshold_ms: u64) -> Self {
        AsyncInstrumenter { guard_depth: 0 }
    }

    pub fn instrument_file(&mut self, file: &mut File) {
        self.guard_depth = 0;
        self.visit_file_mut(file);
    }

    fn is_async_macro(mac: &Macro) -> bool {
        ASYNC_MACROS.iter().any(|name| mac.path.is_ident(name))
            || (!mac.path.segments.is_empty()
                && mac
                    .path
                    .segments
                    .last()
                    .map(|s| ASYNC_MACROS.iter().any(|name| s.ident == name))
                    .unwrap_or(false))
    }

    fn in_guard_scope(&self) -> bool {
        self.guard_depth > 0
    }
}

impl VisitMut for AsyncInstrumenter {
    fn visit_expr_mut(&mut self, expr: &mut Expr) {
        match expr {
            Expr::Await(await_expr) => {
                // First visit the awaited expression to handle nested instrumentation
                self.visit_expr_mut(&mut await_expr.base);

                if !self.in_guard_scope() {
                    return;
                }

                // Transform await expression to use pause/resume pattern
                let l: u32 = await_token_line(await_expr);
                let inner = &await_expr.base;

                // Create a block that pauses guard, awaits, then resets guard
                let new_expr: Expr = parse_quote! {
                    {
                        __guard.end_section(#l);
                        let __result = #inner.await;
                        __guard.start_section(#l);
                        __result
                    }
                };

                *expr = new_expr;
            }
            Expr::Macro(ExprMacro { mac, .. }) => {
                // Check if this is an async macro that internally awaits
                let is_async_macro = Self::is_async_macro(mac);

                if is_async_macro && self.in_guard_scope() {
                    let l: u32 = orig_line(expr);
                    // Clone the original macro to use in the transformation
                    let original_macro = expr.clone();

                    // Wrap the select!/join! using pause/reset
                    let new_expr: Expr = parse_quote! {
                        {
                            __guard.end_section(#l);
                            let __result = #original_macro;
                            __guard.start_section(#l);
                            __result
                        }
                    };

                    *expr = new_expr;
                } else {
                    // Continue visiting other macros
                    visit_mut::visit_expr_mut(self, expr);
                }
            }
            Expr::Continue(cont) => {
                if !self.in_guard_scope() {
                    return;
                }
                let l: u32 = orig_line(cont);
                let label = cont.label.clone();
                let new_expr: Expr = if let Some(label) = label {
                    parse_quote!({
                        __guard.end_section(#l);
                        __guard.start_section(#l);
                        continue #label;
                    })
                } else {
                    parse_quote!({
                        __guard.end_section(#l);
                        __guard.start_section(#l);
                        continue;
                    })
                };
                *expr = new_expr;
            }
            Expr::Break(brk) => {
                if let Some(expr) = brk.expr.as_mut() {
                    self.visit_expr_mut(expr);
                }

                if !self.in_guard_scope() {
                    return;
                }

                let l: u32 = orig_line(brk);
                let label = brk.label.clone();
                let value = brk.expr.as_ref().map(|inner| *inner.clone());
                let new_expr: Expr = match (label, value) {
                    (Some(label), Some(v)) => parse_quote!({
                        __guard.end_section(#l);
                        __guard.start_section(#l);
                        break #label #v;
                    }),
                    (Some(label), None) => parse_quote!({
                        __guard.end_section(#l);
                        __guard.start_section(#l);
                        break #label;
                    }),
                    (None, Some(v)) => parse_quote!({
                        __guard.end_section(#l);
                        __guard.start_section(#l);
                        break #v;
                    }),
                    (None, None) => parse_quote!({
                        __guard.end_section(#l);
                        __guard.start_section(#l);
                        break;
                    }),
                };
                *expr = new_expr;
            }
            Expr::Return(ret) => {
                if let Some(expr) = ret.expr.as_mut() {
                    self.visit_expr_mut(expr);
                }

                if !self.in_guard_scope() {
                    return;
                }

                let l: u32 = orig_line(ret);
                let value = ret.expr.as_ref().map(|inner| *inner.clone());
                let new_expr: Expr = if let Some(v) = value {
                    parse_quote!({
                        __guard.end_section(#l);
                        return #v;
                    })
                } else {
                    parse_quote!({
                        __guard.end_section(#l);
                        return;
                    })
                };
                *expr = new_expr;
            }
            _ => {
                // Continue visiting other expressions
                visit_mut::visit_expr_mut(self, expr);
            }
        }
    }

    fn visit_expr_async_mut(&mut self, async_block: &mut ExprAsync) {
        // Add initial guard at the start of async block
        let l: u32 = orig_line(async_block);
        let init_guard = parse_quote! {
            let mut __guard = crate::__async_profile_guard__::Guard::new(
                concat!(module_path!(), "::async_block"),
                file!(),
                #l
            );
        };

        // Insert guard FIRST before any transformations
        async_block.block.stmts.insert(0, init_guard);

        // Now visit the rest of the block's statements
        // We need to collect indices first to avoid borrowing issues
        let stmt_count = async_block.block.stmts.len();
        self.guard_depth += 1;
        for i in 1..stmt_count {
            let mut stmt = async_block.block.stmts[i].clone();
            self.visit_stmt_mut(&mut stmt);
            async_block.block.stmts[i] = stmt;
        }
        self.guard_depth -= 1;
    }

    fn visit_expr_closure_mut(&mut self, closure: &mut syn::ExprClosure) {
        let prev_depth = self.guard_depth;
        self.guard_depth = 0;
        visit_mut::visit_expr_closure_mut(self, closure);
        self.guard_depth = prev_depth;
    }

    fn visit_expr_for_loop_mut(&mut self, node: &mut ExprForLoop) {
        if self.in_guard_scope() {
            let l: u32 = orig_line(node);
            node.body
                .stmts
                .insert(0, parse_quote! { __guard.checkpoint(#l); });
        }
        visit_mut::visit_expr_for_loop_mut(self, node);
    }

    fn visit_expr_loop_mut(&mut self, node: &mut ExprLoop) {
        if self.in_guard_scope() {
            let l: u32 = orig_line(node);
            node.body
                .stmts
                .insert(0, parse_quote! { __guard.checkpoint(#l); });
        }
        visit_mut::visit_expr_loop_mut(self, node);
    }

    fn visit_expr_while_mut(&mut self, node: &mut ExprWhile) {
        if self.in_guard_scope() {
            let l: u32 = orig_line(node);
            node.body
                .stmts
                .insert(0, parse_quote! { __guard.checkpoint(#l); });
        }
        visit_mut::visit_expr_while_mut(self, node);
    }

    fn visit_impl_item_fn_mut(&mut self, func: &mut ImplItemFn) {
        let prev_depth = self.guard_depth;
        if self.in_guard_scope() {
            self.guard_depth = 0;
        }

        if func.sig.asyncness.is_none() {
            visit_mut::visit_impl_item_fn_mut(self, func);
            self.guard_depth = prev_depth;
            return;
        }

        let func_name = &func.sig.ident;
        let l: u32 = orig_line(&func.block);

        // Add initial guard at the start of async function
        let init_guard = parse_quote! {
            let mut __guard = crate::__async_profile_guard__::Guard::new(
                concat!(module_path!(), "::", stringify!(#func_name)),
                file!(),
                #l
            );
        };

        // Insert guard FIRST before any transformations
        func.block.stmts.insert(0, init_guard);

        // Rebind simple arguments after guard so they drop before guard
        let mut rebinds: Vec<syn::Stmt> = Vec::new();
        for input in func.sig.inputs.iter() {
            match input {
                syn::FnArg::Receiver(_) => {
                    // Cannot shadow `self`
                }
                syn::FnArg::Typed(pat_ty) => match &*pat_ty.pat {
                    syn::Pat::Ident(pat_ident) => {
                        let ident = &pat_ident.ident;
                        if pat_ident.mutability.is_some() {
                            let stmt: syn::Stmt = parse_quote! { let mut #ident = #ident; };
                            rebinds.push(stmt);
                        } else {
                            let stmt: syn::Stmt = parse_quote! { let #ident = #ident; };
                            rebinds.push(stmt);
                        }
                    }
                    syn::Pat::Tuple(pat_tuple) => {
                        let mut rhs_exprs: Vec<syn::Expr> = Vec::new();
                        let mut all_ident = true;
                        for elem in pat_tuple.elems.iter() {
                            if let syn::Pat::Ident(pi) = elem {
                                let id = pi.ident.clone();
                                let expr: syn::Expr = parse_quote! { #id };
                                rhs_exprs.push(expr);
                            } else {
                                all_ident = false;
                                break;
                            }
                        }
                        if all_ident {
                            let lhs_pat: syn::Pat = (*pat_ty.pat).clone();
                            let stmt: syn::Stmt =
                                parse_quote! { let #lhs_pat = ( #(#rhs_exprs),* ); };
                            rebinds.push(stmt);
                        }
                    }
                    _ => {}
                },
            }
        }
        if !rebinds.is_empty() {
            let mut idx = 1usize;
            for stmt in rebinds {
                func.block.stmts.insert(idx, stmt);
                idx += 1;
            }
        }

        // Now visit the rest of the block's statements
        // We need to collect indices first to avoid borrowing issues
        let stmt_count = func.block.stmts.len();
        self.guard_depth += 1;
        for i in 1..stmt_count {
            let mut stmt = func.block.stmts[i].clone();
            self.visit_stmt_mut(&mut stmt);
            func.block.stmts[i] = stmt;
        }
        self.guard_depth -= 1;
        self.guard_depth = prev_depth;
    }

    fn visit_item_fn_mut(&mut self, func: &mut ItemFn) {
        let prev_depth = self.guard_depth;
        if self.in_guard_scope() {
            self.guard_depth = 0;
        }

        if func.sig.asyncness.is_none() {
            visit_mut::visit_item_fn_mut(self, func);
            self.guard_depth = prev_depth;
            return;
        }

        let func_name = &func.sig.ident;
        let l: u32 = orig_line(&func.block);

        // Add initial guard at the start of async function
        let init_guard = parse_quote! {
            let mut __guard = crate::__async_profile_guard__::Guard::new(
                concat!(module_path!(), "::", stringify!(#func_name)),
                file!(),
                #l
            );
        };

        // Insert guard FIRST before any transformations
        func.block.stmts.insert(0, init_guard);

        // Rebind simple arguments after guard so they drop before guard
        // This ensures guard drops last when the async function completes.
        let mut rebinds: Vec<syn::Stmt> = Vec::new();
        for input in func.sig.inputs.iter() {
            match input {
                syn::FnArg::Receiver(_) => {
                    // `self`/`&self`/`&mut self` cannot be shadowed; skip
                }
                syn::FnArg::Typed(pat_ty) => {
                    match &*pat_ty.pat {
                        syn::Pat::Ident(pat_ident) => {
                            let ident = &pat_ident.ident;
                            if pat_ident.mutability.is_some() {
                                let stmt: syn::Stmt = parse_quote! { let mut #ident = #ident; };
                                rebinds.push(stmt);
                            } else {
                                let stmt: syn::Stmt = parse_quote! { let #ident = #ident; };
                                rebinds.push(stmt);
                            }
                        }
                        syn::Pat::Tuple(pat_tuple) => {
                            // Support tuple pattern with ident elements: (a, b): (T, U)
                            let mut rhs_exprs: Vec<syn::Expr> = Vec::new();
                            let mut all_ident = true;
                            for elem in pat_tuple.elems.iter() {
                                if let syn::Pat::Ident(pi) = elem {
                                    let id = pi.ident.clone();
                                    let expr: syn::Expr = parse_quote! { #id };
                                    rhs_exprs.push(expr);
                                } else {
                                    all_ident = false;
                                    break;
                                }
                            }
                            if all_ident {
                                let lhs_pat: syn::Pat = (*pat_ty.pat).clone();
                                let stmt: syn::Stmt =
                                    parse_quote! { let #lhs_pat = ( #(#rhs_exprs),* ); };
                                rebinds.push(stmt);
                            }
                        }
                        _ => {}
                    }
                }
            }
        }
        if !rebinds.is_empty() {
            // Insert after the guard we just added (index 1..1 to splice at position 1)
            let mut idx = 1usize;
            for stmt in rebinds {
                func.block.stmts.insert(idx, stmt);
                idx += 1;
            }
        }

        // Now visit the rest of the block's statements
        // We need to collect indices first to avoid borrowing issues
        let stmt_count = func.block.stmts.len();
        self.guard_depth += 1;
        for i in 1..stmt_count {
            let mut stmt = func.block.stmts[i].clone();
            self.visit_stmt_mut(&mut stmt);
            func.block.stmts[i] = stmt;
        }
        self.guard_depth -= 1;
        self.guard_depth = prev_depth;
    }

    fn visit_stmt_mut(&mut self, stmt: &mut syn::Stmt) {
        if let syn::Stmt::Macro(stmt_macro) = stmt {
            let mac = &stmt_macro.mac;

            // Check if this is an async macro that internally awaits
            let is_async_macro = Self::is_async_macro(mac);

            if is_async_macro && self.in_guard_scope() {
                // Transform the statement macro using pause/resume
                let l: u32 = orig_line(&stmt_macro.mac);
                let original_macro = quote!(#mac);

                // Check if the original macro had a semicolon
                let has_semicolon = stmt_macro.semi_token.is_some();

                // Create statements to pause before and reset after
                if has_semicolon {
                    // Statement with semicolon - transform to multiple statements
                    let wrapped: syn::Stmt = parse_quote! {
                        {
                            __guard.end_section(#l);
                            #original_macro;
                            __guard.start_section(#l);
                        };
                    };
                    *stmt = wrapped;
                } else {
                    // Expression without semicolon (for return values)
                    let wrapped: syn::Stmt = parse_quote! {
                        {
                            __guard.end_section(#l);
                            let __result = #original_macro;
                            __guard.start_section(#l);
                            __result
                        }
                    };
                    *stmt = wrapped;
                }

                return;
            }
        }

        // Continue with default visiting for other statements
        visit_mut::visit_stmt_mut(self, stmt);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use insta::assert_snapshot;
    use quote::quote;

    #[test]
    fn test_async_fn_args_rebound_after_guard() {
        let input = quote! {
            async fn fron(arg: Bla) {}
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        // Guard inserted and argument is rebound after guard
        assert!(output.contains("let mut __guard = crate::__async_profile_guard__::Guard::new"));
        // Ensure rebind happens
        assert!(output.contains("let arg = arg;"));
        // And appears after guard by simple index check
        assert!(output.find("__guard").unwrap() < output.find("let arg = arg;").unwrap());
    }

    #[test]
    fn test_async_fn_mut_args_preserve_mut() {
        let input = quote! {
            async fn foo(mut x: i32, y: i32) {}
        };
        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);
        let output = prettyplease::unparse(&file);
        assert!(output.contains("let mut x = x;"));
        assert!(output.contains("let y = y;"));
    }

    #[test]
    fn test_async_method_skips_self_rebinds_others() {
        let input = quote! {
            struct S;
            impl S {
                async fn m(&mut self, x: i32) {}
            }
        };
        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);
        let output = prettyplease::unparse(&file);
        // Must not try to rebind `self`
        assert!(!output.contains("let self = self;"));
        // But should rebind x
        assert!(output.contains("let x = x;"));
    }

    #[test]
    fn test_async_fn_tuple_args_rebind() {
        let input = quote! {
            async fn t((a, b): (A, B)) {}
        };
        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);
        let output = prettyplease::unparse(&file);
        assert!(output.contains("let (a, b) = (a, b);"));
        // Appears after guard
        assert!(output.find("__guard").unwrap() < output.find("let (a, b) = (a, b);").unwrap());
    }

    #[test]
    fn test_async_fn_tuple_args_rebind_with_mut() {
        let input = quote! {
            async fn t((mut a, b): (A, B)) {}
        };
        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);
        let output = prettyplease::unparse(&file);
        assert!(output.contains("let (mut a, b) = (a, b);"));
    }

    #[test]
    fn test_simple_async_function() {
        let input = quote! {
            async fn process() {
                let data = fetch().await;
                process_data(data);
                save().await;
            }
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output, @r#"
        async fn process() {
            let mut __guard = crate::__async_profile_guard__::Guard::new(
                concat!(module_path!(), "::", stringify!(process)),
                file!(),
                1u32,
            );
            let data = {
                __guard.end_section(1u32);
                let __result = fetch().await;
                __guard.start_section(1u32);
                __result
            };
            process_data(data);
            {
                __guard.end_section(1u32);
                let __result = save().await;
                __guard.start_section(1u32);
                __result
            };
        }
        "#);
    }

    #[test]
    fn test_async_block() {
        let input = quote! {
            fn main() {
                let future = async {
                    let x = foo().await;
                    bar(x);
                };
            }
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output, @r###"
        fn main() {
            let future = async {
                let mut __guard = crate::__async_profile_guard__::Guard::new(
                    concat!(module_path!(), "::async_block"),
                    file!(),
                    1u32,
                );
                let x = {
                    __guard.end_section(1u32);
                    let __result = foo().await;
                    __guard.start_section(1u32);
                    __result
                };
                bar(x);
            };
        }
        "###);
    }

    #[test]
    fn test_multiline_await_uses_await_token_line() {
        let input = r#"async fn foo() {
    let value = compute()
        .await;
}
"#;

        let mut file: File = syn::parse_file(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output, @r###"
        async fn foo() {
            let mut __guard = crate::__async_profile_guard__::Guard::new(
                concat!(module_path!(), "::", stringify!(foo)),
                file!(),
                1u32,
            );
            let value = {
                __guard.end_section(3u32);
                let __result = compute().await;
                __guard.start_section(3u32);
                __result
            };
        }
        "###);
    }

    #[test]
    fn test_non_async_unchanged() {
        let input = quote! {
            fn regular_function() {
                let x = 42;
                println!("{}", x);
            }
        };

        let mut file: File = syn::parse2(input).unwrap();

        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output, @r###"
        fn regular_function() {
            let x = 42;
            println!("{}", x);
        }
        "###);
    }

    #[test]
    fn test_loop_with_await() {
        let input = quote! {
            async fn looper() {
                loop {
                    work();
                    sleep(Duration::from_secs(1)).await;
                    more_work();
                }
            }
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output, @r###"
        async fn looper() {
            let mut __guard = crate::__async_profile_guard__::Guard::new(
                concat!(module_path!(), "::", stringify!(looper)),
                file!(),
                1u32,
            );
            loop {
                __guard.checkpoint(1u32);
                work();
                {
                    __guard.end_section(1u32);
                    let __result = sleep(Duration::from_secs(1)).await;
                    __guard.start_section(1u32);
                    __result
                };
                more_work();
            }
        }
        "###);
    }

    #[test]
    fn test_async_block_comprehensive() {
        let input = quote! {
            fn spawn_tasks() {
                tokio::spawn(async {
                    let data = fetch_data().await;
                    println!("Got data: {}", data);
                    process(data).await;
                    let result = compute().await;
                    save(result).await;
                });
            }
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output);
    }

    #[test]
    fn test_continue_instrumentation() {
        let input = quote! {
            async fn worker() {
                for item in items() {
                    if should_skip(&item) {
                        continue;
                    }
                    process(item).await;
                }
            }
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output);
        assert!(output.contains("__guard.checkpoint("));
        assert!(output.contains("__guard.end_section("));
        assert!(output.contains("__guard.start_section("));
        assert!(output.contains("continue;"));
        assert!(!output.contains("line!()"));
    }

    #[test]
    fn test_break_instrumentation() {
        let input = quote! {
            async fn breaker() -> i32 {
                'outer: loop {
                    if ready() {
                        break 'outer value();
                    }
                    perform().await;
                }
                follow_up().await;
                7
            }
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output);
        assert!(output.contains("__guard.checkpoint("));
        assert!(output.contains("__guard.end_section("));
        assert!(output.contains("__guard.start_section("));
        assert!(output.contains("follow_up"));
        assert!(!output.contains("line!()"));
    }

    #[test]
    fn test_return_instrumentation() {
        let input = quote! {
            async fn maybe_fail() -> Result<(), Error> {
                loop {
                    if fatal() {
                        return Err(Error::Fatal);
                    }
                    step().await;
                }
            }
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output);
        assert!(output.contains("__guard.checkpoint("));
        assert!(output.contains("__guard.end_section("));
        assert!(!output.contains("line!()"));
    }

    #[test]
    fn test_async_move_block() {
        let input = quote! {
            fn transfer_ownership() {
                let data = vec![1, 2, 3];
                let processor = String::from("test");

                tokio::spawn(async move {
                    println!("Processing with: {}", processor);
                    let result = process_vec(data).await;
                    store(result).await;
                    finalize().await;
                });
            }
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output);
    }

    #[test]
    fn test_async_move_closure() {
        let input = quote! {
            fn register() {
                spawn_metrics_loop(async move |this: Arc<Service>| {
                    this.scan().await;
                    finish().await;
                });
            }
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output);
    }

    #[test]
    fn test_async_move_closure_inside_async_fn() {
        let input = quote! {
            async fn monitor() {
                spawn_metrics_loop(async move |this: Arc<Service>| {
                    this.scan().await;
                    finish().await;
                });
            }
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output);
    }

    #[test]
    fn test_complex_real_world_function() {
        // Read the test file
        let content = std::fs::read_to_string("test_fixtures/test_complex_fn.rs")
            .expect("Failed to read test_fixtures/test_complex_fn.rs");

        let mut file: File = syn::parse_file(&content).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);

        // Use snapshot for full output verification
        // This will show us exactly what the transformation produces
        assert_snapshot!(output);
    }

    #[test]
    fn test_manual_fut() {
        // Read the test file
        let content = std::fs::read_to_string("test_fixtures/manual_fut.rs")
            .expect("Failed to read test_fixtures/test_complex_fn.rs");

        let mut file: File = syn::parse_file(&content).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);

        // Use snapshot for full output verification
        // This will show us exactly what the transformation produces
        assert_snapshot!(output);
    }

    #[test]
    fn test_async_move() {
        // Read the test file
        let content = std::fs::read_to_string("test_fixtures/async_move.rs")
            .expect("Failed to read test_fixtures/async_move.rs");

        let mut file: File = syn::parse_file(&content).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);

        // Use snapshot for full output verification
        // This will show us exactly what the transformation produces
        assert_snapshot!(output);
    }

    #[test]
    fn test_blocking_move() {
        // Read the test file
        let content = std::fs::read_to_string("test_fixtures/blocking_move.rs")
            .expect("Failed to read test_fixtures/blocking_move.rs");

        let mut file: File = syn::parse_file(&content).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);

        // Use snapshot for full output verification
        // This will show us exactly what the transformation produces
        assert_snapshot!(output);
    }

    #[test]
    fn test_closure() {
        // Read the test file
        let content = std::fs::read_to_string("test_fixtures/closure.rs")
            .expect("Failed to read test_fixtures/blocking_move.rs");

        let mut file: File = syn::parse_file(&content).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);

        // Use snapshot for full output verification
        // This will show us exactly what the transformation produces
        assert_snapshot!(output);
    }

    #[test]
    fn test_nested_sync_fn_not_instrumented() {
        let content = std::fs::read_to_string("test_fixtures/nested_sync_fn.rs")
            .expect("Failed to read test_fixtures/nested_sync_fn.rs");

        let mut file: File = syn::parse_file(&content).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output);
    }

    #[test]
    fn test_nested_sync_min_not_instrumented() {
        let content = std::fs::read_to_string("test_fixtures/nested_sync_min.rs")
            .expect("Failed to read test_fixtures/nested_sync_min.rs");

        let mut file: File = syn::parse_file(&content).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output);
    }

    #[test]
    fn test_tokio_select() {
        let input = quote! {
            async fn with_select() {
                let mut interval = tokio::time::interval(Duration::from_secs(1));
                let mut rx = channel.subscribe();

                loop {
                    tokio::select! {
                        _ = interval.tick() => {
                            println!("tick");
                        }
                        msg = rx.recv() => {
                            println!("got message: {:?}", msg);
                        }
                    }
                }
            }
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output, @r#"
        async fn with_select() {
            let mut __guard = crate::__async_profile_guard__::Guard::new(
                concat!(module_path!(), "::", stringify!(with_select)),
                file!(),
                1u32,
            );
            let mut interval = tokio::time::interval(Duration::from_secs(1));
            let mut rx = channel.subscribe();
            loop {
                __guard.checkpoint(1u32);
                {
                    __guard.end_section(1u32);
                    let __result = tokio::select! {
                        _ = interval.tick() => { println!("tick"); } msg = rx.recv() => {
                        println!("got message: {:?}", msg); }
                    };
                    __guard.start_section(1u32);
                    __result
                }
            }
        }
        "#);
    }

    #[test]
    fn test_plain_select() {
        let input = quote! {
            async fn with_imported_select() {
                use tokio::select;

                loop {
                    select! {
                        _ = async_operation() => handle_result(),
                    }
                }
            }
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output, @r#"
        async fn with_imported_select() {
            let mut __guard = crate::__async_profile_guard__::Guard::new(
                concat!(module_path!(), "::", stringify!(with_imported_select)),
                file!(),
                1u32,
            );
            use tokio::select;
            loop {
                __guard.checkpoint(1u32);
                {
                    __guard.end_section(1u32);
                    let __result = select! {
                        _ = async_operation() => handle_result(),
                    };
                    __guard.start_section(1u32);
                    __result
                }
            }
        }
        "#);
    }

    #[test]
    fn test_tokio_join() {
        let input = quote! {
            async fn with_join() {
                let fut1 = async_operation1();
                let fut2 = async_operation2();

                let (result1, result2) = tokio::join!(fut1, fut2);
                process_results(result1, result2);
            }
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output, @r###"
        async fn with_join() {
            let mut __guard = crate::__async_profile_guard__::Guard::new(
                concat!(module_path!(), "::", stringify!(with_join)),
                file!(),
                1u32,
            );
            let fut1 = async_operation1();
            let fut2 = async_operation2();
            let (result1, result2) = {
                __guard.end_section(1u32);
                let __result = tokio::join!(fut1, fut2);
                __guard.start_section(1u32);
                __result
            };
            process_results(result1, result2);
        }
        "###);
    }

    #[test]
    fn test_join_as_expression() {
        let input = quote! {
            async fn join_in_expression() {
                let results = tokio::join!(
                    fetch_data(),
                    fetch_metadata(),
                    fetch_config()
                );

                process(results);
            }
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output, @r###"
        async fn join_in_expression() {
            let mut __guard = crate::__async_profile_guard__::Guard::new(
                concat!(module_path!(), "::", stringify!(join_in_expression)),
                file!(),
                1u32,
            );
            let results = {
                __guard.end_section(1u32);
                let __result = tokio::join!(fetch_data(), fetch_metadata(), fetch_config());
                __guard.start_section(1u32);
                __result
            };
            process(results);
        }
        "###);
    }

    #[test]
    fn test_tokio_try_join() {
        let input = quote! {
            async fn with_try_join() -> Result<(), Error> {
                let fut1 = async_operation1();
                let fut2 = async_operation2();

                let (result1, result2) = tokio::try_join!(fut1, fut2)?;
                process_results(result1, result2);
                Ok(())
            }
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output, @r###"
        async fn with_try_join() -> Result<(), Error> {
            let mut __guard = crate::__async_profile_guard__::Guard::new(
                concat!(module_path!(), "::", stringify!(with_try_join)),
                file!(),
                1u32,
            );
            let fut1 = async_operation1();
            let fut2 = async_operation2();
            let (result1, result2) = {
                __guard.end_section(1u32);
                let __result = tokio::try_join!(fut1, fut2);
                __guard.start_section(1u32);
                __result
            }?;
            process_results(result1, result2);
            Ok(())
        }
        "###);
    }

    #[test]
    fn test_futures_join() {
        let input = quote! {
            async fn with_futures_join() {
                use futures::join;

                let (a, b) = join!(async_a(), async_b());
                combine(a, b);
            }
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert_snapshot!(output, @r###"
        async fn with_futures_join() {
            let mut __guard = crate::__async_profile_guard__::Guard::new(
                concat!(module_path!(), "::", stringify!(with_futures_join)),
                file!(),
                1u32,
            );
            use futures::join;
            let (a, b) = {
                __guard.end_section(1u32);
                let __result = join!(async_a(), async_b());
                __guard.start_section(1u32);
                __result
            };
            combine(a, b);
        }
        "###);
    }

    #[test]
    fn test_archives_gc_pattern() {
        // Read the test file with complex async patterns
        let content = std::fs::read_to_string("test_fixtures/archives_gc.rs")
            .expect("Failed to read test_fixtures/archives_gc.rs");

        let mut file: File = syn::parse_file(&content).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);

        // Use snapshot for this complex transformation
        // The snapshot file will be created/updated with cargo insta review
        assert_snapshot!(output);
    }

    #[test]
    fn test_broadcaster_run_with_biased_select() {
        // Read the broadcaster test fixture
        let content = std::fs::read_to_string("test_fixtures/broadcaster_run.rs")
            .expect("Failed to read test_fixtures/broadcaster_run.rs");

        let mut file: File = syn::parse_file(&content).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);

        // Use snapshot for this complex biased select transformation
        assert_snapshot!(output);
    }

    #[test]
    fn test_select_biased_macro() {
        let input = quote! {
            async fn with_select_biased() {
                futures::select_biased! {
                    _ = async_operation() => handle(),
                }
            }
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert!(output.contains("let __result = futures::select_biased!"));
        assert!(output.contains("futures::select_biased!"));
    }

    #[test]
    fn test_race_macro() {
        let input = quote! {
            async fn with_race() {
                let value = futures::future::race!(future_one(), future_two());
                consume(value);
            }
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        assert!(
            output.contains("let __result = futures::future::race!(future_one(), future_two())")
        );
        assert!(output.contains("futures::future::race!(future_one(), future_two())"));
    }

    #[test]
    fn test_wait_with_sleep_pattern() {
        let input = quote! {
            async fn wait_with_sleep(
                tick_rx: &mut TickRx,
                manual_rx: &mut ManualTriggerRx,
                sleep_until: Option<Instant>,
            ) -> Option<GcSource> {
                use futures_util::future::Either;

                let fut = match sleep_until {
                    Some(deadline) => Either::Left(async move {
                        tokio::time::sleep_until(deadline.into()).await;
                        Some(GcSource::Schedule)
                    }),
                    None => Either::Right(async {
                        let res = tick_rx.changed().await;
                        res.is_ok().then_some(GcSource::Schedule)
                    }),
                };

                tokio::select! {
                    res = fut => res,
                    trigger = manual_rx.changed() => {
                        trigger.is_ok().then_some(GcSource::Manual)
                    },
                }
            }
        };

        let mut file: File = syn::parse2(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut file);

        let output = prettyplease::unparse(&file);
        // This test is complex so we just verify key parts
        // The transformation is correct - select! returns value without semicolon
        assert!(output.contains("wait_with_sleep"));
        assert!(output.contains("Either::Left(async move"));
        assert!(output.contains("Either::Right(async"));
        assert!(output.contains("tokio::select!"));
        // Verify it ends with an expression (no semicolon after the select block)
        assert!(output.contains("__result\n    }\n}"));
    }

    #[test]
    fn test_loop() {
        let input = include_str!("../test_fixtures/loop_ex.rs");
        let mut input: File = syn::parse_file(input).unwrap();
        let mut instrumenter = AsyncInstrumenter::new(10);
        instrumenter.instrument_file(&mut input);

        let output = prettyplease::unparse(&input);
        println!("{output}")
    }
}