token.rs

//! LaTeX tokens
//!
//! Splits the source string into tokens. Only math, chemical structures, and verbatim text keep
//! their contents as written, because math is passed to suki and chemical structures are passed
//! to lina as unchanged strings.
//!
//! Positions are counted in characters. [`crate::error`] converts them to lines and columns.

use crate::error::Error;


/// Characters that end a word
const SPECIAL: &[char] = &['\\', '{', '}', '$', '&', '~', '%', '#', '[', ']'];


/// Environments that keep their contents as written
const RAW_ENVIRONMENT: &[&str] = &[
    "verbatim",
    "equation",
    "equation*",
    "displaymath",
    "align",
    "align*",
    "gather",
    "gather*",
];


/// Commands that keep their contents as written, and the number of `{}` groups each one takes
const RAW_COMMAND: &[(&str, usize)] = &[("chemfig", 1), ("chemname", 2)];


/// The kind of a token
#[derive(Debug, Clone, PartialEq)]
pub enum Kind {
    /// A sequence of characters with no whitespace and no special characters
    Word(String),
    /// A break between words
    Space,
    /// An empty line. A paragraph break
    Break,
    /// `\name`, or a single symbol that follows `\`
    Command(String),
    /// `{`
    Open,
    /// `}`
    Close,
    /// `&`. A boundary between columns of a table
    Cell,
    /// `~`. A non-breaking space
    HardSpace,
    /// An argument number written in a command definition, such as `#1`
    Parameter(usize),
    /// Math. The contents are kept as written
    Math {source: String, display: bool},
    /// A chemical structure. Kept as written, including the command
    Chem(String),
    /// A string that is output unchanged. block says whether it is the `verbatim` environment
    Verbatim {text: String, block: bool},
    /// An environment that keeps its contents as written
    Environment {name: String, body: String},
}


/// One token and its position in the source string
#[derive(Debug, Clone, PartialEq)]
pub struct Token {
    pub kind: Kind,
    pub position: usize,
}


/// Splits the source string into tokens
pub fn scan(source: &[char]) -> Result<Vec<Token>, Error> {
    let mut found = Vec::new();
    let mut at = 0;

    while at < source.len() {
        at = match source[at] {
            '%' => skip_comment(source, at, &mut found),
            '\\' => read_escape(source, at, &mut found)?,
            '{' => push_one(&mut found, Kind::Open, at),
            '}' => push_one(&mut found, Kind::Close, at),
            '&' => push_one(&mut found, Kind::Cell, at),
            '~' => push_one(&mut found, Kind::HardSpace, at),
            '[' | ']' => push_one(&mut found, Kind::Word(source[at].to_string()), at),
            '$' => read_dollar(source, at, &mut found)?,
            '#' => read_parameter(source, at, &mut found)?,
            character if character.is_whitespace() => read_space(source, at, &mut found),
            _ => read_word(source, at, &mut found),
        };
    }

    Ok(found)
}


/// Adds a token for one character and returns the next position
fn push_one(found: &mut Vec<Token>, kind: Kind, at: usize) -> usize {
    found.push(Token {kind, position: at});

    at + 1
}


/// Reads a sequence of characters that are neither whitespace nor special characters
fn read_word(source: &[char], at: usize, found: &mut Vec<Token>) -> usize {
    let mut end = at;

    while end < source.len() {
        let character = source[end];

        if character.is_whitespace() || SPECIAL.contains(&character) {
            break;
        }

        end += 1;
    }

    found.push(Token {kind: Kind::Word(source[at..end].iter().collect()), position: at});

    end
}


/// Reads consecutive whitespace. Two or more newlines make a paragraph break
fn read_space(source: &[char], at: usize, found: &mut Vec<Token>) -> usize {
    let mut end = at;
    let mut newline = 0;

    while end < source.len() && source[end].is_whitespace() {
        if source[end] == '\n' {
            newline += 1;
        }

        end += 1;
    }

    let kind = if newline >= 2 {Kind::Break} else {Kind::Space};
    found.push(Token {kind, position: at});

    end
}


