plang-2/plang2_lib/src/token.rs

use std::{fmt::Display, borrow::Cow};

/// Operators
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Op {
    // Addition
    Add,
    Sub,

    // Multiplications
    Mul,
    Div,
    Mod,

    // Assignment
    Assign,

    // Equality
    Eq,
    Neq,
    Gt,
    Lt,
    Ge,
    Le,

    // Bool
    And,
    Or,
    Not,
    Xor,

    Arrow,
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Group {
    Paren,
    Bracket,
    Braces,
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Literal {
    Boolean(bool),
    Int64(i64),
    String(String),
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Keyword {
    Let,
    While,
    Loop,
    If,
    Else,
    Fn,
    Return,
    Void,
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Token {
    Literal(Literal),
    Op(Op),
    Open(Group),
    Close(Group),

    Ident(String),

    Keyword(Keyword),

    Semicolon,
    Colon,
    Comma,
    Dot,
}

pub struct TokenStream {
    tokens: Vec<Token>,
    idx: usize,
}

impl Display for Token {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

        let op: Cow<'static, str> = match self {
            Token::Op(Op::Add) => "+".into(),
            Token::Op(Op::Sub) => "-".into(),

            Token::Op(Op::Mul) => "*".into(),
            Token::Op(Op::Div) => "/".into(),
            Token::Op(Op::Mod) => "%".into(),
            
            Token::Op(Op::Eq) => "==".into(),
            Token::Op(Op::Neq) => "!=".into(),
            Token::Op(Op::Gt) => ">".into(),
            Token::Op(Op::Lt) => "<".into(),
            Token::Op(Op::Ge) => ">=".into(),
            Token::Op(Op::Le) => "<=".into(),

            Token::Op(Op::Assign) => "=".into(),
            Token::Op(Op::Arrow) => "->".into(),
            
            Token::Op(Op::And) => "&&".into(),
            Token::Op(Op::Or) => "||".into(),
            Token::Op(Op::Xor) => "^".into(),
            Token::Op(Op::Not) => "!".into(),

            Token::Open(Group::Paren) => "(".into(),
            Token::Open(Group::Bracket) => "[".into(),
            Token::Open(Group::Braces) => "{".into(),
            Token::Close(Group::Paren) => ")".into(),
            Token::Close(Group::Bracket) => "]".into(),
            Token::Close(Group::Braces) => "}".into(),

            Token::Literal(Literal::Int64(num)) => format!("Int64({})", num).into(),
            Token::Literal(Literal::String(text)) => format!("String({})", text).into(),
            Token::Literal(Literal::Boolean(val)) => format!("Boolean({})", val).into(),

            Token::Ident(ident) => format!("Ident({})", ident).into(),

            Token::Semicolon => ";".into(),
            Token::Colon => ":".into(),
            Token::Comma => ",".into(),
            Token::Dot => ".".into(),

            Token::Keyword(Keyword::Let) => "let".into(),
            Token::Keyword(Keyword::If) => "if".into(),
            Token::Keyword(Keyword::While) => "while".into(),
            Token::Keyword(Keyword::Loop) => "loop".into(),
            Token::Keyword(Keyword::Else) => "else".into(),
            Token::Keyword(Keyword::Fn) => "fn".into(),
            Token::Keyword(Keyword::Return) => "return".into(),
            Token::Keyword(Keyword::Void) => "void".into(),
        };

        write!(f, "{}", op)
    }
}

impl TokenStream {
    pub fn new(tokens: Vec<Token>) -> Self {
        Self { tokens, idx: 0 }
    }

    pub fn as_vec(&self) -> &Vec<Token> {
        &self.tokens
    }

    pub fn curr(&self) -> Option<&Token> {
        self.tokens.get(self.idx)
    }

    pub fn peek(&self) -> Option<&Token> {
        self.tokens.get(self.idx + 1)
    }

    /// Advance to the next token. Sets curr to next and returns the old curr.
    pub fn advance(&mut self) -> Option<&Token> {
        self.idx += 1;
        self.tokens.get(self.idx - 1)
    }
}

impl Display for TokenStream {
    /// Print the TokenStream with autofomatting
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut indent = 0_usize;
        let mut fresh_line = true; 

        for tok in self.tokens.iter() {
            if matches!(tok, Token::Close(Group::Braces)) {
                indent = indent.saturating_sub(1);
                fresh_line = true;
            }

            if fresh_line {
                write!(f, "{}", " ".repeat(indent * 4))?;
                fresh_line = false;
            }

            write!(f, "{} ", tok)?;
            

            match tok {
                Token::Open(Group::Braces) => {
                    writeln!(f)?;
                    indent += 1;
                    fresh_line = true;
                }
                Token::Semicolon | Token::Close(Group::Braces) => {
                    writeln!(f)?;
                    fresh_line = true;
                }
                _ => ()
            }
        }

        Ok(())
    }
}