flake_edit/
walk.rs

1//! AST walking and manipulation 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::input::Input;
16
17pub use context::Context;
18pub use error::WalkerError;
19
20use inputs::{remove_child_with_whitespace, walk_inputs};
21use node::{
22    get_sibling_whitespace, make_toplevel_flake_false_attr, make_toplevel_url_attr, parse_node,
23};
24
25#[derive(Debug, Clone)]
26pub struct Walker {
27    pub root: SyntaxNode,
28    pub inputs: HashMap<String, Input>,
29    pub add_toplevel: bool,
30}
31
32impl<'a> Walker {
33    pub fn new(stream: &'a str) -> Self {
34        let root = Root::parse(stream).syntax();
35        Self {
36            root,
37            inputs: HashMap::new(),
38            add_toplevel: false,
39        }
40    }
41
42    /// Traverse the toplevel `flake.nix` file.
43    /// It should consist of three attribute keys:
44    /// - description
45    /// - inputs
46    /// - outputs
47    pub fn walk(&mut self, change: &Change) -> Result<Option<SyntaxNode>, WalkerError> {
48        let cst = &self.root;
49        if cst.kind() != SyntaxKind::NODE_ROOT {
50            return Err(WalkerError::NotARoot(cst.kind()));
51        }
52        self.walk_toplevel(cst.clone(), None, change)
53    }
54
55    /// Only walk the outputs attribute
56    pub(crate) fn list_outputs(&mut self) -> Result<Outputs, WalkerError> {
57        outputs::list_outputs(&self.root)
58    }
59
60    /// Only change the outputs attribute
61    pub(crate) fn change_outputs(
62        &mut self,
63        change: OutputChange,
64    ) -> Result<Option<SyntaxNode>, WalkerError> {
65        outputs::change_outputs(&self.root, change)
66    }
67
68    /// Traverse the toplevel `flake.nix` file.
69    fn walk_toplevel(
70        &mut self,
71        node: SyntaxNode,
72        ctx: Option<Context>,
73        change: &Change,
74    ) -> Result<Option<SyntaxNode>, WalkerError> {
75        let Some(root) = node.first_child() else {
76            return Ok(None);
77        };
78
79        for toplevel in root.children() {
80            if toplevel.kind() != SyntaxKind::NODE_ATTRPATH_VALUE {
81                return Err(WalkerError::UnexpectedNodeKind {
82                    expected: SyntaxKind::NODE_ATTRPATH_VALUE,
83                    found: toplevel.kind(),
84                });
85            }
86
87            for child in toplevel.children() {
88                let child_str = child.to_string();
89
90                if child_str == "description" {
91                    break;
92                }
93
94                if child_str == "inputs" {
95                    if let Some(result) =
96                        self.handle_inputs_attr(&root, &toplevel, &child, &ctx, change)
97                    {
98                        return Ok(Some(result));
99                    }
100                    continue;
101                }
102
103                if child_str.starts_with("inputs") {
104                    if let Some(result) =
105                        self.handle_inputs_flat(&root, &toplevel, &child, &ctx, change)
106                    {
107                        return Ok(Some(result));
108                    }
109                    continue;
110                }
111
112                if child_str == "outputs"
113                    && let Some(result) =
114                        self.handle_add_at_outputs(&root, &toplevel, &child, change)
115                {
116                    return Ok(Some(result));
117                }
118            }
119        }
120        Ok(None)
121    }
122
123    /// Handle `inputs = { ... }` attribute
124    fn handle_inputs_attr(
125        &mut self,
126        _root: &SyntaxNode,
127        toplevel: &SyntaxNode,
128        child: &SyntaxNode,
129        ctx: &Option<Context>,
130        change: &Change,
131    ) -> Option<SyntaxNode> {
132        let sibling = child.next_sibling()?;
133        let replacement = walk_inputs(&mut self.inputs, sibling.clone(), ctx, change)?;
134
135        let green = toplevel
136            .green()
137            .replace_child(sibling.index(), replacement.green().into());
138        let green = toplevel.replace_with(green);
139        Some(parse_node(&green.to_string()))
140    }
141
142    /// Handle flat-style `inputs.foo.url = "..."` attributes
143    fn handle_inputs_flat(
144        &mut self,
145        root: &SyntaxNode,
146        toplevel: &SyntaxNode,
147        child: &SyntaxNode,
148        ctx: &Option<Context>,
149        change: &Change,
150    ) -> Option<SyntaxNode> {
151        let replacement = walk_inputs(&mut self.inputs, child.clone(), ctx, change)?;
152
153        // If replacement is empty, remove the entire toplevel node
154        if replacement.to_string().is_empty() {
155            return Some(remove_child_with_whitespace(
156                root,
157                toplevel,
158                toplevel.index(),
159            ));
160        }
161
162        let sibling = child.next_sibling()?;
163        let green = toplevel
164            .green()
165            .replace_child(sibling.index(), replacement.green().into());
166        let green = toplevel.replace_with(green);
167        Some(parse_node(&green.to_string()))
168    }
169
170    /// Handle adding inputs when we've reached `outputs` but have no inputs yet
171    fn handle_add_at_outputs(
172        &mut self,
173        root: &SyntaxNode,
174        toplevel: &SyntaxNode,
175        child: &SyntaxNode,
176        change: &Change,
177    ) -> Option<SyntaxNode> {
178        if !self.add_toplevel {
179            return None;
180        }
181
182        let Change::Add {
183            id: Some(id),
184            uri: Some(uri),
185            flake,
186        } = change
187        else {
188            return None;
189        };
190
191        if toplevel.index() == 0 {
192            return None;
193        }
194
195        let addition = make_toplevel_url_attr(id, uri);
196        let insert_pos = toplevel.index() - 1;
197
198        let mut green = root
199            .green()
200            .insert_child(insert_pos, addition.green().into());
201
202        // Add whitespace before the new input
203        if let Some(prev_child) = root.children().find(|c| c.index() == toplevel.index() - 2)
204            && let Some(whitespace) = get_sibling_whitespace(&prev_child)
205        {
206            green = green.insert_child(insert_pos, whitespace.green().into());
207        }
208
209        // Add flake=false if needed
210        if !flake {
211            let no_flake = make_toplevel_flake_false_attr(id);
212            green = green.insert_child(toplevel.index() + 1, no_flake.green().into());
213
214            if let Some(prev_child) = root.children().find(|c| c.index() == toplevel.index() - 2)
215                && let Some(whitespace) = get_sibling_whitespace(&prev_child)
216            {
217                green = green.insert_child(toplevel.index() + 1, whitespace.green().into());
218            }
219        }
220
221        // Preserve whitespace after outputs
222        if let Some(next) = child.next_sibling_or_token()
223            && next.kind() == SyntaxKind::TOKEN_WHITESPACE
224        {
225            let whitespace = parse_node(next.as_token().unwrap().green().text());
226            green = green.insert_child(child.index() + 1, whitespace.green().into());
227        }
228
229        Some(parse_node(&green.to_string()))
230    }
231}
flake_edit/walk.rs

flake_edit/
walk.rs
