amql_mutate/

lib.rs

//! Pure source code mutation operations.
//!
//! All functions are pure: source text + node references in, modified source +
//! updated node references out. No file I/O happens here — callers are
//! responsible for reading/writing files.

mod types;
#[cfg(feature = "wasm")]
mod index;

use serde::{Deserialize, Serialize};
pub use types::{NodeKind, RelativePath};

/// A self-contained reference to a node in a source file.
///
/// Encodes everything needed to locate and re-parse the node:
/// file path + byte range. Each MCP call receives and returns these.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[cfg_attr(feature = "jsonschema", derive(schemars::JsonSchema))]
#[cfg_attr(feature = "ts", derive(ts_rs::TS))]
#[cfg_attr(feature = "flow", derive(flowjs_rs::Flow))]
#[cfg_attr(feature = "ts", ts(export))]
#[cfg_attr(feature = "flow", flow(export))]
pub struct NodeRef {
    /// Relative path to the source file.
    pub file: RelativePath,
    /// Byte offset of the start of this node.
    pub start_byte: usize,
    /// Byte offset of the end of this node.
    pub end_byte: usize,
    /// tree-sitter node kind (e.g. "function_declaration", "class_declaration").
    pub kind: NodeKind,
    /// 1-based start line number.
    pub line: usize,
    /// 0-based start column offset.
    pub column: usize,
    /// 1-based end line number.
    pub end_line: usize,
    /// 0-based end column offset.
    pub end_column: usize,
}

/// Position relative to a target node for insertion.
#[non_exhaustive]
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[cfg_attr(feature = "jsonschema", derive(schemars::JsonSchema))]
#[cfg_attr(feature = "ts", derive(ts_rs::TS))]
#[cfg_attr(feature = "flow", derive(flowjs_rs::Flow))]
#[cfg_attr(feature = "ts", ts(export))]
#[cfg_attr(feature = "flow", flow(export))]
#[serde(rename_all = "lowercase")]
pub enum InsertPosition {
    /// Before the target node.
    Before,
    /// After the target node.
    After,
    /// As the first child inside the target node's body.
    Into,
}

/// Result of a mutation operation.
#[non_exhaustive]
#[derive(Debug, Clone, Serialize)]
#[cfg_attr(feature = "ts", derive(ts_rs::TS))]
#[cfg_attr(feature = "flow", derive(flowjs_rs::Flow))]
#[cfg_attr(feature = "ts", ts(export))]
#[cfg_attr(feature = "flow", flow(export))]
pub struct MutationResult {
    /// The modified source text.
    pub source: String,
    /// Updated node references for nodes affected by the mutation.
    /// After a mutation, previous node refs are stale — use these instead.
    pub affected_nodes: Vec<NodeRef>,
}

/// Result of a node removal — modified source and the extracted node text.
#[derive(Debug, Clone, Serialize)]
#[cfg_attr(feature = "ts", derive(ts_rs::TS))]
#[cfg_attr(feature = "flow", derive(flowjs_rs::Flow))]
#[cfg_attr(feature = "ts", ts(export))]
#[cfg_attr(feature = "flow", flow(export))]
pub struct RemoveResult {
    /// Mutation result containing the modified source.
    pub result: MutationResult,
    /// The text of the removed node.
    pub detached: String,
}

/// Remove a node from the source text, returning the modified source
/// and the detached node's text.
#[must_use = "remove result contains modified source and detached text"]
pub fn remove_node(source: &str, node: &NodeRef) -> Result<RemoveResult, String> {
    validate_range(source, node)?;

    let detached = source[node.start_byte..node.end_byte].to_string();

    // Remove the node and any trailing whitespace/newline
    let end = skip_trailing_whitespace(source, node.end_byte);
    let mut modified = String::with_capacity(source.len());
    modified.push_str(&source[..node.start_byte]);
    modified.push_str(&source[end..]);

    Ok(RemoveResult {
        result: MutationResult {
            source: modified,
            affected_nodes: vec![],
        },
        detached,
    })
}

