koi_parser/

token_tree.rs

use std::{fmt::Display, ops::Range};

use crate::lexer::{Token, lexer};

/// A token tree node - either a single token or a delimited group
#[derive(Debug, PartialEq)]
pub enum TokenTree {
    /// A single token
    Token(TokenTreeToken),
    /// A delimited group of token trees
    Block(TokenTreeBlock),
}

#[derive(Debug, PartialEq)]
pub struct TokenTreeToken {
    pub token: Token,
    pub span: Range<usize>,
}

#[derive(Debug, PartialEq)]
pub struct TokenTreeBlock {
    pub delimiter: Block,
    pub open_span: Range<usize>,
    pub trees: Vec<TokenTree>,
    pub close_span: Range<usize>,
}

impl TokenTreeBlock {
    fn span(&self) -> Range<usize> {
        Range {
            start: self.open_span.start,
            end: self.open_span.end,
        }
    }
}

impl TokenTree {
    pub fn token(&self) -> Option<&TokenTreeToken> {
        match self {
            TokenTree::Token(token) => Some(token),
            TokenTree::Block(..) => None,
        }
    }

    pub fn block(&self) -> Option<&TokenTreeBlock> {
        match self {
            TokenTree::Token(..) => None,
            TokenTree::Block(block) => Some(block),
        }
    }

    pub fn paren_block(&self) -> Option<&TokenTreeBlock> {
        match self {
            TokenTree::Token(..) => None,
            TokenTree::Block(block) => match block.delimiter {
                Block::Parenthesis => Some(block),
                Block::Brace => None,
                Block::Bracket => None,
            },
        }
    }

    pub fn brace_block(&self) -> Option<&TokenTreeBlock> {
        match self {
            TokenTree::Token(..) => None,
            TokenTree::Block(block) => match block.delimiter {
                Block::Parenthesis => None,
                Block::Brace => Some(block),
                Block::Bracket => None,
            },
        }
    }

    pub fn bracket_block(&self) -> Option<&TokenTreeBlock> {
        match self {
            TokenTree::Token(..) => None,
            TokenTree::Block(block) => match block.delimiter {
                Block::Parenthesis => None,
                Block::Brace => None,
                Block::Bracket => Some(block),
            },
        }
    }

    pub fn block_of(&self, delimiter: Block) -> Option<&TokenTreeBlock> {
        match self {
            TokenTree::Token(..) => None,
            TokenTree::Block(block) => {
                if block.delimiter == delimiter {
                    return Some(block);
                } else {
                    return None;
                }
            }
        }
    }

    pub fn span(&self) -> Range<usize> {
        match self {
            TokenTree::Token(token_tree_token) => token_tree_token.span.clone(),
            TokenTree::Block(token_tree_block) => token_tree_block.span(),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Block {
    Parenthesis, // ( )
    Brace,       // { }
    Bracket,     // [ ]
}

impl Display for Block {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Block::Parenthesis => write!(f, "( )"),
            Block::Brace => write!(f, "{{ }}"),
            Block::Bracket => write!(f, "[ ]"),
        }
    }
}

impl Block {
    /// This delimiter as a opening token
    pub fn opening_token(&self) -> Token {
        match self {
            Block::Parenthesis => Token::LParen,
            Block::Brace => Token::LBrace,
            Block::Bracket => Token::LBracket,
        }
    }

