ryo_executor/engine/impls/

method.rs

//! ASTRegApply implementations for AddMethodMutation and RemoveMethodMutation
//!
//! New design: Methods are registered directly on types (Struct/Enum) as Type::method.
//! Impl blocks are file-level constructs added to module_items for code generation.

use ryo_mutations::{AddMethodMutation, MutationResult, RemoveMethodMutation};
use ryo_source::pure::{
    PureBlock, PureExpr, PureFn, PureGenerics, PureImpl, PureImplItem, PureItem, PureParam,
    PureStmt, PureType, PureVis,
};
use ryo_symbol::SymbolKind;

use crate::engine::{ASTMutationContext, ASTRegApply, ModificationType};

// ============================================================================
// Helper functions
// ============================================================================

fn parse_type_simple(s: &str) -> PureType {
    let s = s.trim();

    // Handle reference types: &str, &mut T, &'a T, &[T]
    if let Some(rest) = s.strip_prefix('&') {
        let rest = rest.trim_start();

        // Check for lifetime: &'a T
        if rest.starts_with('\'') {
            // Find end of lifetime
            let lifetime_end = rest.find(|c: char| c.is_whitespace()).unwrap_or(rest.len());
            let lifetime = rest[..lifetime_end].to_string();
            let rest = rest[lifetime_end..].trim_start();

            // Check for mut after lifetime
            if let Some(inner) = rest.strip_prefix("mut ") {
                return PureType::Ref {
                    lifetime: Some(lifetime),
                    is_mut: true,
                    ty: Box::new(parse_type_simple(inner)),
                };
            } else {
                return PureType::Ref {
                    lifetime: Some(lifetime),
                    is_mut: false,
                    ty: Box::new(parse_type_simple(rest)),
                };
            }
        }

        // Check for mut: &mut T
        if let Some(inner) = rest.strip_prefix("mut ") {
            return PureType::Ref {
                lifetime: None,
                is_mut: true,
                ty: Box::new(parse_type_simple(inner)),
            };
        }

        // Simple reference: &T (including &[T])
        return PureType::Ref {
            lifetime: None,
            is_mut: false,
            ty: Box::new(parse_type_simple(rest)),
        };
    }

    // Handle slice types: [T]
    if s.starts_with('[') && s.ends_with(']') {
        let inner = &s[1..s.len() - 1];
        return PureType::Slice(Box::new(parse_type_simple(inner)));
    }

    // Handle tuple types: (), (A,), (A, B), etc.
    if s.starts_with('(') && s.ends_with(')') {
        let inner = s[1..s.len() - 1].trim();
        if inner.is_empty() {
            // Unit type ()
            return PureType::Tuple(vec![]);
        }
        // Parse tuple elements (simple split by comma - may not handle nested types correctly)
        let elements: Vec<PureType> = inner
            .split(',')
            .map(|e| parse_type_simple(e.trim()))
            .collect();
        return PureType::Tuple(elements);
    }

    // Handle Option<T>, Result<T, E>, Vec<T> etc. - keep as path for now
    // (These are valid path types)

    // Default: treat as path
    PureType::Path(s.to_string())
}

// ============================================================================
// ASTRegApply implementations
// ============================================================================

