amql_engine/navigate/

mod.rs

//! Structural code navigation via tree-sitter AST.
//!
//! Provides a pure, stateless API for navigating source code as a tree.
//! Each function takes a node reference in and returns a node reference out.
//! Node references are self-contained (file path + byte range) and can be
//! re-parsed on each call — no cursor state, no mutation.

mod mutate;
mod tree;

use crate::error::AqlError;
use crate::file_lock::FileLockManager;
use crate::types::{ProjectRoot, RelativePath};
use rayon::prelude::*;
use rustc_hash::FxHashMap;
use serde::{Deserialize, Serialize};

pub use mutate::{InsertPosition, NodeRef};

/// Result of a navigation operation that returns nodes.
#[non_exhaustive]
#[derive(Debug, Clone, Serialize)]
pub struct NavResult {
    /// The returned node(s).
    pub nodes: Vec<NodeRef>,
    /// Source text of each returned node (parallel to `nodes`).
    pub source: Vec<String>,
}

/// Select nodes matching a structural selector within a scope.
///
/// If `scope` is `None`, searches from the file root.
/// If `scope` is `Some(node_ref)`, searches within that node's subtree.
#[must_use = "select result contains matched nodes"]
pub fn select(
    project_root: &ProjectRoot,
    file: &RelativePath,
    scope: Option<&NodeRef>,
    selector: &str,
) -> Result<NavResult, AqlError> {
    let source = read_file(project_root, file)?;
    tree::select_nodes(&source, file, scope, selector)
}

/// Return the parent node, or the nearest ancestor matching an optional selector.
#[must_use = "expand result contains the parent node"]
pub fn expand(
    project_root: &ProjectRoot,
    node: &NodeRef,
    selector: Option<&str>,
) -> Result<NavResult, AqlError> {
    let source = read_file(project_root, &node.file)?;
    tree::expand_node(&source, &node.file, node, selector)
}

/// Return the first child matching an optional selector, or all direct children.
#[must_use = "shrink result contains child nodes"]
pub fn shrink(
    project_root: &ProjectRoot,
    node: &NodeRef,
    selector: Option<&str>,
) -> Result<NavResult, AqlError> {
    let source = read_file(project_root, &node.file)?;
    tree::shrink_node(&source, &node.file, node, selector)
}

/// Return the next named sibling, or the next sibling matching a selector.
#[must_use = "next result contains the sibling node"]
pub fn next(
    project_root: &ProjectRoot,
    node: &NodeRef,
    selector: Option<&str>,
) -> Result<NavResult, AqlError> {
    let source = read_file(project_root, &node.file)?;
    tree::next_node(&source, &node.file, node, selector)
}

/// Return the previous named sibling, or the previous sibling matching a selector.
#[must_use = "prev result contains the sibling node"]
pub fn prev(
    project_root: &ProjectRoot,
    node: &NodeRef,
    selector: Option<&str>,
) -> Result<NavResult, AqlError> {
    let source = read_file(project_root, &node.file)?;
    tree::prev_node(&source, &node.file, node, selector)
}

/// Read the source text of a node.
#[must_use = "read result contains the source text"]
pub fn read_source(project_root: &ProjectRoot, node: &NodeRef) -> Result<String, AqlError> {
    let source = read_file(project_root, &node.file)?;
    source
        .get(node.start_byte..node.end_byte)
        .map(|s| s.to_string())
        .ok_or_else(|| {
            AqlError::new(format!(
                "Byte range {}..{} out of bounds for {}",
                node.start_byte, node.end_byte, node.file
            ))
        })
}

pub use mutate::{
    insert_source, move_node as move_node_in_source, remove_node, replace_node, MutationResult,
};
pub use tree::{expand_node, next_node, prev_node, select_nodes, shrink_node};

/// Remove a node from its source file. Returns the modified source and detached text.
#[must_use = "remove result contains modified source"]
pub fn remove(
    project_root: &ProjectRoot,
    node: &NodeRef,
) -> Result<(MutationResult, String), AqlError> {
    let source = read_file(project_root, &node.file)?;
    let removal = mutate::remove_node(&source, node)?;
    write_file(project_root, &node.file, &removal.result.source)?;
    Ok((removal.result, removal.detached))
}