/// Skips a comment
///
/// A comment removes the rest of the line together with the newline. The whitespace at the
/// start of the next line is removed too.
/// If an empty line follows the comment, it is kept as a paragraph break
fn skip_comment(source: &[char], at: usize, found: &mut Vec<Token>) -> usize {
    let mut end = at;

    while end < source.len() && source[end] != '\n' {
        end += 1;
    }

    if end < source.len() {
        end += 1;
    }

    let mut after = end;
    let mut newline = 0;

    while after < source.len() && source[after].is_whitespace() {
        if source[after] == '\n' {
            newline += 1;
        }

        after += 1;
    }

    if newline >= 1 {
        found.push(Token {kind: Kind::Break, position: at});
    }

    after
}


/// Reads math that starts with `$`
fn read_dollar(source: &[char], at: usize, found: &mut Vec<Token>) -> Result<usize, Error> {
    let display = source.get(at + 1) == Some(&'$');
    let open = if display {at + 2} else {at + 1};
    let close: &[char] = if display {&['$', '$']} else {&['$']};
    let (body, end) = read_until(source, open, close)?;
    found.push(Token {kind: Kind::Math {source: body, display}, position: at});

    Ok(end)
}


/// Reads `#1`
fn read_parameter(source: &[char], at: usize, found: &mut Vec<Token>) -> Result<usize, Error> {
    match source.get(at + 1).and_then(|character| character.to_digit(10)) {
        Some(number) if number >= 1 => {
            found.push(Token {kind: Kind::Parameter(number as usize), position: at});

            Ok(at + 2)
        }
        _ => Err(Error::new("# のあとには 1 から 9 の数字が要ります", at)),
    }
}


/// Reads something that starts with `\`
fn read_escape(source: &[char], at: usize, found: &mut Vec<Token>) -> Result<usize, Error> {
    let next = match source.get(at + 1) {
        Some(character) => *character,
        None => return Err(Error::new("文書が \\ で終わっています", at)),
    };

    if next == '(' || next == '[' {
        let display = next == '[';
        let close: &[char] = if display {&['\\', ']']} else {&['\\', ')']};
        let (body, end) = read_until(source, at + 2, close)?;
        found.push(Token {kind: Kind::Math {source: body, display}, position: at});

        return Ok(end);
    }

    if !next.is_alphabetic() {
        found.push(Token {kind: Kind::Command(next.to_string()), position: at});

        return Ok(at + 2);
    }

    let (name, end) = read_name(source, at + 1);

    if name == "verb" || name == "verb*" {
        return read_verb(source, at, end, found);
    }

    if let Some((_, count)) = RAW_COMMAND.iter().find(|(known, _)| *known == name) {
        return read_chem(source, at, end, *count, found);
    }

    if name == "begin" {
        if let Some(after) = read_raw_environment(source, at, end, found)? {
            return Ok(after);
        }
    }

    found.push(Token {kind: Kind::Command(name), position: at});

    Ok(skip_after_command(source, end))
}


/// Reads the name of a command. A `*` at the end is included in the name
fn read_name(source: &[char], at: usize) -> (String, usize) {
    let mut end = at;

    while end < source.len() && source[end].is_alphabetic() {
        end += 1;
    }

    if end < source.len() && source[end] == '*' {
        end += 1;
    }

    (source[at..end].iter().collect(), end)
}


/// Skips whitespace after a command name. An empty line is a paragraph break, so it is kept
fn skip_after_command(source: &[char], at: usize) -> usize {
    let mut end = at;

    while end < source.len() && (source[end] == ' ' || source[end] == '\t') {
        end += 1;
    }

    if source.get(end) != Some(&'\n') {
        return end;
    }

    let mut after = end + 1;
    let mut newline = 0;

    while after < source.len() && source[after].is_whitespace() {
        if source[after] == '\n' {
            newline += 1;
        }

        after += 1;
    }

    if newline == 0 {
        return after;
    }

    end
}


/// Skips whitespace
fn skip_blank(source: &[char], at: usize) -> usize {
    let mut end = at;

    while end < source.len() && source[end].is_whitespace() {
        end += 1;
    }

    end
}


/// Reads `\verb|…|`
fn read_verb(
    source: &[char],
    at: usize,
    after_name: usize,
    found: &mut Vec<Token>,
) -> Result<usize, Error> {
    let mark = match source.get(after_name) {
        Some(character) => *character,
        None => return Err(Error::new("\\verb の区切りの文字がありません", at)),
    };
    let mut end = after_name + 1;

    while end < source.len() && source[end] != mark {
        end += 1;
    }

    if end >= source.len() {
        return Err(Error::new("\\verb が閉じていません", at));
    }

    let text: String = source[after_name + 1..end].iter().collect();
    found.push(Token {kind: Kind::Verbatim {text, block: false}, position: at});

    Ok(end + 1)
}


