flake_edit/
walk.rs

1//! CST walking and mutation for `flake.nix` files.
2
3mod context;
4mod error;
5mod inputs;
6mod node;
7mod outputs;
8
9use std::collections::HashMap;
10
11use rnix::{Root, SyntaxKind, SyntaxNode};
12
13use crate::change::Change;
14use crate::edit::{OutputChange, Outputs};
15use crate::follows::path::follows_idents_prefixed;
16use crate::follows::{AttrPath, Segment, strip_outer_quotes};
17use crate::input::Input;
18
19pub(crate) use context::Context;
20pub use error::WalkerError;
21
22use inputs::walk_inputs;
23use node::{
24    FollowsKind, adjacent_whitespace_index, get_sibling_whitespace, insertion_index_after,
25    last_line_with_newline, make_quoted_string, make_toplevel_flake_false_attr,
26    make_toplevel_url_attr, parse_node, substitute_child,
27};
28
29/// The flake's top-level attribute set.
30///
31/// A flake's root expression is normally a bare attribute set, but it may be
32/// wrapped in `let <bindings> in { ... }`. In that case the inputs and outputs
33/// live in the `in` body, not in the let bindings, so we descend through the
34/// `NODE_LET_IN` to reach the body attrset and walk it exactly as a bare one.
35///
36/// Returns `None` when the root has no expression, or when a `let ... in` body
37/// is not an attribute set (for example a function): such a flake exposes no
38/// inputs or outputs to walk.
39pub(crate) fn flake_attr_set(root: &SyntaxNode) -> Option<SyntaxNode> {
40    let first = root.first_child()?;
41    if first.kind() != SyntaxKind::NODE_LET_IN {
42        return Some(first);
43    }
44    // `let <bindings> in <body>`: the body expression is the last child node,
45    // after the bindings. Only an attrset body carries flake attributes.
46    let body = first.last_child()?;
47    (body.kind() == SyntaxKind::NODE_ATTR_SET).then_some(body)
48}
49
50/// Whether a CST attrpath (idents may carry surrounding `"..."`) matches `expected`
51/// pairwise after unquoting.
52fn idents_match(have: &[String], expected: &[&str]) -> bool {
53    if have.len() != expected.len() {
54        return false;
55    }
56    have.iter()
57        .zip(expected.iter())
58        .all(|(h, e)| strip_outer_quotes(h) == *e)
59}
60
61fn is_flat_inputs_attr_for(idents: &[String], parent_id: &str) -> bool {
62    idents.len() >= 2 && idents[0] == "inputs" && strip_outer_quotes(&idents[1]) == parent_id
63}
64
65fn block_parent_attrset(
66    toplevel: &SyntaxNode,
67    idents: &[String],
68    parent_id: &str,
69) -> Option<SyntaxNode> {
70    if idents.len() != 2 {
71        return None;
72    }
73    if !is_flat_inputs_attr_for(idents, parent_id) {
74        return None;
75    }
76    toplevel
77        .children()
78        .find(|c| c.kind() == SyntaxKind::NODE_ATTR_SET)
79}
80
81/// Same-target hits return the unchanged root. The caller treats `Some` as
82/// "claimed", so the no-op variant must still return `Some` to short-circuit
83/// the surrounding scan.
84fn retarget_existing_flat_follows(
85    attr_set: &SyntaxNode,
86    toplevel: &SyntaxNode,
87    value_node: Option<SyntaxNode>,
88    target: &str,
89) -> Option<SyntaxNode> {
90    let current_target = value_node
91        .as_ref()
92        .map(|v| strip_outer_quotes(&v.to_string()).to_string())
93        .unwrap_or_default();
94
95    if current_target == target {
96        return attr_set.ancestors().last();
97    }
98
99    let value = value_node?;
100    let new_value = make_quoted_string(target);
101    let new_toplevel = substitute_child(toplevel, value.index(), &new_value);
102    let green = attr_set
103        .green()
104        .replace_child(toplevel.index(), new_toplevel.green().into());
105    Some(SyntaxNode::new_root(attr_set.replace_with(green)))
106}
107
108/// Mirrors `ref_child`'s leading newline + indent so the inserted line
109/// reads at the same column as its neighbour. Without normalization, the
110/// raw whitespace token includes any pre-`ref_child` spacing too.
111fn insert_flat_follows_after(
112    attr_set: &SyntaxNode,
113    ref_child: &SyntaxNode,
114    path: &AttrPath,
115    target: &str,
116) -> SyntaxNode {
117    let follows_node = FollowsKind::TopLevelNested { path, target }.emit();
118    let insert_index = insertion_index_after(ref_child);
119
120    let mut green = attr_set
121        .green()
122        .insert_child(insert_index, follows_node.green().into());
123
124    if let Some(whitespace) = get_sibling_whitespace(ref_child) {
125        let ws_str = whitespace.to_string();
126        let ws_node = parse_node(last_line_with_newline(&ws_str));
127        green = green.insert_child(insert_index, ws_node.green().into());
128    }
129
130    SyntaxNode::new_root(attr_set.replace_with(green))
131}
132
133#[derive(Debug, Clone)]
134pub struct Walker {
135    pub(crate) root: SyntaxNode,
136    pub(crate) inputs: HashMap<String, Input>,
137    pub(crate) add_toplevel: bool,
138}
139
140impl<'a> Walker {
141    pub fn new(stream: &'a str) -> Self {
142        let root = Root::parse(stream).syntax();
143        Self::from_root(root)
144    }
145
146    /// Build a walker around an already-parsed root, skipping the rnix parse.
147    /// Lets callers that ran a parse for validation share the result.
148    pub fn from_root(root: SyntaxNode) -> Self {
149        Self {
150            root,
151            inputs: HashMap::new(),
152            add_toplevel: false,
153        }
154    }
155
156    /// Apply `change` to the parsed `flake.nix`, returning the rebuilt root if
157    /// the tree was modified.
158    ///
159    /// Expects the parsed root to be an attrset with `description`, `inputs`, and
160    /// `outputs` keys.
161    pub fn walk(&mut self, change: &Change) -> Result<Option<SyntaxNode>, WalkerError> {
162        let cst = self.root.clone();
163        if cst.kind() != SyntaxKind::NODE_ROOT {
164            return Err(WalkerError::NotARoot);
165        }
166        self.walk_toplevel(cst, None, change)
167    }
168
169    /// List the `outputs` arguments without touching `inputs`.
170    pub(crate) fn list_outputs(&mut self) -> Result<Outputs, WalkerError> {
171        outputs::list_outputs(&self.root)
172    }
173
174    /// Apply an [`OutputChange`] to the `outputs` attribute alone.
175    pub(crate) fn change_outputs(
176        &mut self,
177        change: OutputChange,
178    ) -> Result<Option<SyntaxNode>, WalkerError> {
179        outputs::change_outputs(&self.root, change)
180    }
181
182    /// Walk the top-level attrset, dispatching on `description`/`inputs`/`outputs`.
183    fn walk_toplevel(
184        &mut self,
185        node: SyntaxNode,
186        ctx: Option<Context>,
187        change: &Change,
188    ) -> Result<Option<SyntaxNode>, WalkerError> {
189        let Some(attr_set) = flake_attr_set(&node) else {
190            return Ok(None);
191        };
192
193        for toplevel in attr_set.children() {
194            if toplevel.kind() != SyntaxKind::NODE_ATTRPATH_VALUE {
195                let range = toplevel.text_range();
196                return Err(WalkerError::unexpected_top_level(
197                    &toplevel.to_string(),
198                    range.start().into(),
199                ));
200            }
201
202            // Dispatch on the NODE_ATTRPATH child alone, not on the value.
203            // For `inputs = { ... }` the value is the whole inputs attrset;
204            // stringifying it dominates this walk on large flakes.
205            let Some(attrpath) = toplevel
206                .children()
207                .find(|c| c.kind() == SyntaxKind::NODE_ATTRPATH)
208            else {
209                continue;
210            };
211            let mut path_idents = attrpath.children();
212            let Some(first_ident) = path_idents.next() else {
213                continue;
214            };
215            let has_more_idents = path_idents.next().is_some();
216            let first_text = first_ident.to_string();
217            let first_unquoted = strip_outer_quotes(&first_text);
218
219            if !has_more_idents && first_unquoted == "description" {
220                continue;
221            }
222
223            if first_unquoted == "inputs" {
224                if has_more_idents {
225                    if let Some(result) =
226                        self.handle_inputs_flat(&attr_set, &toplevel, &attrpath, &ctx, change)
227                    {
228                        return Ok(Some(result));
229                    }
230                } else if let Some(result) =
231                    self.handle_inputs_attr(&toplevel, &attrpath, &ctx, change)
232                {
233                    return Ok(Some(result));
234                }
235                continue;
236            }
237
238            if !has_more_idents
239                && first_unquoted == "outputs"
240                && let Some(result) = self.handle_add_at_outputs(&attr_set, &toplevel, change)
241            {
242                return Ok(Some(result));
243            }
244        }
245
246        // Follows on toplevel flat-style inputs (`inputs.X.url = "..."`).
247        if let Change::Follows { input, target } = change {
248            let path = input.path();
249            if path.len() >= 2 {
250                let parent_id = input.input();
251                if self.inputs.contains_key(parent_id.as_str()) {
252                    let target_str = target.to_flake_follows_string();
253                    return self.handle_follows_flat_toplevel(&attr_set, path, &target_str);
254                }
255            }
256        }
257
258        Ok(None)
259    }
260
261    /// Add a follows attribute next to a toplevel flat-style input.
262    ///
263    /// Converts `inputs.crane.url = "github:...";` into:
264    /// ```nix
265    /// inputs.crane.url = "github:...";
266    /// inputs.crane.inputs.nixpkgs.follows = "nixpkgs";
267    /// ```
268    fn handle_follows_flat_toplevel(
269        &self,
270        attr_set: &SyntaxNode,
271        path: &AttrPath,
272        target: &str,
273    ) -> Result<Option<SyntaxNode>, WalkerError> {
274        let parent_id = path.first();
275        let expected_flat = follows_idents_prefixed(path.segments());
276        let mut last_parent_attr: Option<SyntaxNode> = None;
277        let mut block_parent: Option<(SyntaxNode, SyntaxNode)> = None;
278
279        for toplevel in attr_set.children() {
280            if toplevel.kind() != SyntaxKind::NODE_ATTRPATH_VALUE {
281                continue;
282            }
283            let Some(attrpath) = toplevel
284                .children()
285                .find(|c| c.kind() == SyntaxKind::NODE_ATTRPATH)
286            else {
287                continue;
288            };
289            let idents: Vec<String> = attrpath.children().map(|c| c.to_string()).collect();
290
291            if let Some(block_attr_set) =
292                block_parent_attrset(&toplevel, &idents, parent_id.as_str())
293            {
294                block_parent = Some((toplevel.clone(), block_attr_set));
295            }
296
297            if idents_match(&idents, &expected_flat)
298                && let Some(rebuilt) = retarget_existing_flat_follows(
299                    attr_set,
300                    &toplevel,
301                    attrpath.next_sibling(),
302                    target,
303                )
304            {
305                return Ok(Some(rebuilt));
306            }
307
308            if is_flat_inputs_attr_for(&idents, parent_id.as_str()) {
309                last_parent_attr = Some(toplevel.clone());
310            }
311        }
312
313        if let Some((toplevel, block_attr_set)) = block_parent {
314            let rest: Vec<Segment> = path.segments()[1..].to_vec();
315            return self.handle_follows_block_toplevel(
316                attr_set,
317                &toplevel,
318                &block_attr_set,
319                &rest,
320                target,
321            );
322        }
323
324        if let Some(ref_child) = last_parent_attr {
325            return Ok(Some(insert_flat_follows_after(
326                attr_set, &ref_child, path, target,
327            )));
328        }
329
330        Ok(None)
331    }
332
333    fn handle_follows_block_toplevel(
334        &self,
335        attr_set: &SyntaxNode,
336        toplevel: &SyntaxNode,
337        block_attr_set: &SyntaxNode,
338        rest: &[Segment],
339        target: &str,
340    ) -> Result<Option<SyntaxNode>, WalkerError> {
341        let expected_block = follows_idents_prefixed(rest);
342        for attr in block_attr_set.children() {
343            if attr.kind() != SyntaxKind::NODE_ATTRPATH_VALUE {
344                continue;
345            }
346            let Some(attrpath) = attr
347                .children()
348                .find(|c| c.kind() == SyntaxKind::NODE_ATTRPATH)
349            else {
350                continue;
351            };
352            let idents: Vec<String> = attrpath.children().map(|c| c.to_string()).collect();
353
354            if idents_match(&idents, &expected_block) {
355                let value_node = attrpath.next_sibling();
356                let current_target = value_node
357                    .as_ref()
358                    .map(|v| strip_outer_quotes(&v.to_string()).to_string())
359                    .unwrap_or_default();
360
361                if current_target == target {
362                    return Ok(attr_set.ancestors().last());
363                }
364
365                if let Some(value) = value_node {
366                    let new_value = make_quoted_string(target);
367                    let new_attr = substitute_child(&attr, value.index(), &new_value);
368                    let new_block = substitute_child(block_attr_set, attr.index(), &new_attr);
369                    let new_toplevel =
370                        substitute_child(toplevel, block_attr_set.index(), &new_block);
371                    let green = attr_set
372                        .green()
373                        .replace_child(toplevel.index(), new_toplevel.green().into());
374                    return Ok(Some(SyntaxNode::new_root(attr_set.replace_with(green))));
375                }
376            }
377        }
378
379        let follows_node = FollowsKind::BlockNested { rest, target }.emit();
380        let children: Vec<_> = block_attr_set.children().collect();
381        if let Some(last_child) = children.last() {
382            let insert_index = last_child.index() + 1;
383
384            let mut green = block_attr_set
385                .green()
386                .insert_child(insert_index, follows_node.green().into());
387
388            if let Some(whitespace) = get_sibling_whitespace(last_child) {
389                green = green.insert_child(insert_index, whitespace.green().into());
390            }
391
392            let new_block = SyntaxNode::new_root(green);
393            let new_toplevel = substitute_child(toplevel, block_attr_set.index(), &new_block);
394            let green = attr_set
395                .green()
396                .replace_child(toplevel.index(), new_toplevel.green().into());
397            return Ok(Some(SyntaxNode::new_root(attr_set.replace_with(green))));
398        }
399
400        Ok(None)
401    }
402
403    /// Apply `change` to the `inputs = { ... }` attribute.
404    ///
405    /// `toplevel.replace_with()` propagates through `NODE_ATTR_SET` up to `NODE_ROOT`,
406    /// preserving leading comments and trivia.
407    fn handle_inputs_attr(
408        &mut self,
409        toplevel: &SyntaxNode,
410        child: &SyntaxNode,
411        ctx: &Option<Context>,
412        change: &Change,
413    ) -> Option<SyntaxNode> {
414        let sibling = child.next_sibling()?;
415        let replacement = walk_inputs(&mut self.inputs, sibling.clone(), ctx, change)?;
416
417        let green = toplevel
418            .green()
419            .replace_child(sibling.index(), replacement.green().into());
420        let green = toplevel.replace_with(green);
421        Some(SyntaxNode::new_root(green))
422    }
423
424    /// Apply `change` to flat-style `inputs.foo.url = "..."` attributes.
425    ///
426    /// Removals rebuild the parent attrset green and `replace_with()` propagates to
427    /// `NODE_ROOT`. Replacements rely on `toplevel.replace_with()` to propagate.
428    fn handle_inputs_flat(
429        &mut self,
430        attr_set: &SyntaxNode,
431        toplevel: &SyntaxNode,
432        child: &SyntaxNode,
433        ctx: &Option<Context>,
434        change: &Change,
435    ) -> Option<SyntaxNode> {
436        let replacement = walk_inputs(&mut self.inputs, child.clone(), ctx, change)?;
437
438        // Empty replacement means we remove the entire toplevel node and
439        // propagate through attr_set to NODE_ROOT.
440        if replacement.to_string().is_empty() {
441            let element: rnix::SyntaxElement = toplevel.clone().into();
442            let mut green = attr_set.green().remove_child(toplevel.index());
443            if let Some(ws_index) = adjacent_whitespace_index(&element) {
444                green = green.remove_child(ws_index);
445            }
446            return Some(SyntaxNode::new_root(attr_set.replace_with(green)));
447        }
448
449        let sibling = child.next_sibling()?;
450        let green = toplevel
451            .green()
452            .replace_child(sibling.index(), replacement.green().into());
453        let green = toplevel.replace_with(green);
454        Some(SyntaxNode::new_root(green))
455    }
456
457    /// Add a new input just before `outputs` when no `inputs` block exists yet.
458    ///
459    /// Rebuilds the parent attrset green. `replace_with()` propagates to `NODE_ROOT`
460    /// while preserving leading comments.
461    fn handle_add_at_outputs(
462        &mut self,
463        attr_set: &SyntaxNode,
464        toplevel: &SyntaxNode,
465        change: &Change,
466    ) -> Option<SyntaxNode> {
467        if !self.add_toplevel {
468            return None;
469        }
470
471        let Change::Add {
472            id: Some(id),
473            uri: Some(uri),
474            flake,
475        } = change
476        else {
477            return None;
478        };
479        let id = id.input().as_str();
480
481        if toplevel.index() == 0 {
482            return None;
483        }
484
485        // Walk back from `outputs` through tokens to find a whitespace run, then
486        // normalize it to a single newline + indent. Walking through tokens (not
487        // siblings) lets us skip past comments between the last input and `outputs`.
488        let ws_node = {
489            let mut ws: Option<SyntaxNode> = None;
490            let mut cursor = toplevel.prev_sibling_or_token();
491            while let Some(ref tok) = cursor {
492                if tok.kind() == SyntaxKind::TOKEN_WHITESPACE {
493                    let ws_str = tok.to_string();
494                    ws = Some(parse_node(last_line_with_newline(&ws_str)));
495                    break;
496                }
497                cursor = tok.prev_sibling_or_token();
498            }
499            ws
500        };
501
502        let addition = make_toplevel_url_attr(id, uri);
503        let insert_pos = toplevel.index() - 1;
504
505        let mut green = attr_set
506            .green()
507            .insert_child(insert_pos, addition.green().into());
508
509        if let Some(ref ws) = ws_node {
510            green = green.insert_child(insert_pos, ws.green().into());
511        }
512
513        // Append `inputs.<id>.flake = false;` when the new input opts out of flake mode.
514        if !flake {
515            let no_flake = make_toplevel_flake_false_attr(id);
516            green = green.insert_child(toplevel.index() + 1, no_flake.green().into());
517
518            if let Some(ref ws) = ws_node {
519                green = green.insert_child(toplevel.index() + 1, ws.green().into());
520            }
521        }
522
523        Some(SyntaxNode::new_root(attr_set.replace_with(green)))
524    }
525}
526
527#[cfg(test)]
528mod tests {
529    use super::*;
530    use crate::change::{Change, ChangeId};
531    use crate::follows::AttrPath;
532
533    fn apply(flake_text: &str, change: &Change) -> String {
534        let mut walker = Walker::new(flake_text);
535        walker
536            .walk(change)
537            .expect("walker error")
538            .expect("walker did not rewrite the tree")
539            .to_string()
540    }
541
542    fn follows_change(input: &str, target: &str) -> Change {
543        Change::Follows {
544            input: ChangeId::parse(input).unwrap(),
545            target: AttrPath::parse(target).unwrap(),
546        }
547    }
548
549    #[test]
550    fn handle_follows_flat_toplevel_inserts_follows_after_last_parent_attr() {
551        let flake = "{
552  inputs.flake-edit.url = \"github:a-kenji/flake-edit\";
553  inputs.nixpkgs.url = \"github:NixOS/nixpkgs\";
554
555  outputs = { self, ... }: { };
556}
557";
558        let result = apply(flake, &follows_change("flake-edit.nixpkgs", "nixpkgs"));
559        assert_eq!(
560            result,
561            "{
562  inputs.flake-edit.url = \"github:a-kenji/flake-edit\";
563  inputs.flake-edit.inputs.nixpkgs.follows = \"nixpkgs\";
564  inputs.nixpkgs.url = \"github:NixOS/nixpkgs\";
565
566  outputs = { self, ... }: { };
567}
568"
569        );
570    }
571
572    #[test]
573    fn handle_follows_flat_toplevel_retargets_existing_follows() {
574        let flake = "{
575  inputs.flake-edit.url = \"github:a-kenji/flake-edit\";
576  inputs.flake-edit.inputs.nixpkgs.follows = \"old-pkgs\";
577  inputs.nixpkgs.url = \"github:NixOS/nixpkgs\";
578
579  outputs = { self, ... }: { };
580}
581";
582        let result = apply(flake, &follows_change("flake-edit.nixpkgs", "nixpkgs"));
583        assert_eq!(
584            result,
585            "{
586  inputs.flake-edit.url = \"github:a-kenji/flake-edit\";
587  inputs.flake-edit.inputs.nixpkgs.follows = \"nixpkgs\";
588  inputs.nixpkgs.url = \"github:NixOS/nixpkgs\";
589
590  outputs = { self, ... }: { };
591}
592"
593        );
594    }
595
596    #[test]
597    fn handle_follows_flat_toplevel_is_noop_when_target_already_matches() {
598        let flake = "{
599  inputs.flake-edit.url = \"github:a-kenji/flake-edit\";
600  inputs.flake-edit.inputs.nixpkgs.follows = \"nixpkgs\";
601  inputs.nixpkgs.url = \"github:NixOS/nixpkgs\";
602
603  outputs = { self, ... }: { };
604}
605";
606        let result = apply(flake, &follows_change("flake-edit.nixpkgs", "nixpkgs"));
607        assert_eq!(result, flake);
608    }
609
610    #[test]
611    fn handle_follows_flat_toplevel_delegates_to_block_parent_when_present() {
612        let flake = "{
613  inputs.nixpkgs.url = \"github:NixOS/nixpkgs\";
614  inputs.flake-edit = {
615    url = \"github:a-kenji/flake-edit\";
616  };
617
618  outputs = { self, ... }: { };
619}
620";
621        let result = apply(flake, &follows_change("flake-edit.nixpkgs", "nixpkgs"));
622        assert_eq!(
623            result,
624            "{
625  inputs.nixpkgs.url = \"github:NixOS/nixpkgs\";
626  inputs.flake-edit = {
627    url = \"github:a-kenji/flake-edit\";
628    inputs.nixpkgs.follows = \"nixpkgs\";
629  };
630
631  outputs = { self, ... }: { };
632}
633"
634        );
635    }
636
637    #[test]
638    fn is_flat_inputs_attr_for_only_matches_matching_parent_id() {
639        let yes = [
640            "inputs".to_string(),
641            "flake-edit".to_string(),
642            "url".to_string(),
643        ];
644        let no = [
645            "inputs".to_string(),
646            "nixpkgs".to_string(),
647            "url".to_string(),
648        ];
649        assert!(is_flat_inputs_attr_for(&yes, "flake-edit"));
650        assert!(!is_flat_inputs_attr_for(&no, "flake-edit"));
651        // The CST keeps surrounding `"..."` on quoted idents; the comparison
652        // must unquote them, otherwise `"flake-edit" != flake-edit`.
653        let quoted = [
654            "inputs".to_string(),
655            "\"flake-edit\"".to_string(),
656            "url".to_string(),
657        ];
658        assert!(is_flat_inputs_attr_for(&quoted, "flake-edit"));
659    }
660}
flake_edit/walk.rs

flake_edit/
walk.rs