/// New design: Add method directly to Struct/Enum
///
/// Flow:
/// 1. Register method as Type::method in SymbolRegistry
/// 2. Store method AST in ASTRegistry
/// 3. Add/update plain impl block in parent module's module_items
impl ASTRegApply for AddMethodMutation {
    fn apply_to_registry(&self, ctx: &mut ASTMutationContext) -> MutationResult {
        // Verify type exists and is Struct/Enum
        let type_kind = ctx.symbol_registry.kind(self.type_id);
        if !matches!(type_kind, Some(SymbolKind::Struct | SymbolKind::Enum)) {
            return MutationResult {
                mutation_type: "AddMethod".to_string(),
                changes: 0,
                description: format!(
                    "Symbol {} is not a struct or enum (kind: {:?})",
                    self.type_id, type_kind
                ),
            };
        }

        // Get type path
        let type_path = match ctx.symbol_registry.path(self.type_id) {
            Some(path) => path.clone(),
            None => {
                return MutationResult {
                    mutation_type: "AddMethod".to_string(),
                    changes: 0,
                    description: format!("Type {} not found in registry", self.type_id),
                };
            }
        };

        // Build method path: Type::method
        let method_path = match type_path.child(&self.name) {
            Ok(path) => path,
            Err(_) => {
                return MutationResult {
                    mutation_type: "AddMethod".to_string(),
                    changes: 0,
                    description: format!("Failed to create method path for '{}'", self.name),
                };
            }
        };

        // Check if method already exists
        if ctx.symbol_registry.lookup(&method_path).is_some() {
            return MutationResult {
                mutation_type: "AddMethod".to_string(),
                changes: 0,
                description: format!(
                    "Method '{}' already exists on type {}",
                    self.name, type_path
                ),
            };
        }

        // Build parameters
        let mut fn_params = Vec::new();

        // Add self parameter if specified
        if let Some((is_ref, is_mut)) = self.takes_self {
            fn_params.push(PureParam::SelfValue { is_ref, is_mut });
        }

        // Add typed parameters
        for (name, ty) in &self.params {
            fn_params.push(PureParam::Typed {
                name: name.clone(),
                ty: parse_type_simple(ty),
            });
        }

        // Build return type
        let ret = self.return_type.as_ref().map(|ty| parse_type_simple(ty));

        // Create body block
        let body_wrapped = if self.body.trim().starts_with('{') {
            self.body.clone()
        } else {
            format!("{{ {} }}", self.body)
        };

        let body_block = PureBlock {
            stmts: vec![PureStmt::Expr(PureExpr::Other(body_wrapped))],
        };

        // Create the method
        let method_fn = PureFn {
            attrs: Vec::new(),
            vis: if self.is_pub {
                PureVis::Public
            } else {
                PureVis::Private
            },
            is_async: false,
            is_async_inferred: false,
            is_const: false,
            is_unsafe: false,
            name: self.name.clone(),
            generics: PureGenerics::default(),
            params: fn_params,
            ret,
            body: body_block,
            abi: None,
        };

        // Register method in SymbolRegistry + ASTRegistry
        let method_id = match ctx.register_with_ast(
            method_path.clone(),
            SymbolKind::Method,
            PureItem::Fn(method_fn.clone()),
        ) {
            Some(id) => id,
            None => {
                return MutationResult {
                    mutation_type: "AddMethod".to_string(),
                    changes: 0,
                    description: format!("Failed to register method '{}'", method_path),
                };
            }
        };

        // Add/update plain impl block in parent module's module_items
        let type_name = type_path.name().to_string();

        // Get type's generics from its definition
        let type_generics = ctx
            .ast_registry
            .get(self.type_id)
            .and_then(|item| match item {
                PureItem::Struct(s) => Some(s.generics.clone()),
                PureItem::Enum(e) => Some(e.generics.clone()),
                _ => None,
            })
            .unwrap_or_default();

        // Build self_ty with generics for the impl block (e.g., "Foo<T>" instead of "Foo")
        let self_ty_with_generics = if type_generics.params.is_empty() {
            type_name.clone()
        } else {
            use ryo_source::pure::PureGenericParam;
            let param_names: Vec<String> = type_generics
                .params
                .iter()
                .map(|p| match p {
                    PureGenericParam::Type { name, .. } => name.clone(),
                    PureGenericParam::Lifetime { name, .. } => name.clone(),
                    PureGenericParam::Const { name, .. } => name.clone(),
                })
                .collect();
            format!("{}<{}>", type_name, param_names.join(", "))
        };

        if let Some(parent_path) = type_path.parent() {
            if let Some(parent_id) = ctx.symbol_registry.lookup(&parent_path) {
                let mut module_items = ctx
                    .ast_registry
                    .get_module_items(parent_id)
                    .cloned()
                    .unwrap_or_default();

                // Find or create plain impl block for this type
                // Match by base type name (without generics) for existing impl blocks
                let impl_block_index = module_items.iter().position(|item| {
                    if let PureItem::Impl(impl_block) = item {
                        // Match base type name (e.g., "Foo" matches "Foo<T>")
                        let base_self_ty = impl_block
                            .self_ty
                            .split('<')
                            .next()
                            .unwrap_or(&impl_block.self_ty);
                        base_self_ty == type_name && impl_block.trait_.is_none()
                    } else {
                        false
                    }
                });

                if let Some(idx) = impl_block_index {
                    // Add method to existing impl block
                    if let PureItem::Impl(impl_block) = &mut module_items[idx] {
                        impl_block.items.push(PureImplItem::Fn(method_fn));
                    }
                } else {
                    // Create new plain impl block with type's generics
                    let new_impl = PureImpl {
                        attrs: vec![],
                        generics: type_generics,
                        is_unsafe: false,
                        trait_: None,
                        self_ty: self_ty_with_generics,
                        items: vec![PureImplItem::Fn(method_fn)],
                    };
                    module_items.push(PureItem::Impl(new_impl));
                }

                ctx.ast_registry.set_module_items(parent_id, module_items);
            }
        }

        // Emit event
        ctx.emit_modified(method_id, ModificationType::MethodAdded(self.name.clone()));

        MutationResult {
            mutation_type: "AddMethod".to_string(),
            changes: 1,
            description: format!("Added method '{}' to type {}", self.name, type_path),
        }
    }
}