/// Insert source text relative to a target node.
#[must_use = "insert result contains modified source and new node ref"]
pub fn insert_source(
    source: &str,
    file: &RelativePath,
    target: &NodeRef,
    position: InsertPosition,
    new_source: &str,
) -> Result<MutationResult, String> {
    validate_range(source, target)?;

    let (insert_at, prefix, suffix) = match position {
        InsertPosition::Before => {
            let indent = detect_indent(source, target.start_byte);
            (target.start_byte, String::new(), format!("\n{indent}"))
        }
        InsertPosition::After => {
            let indent = detect_indent(source, target.start_byte);
            (target.end_byte, format!("\n{indent}"), String::new())
        }
        InsertPosition::Into => {
            // Insert as last child inside the node's body
            // Find the closing brace/bracket
            let body_end = find_body_end(source, target);
            let indent = detect_indent(source, target.start_byte);
            let child_indent = format!("{indent}    ");
            (body_end, format!("\n{child_indent}"), String::new())
        }
    };

    let inserted = format!("{prefix}{new_source}{suffix}");
    let inserted_len = inserted.len();

    let mut result = String::with_capacity(source.len() + inserted_len);
    result.push_str(&source[..insert_at]);
    result.push_str(&inserted);
    result.push_str(&source[insert_at..]);

    // Build a node ref for the inserted content
    let new_start = insert_at + prefix.len();
    let new_end = new_start + new_source.len();
    let new_ref = build_node_ref_from_range(&result, file, new_start, new_end);

    Ok(MutationResult {
        source: result,
        affected_nodes: vec![new_ref],
    })
}

/// Replace a node's source text with new content.
#[must_use = "replace result contains modified source and new node ref"]
pub fn replace_node(
    source: &str,
    file: &RelativePath,
    node: &NodeRef,
    new_source: &str,
) -> Result<MutationResult, String> {
    validate_range(source, node)?;

    let mut result =
        String::with_capacity(source.len() - (node.end_byte - node.start_byte) + new_source.len());
    result.push_str(&source[..node.start_byte]);
    result.push_str(new_source);
    result.push_str(&source[node.end_byte..]);

    let new_end = node.start_byte + new_source.len();
    let new_ref = build_node_ref_from_range(&result, file, node.start_byte, new_end);

    Ok(MutationResult {
        source: result,
        affected_nodes: vec![new_ref],
    })
}