/// Reads `\chemfig{…}` or `\chemname{…}{…}` as written, including the command
fn read_chem(
    source: &[char],
    at: usize,
    after_name: usize,
    count: usize,
    found: &mut Vec<Token>,
) -> Result<usize, Error> {
    let mut end = after_name;

    for _ in 0..count {
        end = skip_blank(source, end);

        if source.get(end) != Some(&'{') {
            return Err(Error::new("化学式に { がありません", end));
        }

        end = group_end(source, end)?;
    }

    let text: String = source[at..end].iter().collect();
    found.push(Token {kind: Kind::Chem(text), position: at});

    Ok(end)
}


/// If `\begin{…}` starts an environment that keeps its contents as written, reads the whole
/// environment. Otherwise does nothing
fn read_raw_environment(
    source: &[char],
    at: usize,
    after_name: usize,
    found: &mut Vec<Token>,
) -> Result<Option<usize>, Error> {
    let open = skip_blank(source, after_name);

    if source.get(open) != Some(&'{') {
        return Ok(None);
    }

    let close = group_end(source, open)?;
    let name: String = source[open + 1..close - 1].iter().collect();

    if !RAW_ENVIRONMENT.contains(&name.as_str()) {
        return Ok(None);
    }

    let needle: Vec<char> = format!("\\end{{{name}}}").chars().collect();
    let stop = match find_sequence(source, close, &needle) {
        Some(stop) => stop,
        None => return Err(Error::new(format!("\\end{{{name}}} がありません"), at)),
    };
    let body = trim_body(&source[close..stop]);
    let kind = if name == "verbatim" {
        Kind::Verbatim {text: body, block: true}
    } else {
        Kind::Environment {name, body}
    };
    found.push(Token {kind, position: at});

    Ok(Some(stop + needle.len()))
}


/// Drops the newlines on the opening and closing lines from the contents of an environment
fn trim_body(body: &[char]) -> String {
    let mut start = 0;
    let mut end = body.len();

    while start < end && (body[start] == ' ' || body[start] == '\t' || body[start] == '\r') {
        start += 1;
    }

    if start < end && body[start] == '\n' {
        start += 1;
    }

    while end > start && (body[end - 1] == ' ' || body[end - 1] == '\t') {
        end -= 1;
    }

    if end > start && body[end - 1] == '\n' {
        end -= 1;
    }

    if end > start && body[end - 1] == '\r' {
        end -= 1;
    }

    body[start..end].iter().collect()
}


/// Given the position of a `{`, returns the position after the matching `}`
fn group_end(source: &[char], at: usize) -> Result<usize, Error> {
    let mut depth = 0;
    let mut end = at;

    while end < source.len() {
        match source[end] {
            '\\' => end += 1,
            '{' => depth += 1,
            '}' => {
                depth -= 1;

                if depth == 0 {
                    return Ok(end + 1);
                }
            }
            _ => {}
        }

        end += 1;
    }

    Err(Error::new("{ が閉じていません", at))
}


/// Reads until the closing sequence appears. The character after a `\` is skipped
fn read_until(source: &[char], from: usize, needle: &[char]) -> Result<(String, usize), Error> {
    let mut at = from;

    while at < source.len() {
        if at + needle.len() <= source.len() && source[at..at + needle.len()] == *needle {
            return Ok((source[from..at].iter().collect(), at + needle.len()));
        }

        at += if source[at] == '\\' {2} else {1};
    }

    let text: String = needle.iter().collect();

    Err(Error::new(format!("{text} が来ないまま文書が終わりました"), from))
}