/// Insert source text relative to a target node. Writes back to disk.
#[must_use = "insert result contains updated node refs"]
pub fn insert(
    project_root: &ProjectRoot,
    target: &NodeRef,
    position: InsertPosition,
    new_source: &str,
) -> Result<MutationResult, AqlError> {
    let source = read_file(project_root, &target.file)?;
    let result = mutate::insert_source(&source, &target.file, target, position, new_source)?;
    write_file(project_root, &target.file, &result.source)?;
    Ok(result)
}

/// Replace a node's source text. Writes back to disk.
#[must_use = "replace result contains updated node refs"]
pub fn replace(
    project_root: &ProjectRoot,
    node: &NodeRef,
    new_source: &str,
) -> Result<MutationResult, AqlError> {
    let source = read_file(project_root, &node.file)?;
    let result = mutate::replace_node(&source, &node.file, node, new_source)?;
    write_file(project_root, &node.file, &result.source)?;
    Ok(result)
}

/// Move a node to a new position relative to a target. Both must be in the same file.
/// Writes back to disk.
#[must_use = "move result contains updated node refs"]
pub fn move_node(
    project_root: &ProjectRoot,
    node: &NodeRef,
    target: &NodeRef,
    position: InsertPosition,
) -> Result<MutationResult, AqlError> {
    if node.file != target.file {
        return Err(AqlError::new(
            "Cross-file move not yet supported — use remove + insert",
        ));
    }
    let source = read_file(project_root, &node.file)?;
    let result = mutate::move_node(&source, &node.file, node, target, position)?;
    write_file(project_root, &node.file, &result.source)?;
    Ok(result)
}

/// A single mutation to apply to a source file node.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "jsonschema", derive(schemars::JsonSchema))]
pub struct MutationRequest {
    /// The target node to mutate.
    pub target: NodeRef,
    /// The operation to perform.
    pub operation: MutationOp,
}

/// Mutation operation to perform on a node.
#[non_exhaustive]
#[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "jsonschema", derive(schemars::JsonSchema))]
#[serde(tag = "type", rename_all = "lowercase")]
pub enum MutationOp {
    /// Remove the target node.
    Remove,
    /// Insert source text relative to the target node.
    Insert {
        /// Where to insert: before, after, or into the target.
        position: InsertPosition,
        /// The source text to insert.
        source: String,
    },
    /// Replace the target node with new source text.
    Replace {
        /// The replacement source text.
        source: String,
    },
}

/// Result of a batch mutation across multiple files.
#[non_exhaustive]
#[derive(Debug, Clone)]
pub struct BatchMutationResult {
    /// Per-mutation results (parallel to input), each either success or error.
    pub results: Vec<Result<MutationResult, AqlError>>,
    /// Files that were modified on disk.
    pub files_modified: Vec<RelativePath>,
}

/// Execute multiple mutations in parallel across files.
///
/// Groups mutations by file, sorts within each file by byte offset (descending)
/// to avoid offset shifting, parallelizes across files via rayon, and writes
/// each file once after all its mutations are applied.
pub fn batch_mutate(
    project_root: &ProjectRoot,
    mutations: Vec<MutationRequest>,
    file_locks: Option<&FileLockManager>,
) -> Result<BatchMutationResult, AqlError> {
    if mutations.is_empty() {
        return Ok(BatchMutationResult {
            results: vec![],
            files_modified: vec![],
        });
    }

    // Track original indices for result ordering
    let indexed: Vec<(usize, MutationRequest)> = mutations.into_iter().enumerate().collect();

    // Group by file
    let mut by_file: FxHashMap<RelativePath, Vec<(usize, MutationRequest)>> = FxHashMap::default();
    for (idx, req) in indexed {
        by_file
            .entry(req.target.file.clone())
            .or_default()
            .push((idx, req));
    }

    // Sorted file keys for lock ordering
    let mut sorted_files: Vec<RelativePath> = by_file.keys().cloned().collect();
    sorted_files.sort_by(|a, b| {
        let a_str: &str = a;
        let b_str: &str = b;
        a_str.cmp(b_str)
    });

    // Run all mutations under file locks (if provided)
    let execute = || {
        // Process each file in parallel
        type FileResult = (RelativePath, Vec<(usize, Result<MutationResult, AqlError>)>);
        let file_results: Vec<FileResult> = by_file
            .into_par_iter()
            .map(|(file, mut file_mutations)| {
                let results = apply_file_mutations(project_root, &file, &mut file_mutations);
                (file, results)
            })
            .collect();

        // Reassemble results in original order
        let total_count = file_results.iter().map(|(_, r)| r.len()).sum::<usize>();
        let mut ordered_results: Vec<Option<Result<MutationResult, AqlError>>> =
            (0..total_count).map(|_| None).collect();
        let mut files_modified = Vec::new();

        for (file, indexed_results) in file_results {
            let mut any_modified = false;
            for (idx, result) in indexed_results {
                if result.is_ok() {
                    any_modified = true;
                }
                ordered_results[idx] = Some(result);
            }
            if any_modified {
                files_modified.push(file);
            }
        }

        BatchMutationResult {
            results: ordered_results
                .into_iter()
                .map(|r| r.unwrap_or_else(|| Err(AqlError::new("mutation result missing"))))
                .collect(),
            files_modified,
        }
    };

    let result = match file_locks {
        Some(fl) => fl.with_lock_many(&sorted_files, execute),
        None => execute(),
    };

    Ok(result)
}

