AI Flow: Reverse Engineering

//! A small S-expression tree for KiCad files: parse a footprint or symbol, change it, write it back.
#[derive(Clone, Debug, PartialEq)]
pub enum N { A(String), L(Vec<N>) }

pub fn parse(s: &str) -> Result<N, String> {
    let b = s.as_bytes();
    let mut i = 0;
    fn ws(b: &[u8], i: &mut usize) { while *i < b.len() && (b[*i] as char).is_whitespace() { *i += 1; } }
    fn node(b: &[u8], i: &mut usize) -> Result<N, String> {
        ws(b, i);
        if *i >= b.len() { return Err("unexpected end".into()); }
        if b[*i] == b'(' {
            *i += 1; let mut v = vec![];
            loop { ws(b, i); if *i >= b.len() { return Err("unclosed list".into()); } if b[*i] == b')' { *i += 1; return Ok(N::L(v)); } v.push(node(b, i)?); }
        }
        if b[*i] == b'"' {
            let st = *i; *i += 1;
            while *i < b.len() && b[*i] != b'"' { if b[*i] == b'\\' { *i += 1; } *i += 1; }
            *i += 1; return Ok(N::A(String::from_utf8_lossy(&b[st..*i]).to_string()));
        }
        let st = *i;
        while *i < b.len() && !(b[*i] as char).is_whitespace() && b[*i] != b'(' && b[*i] != b')' { *i += 1; }
        Ok(N::A(String::from_utf8_lossy(&b[st..*i]).to_string()))
    }
    node(b, &mut i)
}

impl N {
    pub fn head(&self) -> &str { match self { N::L(v) => match v.first() { Some(N::A(s)) => s, _ => "" }, _ => "" } }
    pub fn items(&self) -> &[N] { match self { N::L(v) => v, _ => &[] } }
    pub fn items_mut(&mut self) -> &mut Vec<N> { match self { N::L(v) => v, _ => panic!("not a list") } }
    pub fn child(&self, h: &str) -> Option<&N> { self.items().iter().find(|n| n.head() == h) }
    pub fn child_mut(&mut self, h: &str) -> Option<&mut N> { self.items_mut().iter_mut().find(|n| n.head() == h) }
    pub fn children<'a>(&'a self, h: &'a str) -> impl Iterator<Item = &'a N> + 'a { self.items().iter().filter(move |n| n.head() == h) }
    /// The i-th atom after the head, unquoted.
    pub fn arg(&self, i: usize) -> Option<String> { match self.items().get(i + 1) { Some(N::A(s)) => Some(unq(s)), _ => None } }
    pub fn num(&self, i: usize) -> Option<f64> { self.arg(i).and_then(|s| s.parse().ok()) }
    pub fn write(&self, out: &mut String, depth: usize) {
        match self {
            N::A(s) => out.push_str(s),
            N::L(v) => {
                let nested = v.iter().any(|n| matches!(n, N::L(_)));
                out.push('(');
                for (k, n) in v.iter().enumerate() {
                    if k > 0 { if nested && matches!(n, N::L(_)) { out.push('\n'); for _ in 0..=depth { out.push('\t'); } } else { out.push(' '); } }
                    n.write(out, depth + 1);
                }
                out.push(')');
            }
        }
    }
    pub fn to_string(&self) -> String { let mut s = String::new(); self.write(&mut s, 0); s }
}

pub fn unq(s: &str) -> String { if s.starts_with('"') && s.ends_with('"') && s.len() >= 2 { s[1..s.len() - 1].replace("\\\"", "\"") } else { s.to_string() } }
pub fn q(s: &str) -> N { N::A(format!("\"{}\"", s.replace('"', "\\\""))) }
pub fn a(s: impl ToString) -> N { N::A(s.to_string()) }
pub fn l(v: Vec<N>) -> N { N::L(v) }
pub fn f(x: f64) -> N { let s = format!("{:.4}", x); let s = s.trim_end_matches('0').trim_end_matches('.').to_string(); a(if s == "-0" { "0".into() } else { s }) }

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn round_trip() {
        let t = r#"(footprint "R_0603" (layer "F.Cu") (pad "1" smd roundrect (at -0.825 0) (size 0.8 0.95) (layers "F.Cu" "F.Mask")))"#;
        let n = parse(t).unwrap();
        assert_eq!(n.head(), "footprint");
        assert_eq!(n.arg(0).unwrap(), "R_0603");
        let p = n.child("pad").unwrap();
        assert_eq!(p.child("at").unwrap().num(0), Some(-0.825));
        assert_eq!(parse(&n.to_string()).unwrap(), n);
    }
}