/// New design: Remove method directly from Struct/Enum
impl ASTRegApply for RemoveMethodMutation {
    fn apply_to_registry(&self, ctx: &mut ASTMutationContext) -> MutationResult {
        // Verify method exists
        if ctx.symbol_registry.kind(self.method_id) != Some(SymbolKind::Method) {
            return MutationResult {
                mutation_type: "RemoveMethod".to_string(),
                changes: 0,
                description: format!("Symbol {} is not a method", self.method_id),
            };
        }

        // Get method path
        let method_path = match ctx.symbol_registry.path(self.method_id) {
            Some(path) => path.clone(),
            None => {
                return MutationResult {
                    mutation_type: "RemoveMethod".to_string(),
                    changes: 0,
                    description: format!("Method {} not found", self.method_id),
                };
            }
        };

        let method_name = method_path.name().to_string();

        // Remove from impl block in module_items AND ASTRegistry
        // Handles both plain impl (Type::method) and trait impl (<impl Trait for Type>::method)
        if let Some(type_path) = method_path.parent() {
            let type_name = type_path.name().to_string();

            // Determine if this is a trait impl or plain impl
            let is_trait_impl = type_name.starts_with("<impl ");

            let (impl_path, type_name_for_match) = if is_trait_impl {
                // Trait impl: type_path IS the impl path (<impl Trait for Type>)
                // Extract type name from "<impl Trait for Type>"
                let type_name_extracted = if let Some(for_pos) = type_name.find(" for ") {
                    let after_for = &type_name[for_pos + 5..];
                    after_for.trim_end_matches('>').trim().to_string()
                } else {
                    type_name.clone()
                };
                (Some(type_path.clone()), type_name_extracted)
            } else {
                // Plain impl: construct impl path <impl Type>
                if let Some(parent_path) = type_path.parent() {
                    let impl_name = format!("<impl {}>", type_name);
                    let impl_path = parent_path.child(&impl_name).ok();
                    (impl_path, type_name)
                } else {
                    (None, type_name)
                }
            };

            // Get parent module
            // Both plain impl and trait impl: parent is type_path.parent()
            // For plain impl: Type::method -> Type -> module
            // For trait impl: <impl Trait for Type>::method -> <impl Trait for Type> -> module
            let parent_id = type_path
                .parent()
                .and_then(|p| ctx.symbol_registry.lookup(&p));

            // Update module_items
            if let Some(parent_id) = parent_id {
                if let Some(module_items) = ctx.ast_registry.get_module_items_mut(parent_id) {
                    for item in module_items.iter_mut() {
                        if let PureItem::Impl(impl_block) = item {
                            // Match by self_ty and trait status
                            let matches = if is_trait_impl {
                                impl_block.trait_.is_some()
                                    && impl_block.self_ty == type_name_for_match
                            } else {
                                impl_block.trait_.is_none()
                                    && impl_block.self_ty == type_name_for_match
                            };

                            if matches {
                                impl_block.items.retain(|impl_item| {
                                    if let PureImplItem::Fn(f) = impl_item {
                                        f.name != method_name
                                    } else {
                                        true
                                    }
                                });
                            }
                        }
                    }
                }
            }

            // Update ASTRegistry impl block
            if let Some(impl_path) = impl_path {
                if let Some(impl_id) = ctx.symbol_registry.lookup(&impl_path) {
                    if let Some(PureItem::Impl(impl_block)) = ctx.ast_registry.get_mut(impl_id) {
                        impl_block.items.retain(|impl_item| {
                            if let PureImplItem::Fn(f) = impl_item {
                                f.name != method_name
                            } else {
                                true
                            }
                        });
                    }
                }
            }
        }

        // Remove from SymbolRegistry and ASTRegistry manually
        // (can't use ctx.remove_symbol because it won't remove methods from impl blocks)
        ctx.symbol_registry.remove(self.method_id);
        ctx.ast_registry.remove(self.method_id);
        ctx.emit_removed(method_path.clone());

        MutationResult {
            mutation_type: "RemoveMethod".to_string(),
            changes: 1,
            description: format!("Removed method '{}'", method_path),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::engine::{multi_file_dumper, ASTMutationEngine};
    use ryo_analysis::testing::ContextBuilder;
    use ryo_symbol::WorkspaceFilePath;

    // =========================================================================
    // TDD: New design tests - Methods registered directly on type
    // =========================================================================

    /// Test AddMethodMutation with new design:
    /// - Methods are registered as Type::method
    /// - Impl block is added to module_items for file generation
    #[test]
    fn test_add_method_to_struct_new_design() {
        // Setup: Struct without impl block
        let mut ctx = ContextBuilder::new()
            .with_file(
                "src/lib.rs",
                r#"pub mod user;
"#,
            )
            .with_file(
                "src/user.rs",
                r#"pub struct User {
    pub name: String,
}
"#,
            )
            .build();

        // Find the User struct
        let user_path = ryo_symbol::SymbolPath::parse("test_crate::user::User").unwrap();
        let user_id = ctx.registry.lookup(&user_path).expect("User not found");

        // Add method "new" to User
        let mutation = AddMethodMutation::new(user_id, "new")
            .public()
            .with_params(vec![("name".to_string(), "String".to_string())])
            .with_return_type("Self")
            .with_body("Self { name }");

        let result = ASTMutationEngine::execute_ast_reg(&mutation, &mut ctx);
        assert!(result.has_changes(), "Expected changes");

        // Verify 1: Method registered as Type::method
        let method_path = ryo_symbol::SymbolPath::parse("test_crate::user::User::new").unwrap();
        let method_id = ctx
            .registry
            .lookup(&method_path)
            .expect("Method should be registered");
        assert_eq!(
            ctx.registry.kind(method_id),
            Some(SymbolKind::Method),
            "Method should have SymbolKind::Method"
        );

        // Verify 2: Method AST exists
        let method_ast = ctx.ast_registry.get(method_id);
        assert!(method_ast.is_some(), "Method AST should exist");
        assert!(
            matches!(method_ast, Some(PureItem::Fn(_))),
            "Method AST should be PureItem::Fn"
        );

        // Verify 3: Impl block added to module_items
        let user_module_path = ryo_symbol::SymbolPath::parse("test_crate::user").unwrap();
        let user_module_id = ctx
            .registry
            .lookup(&user_module_path)
            .expect("Module not found");
        let module_items = ctx
            .ast_registry
            .get_module_items(user_module_id)
            .expect("Module should have items");

        let has_impl_block = module_items.iter().any(|item| {
            if let PureItem::Impl(impl_block) = item {
                impl_block.self_ty == "User"
            } else {
                false
            }
        });
        assert!(has_impl_block, "Module should contain impl block for User");

        // Verify 4: File generation includes impl block
        let files = multi_file_dumper().dump_all(&ctx).unwrap();
        let user_file_path =
            WorkspaceFilePath::new_for_test("src/user.rs", ctx.workspace_root(), "test_crate");
        let user_content = files.get(&user_file_path).expect("user.rs should exist");

        assert!(
            user_content.contains("impl User"),
            "File should contain impl block. Got:\n{}",
            user_content
        );
        assert!(
            user_content.contains("pub fn new(name: String) -> Self"),
            "File should contain method signature. Got:\n{}",
            user_content
        );
    }

    /// Test AddMethodMutation for a generic struct.
    /// Verifies that impl block is generated with correct generics.
    #[test]
    fn test_add_method_to_generic_struct() {
        // Setup: Generic struct without impl block
        let mut ctx = ContextBuilder::new()
            .with_file(
                "src/lib.rs",
                r#"pub mod service;
"#,
            )
            .with_file(
                "src/service.rs",
                r#"pub trait Repository {
    fn find(&self, id: u64) -> Option<String>;
}

pub struct Service<R: Repository> {
    repository: R,
}
"#,
            )
            .build();

        // Find the Service struct
        let service_path = ryo_symbol::SymbolPath::parse("test_crate::service::Service").unwrap();
        let service_id = ctx
            .registry
            .lookup(&service_path)
            .expect("Service not found");

        // Add method "new" to Service
        let mutation = AddMethodMutation::new(service_id, "new")
            .public()
            .with_params(vec![("repository".to_string(), "R".to_string())])
            .with_return_type("Self")
            .with_body("Self { repository }");

        let result = ASTMutationEngine::execute_ast_reg(&mutation, &mut ctx);
        assert!(result.has_changes(), "Expected changes");

        // Verify: File generation includes impl block with generics
        let files = multi_file_dumper().dump_all(&ctx).unwrap();
        let service_file_path =
            WorkspaceFilePath::new_for_test("src/service.rs", ctx.workspace_root(), "test_crate");
        let service_content = files
            .get(&service_file_path)
            .expect("service.rs should exist");

        // The impl block should have the generic parameter
        assert!(
            service_content.contains("impl<R: Repository> Service<R>")
                || service_content.contains("impl<R : Repository> Service<R>"),
            "File should contain impl block with generics. Got:\n{}",
            service_content
        );
        assert!(
            service_content.contains("pub fn new(repository: R) -> Self"),
            "File should contain method signature. Got:\n{}",
            service_content
        );
    }