/// Move a node to a new position relative to a target.
/// Both nodes must be in the same file (same source text).
#[must_use = "move result contains modified source"]
pub fn move_node(
    source: &str,
    file: &RelativePath,
    node: &NodeRef,
    target: &NodeRef,
    position: InsertPosition,
) -> Result<MutationResult, String> {
    validate_range(source, node)?;
    validate_range(source, target)?;

    // Extract the node text first
    let node_text = source[node.start_byte..node.end_byte].to_string();

    // Remove first, then insert. Order matters for byte offsets.
    // If node is before target, removing shifts target left.
    // If node is after target, removing doesn't affect target.
    if node.start_byte < target.start_byte {
        // Remove first (shifts target)
        let removal = remove_node(source, node)?;
        let removed_bytes = (skip_trailing_whitespace(source, node.end_byte)) - node.start_byte;
        let adjusted_target = NodeRef {
            start_byte: target.start_byte - removed_bytes,
            end_byte: target.end_byte - removed_bytes,
            ..target.clone()
        };
        insert_source(
            &removal.result.source,
            file,
            &adjusted_target,
            position,
            &node_text,
        )
    } else {
        // Insert first (doesn't shift node)
        let insert_result = insert_source(source, file, target, position, &node_text)?;
        let inserted_bytes = insert_result.source.len() - source.len();
        let adjusted_node = NodeRef {
            start_byte: node.start_byte + inserted_bytes,
            end_byte: node.end_byte + inserted_bytes,
            ..node.clone()
        };
        let removal = remove_node(&insert_result.source, &adjusted_node)?;
        Ok(removal.result)
    }
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

fn validate_range(source: &str, node: &NodeRef) -> Result<(), String> {
    if node.end_byte > source.len() {
        return Err(format!(
            "Byte range {}..{} out of bounds for source (len={})",
            node.start_byte,
            node.end_byte,
            source.len()
        ));
    }
    if node.start_byte > node.end_byte {
        return Err(format!(
            "Invalid byte range: start {} > end {}",
            node.start_byte, node.end_byte
        ));
    }
    Ok(())
}

/// Skip trailing whitespace and a single newline after a removed node.
fn skip_trailing_whitespace(source: &str, from: usize) -> usize {
    let bytes = source.as_bytes();
    let mut pos = from;
    // Skip spaces/tabs
    while pos < bytes.len() && (bytes[pos] == b' ' || bytes[pos] == b'\t') {
        pos += 1;
    }
    // Skip one newline
    if pos < bytes.len() && bytes[pos] == b'\n' {
        pos += 1;
    } else if pos + 1 < bytes.len() && bytes[pos] == b'\r' && bytes[pos + 1] == b'\n' {
        pos += 2;
    }
    pos
}

/// Detect the indentation of the line containing `byte_offset`.
fn detect_indent(source: &str, byte_offset: usize) -> String {
    let before = &source[..byte_offset];
    let line_start = before.rfind('\n').map(|i| i + 1).unwrap_or(0);
    let line = &source[line_start..byte_offset];
    let indent_len = line.len() - line.trim_start().len();
    line[..indent_len].to_string()
}

/// Find the byte offset just before the closing brace of a node's body.
fn find_body_end(source: &str, node: &NodeRef) -> usize {
    // Walk backwards from node.end_byte to find closing brace/bracket
    let bytes = source.as_bytes();
    let mut pos = node.end_byte;
    while pos > node.start_byte {
        pos -= 1;
        if bytes[pos] == b'}' || bytes[pos] == b']' || bytes[pos] == b')' {
            return pos;
        }
    }
    // Fallback: insert before end
    node.end_byte
}

/// Build a NodeRef from a byte range in source text.
/// Computes line/column from the byte offset.
fn build_node_ref_from_range(
    source: &str,
    file: &RelativePath,
    start: usize,
    end: usize,
) -> NodeRef {
    let (line, column) = byte_to_line_col(source, start);
    let (end_line, end_column) = byte_to_line_col(source, end);
    NodeRef {
        file: file.clone(),
        start_byte: start,
        end_byte: end,
        kind: NodeKind::from("inserted"),
        line,
        column,
        end_line,
        end_column,
    }
}

/// Convert a byte offset to 1-based line and 0-based column.
fn byte_to_line_col(source: &str, byte_offset: usize) -> (usize, usize) {
    let before = &source[..byte_offset.min(source.len())];
    let line = before.matches('\n').count() + 1;
    let col = before.len() - before.rfind('\n').map(|i| i + 1).unwrap_or(0);
    (line, col)
}

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

    fn test_file() -> RelativePath {
        RelativePath::from("test.ts")
    }

    const SOURCE: &str = r#"function greet(name: string): string {
    return `Hello, ${name}!`;
}

async function fetchUser(id: number): Promise<User> {
    const response = await fetch(`/api/users/${id}`);
    return response.json();
}