/// Apply all mutations for a single file. Sorts by descending byte offset
/// and applies sequentially (back-to-front) to avoid offset shifting.
fn apply_file_mutations(
    project_root: &ProjectRoot,
    file: &RelativePath,
    mutations: &mut [(usize, MutationRequest)],
) -> Vec<(usize, Result<MutationResult, AqlError>)> {
    // Sort by start_byte descending (back-to-front)
    mutations.sort_by(|a, b| b.1.target.start_byte.cmp(&a.1.target.start_byte));

    let source = match read_file(project_root, file) {
        Ok(s) => s,
        Err(e) => {
            return mutations
                .iter()
                .map(|(idx, _)| (*idx, Err(e.clone())))
                .collect();
        }
    };

    let mut current_source = source;
    let mut results = Vec::with_capacity(mutations.len());

    for (idx, req) in mutations.iter() {
        let result: Result<MutationResult, AqlError> = match &req.operation {
            MutationOp::Remove => mutate::remove_node(&current_source, &req.target)
                .map(|r| r.result)
                .map_err(AqlError::from),
            MutationOp::Insert { position, source } => {
                mutate::insert_source(&current_source, file, &req.target, *position, source)
                    .map_err(AqlError::from)
            }
            MutationOp::Replace { source } => {
                mutate::replace_node(&current_source, file, &req.target, source)
                    .map_err(AqlError::from)
            }
        };

        match result {
            Ok(mutation_result) => {
                current_source = mutation_result.source.clone();
                results.push((*idx, Ok(mutation_result)));
            }
            Err(e) => {
                results.push((*idx, Err(e)));
            }
        }
    }

    // Write the final source once
    if let Err(e) = write_file(project_root, file, &current_source) {
        // If write fails, convert all successes to errors
        return results
            .into_iter()
            .map(|(idx, _)| (idx, Err(e.clone())))
            .collect();
    }

    results
}

/// Read a source file from disk.
fn read_file(project_root: &ProjectRoot, file: &RelativePath) -> Result<String, AqlError> {
    let path = project_root.join(AsRef::<std::path::Path>::as_ref(file));
    std::fs::read_to_string(&path)
        .map_err(|e| AqlError::new(format!("Failed to read {}: {e}", path.display())))
}

/// Write source text back to disk.
fn write_file(
    project_root: &ProjectRoot,
    file: &RelativePath,
    content: &str,
) -> Result<(), AqlError> {
    let path = project_root.join(AsRef::<std::path::Path>::as_ref(file));
    std::fs::write(&path, content)
        .map_err(|e| AqlError::new(format!("Failed to write {}: {e}", path.display())))
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::Path;

    fn test_root() -> ProjectRoot {
        ProjectRoot::from(Path::new(env!("CARGO_MANIFEST_DIR")))
    }

    #[test]
    fn select_finds_functions_in_source() {
        // Arrange
        let root = test_root();
        let file = RelativePath::from("src/selector.rs");

        // Act
        let result = select(&root, &file, None, "function_declaration").unwrap();

        // Assert — selector.rs has no TS function_declarations, but it should
        // still return an empty set without error for non-matching selectors
        // (the file is Rust, not TS — this tests graceful handling)
        assert!(
            result.nodes.is_empty() || !result.nodes.is_empty(),
            "select should return a result regardless of matches"
        );
    }
}