flake_edit/

walk.rs

//! CST walking and mutation for `flake.nix` files.

mod context;
mod error;
mod inputs;
mod node;
mod outputs;
pub(crate) mod toggle;

use std::collections::HashMap;

use rnix::{Root, SyntaxKind, SyntaxNode};

use crate::change::Change;
use crate::edit::{OutputChange, Outputs};
use crate::follows::path::follows_idents_prefixed;
use crate::follows::{AttrPath, Segment, strip_outer_quotes};
use crate::input::Input;

pub(crate) use context::Context;
pub use error::WalkerError;

use inputs::walk_inputs;
use node::{
    FollowsKind, adjacent_whitespace_index, get_sibling_whitespace, insertion_index_after,
    last_line_with_newline, make_quoted_string, make_toplevel_flake_false_attr,
    make_toplevel_url_attr, parse_node, substitute_child,
};

/// The flake's top-level attribute set.
///
/// A flake's root expression is normally a bare attribute set, but it may be
/// wrapped in `let <bindings> in { ... }`. In that case the inputs and outputs
/// live in the `in` body, not in the let bindings, so we descend through the
/// `NODE_LET_IN` to reach the body attrset and walk it exactly as a bare one.
///
/// Returns `None` when the root has no expression, or when a `let ... in` body
/// is not an attribute set (for example a function): such a flake exposes no
/// inputs or outputs to walk.
pub(crate) fn flake_attr_set(root: &SyntaxNode) -> Option<SyntaxNode> {
    let first = root.first_child()?;
    if first.kind() != SyntaxKind::NODE_LET_IN {
        return Some(first);
    }
    // `let <bindings> in <body>`: the body expression is the last child node,
    // after the bindings. Only an attrset body carries flake attributes.
    let body = first.last_child()?;
    (body.kind() == SyntaxKind::NODE_ATTR_SET).then_some(body)
}

/// Whether a CST attrpath (idents may carry surrounding `"..."`) matches `expected`
/// pairwise after unquoting.
fn idents_match(have: &[String], expected: &[&str]) -> bool {
    if have.len() != expected.len() {
        return false;
    }
    have.iter()
        .zip(expected.iter())
        .all(|(h, e)| strip_outer_quotes(h) == *e)
}

fn is_flat_inputs_attr_for(idents: &[String], parent_id: &str) -> bool {
    idents.len() >= 2 && idents[0] == "inputs" && strip_outer_quotes(&idents[1]) == parent_id
}

fn block_parent_attrset(
    toplevel: &SyntaxNode,
    idents: &[String],
    parent_id: &str,
) -> Option<SyntaxNode> {
    if idents.len() != 2 {
        return None;
    }
    if !is_flat_inputs_attr_for(idents, parent_id) {
        return None;
    }
    toplevel
        .children()
        .find(|c| c.kind() == SyntaxKind::NODE_ATTR_SET)
}

/// Same-target hits return the unchanged root. The caller treats `Some` as
/// "claimed", so the no-op variant must still return `Some` to short-circuit
/// the surrounding scan.
fn retarget_existing_flat_follows(
    attr_set: &SyntaxNode,
    toplevel: &SyntaxNode,
    value_node: Option<SyntaxNode>,
    target: &str,
) -> Option<SyntaxNode> {
    let current_target = value_node
        .as_ref()
        .map(|v| strip_outer_quotes(&v.to_string()).to_string())
        .unwrap_or_default();

    if current_target == target {
        return attr_set.ancestors().last();
    }

    let value = value_node?;
    let new_value = make_quoted_string(target);
    let new_toplevel = substitute_child(toplevel, value.index(), &new_value);
    let green = attr_set
        .green()
        .replace_child(toplevel.index(), new_toplevel.green().into());
    Some(SyntaxNode::new_root(attr_set.replace_with(green)))
}

/// Mirrors `ref_child`'s leading newline + indent so the inserted line
/// reads at the same column as its neighbour. Without normalization, the
/// raw whitespace token includes any pre-`ref_child` spacing too.
fn insert_flat_follows_after(
    attr_set: &SyntaxNode,
    ref_child: &SyntaxNode,
    path: &AttrPath,
    target: &str,
) -> SyntaxNode {
    let follows_node = FollowsKind::TopLevelNested { path, target }.emit();
    let insert_index = insertion_index_after(ref_child);

    let mut green = attr_set
        .green()
        .insert_child(insert_index, follows_node.green().into());

    if let Some(whitespace) = get_sibling_whitespace(ref_child) {
        let ws_str = whitespace.to_string();
        let ws_node = parse_node(last_line_with_newline(&ws_str));
        green = green.insert_child(insert_index, ws_node.green().into());
    }

    SyntaxNode::new_root(attr_set.replace_with(green))
}