    /// Test adding method to a generic struct that was created via AddItem.
    /// This simulates the E2E scenario where struct and method are added in same execution.
    #[test]
    fn test_add_method_to_generic_struct_via_add_item() {
        use ryo_mutations::AddItemMutation;

        // Setup: Empty module
        let mut ctx = ContextBuilder::new()
            .with_file(
                "src/lib.rs",
                r#"pub mod repository;
pub mod service;
"#,
            )
            .with_file(
                "src/repository.rs",
                r#"pub trait InventoryRepository {
    fn find(&self, id: u64) -> Option<String>;
}
"#,
            )
            .with_file("src/service.rs", "//! Service module\n")
            .build();

        // Step 1: Add generic struct via AddItem
        let service_mod_path = ryo_symbol::SymbolPath::parse("test_crate::service").unwrap();
        let service_mod_id = ctx
            .registry
            .lookup(&service_mod_path)
            .expect("service module not found");

        let add_struct_mutation = AddItemMutation::new(
            service_mod_id,
            r#"pub struct InventoryService<R: crate::repository::InventoryRepository> {
    repository: R,
}"#
            .to_string(),
        );

        let result = ASTMutationEngine::execute_ast_reg(&add_struct_mutation, &mut ctx);
        assert!(result.has_changes(), "AddItem should add struct");

        // Step 2: Find the newly added struct
        let struct_path =
            ryo_symbol::SymbolPath::parse("test_crate::service::InventoryService").unwrap();
        let struct_id = ctx
            .registry
            .lookup(&struct_path)
            .expect("InventoryService not found after AddItem");

        // Debug: Verify struct generics are stored
        let struct_ast = ctx.ast_registry.get(struct_id);
        assert!(struct_ast.is_some(), "Struct AST should exist");
        if let Some(PureItem::Struct(s)) = struct_ast {
            assert!(
                !s.generics.params.is_empty(),
                "Struct should have generic params. Got: {:?}",
                s.generics
            );
        } else {
            panic!("Expected PureItem::Struct");
        }

        // Step 3: Add method via AddMethod
        let add_method_mutation = AddMethodMutation::new(struct_id, "new")
            .public()
            .with_params(vec![("repository".to_string(), "R".to_string())])
            .with_return_type("Self")
            .with_body("Self { repository }");

        let result = ASTMutationEngine::execute_ast_reg(&add_method_mutation, &mut ctx);
        assert!(result.has_changes(), "AddMethod should add method");

        // Verify: File generation includes impl block with generics
        let files = multi_file_dumper().dump_all(&ctx).unwrap();
        let service_file_path =
            WorkspaceFilePath::new_for_test("src/service.rs", ctx.workspace_root(), "test_crate");
        let service_content = files
            .get(&service_file_path)
            .expect("service.rs should exist");

        // The impl block should have the generic parameter with trait bound
        assert!(
            service_content
                .contains("impl<R: crate::repository::InventoryRepository> InventoryService<R>")
                || service_content.contains(
                    "impl<R : crate :: repository :: InventoryRepository> InventoryService<R>"
                )
                || service_content.contains(
                    "impl<R: crate :: repository :: InventoryRepository> InventoryService < R >"
                ),
            "File should contain impl block with generics. Got:\n{}",
            service_content
        );
    }

    /// Test RemoveMethodMutation with new design:
    /// - Method is removed from SymbolRegistry (Type::method)
    /// - Method AST is removed from ASTRegistry
    /// - Method is removed from impl block in module_items
    #[test]
    fn test_remove_method_from_struct_new_design() {
        // Setup: Struct with a method
        let mut ctx = ContextBuilder::new()
            .with_file(
                "src/lib.rs",
                r#"pub mod user;
"#,
            )
            .with_file(
                "src/user.rs",
                r#"pub struct User {
    pub name: String,
}

impl User {
    pub fn new(name: String) -> Self {
        Self { name }
    }

    pub fn get_name(&self) -> &str {
        &self.name
    }
}
"#,
            )
            .build();

        // Verify initial state: both methods exist
        let new_method_path = ryo_symbol::SymbolPath::parse("test_crate::user::User::new").unwrap();
        let get_name_path =
            ryo_symbol::SymbolPath::parse("test_crate::user::User::get_name").unwrap();

        assert!(
            ctx.registry.lookup(&new_method_path).is_some(),
            "Method 'new' should exist initially"
        );
        assert!(
            ctx.registry.lookup(&get_name_path).is_some(),
            "Method 'get_name' should exist initially"
        );

        // Remove method "get_name"
        let get_name_id = ctx
            .registry
            .lookup(&get_name_path)
            .expect("Method not found");
        let mutation = RemoveMethodMutation::new(get_name_id);

        let result = ASTMutationEngine::execute_ast_reg(&mutation, &mut ctx);
        assert!(result.has_changes(), "Expected changes");

        // Verify 1: Method removed from SymbolRegistry
        assert!(
            ctx.registry.lookup(&get_name_path).is_none(),
            "Method 'get_name' should be removed from SymbolRegistry"
        );

        // Verify 2: Method AST removed from ASTRegistry
        assert!(
            ctx.ast_registry.get(get_name_id).is_none(),
            "Method AST should be removed from ASTRegistry"
        );

        // Verify 3: Method removed from impl block in module_items
        let user_module_path = ryo_symbol::SymbolPath::parse("test_crate::user").unwrap();
        let user_module_id = ctx
            .registry
            .lookup(&user_module_path)
            .expect("Module not found");
        let module_items = ctx
            .ast_registry
            .get_module_items(user_module_id)
            .expect("Module should have items");

        // Check impl block still exists but without get_name
        let impl_block = module_items.iter().find_map(|item| {
            if let PureItem::Impl(impl_block) = item {
                if impl_block.self_ty == "User" {
                    Some(impl_block)
                } else {
                    None
                }
            } else {
                None
            }
        });

        assert!(impl_block.is_some(), "Impl block should still exist");
        let impl_block = impl_block.unwrap();

        // Should have only "new" method, not "get_name"
        let method_names: Vec<String> = impl_block
            .items
            .iter()
            .filter_map(|item| {
                if let PureImplItem::Fn(f) = item {
                    Some(f.name.clone())
                } else {
                    None
                }
            })
            .collect();

        assert!(
            method_names.contains(&"new".to_string()),
            "Method 'new' should still exist"
        );
        assert!(
            !method_names.contains(&"get_name".to_string()),
            "Method 'get_name' should be removed from impl block"
        );

        // Verify 4: File generation doesn't include removed method
        let files = multi_file_dumper().dump_all(&ctx).unwrap();
        let user_file_path =
            WorkspaceFilePath::new_for_test("src/user.rs", ctx.workspace_root(), "test_crate");
        let user_content = files.get(&user_file_path).expect("user.rs should exist");

        assert!(
            user_content.contains("impl User"),
            "File should contain impl block. Got:\n{}",
            user_content
        );
        assert!(
            user_content.contains("fn new("),
            "File should contain 'new' method. Got:\n{}",
            user_content
        );
        assert!(
            !user_content.contains("fn get_name("),
            "File should NOT contain 'get_name' method. Got:\n{}",
            user_content
        );
    }

    /// Test removing multiple methods sequentially
    #[test]
    fn test_remove_multiple_methods() {
        // Setup: Struct with 3 methods
        let mut ctx = ContextBuilder::new()
            .with_file(
                "src/lib.rs",
                r#"pub mod calc;
"#,
            )
            .with_file(
                "src/calc.rs",
                r#"pub struct Calculator {
    value: i32,
}

impl Calculator {
    pub fn new(value: i32) -> Self {
        Self { value }
    }

    pub fn add(&mut self, x: i32) {
        self.value += x;
    }

    pub fn multiply(&mut self, x: i32) {
        self.value *= x;
    }

    pub fn get_value(&self) -> i32 {
        self.value
    }
}
"#,
            )
            .build();

        // Remove "add" method
        let add_path = ryo_symbol::SymbolPath::parse("test_crate::calc::Calculator::add").unwrap();
        let add_id = ctx
            .registry
            .lookup(&add_path)
            .expect("Method 'add' not found");
        let mutation1 = RemoveMethodMutation::new(add_id);
        let result1 = ASTMutationEngine::execute_ast_reg(&mutation1, &mut ctx);
        assert!(result1.has_changes(), "First removal should succeed");

        // Remove "multiply" method
        let multiply_path =
            ryo_symbol::SymbolPath::parse("test_crate::calc::Calculator::multiply").unwrap();
        let multiply_id = ctx
            .registry
            .lookup(&multiply_path)
            .expect("Method 'multiply' not found");
        let mutation2 = RemoveMethodMutation::new(multiply_id);
        let result2 = ASTMutationEngine::execute_ast_reg(&mutation2, &mut ctx);
        assert!(result2.has_changes(), "Second removal should succeed");

        // Verify: Only "new" and "get_value" remain
        assert!(
            ctx.registry
                .lookup(
                    &ryo_symbol::SymbolPath::parse("test_crate::calc::Calculator::new").unwrap()
                )
                .is_some(),
            "Method 'new' should still exist"
        );
        assert!(
            ctx.registry
                .lookup(
                    &ryo_symbol::SymbolPath::parse("test_crate::calc::Calculator::get_value")
                        .unwrap()
                )
                .is_some(),
            "Method 'get_value' should still exist"
        );
        assert!(
            ctx.registry.lookup(&add_path).is_none(),
            "Method 'add' should be removed"
        );
        assert!(
            ctx.registry.lookup(&multiply_path).is_none(),
            "Method 'multiply' should be removed"
        );

        // Verify file output
        let files = multi_file_dumper().dump_all(&ctx).unwrap();
        let calc_path =
            WorkspaceFilePath::new_for_test("src/calc.rs", ctx.workspace_root(), "test_crate");
        let calc_content = files.get(&calc_path).expect("calc.rs should exist");

        assert!(
            calc_content.contains("fn new("),
            "Should contain 'new' method"
        );
        assert!(
            calc_content.contains("fn get_value("),
            "Should contain 'get_value' method"
        );
        assert!(
            !calc_content.contains("fn add("),
            "Should NOT contain 'add' method. Got:\n{}",
            calc_content
        );
        assert!(
            !calc_content.contains("fn multiply("),
            "Should NOT contain 'multiply' method. Got:\n{}",
            calc_content
        );
    }

    /// Test removing method from trait impl
    #[test]
    fn test_remove_method_from_trait_impl() {
        // Setup: Struct with trait impl (using simple trait name without ::)
        let mut ctx = ContextBuilder::new()
            .with_file(
                "src/lib.rs",
                r#"pub mod shapes;
"#,
            )
            .with_file(
                "src/shapes.rs",
                r#"pub trait Drawable {
    fn draw(&self) -> String;
    fn color(&self) -> String;
}

pub struct Circle {
    pub radius: f64,
}

impl Drawable for Circle {
    fn draw(&self) -> String {
        format!("Circle with radius {}", self.radius)
    }

    fn color(&self) -> String {
        "red".to_string()
    }
}
"#,
            )
            .build();

        // Find trait impl method: <impl Drawable for Circle>::color
        let color_path =
            ryo_symbol::SymbolPath::parse("test_crate::shapes::<impl Drawable for Circle>::color")
                .unwrap();
        let color_id = ctx
            .registry
            .lookup(&color_path)
            .expect("Method 'color' not found in trait impl");

        // Remove method "color"
        let mutation = RemoveMethodMutation::new(color_id);
        let result = ASTMutationEngine::execute_ast_reg(&mutation, &mut ctx);
        assert!(result.has_changes(), "Removal should succeed");

        // Verify: Method removed from SymbolRegistry
        assert!(
            ctx.registry.lookup(&color_path).is_none(),
            "Method 'color' should be removed from SymbolRegistry"
        );

        // Verify: "draw" method still exists
        let draw_path =
            ryo_symbol::SymbolPath::parse("test_crate::shapes::<impl Drawable for Circle>::draw")
                .unwrap();
        assert!(
            ctx.registry.lookup(&draw_path).is_some(),
            "Method 'draw' should still exist"
        );

        // Verify: trait impl block still exists with only "draw"
        let shapes_module_path = ryo_symbol::SymbolPath::parse("test_crate::shapes").unwrap();
        let shapes_module_id = ctx
            .registry
            .lookup(&shapes_module_path)
            .expect("Module not found");
        let module_items = ctx
            .ast_registry
            .get_module_items(shapes_module_id)
            .expect("Module should have items");

        // Check trait impl block exists with only "draw"
        let trait_impl = module_items.iter().find_map(|item| {
            if let PureItem::Impl(impl_block) = item {
                if impl_block.trait_.is_some() && impl_block.self_ty == "Circle" {
                    Some(impl_block)
                } else {
                    None
                }
            } else {
                None
            }
        });

        assert!(trait_impl.is_some(), "Trait impl block should exist");
        let trait_impl = trait_impl.unwrap();
        assert_eq!(
            trait_impl.items.len(),
            1,
            "Trait impl should have 1 method after removal"
        );

        // Verify file output
        let files = multi_file_dumper().dump_all(&ctx).unwrap();
        let shapes_path =
            WorkspaceFilePath::new_for_test("src/shapes.rs", ctx.workspace_root(), "test_crate");
        let shapes_content = files.get(&shapes_path).expect("shapes.rs should exist");

        assert!(
            shapes_content.contains("impl Drawable for Circle"),
            "Should contain trait impl block. Got:\n{}",
            shapes_content
        );
        assert!(
            shapes_content.contains("fn draw("),
            "Should contain 'draw' method. Got:\n{}",
            shapes_content
        );
        // Check that color method implementation is removed (trait definition should still have it)
        assert!(
            !shapes_content.contains("\"red\""),
            "Should NOT contain 'color' method implementation. Got:\n{}",
            shapes_content
        );
    }
}