ryo_executor/engine/impls/

organize_imports.rs

//! V2 ASTRegApply implementation for OrganizeImportsMutation
//!
//! # Implementation Notes
//!
//! OrganizeImports operates on module_items stored in ASTRegistry.
//! This enables access to use statements which are not individual symbols.
//!
//! The mutation:
//! 1. Iterates through all modules with stored module_items
//! 2. Extracts use statements from each module
//! 3. Applies deduplication and/or merging based on options
//! 4. Updates the module_items with reorganized uses

use ryo_analysis::SymbolKind;
use ryo_mutations::idiom::OrganizeImportsMutation;
use ryo_mutations::{Mutation, MutationResult};
use ryo_source::pure::{PureItem, PureUse, PureUseTree, PureVis};
use std::collections::BTreeMap;

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

impl ASTRegApply for OrganizeImportsMutation {
    fn apply_to_registry(&self, ctx: &mut ASTMutationContext) -> MutationResult {
        let mut total_changes = 0;

        // Collect module IDs that have module_items
        // When target_module is set, restrict to that module only (no leak to
        // unrelated modules in the same workspace).
        let module_ids: Vec<_> = if let Some(target) = self.target_module {
            if matches!(ctx.symbol_registry.kind(target), Some(SymbolKind::Mod))
                && ctx.ast_registry.has_module_items(target)
            {
                vec![target]
            } else {
                vec![]
            }
        } else {
            ctx.symbol_registry
                .iter()
                .filter(|(id, _)| {
                    matches!(ctx.symbol_registry.kind(*id), Some(SymbolKind::Mod))
                        && ctx.ast_registry.has_module_items(*id)
                })
                .map(|(id, _)| id)
                .collect()
        };

        let module_count = module_ids.len();

        // Process each module
        for module_id in module_ids {
            if let Some(module_items) = ctx.ast_registry.get_module_items_mut(module_id) {
                let changes =
                    organize_module_items(module_items, self.deduplicate, self.merge_groups);
                if changes > 0 {
                    // Emit event so sync_files_and_rebuild regenerates this module's file
                    ctx.emit_modified(
                        module_id,
                        crate::engine::events::ModificationType::Other(
                            "imports organized".to_string(),
                        ),
                    );
                }
                total_changes += changes;
            }
        }

        MutationResult {
            mutation_type: self.mutation_type().to_string(),
            changes: total_changes,
            description: if total_changes > 0 {
                format!("Organized imports in {} module(s)", module_count)
            } else {
                "No imports to organize".to_string()
            },
        }
    }
}

/// Organize use statements within module items
///
/// Returns number of changes made.
fn organize_module_items(
    items: &mut Vec<PureItem>,
    deduplicate: bool,
    merge_groups: bool,
) -> usize {
    // Extract use statements
    let mut uses: Vec<PureUse> = Vec::new();
    let mut other_items: Vec<PureItem> = Vec::new();

    for item in items.drain(..) {
        if let PureItem::Use(u) = item {
            uses.push(u);
        } else {
            other_items.push(item);
        }
    }

    let original_count = uses.len();

    if uses.is_empty() {
        *items = other_items;
        return 0;
    }

    // Deduplicate if requested
    if deduplicate {
        let mut seen = Vec::new();
        uses.retain(|u| {
            let dominated = seen
                .iter()
                .any(|existing: &PureUseTree| trees_equal(existing, &u.tree));
            if !dominated {
                seen.push(u.tree.clone());
                true
            } else {
                false
            }
        });
    }

    // Merge groups if requested
    if merge_groups {
        // Group by visibility first
        let mut pub_uses: Vec<PureUseTree> = Vec::new();
        let mut priv_uses: Vec<PureUseTree> = Vec::new();

        for u in uses {
            match u.vis {
                PureVis::Private => priv_uses.push(u.tree),
                _ => pub_uses.push(u.tree),
            }
        }

        // Merge each group
        let merged_priv = merge_trees(priv_uses);
        let merged_pub = merge_trees(pub_uses);

        uses = merged_priv
            .into_iter()
            .map(|tree| PureUse {
                vis: PureVis::Private,
                tree,
            })
            .chain(merged_pub.into_iter().map(|tree| PureUse {
                vis: PureVis::Public,
                tree,
            }))
            .collect();
    }

    // Sort uses for deterministic output
    uses.sort_by(use_cmp);

    // Calculate changes
    let changes = if uses.len() != original_count {
        original_count.saturating_sub(uses.len())
    } else {
        0
    };

    // Rebuild items: uses first, then other items
    items.clear();
    items.extend(uses.into_iter().map(PureItem::Use));
    items.extend(other_items);

    changes
}

/// Check if two use trees are structurally equal
fn trees_equal(a: &PureUseTree, b: &PureUseTree) -> bool {
    match (a, b) {
        (PureUseTree::Path { path: pa, tree: ta }, PureUseTree::Path { path: pb, tree: tb }) => {
            pa == pb && trees_equal(ta, tb)
        }
        (PureUseTree::Name(na), PureUseTree::Name(nb)) => na == nb,
        (
            PureUseTree::Rename {
                name: na,
                rename: ra,
            },
            PureUseTree::Rename {
                name: nb,
                rename: rb,
            },
        ) => na == nb && ra == rb,
        (PureUseTree::Glob, PureUseTree::Glob) => true,
        (PureUseTree::Group(ga), PureUseTree::Group(gb)) => {
            ga.len() == gb.len() && ga.iter().zip(gb.iter()).all(|(a, b)| trees_equal(a, b))
        }
        _ => false,
    }
}