#[derive(Debug, Clone)]
pub struct Walker {
    pub(crate) root: SyntaxNode,
    pub(crate) inputs: HashMap<String, Input>,
    pub(crate) add_toplevel: bool,
}

impl<'a> Walker {
    pub fn new(stream: &'a str) -> Self {
        let root = Root::parse(stream).syntax();
        Self::from_root(root)
    }

    /// Build a walker around an already-parsed root, skipping the rnix parse.
    /// Lets callers that ran a parse for validation share the result.
    pub fn from_root(root: SyntaxNode) -> Self {
        Self {
            root,
            inputs: HashMap::new(),
            add_toplevel: false,
        }
    }

    /// Apply `change` to the parsed `flake.nix`, returning the rebuilt root if
    /// the tree was modified.
    ///
    /// Expects the parsed root to be an attrset with `description`, `inputs`, and
    /// `outputs` keys.
    pub fn walk(&mut self, change: &Change) -> Result<Option<SyntaxNode>, WalkerError> {
        let cst = self.root.clone();
        if cst.kind() != SyntaxKind::NODE_ROOT {
            return Err(WalkerError::NotARoot);
        }
        self.walk_toplevel(cst, None, change)
    }

    /// List the `outputs` arguments without touching `inputs`.
    pub(crate) fn list_outputs(&mut self) -> Result<Outputs, WalkerError> {
        outputs::list_outputs(&self.root)
    }

    /// Apply an [`OutputChange`] to the `outputs` attribute alone.
    pub(crate) fn change_outputs(
        &mut self,
        change: OutputChange,
    ) -> Result<Option<SyntaxNode>, WalkerError> {
        outputs::change_outputs(&self.root, change)
    }

    /// Walk the top-level attrset, dispatching on `description`/`inputs`/`outputs`.
    fn walk_toplevel(
        &mut self,
        node: SyntaxNode,
        ctx: Option<Context>,
        change: &Change,
    ) -> Result<Option<SyntaxNode>, WalkerError> {
        let Some(attr_set) = flake_attr_set(&node) else {
            return Ok(None);
        };

        for toplevel in attr_set.children() {
            if toplevel.kind() != SyntaxKind::NODE_ATTRPATH_VALUE {
                let range = toplevel.text_range();
                return Err(WalkerError::unexpected_top_level(
                    &toplevel.to_string(),
                    range.start().into(),
                ));
            }

            // Dispatch on the NODE_ATTRPATH child alone, not on the value.
            // For `inputs = { ... }` the value is the whole inputs attrset;
            // stringifying it dominates this walk on large flakes.
            let Some(attrpath) = toplevel
                .children()
                .find(|c| c.kind() == SyntaxKind::NODE_ATTRPATH)
            else {
                continue;
            };
            let mut path_idents = attrpath.children();
            let Some(first_ident) = path_idents.next() else {
                continue;
            };
            let has_more_idents = path_idents.next().is_some();
            let first_text = first_ident.to_string();
            let first_unquoted = strip_outer_quotes(&first_text);

            if !has_more_idents && first_unquoted == "description" {
                continue;
            }

            if first_unquoted == "inputs" {
                if has_more_idents {
                    if let Some(result) =
                        self.handle_inputs_flat(&attr_set, &toplevel, &attrpath, &ctx, change)
                    {
                        return Ok(Some(result));
                    }
                } else if let Some(result) =
                    self.handle_inputs_attr(&toplevel, &attrpath, &ctx, change)
                {
                    return Ok(Some(result));
                }
                continue;
            }

            if !has_more_idents
                && first_unquoted == "outputs"
                && let Some(result) = self.handle_add_at_outputs(&attr_set, &toplevel, change)
            {
                return Ok(Some(result));
            }
        }

        // Follows on toplevel flat-style inputs (`inputs.X.url = "..."`).
        if let Change::Follows { input, target } = change {
            let path = input.path();
            if path.len() >= 2 {
                let parent_id = input.input();
                if self.inputs.contains_key(parent_id.as_str()) {
                    let target_str = target.to_flake_follows_string();
                    return self.handle_follows_flat_toplevel(&attr_set, path, &target_str);
                }
            }
        }

        Ok(None)
    }

    /// Add a follows attribute next to a toplevel flat-style input.
    ///
    /// Converts `inputs.crane.url = "github:...";` into:
    /// ```nix
    /// inputs.crane.url = "github:...";
    /// inputs.crane.inputs.nixpkgs.follows = "nixpkgs";
    /// ```
    fn handle_follows_flat_toplevel(
        &self,
        attr_set: &SyntaxNode,
        path: &AttrPath,
        target: &str,
    ) -> Result<Option<SyntaxNode>, WalkerError> {
        let parent_id = path.first();
        let expected_flat = follows_idents_prefixed(path.segments());
        let mut last_parent_attr: Option<SyntaxNode> = None;
        let mut block_parent: Option<(SyntaxNode, SyntaxNode)> = None;

        for toplevel in attr_set.children() {
            if toplevel.kind() != SyntaxKind::NODE_ATTRPATH_VALUE {
                continue;
            }
            let Some(attrpath) = toplevel
                .children()
                .find(|c| c.kind() == SyntaxKind::NODE_ATTRPATH)
            else {
                continue;
            };
            let idents: Vec<String> = attrpath.children().map(|c| c.to_string()).collect();

            if let Some(block_attr_set) =
                block_parent_attrset(&toplevel, &idents, parent_id.as_str())
            {
                block_parent = Some((toplevel.clone(), block_attr_set));
            }

            if idents_match(&idents, &expected_flat)
                && let Some(rebuilt) = retarget_existing_flat_follows(
                    attr_set,
                    &toplevel,
                    attrpath.next_sibling(),
                    target,
                )
            {
                return Ok(Some(rebuilt));
            }

            if is_flat_inputs_attr_for(&idents, parent_id.as_str()) {
                last_parent_attr = Some(toplevel.clone());
            }
        }

        if let Some((toplevel, block_attr_set)) = block_parent {
            let rest: Vec<Segment> = path.segments()[1..].to_vec();
            return self.handle_follows_block_toplevel(
                attr_set,
                &toplevel,
                &block_attr_set,
                &rest,
                target,
            );
        }

        if let Some(ref_child) = last_parent_attr {
            return Ok(Some(insert_flat_follows_after(
                attr_set, &ref_child, path, target,
            )));
        }

        Ok(None)
    }

    fn handle_follows_block_toplevel(
        &self,
        attr_set: &SyntaxNode,
        toplevel: &SyntaxNode,
        block_attr_set: &SyntaxNode,
        rest: &[Segment],
        target: &str,
    ) -> Result<Option<SyntaxNode>, WalkerError> {
        let expected_block = follows_idents_prefixed(rest);
        for attr in block_attr_set.children() {
            if attr.kind() != SyntaxKind::NODE_ATTRPATH_VALUE {
                continue;
            }
            let Some(attrpath) = attr
                .children()
                .find(|c| c.kind() == SyntaxKind::NODE_ATTRPATH)
            else {
                continue;
            };
            let idents: Vec<String> = attrpath.children().map(|c| c.to_string()).collect();

            if idents_match(&idents, &expected_block) {
                let value_node = attrpath.next_sibling();
                let current_target = value_node
                    .as_ref()
                    .map(|v| strip_outer_quotes(&v.to_string()).to_string())
                    .unwrap_or_default();

                if current_target == target {
                    return Ok(attr_set.ancestors().last());
                }

                if let Some(value) = value_node {
                    let new_value = make_quoted_string(target);
                    let new_attr = substitute_child(&attr, value.index(), &new_value);
                    let new_block = substitute_child(block_attr_set, attr.index(), &new_attr);
                    let new_toplevel =
                        substitute_child(toplevel, block_attr_set.index(), &new_block);
                    let green = attr_set
                        .green()
                        .replace_child(toplevel.index(), new_toplevel.green().into());
                    return Ok(Some(SyntaxNode::new_root(attr_set.replace_with(green))));
                }
            }
        }

        let follows_node = FollowsKind::BlockNested { rest, target }.emit();
        let children: Vec<_> = block_attr_set.children().collect();
        if let Some(last_child) = children.last() {
            let insert_index = last_child.index() + 1;

            let mut green = block_attr_set
                .green()
                .insert_child(insert_index, follows_node.green().into());

            if let Some(whitespace) = get_sibling_whitespace(last_child) {
                green = green.insert_child(insert_index, whitespace.green().into());
            }

            let new_block = SyntaxNode::new_root(green);
            let new_toplevel = substitute_child(toplevel, block_attr_set.index(), &new_block);
            let green = attr_set
                .green()
                .replace_child(toplevel.index(), new_toplevel.green().into());
            return Ok(Some(SyntaxNode::new_root(attr_set.replace_with(green))));
        }

        Ok(None)
    }

    /// Apply `change` to the `inputs = { ... }` attribute.
    ///
    /// `toplevel.replace_with()` propagates through `NODE_ATTR_SET` up to `NODE_ROOT`,
    /// preserving leading comments and trivia.
    fn handle_inputs_attr(
        &mut self,
        toplevel: &SyntaxNode,
        child: &SyntaxNode,
        ctx: &Option<Context>,
        change: &Change,
    ) -> Option<SyntaxNode> {
        let sibling = child.next_sibling()?;
        let replacement = walk_inputs(&mut self.inputs, sibling.clone(), ctx, change)?;

        let green = toplevel
            .green()
            .replace_child(sibling.index(), replacement.green().into());
        let green = toplevel.replace_with(green);
        Some(SyntaxNode::new_root(green))
    }

    /// Apply `change` to flat-style `inputs.foo.url = "..."` attributes.
    ///
    /// Removals rebuild the parent attrset green and `replace_with()` propagates to
    /// `NODE_ROOT`. Replacements rely on `toplevel.replace_with()` to propagate.
    fn handle_inputs_flat(
        &mut self,
        attr_set: &SyntaxNode,
        toplevel: &SyntaxNode,
        child: &SyntaxNode,
        ctx: &Option<Context>,
        change: &Change,
    ) -> Option<SyntaxNode> {
        let replacement = walk_inputs(&mut self.inputs, child.clone(), ctx, change)?;

        // Empty replacement means we remove the entire toplevel node and
        // propagate through attr_set to NODE_ROOT.
        if replacement.to_string().is_empty() {
            let element: rnix::SyntaxElement = toplevel.clone().into();
            let mut green = attr_set.green().remove_child(toplevel.index());
            if let Some(ws_index) = adjacent_whitespace_index(&element) {
                green = green.remove_child(ws_index);
            }
            return Some(SyntaxNode::new_root(attr_set.replace_with(green)));
        }

        let sibling = child.next_sibling()?;
        let green = toplevel
            .green()
            .replace_child(sibling.index(), replacement.green().into());
        let green = toplevel.replace_with(green);
        Some(SyntaxNode::new_root(green))
    }

    /// Add a new input just before `outputs` when no `inputs` block exists yet.
    ///
    /// Rebuilds the parent attrset green. `replace_with()` propagates to `NODE_ROOT`
    /// while preserving leading comments.
    fn handle_add_at_outputs(
        &mut self,
        attr_set: &SyntaxNode,
        toplevel: &SyntaxNode,
        change: &Change,
    ) -> Option<SyntaxNode> {
        if !self.add_toplevel {
            return None;
        }

        let Change::Add {
            id: Some(id),
            uri: Some(uri),
            flake,
        } = change
        else {
            return None;
        };
        let id = id.input().as_str();

        if toplevel.index() == 0 {
            return None;
        }

        // Walk back from `outputs` through tokens to find a whitespace run, then
        // normalize it to a single newline + indent. Walking through tokens (not
        // siblings) lets us skip past comments between the last input and `outputs`.
        let ws_node = {
            let mut ws: Option<SyntaxNode> = None;
            let mut cursor = toplevel.prev_sibling_or_token();
            while let Some(ref tok) = cursor {
                if tok.kind() == SyntaxKind::TOKEN_WHITESPACE {
                    let ws_str = tok.to_string();
                    ws = Some(parse_node(last_line_with_newline(&ws_str)));
                    break;
                }
                cursor = tok.prev_sibling_or_token();
            }
            ws
        };

        let addition = make_toplevel_url_attr(id, uri);
        let insert_pos = toplevel.index() - 1;

        let mut green = attr_set
            .green()
            .insert_child(insert_pos, addition.green().into());

        if let Some(ref ws) = ws_node {
            green = green.insert_child(insert_pos, ws.green().into());
        }

        // Append `inputs.<id>.flake = false;` when the new input opts out of flake mode.
        if !flake {
            let no_flake = make_toplevel_flake_false_attr(id);
            green = green.insert_child(toplevel.index() + 1, no_flake.green().into());

            if let Some(ref ws) = ws_node {
                green = green.insert_child(toplevel.index() + 1, ws.green().into());
            }
        }

        Some(SyntaxNode::new_root(attr_set.replace_with(green)))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::change::{Change, ChangeId};
    use crate::follows::AttrPath;

    fn apply(flake_text: &str, change: &Change) -> String {
        let mut walker = Walker::new(flake_text);
        walker
            .walk(change)
            .expect("walker error")
            .expect("walker did not rewrite the tree")
            .to_string()
    }

    fn follows_change(input: &str, target: &str) -> Change {
        Change::Follows {
            input: ChangeId::parse(input).unwrap(),
            target: AttrPath::parse(target).unwrap(),
        }
    }

    #[test]
    fn handle_follows_flat_toplevel_inserts_follows_after_last_parent_attr() {
        let flake = "{
  inputs.flake-edit.url = \"github:a-kenji/flake-edit\";
  inputs.nixpkgs.url = \"github:NixOS/nixpkgs\";

  outputs = { self, ... }: { };
}
";
        let result = apply(flake, &follows_change("flake-edit.nixpkgs", "nixpkgs"));
        assert_eq!(
            result,
            "{
  inputs.flake-edit.url = \"github:a-kenji/flake-edit\";
  inputs.flake-edit.inputs.nixpkgs.follows = \"nixpkgs\";
  inputs.nixpkgs.url = \"github:NixOS/nixpkgs\";

  outputs = { self, ... }: { };
}
"
        );
    }

    #[test]
    fn handle_follows_flat_toplevel_retargets_existing_follows() {
        let flake = "{
  inputs.flake-edit.url = \"github:a-kenji/flake-edit\";
  inputs.flake-edit.inputs.nixpkgs.follows = \"old-pkgs\";
  inputs.nixpkgs.url = \"github:NixOS/nixpkgs\";

  outputs = { self, ... }: { };
}
";
        let result = apply(flake, &follows_change("flake-edit.nixpkgs", "nixpkgs"));
        assert_eq!(
            result,
            "{
  inputs.flake-edit.url = \"github:a-kenji/flake-edit\";
  inputs.flake-edit.inputs.nixpkgs.follows = \"nixpkgs\";
  inputs.nixpkgs.url = \"github:NixOS/nixpkgs\";

  outputs = { self, ... }: { };
}
"
        );
    }

    #[test]
    fn handle_follows_flat_toplevel_is_noop_when_target_already_matches() {
        let flake = "{
  inputs.flake-edit.url = \"github:a-kenji/flake-edit\";
  inputs.flake-edit.inputs.nixpkgs.follows = \"nixpkgs\";
  inputs.nixpkgs.url = \"github:NixOS/nixpkgs\";

  outputs = { self, ... }: { };
}
";
        let result = apply(flake, &follows_change("flake-edit.nixpkgs", "nixpkgs"));
        assert_eq!(result, flake);
    }

    #[test]
    fn handle_follows_flat_toplevel_delegates_to_block_parent_when_present() {
        let flake = "{
  inputs.nixpkgs.url = \"github:NixOS/nixpkgs\";
  inputs.flake-edit = {
    url = \"github:a-kenji/flake-edit\";
  };

  outputs = { self, ... }: { };
}
";
        let result = apply(flake, &follows_change("flake-edit.nixpkgs", "nixpkgs"));
        assert_eq!(
            result,
            "{
  inputs.nixpkgs.url = \"github:NixOS/nixpkgs\";
  inputs.flake-edit = {
    url = \"github:a-kenji/flake-edit\";
    inputs.nixpkgs.follows = \"nixpkgs\";
  };

  outputs = { self, ... }: { };
}
"
        );
    }

    #[test]
    fn is_flat_inputs_attr_for_only_matches_matching_parent_id() {
        let yes = [
            "inputs".to_string(),
            "flake-edit".to_string(),
            "url".to_string(),
        ];
        let no = [
            "inputs".to_string(),
            "nixpkgs".to_string(),
            "url".to_string(),
        ];
        assert!(is_flat_inputs_attr_for(&yes, "flake-edit"));
        assert!(!is_flat_inputs_attr_for(&no, "flake-edit"));
        // The CST keeps surrounding `"..."` on quoted idents; the comparison
        // must unquote them, otherwise `"flake-edit" != flake-edit`.
        let quoted = [
            "inputs".to_string(),
            "\"flake-edit\"".to_string(),
            "url".to_string(),
        ];
        assert!(is_flat_inputs_attr_for(&quoted, "flake-edit"));
    }
}