    /// This delimiter as a closing token
    pub fn closing_token(&self) -> Token {
        match self {
            Block::Parenthesis => Token::RParen,
            Block::Brace => Token::RBrace,
            Block::Bracket => Token::RBracket,
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum ParseError {
    /// EOF reached before closing delimiter
    UnexpectedEOF {
        expected: Block,
        expected_span: Range<usize>,
    },
    /// A closing delimiter found before any opening delimiter
    UnexpectedClosing {
        token: Block,
        span: Range<usize>,
    },
    /// Currently looking for closing delimiter but found another closing or opening delimiter
    MismatchedDelimiter {
        expected: Block,
        expected_span: Range<usize>,
        found: Block,
        found_span: Range<usize>,
    },
}

pub struct TokenTreeParser {
    tokens: Vec<(Token, Range<usize>)>,
    pos: usize,
}

impl TokenTreeParser {
    pub fn new(tokens: Vec<(Token, Range<usize>)>) -> Self {
        Self { tokens, pos: 0 }
    }

    /// Parse the entire token stream into a flat list of token trees
    pub fn parse(&mut self) -> Result<Vec<TokenTree>, ParseError> {
        let mut trees = Vec::new();

        while !self.is_at_end() {
            trees.push(self.parse_tree()?);
        }

        Ok(trees)
    }

    /// Parse a single token tree
    fn parse_tree(&mut self) -> Result<TokenTree, ParseError> {
        let (token, span) = self.current();

        match token {
            Token::LParen => self.parse_delimited(Block::Parenthesis),
            Token::LBrace => self.parse_delimited(Block::Brace),
            Token::RParen => Err(ParseError::UnexpectedClosing {
                token: Block::Parenthesis,
                span: span.clone(),
            }),
            Token::RBrace => Err(ParseError::UnexpectedClosing {
                token: Block::Brace,
                span: span.clone(),
            }),
            Token::RBracket => Err(ParseError::UnexpectedClosing {
                token: Block::Bracket,
                span: span.clone(),
            }),
            Token::EOF => unreachable!(),
            _ => {
                let tree = TokenTree::Token(TokenTreeToken {
                    token: token.clone(),
                    span: span.clone(),
                });
                self.advance();
                Ok(tree)
            }
        }
    }

    /// Parse a delimited group
    fn parse_delimited(&mut self, delimiter: Block) -> Result<TokenTree, ParseError> {
        let open_span = self.current().1.clone();
        self.advance(); // consume opening delimiter
        let closing_token = delimiter.closing_token();
        let mut trees = Vec::new();

        loop {
            let (token, span) = self.current();
            if token == &closing_token {
                let close_span = span.clone();
                self.advance(); // consume closing delimiter
                return Ok(TokenTree::Block(TokenTreeBlock {
                    delimiter,
                    open_span,
                    trees,
                    close_span,
                }));
            }

            match token {
                Token::EOF => {
                    return Err(ParseError::UnexpectedEOF {
                        expected: delimiter,
                        expected_span: open_span,
                    });
                }
                Token::RParen => {
                    return Err(ParseError::MismatchedDelimiter {
                        expected: delimiter,
                        expected_span: open_span,
                        found: Block::Parenthesis,
                        found_span: span.clone(),
                    });
                }
                Token::RBrace => {
                    return Err(ParseError::MismatchedDelimiter {
                        expected: delimiter,
                        expected_span: open_span,
                        found: Block::Brace,
                        found_span: span.clone(),
                    });
                }
                Token::RBracket => {
                    return Err(ParseError::MismatchedDelimiter {
                        expected: delimiter,
                        expected_span: open_span,
                        found: Block::Bracket,
                        found_span: span.clone(),
                    });
                }
                _ => {
                    trees.push(self.parse_tree()?);
                }
            }
        }
    }

    fn current(&self) -> &(Token, Range<usize>) {
        &self.tokens[self.pos]
    }

    fn advance(&mut self) {
        if self.pos < self.tokens.len() - 1 {
            self.pos += 1;
        }
    }

    fn is_at_end(&self) -> bool {
        matches!(self.current().0, Token::EOF)
    }

    pub fn token_trees_to_string(trees: &[TokenTree]) -> String {
        let mut result = String::new();
        write_token_trees_internal(&mut result, &trees.iter().collect::<Vec<&TokenTree>>(), 0);
        result
    }
}

impl Display for TokenTree {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut string = String::new();
        write_token_trees_internal(&mut string, &[self], 0);
        write!(f, "{}", string)
    }
}
fn write_token_trees_internal(string: &mut String, trees: &[&TokenTree], indent: usize) {
    for tree in trees {
        match tree {
            TokenTree::Token(TokenTreeToken { token, span }) => {
                string.push_str(&format!(
                    "{:indent$}{:?} @ {:?}\n",
                    "",
                    token,
                    span,
                    indent = indent
                ));
            }
            TokenTree::Block(TokenTreeBlock {
                delimiter,
                open_span,
                trees,
                close_span,
            }) => {
                string.push_str(&format!(
                    "{:indent$}{:?} {{ @ {:?}\n",
                    "",
                    delimiter,
                    open_span,
                    indent = indent
                ));
                let tree_refs: Vec<&TokenTree> = trees.iter().collect();
                write_token_trees_internal(string, &tree_refs, indent + 2);
                string.push_str(&format!(
                    "{:indent$}}} @ {:?}\n",
                    "",
                    close_span,
                    indent = indent
                ));
            }
        };
    }
}

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

    fn parse_token_trees(source: &str) -> Result<Vec<TokenTree>, ParseError> {
        let mut parser = TokenTreeParser::new(lexer(source).collect());
        parser.parse()
    }

    #[test]
    fn test_simple_tokens() {
        let source = "let x = 5;";
        let trees = parse_token_trees(source).unwrap();
        assert_eq!(trees.len(), 5); // let, x, =, 5, ;
    }

    #[test]
    fn test_parentheses() {
        let source = "fn add(x, y)";
        let trees = parse_token_trees(source).unwrap();
        assert_eq!(trees.len(), 3); // fn, add, (x, y)

        match &trees[2] {
            TokenTree::Block(TokenTreeBlock {
                delimiter, trees, ..
            }) => {
                assert_eq!(*delimiter, Block::Parenthesis);
                assert_eq!(trees.len(), 3); // x, comma, y
            }
            _ => panic!("Expected delimited group"),
        }
    }

    #[test]
    fn test_nested_delimiters() {
        let source = "fn foo() { let x = (1 + 2); }";
        let trees = parse_token_trees(source).unwrap();

        // Should have: fn, foo, (), {}
        assert_eq!(trees.len(), 4);

        // Check the body
        match &trees[3] {
            TokenTree::Block(TokenTreeBlock {
                delimiter, trees, ..
            }) => {
                assert_eq!(*delimiter, Block::Brace);
                // Should contain: let, x, =, (1 + 2), ;
                assert_eq!(trees.len(), 5);
            }
            _ => panic!("Expected brace group"),
        }
    }

    #[test]
    fn test_mismatched_delimiter() {
        let source = "fn foo(] { }";
        let result = parse_token_trees(source);
        assert!(result.is_err());

        match result.unwrap_err() {
            ParseError::MismatchedDelimiter { .. } => {}
            e => panic!("Expected MismatchedDelimiter, got {:?}", e),
        }
    }

    #[test]
    fn test_unclosed_delimiter() {
        let source = "fn foo() { let x = 5;";
        let result = parse_token_trees(source);
        assert!(result.is_err());

        match result.unwrap_err() {
            ParseError::UnexpectedEOF { .. } => {}
            e => panic!("Expected UnexpectedEOF, got {:?}", e),
        }
    }

    #[test]
    fn test_unexpected_closing() {
        let source = "let x = 5 }";
        let result = parse_token_trees(source);
        assert!(result.is_err());

        match result.unwrap_err() {
            ParseError::UnexpectedClosing { .. } => {}
            e => panic!("Expected UnexpectedClosing, got {:?}", e),
        }
    }

    #[test]
    fn test_complex_nesting() {
        let source = r#"
            fn calculate() {
                let result = add(mul(2, 3), sub(10, 5));
                return result;
            }
        "#;
        let trees = parse_token_trees(source).unwrap();
        println!("{}", TokenTreeParser::token_trees_to_string(&trees));
    }
}