Skip to main content

seahorse_dev/core/
util.rs

1use owo_colors::OwoColorize;
2use rustpython_parser::ast::Location as SrcLocation;
3use std::{
4    collections::{BTreeMap, HashMap},
5    error, fmt,
6    rc::Rc,
7};
8
9/// Match and extract against a single pattern, panicking if no match.
10#[macro_export]
11macro_rules! match1 {
12    ($obj:expr, $var:pat => $up:expr) => {
13        match $obj {
14            $var => $up,
15            obj => panic!("match1 failed: received {:?}", obj),
16        }
17    };
18}
19
20#[derive(Debug, Clone)]
21pub struct CoreError {
22    message: (String, String),
23    src: Option<Rc<String>>,
24    loc: Option<SrcLocation>,
25}
26
27impl CoreError {
28    pub fn located(self, loc: Location) -> Self {
29        Self {
30            loc: Some(loc.loc),
31            src: Some(loc.src),
32            ..self
33        }
34    }
35
36    pub fn updated(self, loc: &Location) -> Self {
37        Self {
38            loc: Some(self.loc.unwrap_or_else(|| loc.loc.clone())),
39            src: Some(self.src.unwrap_or_else(|| loc.src.clone())),
40            ..self
41        }
42    }
43
44    pub fn with_loc(self, loc: SrcLocation) -> Self {
45        Self {
46            loc: Some(loc),
47            ..self
48        }
49    }
50
51    pub fn with_src(self, src: Rc<String>) -> Self {
52        Self {
53            src: Some(src),
54            ..self
55        }
56    }
57
58    pub fn make_raw<H, F>(header: H, footer: F) -> Self
59    where
60        H: ToString,
61        F: ToString,
62    {
63        Self {
64            message: (header.to_string(), footer.to_string()),
65            src: None,
66            loc: None,
67        }
68    }
69}
70
71impl fmt::Display for CoreError {
72    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
73        match self {
74            Self {
75                message: (header, footer),
76                src: Some(src),
77                loc: Some(loc),
78            } => {
79                let line = format!("{} |", loc.row());
80                let context = src.lines().nth(loc.row() - 1).map(|s| s.to_string());
81
82                if footer.len() > 0 {
83                    write!(
84                        f,
85                        "{}: {}.\n{} {}\n{: >col$}\n{}\n",
86                        "Error".red().bold(),
87                        header.bold(),
88                        line,
89                        context.unwrap(),
90                        "^".bold(),
91                        footer,
92                        col = line.len() + 1 + loc.column()
93                    )
94                } else {
95                    write!(
96                        f,
97                        "{}: {}.\n{} {}\n{: >col$}\n",
98                        "Error".red().bold(),
99                        header.bold(),
100                        line,
101                        context.unwrap(),
102                        "^".bold(),
103                        col = line.len() + 1 + loc.column()
104                    )
105                }
106            }
107            Self {
108                message: (header, footer),
109                ..
110            } => {
111                write!(f, "Error: {}.\n{}", header, footer)
112            }
113        }
114    }
115}
116
117impl error::Error for CoreError {}
118
119pub type CResult<T> = Result<T, CoreError>;
120
121/// Generic trait that performs a more general version of TryInto.
122///
123/// Fallibly performs a pass turning T into a U, using self as an object with necessary context.
124/// Self is taken as a mutable reference, since the transformation might accumulate some information
125/// in self.
126pub trait TryPass<T, U>
127where
128    Self: Sized,
129    T: Sized,
130    U: Sized,
131{
132    fn try_pass(&mut self, t: T) -> CResult<U>;
133}
134
135#[derive(Debug, Clone)]
136pub struct Location {
137    pub src: Rc<String>,
138    pub loc: SrcLocation,
139}
140
141impl Location {
142    pub fn new(src: &Rc<String>, loc: SrcLocation) -> Self {
143        Self {
144            src: src.clone(),
145            loc,
146        }
147    }
148}
149
150/// Generic type to encapsulate syntax elements with a source code location.
151///
152/// Some elements always carry location information with them. These are newtyped as an alias for
153/// Located<T>, then by convention given the suffix "Obj".
154#[derive(Clone, Debug)]
155pub struct Located<T>(pub Location, pub T);
156
157/// Generic type for a string-indexed tree.
158#[derive(Clone, Debug)]
159pub enum Tree<T> {
160    Node(HashMap<String, Tree<T>>),
161    Leaf(T),
162}
163
164impl<T> Tree<T> {
165    /// Map every leaf element with a function, consuming this tree and returning a new one with
166    /// the same structure.
167    pub fn map<U, F>(self, f: F) -> Tree<U>
168    where
169        F: Fn(T) -> U,
170    {
171        self._map(Rc::new(f))
172    }
173
174    fn _map<U, F>(self, f: Rc<F>) -> Tree<U>
175    where
176        F: Fn(T) -> U,
177    {
178        match self {
179            Self::Node(node) => Tree::Node(HashMap::from_iter(
180                node.into_iter().map(|(k, v)| (k, v._map(f.clone()))),
181            )),
182            Self::Leaf(leaf) => Tree::Leaf((*f)(leaf)),
183        }
184    }
185
186    /// Like map, but gives you an argument containing the leaf's path.
187    pub fn map_with_path<U, F>(self, f: F) -> Tree<U>
188    where
189        F: Fn(T, &Vec<String>) -> U,
190    {
191        self._map_with_path(Rc::new(f), &mut vec![])
192    }
193
194    pub fn _map_with_path<U, F>(self, f: Rc<F>, path: &mut Vec<String>) -> Tree<U>
195    where
196        F: Fn(T, &Vec<String>) -> U,
197    {
198        match self {
199            Self::Node(node) => {
200                let mut node_ = HashMap::new();
201                for (k, v) in node.into_iter() {
202                    path.push(k.clone());
203                    node_.insert(k, v._map_with_path(f.clone(), path));
204                    path.pop();
205                }
206
207                Tree::Node(node_)
208            }
209            Self::Leaf(leaf) => Tree::Leaf((*f)(leaf, path)),
210        }
211    }
212
213    pub fn zip<U>(self, tree: Tree<U>) -> Tree<(T, U)> {
214        match (self, tree) {
215            (Self::Node(node), Tree::Node(mut node_)) => {
216                let mut zipped = HashMap::new();
217                for (key, val) in node.into_iter() {
218                    let val_ = node_.remove(&key).unwrap();
219                    zipped.insert(key, val.zip(val_));
220                }
221                Tree::Node(zipped)
222            }
223            (Self::Leaf(leaf), Tree::Leaf(leaf_)) => Tree::Leaf((leaf, leaf_)),
224            _ => panic!(),
225        }
226    }
227
228    /// Get a reference to a node in the tree at the given `path`.
229    pub fn get<'a>(&'a self, path: &Vec<String>) -> Option<&'a Self> {
230        let mut tree = self;
231
232        for part in path.iter() {
233            match tree {
234                Self::Node(node) => match node.get(part) {
235                    Some(tree_) => {
236                        tree = tree_;
237                    }
238                    None => {
239                        return None;
240                    }
241                },
242                _ => {
243                    return None;
244                }
245            }
246        }
247
248        return Some(tree);
249    }
250
251    /// Get a reference to a leaf in the tree at the given `path`.
252    pub fn get_leaf<'a>(&'a self, path: &Vec<String>) -> Option<&'a T> {
253        match self.get(path) {
254            Some(Self::Leaf(leaf)) => Some(leaf),
255            _ => None,
256        }
257    }
258
259    /// Get a mutable reference to a node in the tree at the given `path`.
260    pub fn get_mut<'a>(&'a mut self, path: &Vec<String>) -> Option<&'a mut Self> {
261        let mut tree = self;
262
263        for part in path.iter() {
264            match tree {
265                Self::Node(node) => match node.get_mut(part) {
266                    Some(tree_) => {
267                        tree = tree_;
268                    }
269                    _ => {
270                        return None;
271                    }
272                },
273                _ => {
274                    return None;
275                }
276            }
277        }
278
279        return Some(tree);
280    }
281
282    /// Insert a subtree into the tree at the given path.
283    ///
284    /// Returns whether the operation completed successfully - will fail if the insertion would have
285    /// to override a preexisting part of the tree.
286    pub fn insert(&mut self, mut path: Vec<String>, tree: Tree<T>, allow_overwrite: bool) -> bool {
287        let last = match path.pop() {
288            Some(last) => last,
289            _ => {
290                return false;
291            }
292        };
293
294        return self._insert(path.into_iter(), last, tree, allow_overwrite);
295    }
296
297    fn _insert<I: Iterator<Item = String>>(
298        &mut self,
299        mut path: I,
300        last: String,
301        tree: Self,
302        allow_overwrite: bool,
303    ) -> bool {
304        match self {
305            Self::Node(node) => match path.next() {
306                Some(part) => {
307                    if !node.contains_key(&part) {
308                        node.insert(part.clone(), Tree::Node(HashMap::new()));
309                    }
310                    let next = node.get_mut(&part).unwrap();
311
312                    next._insert(path, last, tree, allow_overwrite)
313                }
314                None => {
315                    if !allow_overwrite && node.contains_key(&last) {
316                        return false;
317                    }
318
319                    node.insert(last, tree);
320                    true
321                }
322            },
323            _ => false,
324        }
325    }
326
327    /// Return whether the tree contains a leaf at the given `path`.
328    pub fn contains_leaf(&self, path: &Vec<String>) -> bool {
329        match self.get(path) {
330            Some(Self::Leaf(..)) => true,
331            _ => false,
332        }
333    }
334
335    /// Remove the subtree at the given `path`. Returns the removed subtree, if one exists.
336    pub fn remove(&mut self, path: &Vec<String>) -> Option<Self> {
337        if path.len() == 0 {
338            return None;
339        }
340
341        let mut tree = self;
342        for part in path.iter().take(path.len() - 1) {
343            match tree {
344                Self::Leaf(..) => {
345                    return None;
346                }
347                Self::Node(node) => {
348                    if !node.contains_key(part) {
349                        return None;
350                    }
351
352                    tree = node.get_mut(part).unwrap();
353                }
354            }
355        }
356
357        return match tree {
358            Self::Leaf(..) => None,
359            Self::Node(node) => node.remove(path.last().unwrap()),
360        };
361    }
362
363    /// Return whether this tree is "dead" (has no leaves).
364    pub fn is_dead(&self) -> bool {
365        match self {
366            Self::Leaf(..) => false,
367            Self::Node(node) => {
368                if node.len() == 0 {
369                    true
370                } else {
371                    node.values().all(|tree| tree.is_dead())
372                }
373            }
374        }
375    }
376}
377
378impl<T, E> Tree<Result<T, E>> {
379    /// Turns a Tree<Result<T, E>> into a Result<Tree<T>, Vec<E>>.
380    ///
381    /// Only one error is returned. Multiple may have been generated, but whichever one is reached
382    /// first gets returned.
383    pub fn transpose(self) -> Result<Tree<T>, E> {
384        match self {
385            Self::Leaf(Ok(t)) => Ok(Tree::Leaf(t)),
386            Self::Leaf(Err(e)) => Err(e),
387            Self::Node(node) => {
388                let mut node_ = HashMap::new();
389                for (k, v) in node.into_iter() {
390                    let v = v.transpose()?;
391                    node_.insert(k, v);
392                }
393
394                Ok(Tree::Node(node_))
395            }
396        }
397    }
398}
399
400impl<T> Tree<Option<T>> {
401    /// Turn a tree of Option<T> into a tree of T by pruning `None` leaves.
402    pub fn prune(self) -> Tree<T> {
403        match self {
404            Self::Node(node) => {
405                let mut node_ = HashMap::new();
406                for (key, val) in node.into_iter() {
407                    match val {
408                        node @ Self::Node(..) => {
409                            node_.insert(key, node.prune());
410                        }
411                        Self::Leaf(Some(leaf)) => {
412                            node_.insert(key, Tree::Leaf(leaf));
413                        }
414                        _ => {}
415                    }
416                }
417
418                Tree::Node(node_)
419            }
420            _ => panic!("Cannot prune a leaf"),
421        }
422    }
423}
424
425impl<T> Tree<HashMap<String, T>> {
426    /// Get the "leaf of a leaf" from a tree.
427    pub fn get_leaf_ext<'a>(&'a self, path: &Vec<String>) -> Option<&'a T> {
428        let mut path = path.clone();
429        let name = path.pop()?;
430
431        return self.get_leaf(&path).and_then(|leaf| leaf.get(&name));
432    }
433}
434
435impl<T> Tree<BTreeMap<String, T>> {
436    /// Get the "leaf of a leaf" from a tree.
437    pub fn get_leaf_ext<'a>(&'a self, path: &Vec<String>) -> Option<&'a T> {
438        let mut path = path.clone();
439        let name = path.pop()?;
440
441        return self.get_leaf(&path).and_then(|leaf| leaf.get(&name));
442    }
443}
444
445/// Helper function to wait for 1 second, in case debug output is going by too quickly.
446pub fn wait() {
447    std::thread::sleep(std::time::Duration::from_secs(1));
448}