nix_editor/
write.rs

1use std::collections::HashMap;
2
3use crate::parse::{findattr, getcfgbase, getkey};
4use thiserror::Error;
5use rnix::{self, SyntaxKind, SyntaxNode};
6
7#[derive(Error, Debug)]
8pub enum WriteError {
9    #[error("Error while parsing.")]
10    ParseError,
11    #[error("No attributes.")]
12    NoAttr,
13    #[error("Error with array.")]
14    ArrayError,
15    #[error("Writing value to attribute set.")]
16    WriteValueToSet,
17}
18
19pub fn write(f: &str, query: &str, val: &str) -> Result<String, WriteError> {
20    let ast = rnix::Root::parse(f);
21    let configbase = match getcfgbase(&ast.syntax()) {
22        Some(x) => x,
23        None => {
24            return Err(WriteError::ParseError);
25        }
26    };
27    if val.trim_start().starts_with('{') && val.trim_end().ends_with('}'){
28        if let Some(x) = getcfgbase(&rnix::Root::parse(val).syntax()) {
29            if x.kind() == SyntaxKind::NODE_ATTR_SET {
30                return addattrval(f, &configbase, query, &x);
31            }
32        }
33    }
34    
35    let outnode = match findattr(&configbase, query) {
36        Some(x) => {
37            if let Some(n) = x.children().last() {
38                if n.kind() == SyntaxKind::NODE_ATTR_SET {
39                    return Err(WriteError::WriteValueToSet);
40                }
41            }
42            modvalue(&x, val).unwrap()
43        }
44        None => {
45            let mut y = query.split('.').collect::<Vec<_>>();
46            y.pop();
47            let x = findattrset(&configbase, &y.join("."), 0);
48            match x {
49                Some((base, v, spaces)) => addvalue(
50                    &base,
51                    &format!("{}{}", " ".repeat(spaces), &query[v.len() + 1..]),
52                    val,
53                ),
54                None => addvalue(&configbase, query, val),
55            }
56        }
57    };
58    Ok(outnode.to_string())
59}
60
61fn addvalue(configbase: &SyntaxNode, query: &str, val: &str) -> SyntaxNode {
62    let mut index = configbase.green().children().len() - 2;
63    // To find a better index for insertion, first find a matching node, then find the next newline token, after that, insert
64    if let Some(x) = matchval(configbase, query, query.split('.').count()) {
65        let i = configbase
66            .green()
67            .children()
68            .position(|y| match y.into_node() {
69                Some(y) => y.to_owned() == x.green().into_owned(),
70                None => false,
71            })
72            .unwrap();
73        let configgreen = configbase.green().to_owned();
74        let configafter = &configgreen.children().collect::<Vec<_>>()[i..];
75        for child in configafter {
76            if let Some(x) = child.as_token() {
77                if x.text().contains('\n') {
78                    let cas = configafter.to_vec();
79                    index = i + cas
80                        .iter()
81                        .position(|y| match y.as_token() {
82                            Some(t) => t == x,
83                            None => false,
84                        })
85                        .unwrap();
86                    break;
87                }
88            }
89        }
90    }
91    let input = rnix::Root::parse(format!("\n  {} = {};", &query, &val).as_str())
92        .syntax();
93    let input = input
94        .green()
95        .to_owned();
96    if index == 0 {
97        index += 1;
98    };
99    let new = configbase
100        .green()
101        .insert_child(index, rnix::NodeOrToken::Node(input.into_owned()));
102    let replace = configbase.replace_with(new);
103    rnix::Root::parse(&replace.to_string()).syntax()
104}
105
106// Currently indentation is badly done by inserting spaces, it should check the spaces of the previous attr instead
107fn findattrset(
108    configbase: &SyntaxNode,
109    name: &str,
110    spaces: usize,
111) -> Option<(SyntaxNode, String, usize)> {
112    for child in configbase.children() {
113        if child.kind() == SyntaxKind::NODE_ATTRPATH_VALUE {
114            // Now we have to read all the indent values from the key
115            for subchild in child.children() {
116                if subchild.kind() == SyntaxKind::NODE_ATTRPATH {
117                    // We have a key, now we need to check if it's the one we're looking for
118                    let key = getkey(&subchild);
119                    let qkey = name
120                        .split('.')
121                        .map(|s| s.to_string())
122                        .collect::<Vec<String>>();
123                    if qkey == key {
124                        // We have key, now lets find the attrset
125                        for possibleset in child.children() {
126                            if possibleset.kind() == SyntaxKind::NODE_ATTR_SET {
127                                return Some((possibleset, name.to_string(), spaces + 2));
128                            }
129                        }
130                        return None;
131                    } else if qkey.len() > key.len() {
132                        // We have a subkey, so we need to recurse
133                        if key == qkey[0..key.len()] {
134                            // We have a subkey, so we need to recurse
135                            let subkey = &qkey[key.len()..].join(".").to_string();
136                            let newbase = getcfgbase(&child).unwrap();
137                            let subattr = findattrset(&newbase, subkey, spaces + 2);
138                            if let Some((node, _, spaces)) = subattr {
139                                return Some((node, name.to_string(), spaces));
140                            }
141                        }
142                    }
143                }
144            }
145        }
146    }
147    None
148}
149
150fn matchval(configbase: &SyntaxNode, query: &str, acc: usize) -> Option<SyntaxNode> {
151    let qvec = &query
152        .split('.')
153        .map(|s| s.to_string())
154        .collect::<Vec<String>>();
155    let q = &qvec[..acc];
156    for child in configbase.children() {
157        if child.kind() == SyntaxKind::NODE_ATTRPATH_VALUE {
158            for subchild in child.children() {
159                if subchild.kind() == SyntaxKind::NODE_ATTRPATH {
160                    let key = getkey(&subchild);
161                    if key.len() >= q.len() && &key[..q.len()] == q {
162                        return Some(child);
163                    }
164                }
165            }
166        }
167    }
168    if acc == 1 {
169        None
170    } else {
171        matchval(configbase, query, acc - 1)
172    }
173}
174
175fn modvalue(node: &SyntaxNode, val: &str) -> Option<SyntaxNode> {
176    // First find the IDENT node
177    for child in node.children() {
178        if child.kind() != SyntaxKind::NODE_ATTRPATH {
179            let c = &child;
180            let input = val.to_string();
181            let rep = &rnix::Root::parse(&input)
182                .syntax()
183                .children()
184                .collect::<Vec<SyntaxNode>>()[0];
185            let index = node
186                .green()
187                .children()
188                .position(|y| match y.into_node() {
189                    Some(y) => y.to_owned() == c.green().into_owned(),
190                    None => false,
191                })
192                .unwrap();
193            let replaced = node
194                .green()
195                .replace_child(index, rnix::NodeOrToken::Node(rep.green().into_owned()));
196            let out = node.replace_with(replaced);
197            let rnode = rnix::Root::parse(&out.to_string()).syntax();
198            return Some(rnode);
199        }
200    }
201    None
202}
203
204// Add an attribute to the config
205fn addattrval(
206    f: &str,
207    configbase: &SyntaxNode,
208    query: &str,
209    val: &SyntaxNode,
210) -> Result<String, WriteError> {
211    let mut attrmap = HashMap::new();
212    buildattrvec(val, vec![], &mut attrmap);
213    let mut file = f.to_string();
214
215    if attrmap.iter().any(|(key, _)| findattr(configbase, &format!("{}.{}", query, key)).is_some()) {
216        for (key, val) in attrmap {
217            match write(&file, &format!("{}.{}", query, key), &val) {
218                Ok(x) => {
219                    file = x
220                },
221                Err(e) => return Err(e),
222            }
223        }
224    } else if let Some(c) = getcfgbase(&rnix::Root::parse(&file).syntax()) {
225        file = addvalue(&c, query, &val.to_string()).to_string();
226    }    
227    Ok(file)
228}
229
230fn buildattrvec(val: &SyntaxNode, prefix: Vec<String>, map: &mut HashMap<String, String>) {
231    for child in val.children() {
232        if child.kind() == SyntaxKind::NODE_ATTRPATH_VALUE {
233            if let Some(subchild) = child.children().last() {
234                if subchild.kind() == SyntaxKind::NODE_ATTR_SET {
235                    for c in child.children() {
236                        if c.kind() == SyntaxKind::NODE_ATTRPATH {
237                            let key = getkey(&c);
238                            let mut newprefix = prefix.clone();
239                            newprefix.append(&mut key.clone());
240                            buildattrvec(&subchild, newprefix, map);
241                            break;
242                        }
243                    }
244                } else {
245                    for c in child.children() {
246                        if c.kind() == SyntaxKind::NODE_ATTRPATH {
247                            let key = getkey(&c);
248                            let mut newprefix = prefix.clone();
249                            newprefix.append(&mut key.clone());
250                            map.insert(newprefix.join("."), subchild.to_string());
251                        }
252                    }
253                }
254            }
255        }
256    }
257}
258
259pub fn addtoarr(f: &str, query: &str, items: Vec<String>) -> Result<String, WriteError> {
260    let ast = rnix::Root::parse(f);
261    let configbase = match getcfgbase(&ast.syntax()) {
262        Some(x) => x,
263        None => return Err(WriteError::ParseError),
264    };
265    let outnode = match findattr(&configbase, query) {
266        Some(x) => match addtoarr_aux(&x, items) {
267            Some(x) => x,
268            None => return Err(WriteError::ArrayError),
269        },
270        // If no arrtibute is found, create a new one
271        None => {
272            let newval = addvalue(&configbase, query, "[\n  ]");
273            return addtoarr(&newval.to_string(), query, items);
274        }
275    };
276    Ok(outnode.to_string())
277}
278
279fn addtoarr_aux(node: &SyntaxNode, items: Vec<String>) -> Option<SyntaxNode> {
280    for child in node.children() {
281        if child.kind() == rnix::SyntaxKind::NODE_WITH {
282            return addtoarr_aux(&child, items);
283        }
284        if child.kind() == SyntaxKind::NODE_LIST {
285            let mut green = child.green().into_owned();
286
287            for elem in items {
288                let mut i = 0;
289                for c in green.children() {
290                    if c.to_string() == "]" {
291                        if green.children().collect::<Vec<_>>()[i - 1]
292                            .as_token()
293                            .unwrap()
294                            .to_string()
295                            .contains('\n')
296                        {
297                            i -= 1;
298                        }
299                        green = green.insert_child(
300                            i,
301                            rnix::NodeOrToken::Node(
302                                rnix::Root::parse(&format!("\n{}{}", " ".repeat(4), elem))
303                                    .syntax()
304                                    .green()
305                                    .into_owned(),
306                            ),
307                        );
308                        break;
309                    }
310                    i += 1;
311                }
312            }
313
314            let index = match node.green().children().position(|x| match x.into_node() {
315                Some(x) => x.to_owned() == child.green().into_owned(),
316                None => false,
317            }) {
318                Some(x) => x,
319                None => return None,
320            };
321
322            let replace = node
323                .green()
324                .replace_child(index, rnix::NodeOrToken::Node(green));
325            let out = node.replace_with(replace);
326            let output = rnix::Root::parse(&out.to_string()).syntax();
327            return Some(output);
328        }
329    }
330    None
331}
332
333pub fn rmarr(f: &str, query: &str, items: Vec<String>) -> Result<String, WriteError> {
334    let ast = rnix::Root::parse(f);
335    let configbase = match getcfgbase(&ast.syntax()) {
336        Some(x) => x,
337        None => return Err(WriteError::ParseError),
338    };
339    let outnode = match findattr(&configbase, query) {
340        Some(x) => match rmarr_aux(&x, items) {
341            Some(x) => x,
342            None => return Err(WriteError::ArrayError),
343        },
344        None => return Err(WriteError::NoAttr),
345    };
346    Ok(outnode.to_string())
347}
348
349fn rmarr_aux(node: &SyntaxNode, items: Vec<String>) -> Option<SyntaxNode> {
350    for child in node.children() {
351        if child.kind() == rnix::SyntaxKind::NODE_WITH {
352            return rmarr_aux(&child, items);
353        }
354        if child.kind() == SyntaxKind::NODE_LIST {
355            let green = child.green().into_owned();
356            let mut idx = vec![];
357            for elem in green.children() {
358                if elem.as_node() != None && items.contains(&elem.to_string()) {
359                    let index = match green.children().position(|x| match x.into_node() {
360                        Some(x) => {
361                            if let Some(y) = elem.as_node() {
362                                x.eq(y)
363                            } else {
364                                false
365                            }
366                        }
367                        None => false,
368                    }) {
369                        Some(x) => x,
370                        None => return None,
371                    };
372                    idx.push(index)
373                }
374            }
375            let mut acc = 0;
376            let mut replace = green;
377
378            for i in idx {
379                replace = replace.remove_child(i - acc);
380                let mut v = vec![];
381                for c in replace.children() {
382                    v.push(c);
383                }
384                if let Some(x) = v.get(i - acc - 1).unwrap().as_token() {
385                    if x.to_string().contains('\n') {
386                        replace = replace.remove_child(i - acc - 1);
387                        acc += 1;
388                    }
389                }
390                acc += 1;
391            }
392            let out = child.replace_with(replace);
393
394            let output = rnix::Root::parse(&out.to_string()).syntax();
395            return Some(output);
396        }
397    }
398    None
399}
400
401pub fn deref(f: &str, query: &str) -> Result<String, WriteError> {
402    let ast = rnix::Root::parse(f);
403    let configbase = match getcfgbase(&ast.syntax()) {
404        Some(x) => x,
405        None => return Err(WriteError::ParseError),
406    };
407    let outnode = match deref_aux(&configbase, query) {
408        Some(x) => x,
409        None => return Err(WriteError::NoAttr),
410    };
411    Ok(outnode.to_string())
412}
413
414fn deref_aux(configbase: &SyntaxNode, name: &str) -> Option<SyntaxNode> {
415    for child in configbase.children() {
416        if child.kind() == SyntaxKind::NODE_ATTRPATH_VALUE {
417            // Now we have to read all the indent values from the key
418            for subchild in child.children() {
419                if subchild.kind() == SyntaxKind::NODE_ATTRPATH {
420                    // We have a key, now we need to check if it's the one we're looking for
421                    let key = getkey(&subchild);
422                    let qkey = name
423                        .split('.')
424                        .map(|s| s.to_string())
425                        .collect::<Vec<String>>();
426                    if qkey == key {
427                        let index =
428                            match configbase
429                                .green()
430                                .children()
431                                .position(|x| match x.into_node() {
432                                    Some(x) => x.to_owned() == child.green().into_owned(),
433                                    None => false,
434                                }) {
435                                Some(x) => x,
436                                None => return None,
437                            };
438                        let mut del = configbase.green().remove_child(index);
439
440                        // Remove leading newline if it still exists
441                        if del.children().collect::<Vec<_>>()[index]
442                            .to_string()
443                            .contains('\n')
444                        {
445                            del = del.remove_child(index);
446                        }
447                        let out = configbase.replace_with(del);
448                        return Some(rnix::Root::parse(&out.to_string()).syntax());
449                    } else if qkey.len() > key.len() {
450                        // We have a subkey, so we need to recurse
451                        if key == qkey[0..key.len()] {
452                            // We have a subkey, so we need to recurse
453                            let subkey = &qkey[key.len()..].join(".").to_string();
454                            let newbase = getcfgbase(&child).unwrap();
455                            let subattr = deref_aux(&newbase, subkey);
456                            if let Some(s) = subattr {
457                                return Some(s);
458                            }
459                        }
460                    }
461                }
462            }
463        }
464    }
465    None
466}
nix_editor/write.rs

nix_editor/
write.rs