const MAX_RETRIES = 3;
"#;

    fn make_ref(start: usize, end: usize) -> NodeRef {
        NodeRef {
            file: test_file(),
            start_byte: start,
            end_byte: end,
            kind: NodeKind::from("test"),
            line: 1,
            column: 0,
            end_line: 1,
            end_column: 0,
        }
    }

    #[test]
    fn remove_extracts_node() {
        // Arrange
        let greet_start = SOURCE.find("function greet").unwrap();
        let greet_end = SOURCE.find("}\n\nasync").unwrap() + 1;
        let node = make_ref(greet_start, greet_end);

        // Act
        let removal = remove_node(SOURCE, &node).unwrap();
        let result = removal.result;
        let detached = removal.detached;

        // Assert
        assert!(
            detached.contains("function greet"),
            "detached should contain the function"
        );
        assert!(
            !result.source.contains("function greet"),
            "result should not contain the removed function"
        );
        assert!(
            result.source.contains("async function fetchUser"),
            "result should still contain fetchUser"
        );
    }

    #[test]
    fn insert_after_adds_text() {
        // Arrange
        let greet_start = SOURCE.find("function greet").unwrap();
        let greet_end = SOURCE.find("}\n\nasync").unwrap() + 1;
        let target = make_ref(greet_start, greet_end);
        let new_fn = "function goodbye(): void {\n    console.log('bye');\n}";

        // Act
        let result =
            insert_source(SOURCE, &test_file(), &target, InsertPosition::After, new_fn).unwrap();

        // Assert
        assert!(
            result.source.contains(new_fn),
            "result should contain the new function"
        );
        let greet_pos = result.source.find("function greet").unwrap();
        let goodbye_pos = result.source.find("function goodbye").unwrap();
        assert!(goodbye_pos > greet_pos, "goodbye should come after greet");
    }

    #[test]
    fn insert_before_adds_text() {
        // Arrange
        let fetch_start = SOURCE.find("async function fetchUser").unwrap();
        let fetch_end = SOURCE.find("}\n\nconst").unwrap() + 1;
        let target = make_ref(fetch_start, fetch_end);
        let new_fn = "function middleware(): void {}";

        // Act
        let result = insert_source(
            SOURCE,
            &test_file(),
            &target,
            InsertPosition::Before,
            new_fn,
        )
        .unwrap();

        // Assert
        let middleware_pos = result.source.find("function middleware").unwrap();
        let fetch_pos = result.source.find("async function fetchUser").unwrap();
        assert!(
            middleware_pos < fetch_pos,
            "middleware should come before fetchUser"
        );
    }

    #[test]
    fn replace_swaps_content() {
        // Arrange
        let max_start = SOURCE.find("const MAX_RETRIES = 3;").unwrap();
        let max_end = max_start + "const MAX_RETRIES = 3;".len();
        let node = make_ref(max_start, max_end);

        // Act
        let result = replace_node(SOURCE, &test_file(), &node, "const MAX_RETRIES = 5;").unwrap();

        // Assert
        assert!(
            result.source.contains("const MAX_RETRIES = 5;"),
            "should contain new value"
        );
        assert!(
            !result.source.contains("const MAX_RETRIES = 3;"),
            "should not contain old value"
        );
    }

    #[test]
    fn move_node_forward() {
        // Arrange — move greet after fetchUser
        let greet_start = SOURCE.find("function greet").unwrap();
        let greet_end = SOURCE.find("}\n\nasync").unwrap() + 1;
        let greet = make_ref(greet_start, greet_end);

        let fetch_start = SOURCE.find("async function fetchUser").unwrap();
        let fetch_end = SOURCE.find("}\n\nconst").unwrap() + 1;
        let fetch = make_ref(fetch_start, fetch_end);

        // Act
        let result =
            move_node(SOURCE, &test_file(), &greet, &fetch, InsertPosition::After).unwrap();

        // Assert
        let fetch_pos = result.source.find("async function fetchUser").unwrap();
        let greet_pos = result.source.find("function greet").unwrap();
        assert!(
            greet_pos > fetch_pos,
            "greet should now come after fetchUser"
        );
    }

    #[test]
    fn move_node_backward() {
        // Arrange — move MAX_RETRIES before greet
        let max_start = SOURCE.find("const MAX_RETRIES = 3;").unwrap();
        let max_end = max_start + "const MAX_RETRIES = 3;".len();
        let max_node = make_ref(max_start, max_end);

        let greet_start = SOURCE.find("function greet").unwrap();
        let greet_end = SOURCE.find("}\n\nasync").unwrap() + 1;
        let greet = make_ref(greet_start, greet_end);

        // Act
        let result = move_node(
            SOURCE,
            &test_file(),
            &max_node,
            &greet,
            InsertPosition::Before,
        )
        .unwrap();

        // Assert
        let max_pos = result.source.find("const MAX_RETRIES = 3;").unwrap();
        let greet_pos = result.source.find("function greet").unwrap();
        assert!(
            max_pos < greet_pos,
            "MAX_RETRIES should now come before greet"
        );
    }

    #[test]
    fn out_of_bounds_returns_error() {
        // Arrange
        let bad_ref = make_ref(0, SOURCE.len() + 100);

        // Act
        let result = remove_node(SOURCE, &bad_ref);

        // Assert
        assert!(result.is_err(), "should error on out-of-bounds range");
    }

    #[test]
    fn byte_to_line_col_correct() {
        // Arrange
        let src = "line1\nline2\nline3";

        // Act and Assert
        assert_eq!(byte_to_line_col(src, 0), (1, 0), "start of file");
        assert_eq!(byte_to_line_col(src, 6), (2, 0), "start of line 2");
        assert_eq!(byte_to_line_col(src, 8), (2, 2), "col 2 of line 2");
    }
}