/// Finds the first position where the sequence appears. The character after `\` is not skipped
fn find_sequence(source: &[char], from: usize, needle: &[char]) -> Option<usize> {
    if needle.is_empty() || source.len() < needle.len() {
        return None;
    }

    for at in from..=source.len() - needle.len() {
        if source[at..at + needle.len()] == *needle {
            return Some(at);
        }
    }

    None
}


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

    fn scan_text(source: &str) -> Vec<Kind> {
        let characters: Vec<char> = source.chars().collect();

        scan(&characters).unwrap().into_iter().map(|token| token.kind).collect()
    }

    #[test]
    fn words_and_spaces_come_apart() {
        assert_eq!(
            scan_text("ab cd"),
            vec![Kind::Word("ab".into()), Kind::Space, Kind::Word("cd".into())]
        );
    }

    #[test]
    fn an_empty_line_is_a_paragraph_break() {
        assert_eq!(
            scan_text("a\n\nb"),
            vec![Kind::Word("a".into()), Kind::Break, Kind::Word("b".into())]
        );
    }

    #[test]
    fn one_newline_is_only_a_space() {
        assert_eq!(
            scan_text("a\nb"),
            vec![Kind::Word("a".into()), Kind::Space, Kind::Word("b".into())]
        );
    }

    #[test]
    fn a_command_swallows_the_space_after_it() {
        assert_eq!(
            scan_text("\\LaTeX  is"),
            vec![Kind::Command("LaTeX".into()), Kind::Word("is".into())]
        );
    }

    #[test]
    fn a_starred_command_keeps_its_star() {
        assert_eq!(scan_text("\\section*"), vec![Kind::Command("section*".into())]);
    }

    #[test]
    fn a_symbol_after_a_backslash_is_its_own_command() {
        assert_eq!(
            scan_text("\\% \\{"),
            vec![Kind::Command("%".into()), Kind::Space, Kind::Command("{".into())]
        );
    }

    #[test]
    fn a_comment_reaches_to_the_end_of_the_line() {
        assert_eq!(
            scan_text("a% ここは消える\nb"),
            vec![Kind::Word("a".into()), Kind::Word("b".into())]
        );
    }

    #[test]
    fn a_comment_before_an_empty_line_still_breaks_the_paragraph() {
        assert_eq!(
            scan_text("a% 注\n\nb"),
            vec![Kind::Word("a".into()), Kind::Break, Kind::Word("b".into())]
        );
    }

    #[test]
    fn math_keeps_its_source() {
        assert_eq!(
            scan_text("$x^2$"),
            vec![Kind::Math {source: "x^2".into(), display: false}]
        );
        assert_eq!(
            scan_text("$$x^2$$"),
            vec![Kind::Math {source: "x^2".into(), display: true}]
        );
        assert_eq!(
            scan_text("\\[x\\]"),
            vec![Kind::Math {source: "x".into(), display: true}]
        );
    }

    #[test]
    fn a_row_separator_inside_math_is_not_a_closing_mark() {
        assert_eq!(
            scan_text("\\[a \\\\ b\\]"),
            vec![Kind::Math {source: "a \\\\ b".into(), display: true}]
        );
    }

    #[test]
    fn a_raw_environment_keeps_its_body() {
        assert_eq!(
            scan_text("\\begin{equation}\nx = 1\n\\end{equation}"),
            vec![Kind::Environment {name: "equation".into(), body: "x = 1".into()}]
        );
    }

    #[test]
    fn verbatim_keeps_everything() {
        assert_eq!(
            scan_text("\\begin{verbatim}\n\\foo{}\n\\end{verbatim}"),
            vec![Kind::Verbatim {text: "\\foo{}".into(), block: true}]
        );
        assert_eq!(
            scan_text("\\verb|a\\b|"),
            vec![Kind::Verbatim {text: "a\\b".into(), block: false}]
        );
    }

    #[test]
    fn an_ordinary_environment_is_left_to_the_parser() {
        assert_eq!(
            scan_text("\\begin{itemize}"),
            vec![
                Kind::Command("begin".into()),
                Kind::Open,
                Kind::Word("itemize".into()),
                Kind::Close,
            ]
        );
    }

    #[test]
    fn a_molecule_is_kept_with_its_command() {
        assert_eq!(
            scan_text("\\chemfig{H-O-H}"),
            vec![Kind::Chem("\\chemfig{H-O-H}".into())]
        );
        assert_eq!(
            scan_text("\\chemname{\\chemfig{O=C=O}}{二酸化炭素}"),
            vec![Kind::Chem("\\chemname{\\chemfig{O=C=O}}{二酸化炭素}".into())]
        );
    }

    #[test]
    fn a_parameter_carries_its_number() {
        assert_eq!(scan_text("#1"), vec![Kind::Parameter(1)]);
    }

    #[test]
    fn an_unclosed_group_is_an_error() {
        let source: Vec<char> = "\\chemfig{H-O".chars().collect();
        assert!(scan(&source).is_err());
    }
}