lib.rs

use anyhow::{anyhow, Context, Error};
use std::fmt;

mod ffi;

#[derive(Clone, Debug)]
pub enum Sexp {
    Atom(String),
    List(Vec<Sexp>),
    Nil,
}

impl Sexp {
    pub fn of_str(input: &str) -> Result<Sexp, Error> {
        let mut parser = ffi::parser();
        let tree = parser
            .parse(input, None)
            .context("Could not parse anything")?;
        let root = tree.root_node();

        let mut walker = root.walk();
        walker.goto_first_child(); // we skip the top-level `sexp` node
        Sexp::build_tree(walker.node(), input.as_bytes())
    }

    fn build_tree(root: tree_sitter::Node, bytes: &[u8]) -> Result<Sexp, Error> {
        match root.kind() {
            "atom" => {
                let text = root.utf8_text(&bytes)?.to_string();
                Ok(Sexp::Atom(text))
            }
            "list" => {
                let mut walker = root.walk();
                walker.goto_first_child();
                let mut children = vec![];
                while walker.goto_next_sibling() {
                    let child = walker.node();
                    children.push(Sexp::build_tree(child, bytes)?);
                }
                Ok(Sexp::List(children))
            }
            ")" => Ok(Sexp::Nil),
            kind => Err(anyhow!("Unknown node kind {:?}", kind)),
        }
    }

    pub fn size(&self) -> usize {
        match self {
            Sexp::Nil => 0,
            Sexp::Atom(s) => s.len(),
            Sexp::List(parts) => {
                let mut s = 0;
                for p in parts {
                    s += p.size() + 1;
                }
                s
            }
        }
    }
}

#[derive(Debug)]
pub struct PrettyPrinter {
    max_width: u8,
    current_width: u8,
    padding: u8,
    indent_size: u8,
    is_root: bool,
}

impl PrettyPrinter {
    pub fn new() -> PrettyPrinter {
        PrettyPrinter {
            is_root: true,
            max_width: 40,
            current_width: 0,
            padding: 0,
            indent_size: 2,
        }
    }

    pub fn pp(&mut self, sexp: &Sexp, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match sexp {
            Sexp::Atom(atom) => {
                if self.current_width + (atom.len() as u8) > self.max_width {
                    self.current_width = 0;
                    self.padding += self.indent_size;
                    write!(fmt, "\n")?;
                    for _ in 0..self.padding {
                        write!(fmt, " ")?
                    }
                } else {
                    self.current_width += atom.len() as u8;
                }
                write!(fmt, "{}", atom)
            }
            Sexp::Nil => {
                self.padding = 0;
                Ok(())
            }
            Sexp::List(parts) if parts.len() > 0 => {
                let next_term_width = parts[0].size() as u8;
                let term_overflows = next_term_width > self.max_width;

                if term_overflows {
                    self.current_width = self.padding;
                    write!(fmt, "\n")?;
                    for _ in 0..self.padding {
                        write!(fmt, " ")?
                    }
                }

                write!(fmt, "(")?;
                self.pp(&parts[0], fmt)?;

                if self.is_root {
                    write!(fmt, "\n")?;
                    for _ in 0..self.padding {
                        write!(fmt, " ")?
                    }
                    self.is_root = false;
                }

                for p in parts[1..].iter() {
                    match p {
                        Sexp::Nil => self.pp(&p, fmt)?,
                        _ => {
                            write!(fmt, " ")?;
                            self.padding += self.indent_size;
                            self.pp(&p, fmt)?;
                            if self.padding < self.indent_size {
                                self.padding = 0;
                            } else {
                                self.padding -= self.indent_size;
                            }
                        }
                    }
                }
                write!(fmt, ")")
            }
            Sexp::List(_) => write!(fmt, "()"),
        }
    }
}

impl fmt::Display for Sexp {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        PrettyPrinter::new().pp(self, fmt)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_single_expr() {
        assert_eq!(
            Sexp::of_str(&"source_file").unwrap().to_string(),
            r#"source_file"#.trim().to_string()
        );
    }

    #[test]
    fn test_sibling_expr() {
        assert_eq!(
            Sexp::of_str(&"(source file)").unwrap().to_string(),
            r#"(source file)"#.trim().to_string()
        );
        assert_eq!(
            Sexp::of_str(&"(source file tree)").unwrap().to_string(),
            r#"(source file tree)"#.trim().to_string()
        );
    }

    #[test]
    fn test_nested_expr() {
        assert_eq!(
            Sexp::of_str(&"(source (file))").unwrap().to_string(),
            r#"(source (file))"#.trim().to_string()
        );
    }

    #[test]
    fn test_nested_sibling_expr() {
        assert_eq!(
            Sexp::of_str(&"(source (file tree))").unwrap().to_string(),
            r#"(source (file tree))"#.trim().to_string()
        );
    }

    #[test]
    fn test_field_expr() {
        assert_eq!(
            Sexp::of_str(&"(source file: test)").unwrap().to_string(),
            r#"(source file: test)"#.trim().to_string()
        );
    }

    #[test]
    fn test_real_life_expr() {
        println!(
            "{}",
            Sexp::of_str(include_str!("./big_fixture.sexp")).unwrap()
        );
        assert_eq!(
            Sexp::of_str(include_str!("./big_fixture.sexp"))
                .unwrap()
                .to_string(),
            include_str!("./big_fixture.sexp").trim().to_string()
        );
    }

    #[test]
    fn test_pretty_printing_expr() {
        let sexp = Sexp::of_str(
            &"(source_file
  (expression
    (function_call
      (qualified_function_name
        (expression
          (term
            (atom
              (unquoted_atom))))
        (atom
          (unquoted_atom)))
      (expression
        (term
          (integer)))
      (expression
        (function_call
          (qualified_function_name
            (expression
              (term
                (atom
                  (unquoted_atom))))
            (atom
              (unquoted_atom)))
          (expression
            (term
              (integer)))
          (expression
            (term
              (integer)))
          (expression
            (term
              (integer))))))))
",
        )
        .unwrap();
        println!("{}", sexp);
        assert_eq!(
            sexp.to_string(),
            r#"(source_file
  (expression
    (function_call
      (qualified_function_name
        (expression (term (atom (unquoted_atom))))
        (atom (unquoted_atom)))
      (expression (term (integer)))
      (expression
        (function_call
          (qualified_function_name
            (expression (term (atom (unquoted_atom))))
            (atom (unquoted_atom)))
          (expression (term (integer)))
          (expression (term (integer)))
          (expression (term (integer))))))))"#
                .to_string()
        );
    }
}