/// Merge use trees by common prefix
fn merge_trees(trees: Vec<PureUseTree>) -> Vec<PureUseTree> {
    if trees.is_empty() {
        return vec![];
    }

    // Group by first path segment
    let mut groups: BTreeMap<String, Vec<PureUseTree>> = BTreeMap::new();
    let mut ungroupable: Vec<PureUseTree> = Vec::new();

    for tree in trees {
        match &tree {
            PureUseTree::Path {
                path,
                tree: subtree,
            } => {
                groups
                    .entry(path.clone())
                    .or_default()
                    .push((**subtree).clone());
            }
            _ => ungroupable.push(tree),
        }
    }

    // Build merged trees
    let mut result = ungroupable;
    for (path, subtrees) in groups {
        if subtrees.len() == 1 {
            result.push(PureUseTree::Path {
                path,
                tree: Box::new(
                    subtrees
                        .into_iter()
                        .next()
                        .expect("len() == 1 guard above guarantees at least one element"),
                ),
            });
        } else {
            // Recursively merge subtrees
            let merged = merge_trees(subtrees);
            result.push(PureUseTree::Path {
                path,
                tree: Box::new(PureUseTree::Group(merged)),
            });
        }
    }

    result
}

/// Compare use statements for sorting
fn use_cmp(a: &PureUse, b: &PureUse) -> std::cmp::Ordering {
    // pub uses come after private uses
    let vis_cmp = match (&a.vis, &b.vis) {
        (PureVis::Private, PureVis::Private) => std::cmp::Ordering::Equal,
        (PureVis::Private, _) => std::cmp::Ordering::Less,
        (_, PureVis::Private) => std::cmp::Ordering::Greater,
        _ => std::cmp::Ordering::Equal,
    };

    if vis_cmp != std::cmp::Ordering::Equal {
        return vis_cmp;
    }

    // Then sort by tree string representation
    tree_cmp(&a.tree, &b.tree)
}

fn tree_cmp(a: &PureUseTree, b: &PureUseTree) -> std::cmp::Ordering {
    let a_str = tree_to_string(a);
    let b_str = tree_to_string(b);
    a_str.cmp(&b_str)
}

fn tree_to_string(tree: &PureUseTree) -> String {
    match tree {
        PureUseTree::Path { path, tree } => format!("{}::{}", path, tree_to_string(tree)),
        PureUseTree::Name(name) => name.clone(),
        PureUseTree::Rename { name, rename } => format!("{} as {}", name, rename),
        PureUseTree::Glob => "*".to_string(),
        PureUseTree::Group(trees) => {
            let parts: Vec<_> = trees.iter().map(tree_to_string).collect();
            format!("{{{}}}", parts.join(", "))
        }
    }
}

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

    #[test]
    fn test_v2_organize_imports_target_module_non_mod_id_no_leak() {
        // Boundary regression: when target_module is set to a non-Mod kind id
        // (e.g. a Function id via scope_spec_to_function defensive setting),
        // the impl must not iterate any module. Pre-fix the converter discarded
        // module_id via `..` and the impl iterated every module unconditionally,
        // leaking import reorganization across the workspace.
        let mut ctx = ContextBuilder::new()
            .with_file(
                "src/lib.rs",
                r#"
use std::io;
use std::io;

fn foo() {}
"#,
            )
            .build();

        let foo_id = ctx
            .registry
            .iter()
            .find(|(id, path)| {
                matches!(ctx.registry.kind(*id), Some(SymbolKind::Function)) && path.name() == "foo"
            })
            .map(|(id, _)| id)
            .expect("foo function not found");

        // Pass a Function id as target_module — must result in 0 changes
        // because the impl filters to SymbolKind::Mod only.
        let mutation = OrganizeImportsMutation::new().in_module(foo_id);
        let result = ASTMutationEngine::execute_ast_reg(&mutation, &mut ctx);

        assert_eq!(
            result.result.changes, 0,
            "non-Mod target_module must not leak"
        );
    }

    fn make_use(path: &str) -> PureUse {
        let parts: Vec<&str> = path.split("::").collect();
        let tree = parts
            .iter()
            .rev()
            .fold(None, |acc, part| {
                Some(match acc {
                    None => PureUseTree::Name(part.to_string()),
                    Some(subtree) => PureUseTree::Path {
                        path: part.to_string(),
                        tree: Box::new(subtree),
                    },
                })
            })
            .unwrap();

        PureUse {
            vis: PureVis::Private,
            tree,
        }
    }

    #[test]
    fn test_organize_deduplicates() {
        let mut items = vec![
            PureItem::Use(make_use("std::io")),
            PureItem::Use(make_use("std::io")),
            PureItem::Use(make_use("std::fs")),
        ];

        let changes = organize_module_items(&mut items, true, false);

        assert_eq!(changes, 1);
        assert_eq!(
            items
                .iter()
                .filter(|i| matches!(i, PureItem::Use(_)))
                .count(),
            2
        );
    }

    #[test]
    fn test_organize_no_changes_when_unique() {
        let mut items = vec![
            PureItem::Use(make_use("std::io")),
            PureItem::Use(make_use("std::fs")),
        ];

        let changes = organize_module_items(&mut items, true, false);

        assert_eq!(changes, 0);
    }

    #[test]
    fn test_trees_equal() {
        let tree1 = PureUseTree::Name("foo".to_string());
        let tree2 = PureUseTree::Name("foo".to_string());

        assert!(trees_equal(&tree1, &tree2));
    }

    #[test]
    fn test_trees_not_equal_different_names() {
        let tree1 = PureUseTree::Name("foo".to_string());
        let tree2 = PureUseTree::Name("bar".to_string());

        assert!(!trees_equal(&tree1, &tree2));
    }
}