Board to Molecule

//! KiCad s-expression parser that keeps byte spans.
//!
//! Writers patch the original text by span so everything they do not touch stays
//! byte-identical; readers walk the tree for the preview model.

use anyhow::{bail, Result};

#[derive(Debug, Clone)]
pub enum Item {
    Atom(String),
    Str(String),
    List(Node),
}

#[derive(Debug, Clone)]
pub struct Node {
    pub items: Vec<Item>,
    /// Byte range of the whole list, from '(' through ')'.
    pub start: usize,
    pub end: usize,
}

pub fn parse(text: &str) -> Result<Node> {
    let b = text.as_bytes();
    let mut stack: Vec<Node> = vec![];
    let mut root: Option<Node> = None;
    let mut i = 0;
    while i < b.len() {
        match b[i] {
            b'(' => {
                stack.push(Node { items: vec![], start: i, end: 0 });
                i += 1;
            }
            b')' => {
                let Some(mut n) = stack.pop() else { bail!("unbalanced ')' at byte {i}") };
                n.end = i + 1;
                match stack.last_mut() {
                    Some(p) => p.items.push(Item::List(n)),
                    None => {
                        root = Some(n);
                        break;
                    }
                }
                i += 1;
            }
            b'"' => {
                let mut s = String::new();
                let mut j = i + 1;
                let mut seg = j;
                while j < b.len() && b[j] != b'"' {
                    if b[j] == b'\\' && j + 1 < b.len() {
                        s.push_str(&text[seg..j]);
                        s.push(b[j + 1] as char);
                        j += 2;
                        seg = j;
                    } else {
                        j += 1;
                    }
                }
                s.push_str(&text[seg..j.min(b.len())]);
                if let Some(p) = stack.last_mut() {
                    p.items.push(Item::Str(s));
                }
                i = j + 1;
            }
            c if c.is_ascii_whitespace() => i += 1,
            _ => {
                let j = b[i..]
                    .iter()
                    .position(|&c| c == b'(' || c == b')' || c == b'"' || c.is_ascii_whitespace())
                    .map_or(b.len(), |p| i + p);
                if let Some(p) = stack.last_mut() {
                    p.items.push(Item::Atom(text[i..j].to_string()));
                }
                i = j;
            }
        }
    }
    root.ok_or_else(|| anyhow::anyhow!("no s-expression found"))
}

impl Node {
    pub fn head(&self) -> &str {
        match self.items.first() {
            Some(Item::Atom(s)) => s,
            _ => "",
        }
    }

    pub fn lists(&self) -> impl Iterator<Item = &Node> {
        self.items.iter().filter_map(|i| if let Item::List(n) = i { Some(n) } else { None })
    }

    pub fn all<'a>(&'a self, name: &'a str) -> impl Iterator<Item = &'a Node> + 'a {
        self.lists().filter(move |n| n.head() == name)
    }

    pub fn get(&self, name: &str) -> Option<&Node> {
        self.lists().find(|n| n.head() == name)
    }

    /// The i-th value after the head, as text (atom or string).
    pub fn arg(&self, i: usize) -> Option<&str> {
        match self.items.get(i + 1)? {
            Item::Atom(s) | Item::Str(s) => Some(s),
            _ => None,
        }
    }

    pub fn nums(&self) -> Vec<f64> {
        self.items.iter().skip(1).filter_map(|i| if let Item::Atom(s) = i { s.parse().ok() } else { None }).collect()
    }

    /// (name x y ...) child as a point.
    pub fn xy(&self, name: &str) -> Option<(f64, f64)> {
        let v = self.get(name)?.nums();
        (v.len() >= 2).then(|| (v[0], v[1]))
    }

    pub fn num(&self, name: &str) -> Option<f64> {
        self.get(name)?.nums().first().copied()
    }

    pub fn layer(&self) -> &str {
        self.get("layer").and_then(|n| n.arg(0)).unwrap_or("")
    }

    pub fn layers(&self) -> Vec<&str> {
        match self.get("layers") {
            Some(n) => n.items.iter().skip(1).filter_map(|i| match i {
                Item::Atom(s) | Item::Str(s) => Some(s.as_str()),
                _ => None,
            }).collect(),
            None => vec![],
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn spans_and_strings() {
        let t = r#"(kicad_pcb (version 20241229) (gr_line (start 1 2) (end 3 4) (layer "Edge.Cuts")) (x "a \"q\" (b)"))"#;
        let n = parse(t).unwrap();
        assert_eq!(n.head(), "kicad_pcb");
        let g = n.get("gr_line").unwrap();
        assert_eq!(&t[g.start..g.end], r#"(gr_line (start 1 2) (end 3 4) (layer "Edge.Cuts"))"#);
        assert_eq!(g.layer(), "Edge.Cuts");
        assert_eq!(g.xy("end"), Some((3.0, 4.0)));
        assert_eq!(n.get("x").unwrap().arg(0), Some("a \"q\" (b)"));
    }
}