expr eval

use std::f64::consts::PI;
use std::io;
use std::io::Write;
use std::str::Chars;

#[derive(Debug)]
struct Tokenizer<'t> {
    iterator: Chars<'t>,
}

impl<'t> Tokenizer<'t> {
    pub fn new(source: &'t str) -> Self {
        Self {
            iterator: source.chars(),
        }
    }

    pub fn next_token(&mut self) -> Result<Option<Token>, String> {
        let Some(next) = self.iterator.next() else {
            return Ok(None);
        };

        match next {
            '(' => Ok(Some(Token::BrOpen)),
            ')' => Ok(Some(Token::BrClose)),
            '^' => Ok(Some(Token::Op(MathOp::Pow))),
            '+' => Ok(Some(Token::Op(MathOp::Add))),
            '-' => Ok(Some(Token::Op(MathOp::Sub))),
            '*' => Ok(Some(Token::Op(MathOp::Mul))),
            '/' => Ok(Some(Token::Op(MathOp::Div))),
            'π' => Ok(Some(Token::Num(PI))),
            ch if ch.is_ascii_digit() || ch == '.' => {
                let mut num = String::from(ch);
                let mut it2 = self.iterator.clone();

                while let Some(next_ch) = it2.next() {
                    if !next_ch.is_ascii_digit() && next_ch != '.' {
                        break;
                    }

                    num.push(next_ch);
                    self.iterator.next();
                }

                Ok(Some(Token::Num(num.parse().map_err(|e| {
                    format!("Encountered error while parsing number: {e}")
                })?)))
            }
            ' ' => self.next_token(),
            _ => Err("Encountered unexpected character".to_string()),
        }
    }
}

/// Parser grammar:
///
/// ```
/// expr = term;
/// term = factor (("+" | "-") factor)*;
/// factor = power (("*" | "/") power)*;
/// power = unary ("^" unary)*;
/// unary = "-"? (unary | primary);
/// primary = Num | "(" expr ")";
/// ```
#[derive(Debug)]
struct Parser<'ecx> {
    tokenizer: Tokenizer<'ecx>,
    current: Option<Token>,
}

impl<'ecx> Parser<'ecx> {
    pub fn new(mut tokenizer: Tokenizer<'ecx>) -> Result<Option<Self>, String> {
        let current = tokenizer.next_token()?;

        Ok(Some(Self {
            tokenizer,
            current,
        }))
    }

    pub fn parse(&mut self) -> Result<AST, String> {
        self.term()
    }

    fn term(&mut self) -> Result<AST, String> {
        let mut expr = self.factor()?;

        while self.current == Some(Token::Op(MathOp::Add)) || self.current == Some(Token::Op(MathOp::Sub)) {
            let Token::Op(op) = self.current.unwrap() else {
                unreachable!()
            };

            self.next()?;
            expr = AST::Op(Operation {
                lhs: Box::new(expr),
                rhs: Box::new(self.factor()?),
                op,
            })
        }

        Ok(expr)
    }

    fn factor(&mut self) -> Result<AST, String> {
        let mut expr = self.power()?;

        while self.current == Some(Token::Op(MathOp::Mul)) || self.current == Some(Token::Op(MathOp::Div)) {
            let Token::Op(op) = self.current.unwrap() else {
                unreachable!()
            };

            self.next()?;
            expr = AST::Op(Operation {
                lhs: Box::new(expr),
                rhs: Box::new(self.power()?),
                op,
            })
        }

        Ok(expr)
    }

    fn power(&mut self) -> Result<AST, String> {
        let mut expr = self.unary()?;

        while self.current == Some(Token::Op(MathOp::Pow)) {
            self.next()?;
            expr = AST::Op(Operation {
                lhs: Box::new(expr),
                rhs: Box::new(self.unary()?),
                op: MathOp::Pow,
            })
        }

        Ok(expr)
    }

    fn unary(&mut self) -> Result<AST, String> {
        if matches!(self.current, Some(Token::Op(MathOp::Sub))) {
            let _ = self.next()?;

            let operand = self.unary()?;
            Ok(AST::Neg(Box::new(operand)))
        } else {
            self.primary()
        }
    }

    fn primary(&mut self) -> Result<AST, String> {
        match self.current {
            Some(Token::Num(num)) => {
                self.next()?;
                Ok(AST::Num(num))
            },
            Some(Token::BrOpen) => {
                // Consume "("
                let _ = self.next()?;

                // Parse inner expression
                let parsed = self.parse();

                // Consume ")"
                let _ = self.expect(Token::BrClose)?;

                parsed
            }
            Some(_) => Err("Unexpected token".to_string()),
            None => Err("Unexpected EOF".to_string())
        }
    }

    fn expect(&mut self, expected: Token) -> Result<Token, String> {
        if !matches!(self.current, Some(_expected)) {
            Err(format!("Expected {expected:?}, found {:?}", self.current))
        } else {
            Ok(self.current.unwrap())
        }
    }

    fn next(&mut self) -> Result<Option<Token>, String> {
        self.current = self.tokenizer.next_token()?;
        Ok(self.current)
    }
}

#[derive(Debug)]
enum AST {
    Num(f64),
    Op(Operation),
    Neg(Box<AST>),
}

#[derive(Debug)]
struct Operation {
    lhs: Box<AST>,
    rhs: Box<AST>,
    op: MathOp,
}

#[derive(Debug, Copy, Clone, PartialEq)]
enum Token {
    Num(f64),
    BrOpen,
    BrClose,
    Op(MathOp),
}

#[derive(Debug, Copy, Clone, PartialEq)]
enum MathOp {
    Add,
    Sub,
    Mul,
    Div,
    Pow,
}

fn eval_ast(ast: &AST) -> Result<f64, String> {
    match ast {
        AST::Num(num) => Ok(*num),
        AST::Op(operation) => {
            let lhs = eval_ast(&operation.lhs)?;
            let rhs = eval_ast(&operation.rhs)?;

            Ok(match operation.op {
                MathOp::Add => lhs + rhs,
                MathOp::Sub => lhs - rhs,
                MathOp::Mul => lhs * rhs,
                MathOp::Div => lhs / rhs,
                MathOp::Pow => lhs.powf(rhs)
            })
        }
        AST::Neg(expr) => Ok(-eval_ast(expr)?)
    }
}

fn process_expr(expr: &str) {
    let tokenizer = Tokenizer::new(&expr);
    let mut parser = match Parser::new(tokenizer) {
        Ok(Some(parser)) => parser,
        Ok(None) => return,
        Err(e) => {
            eprintln!("Err: {e}");
            return;
        }
    };

    let ast = match parser.parse() {
        Ok(ast) => ast,
        Err(e) => {
            eprintln!("Err: {e}");
            return;
        }
    };

    match eval_ast(&ast) {
        Ok(res) => println!("{res}"),
        Err(e) => {
            eprintln!("Err: {e}");
            return;
        }
    }
}

fn take_input() -> String {
    print!("> ");
    io::stdout().flush().unwrap();
    let mut buf = String::new();
    io::stdin().read_line(&mut buf).unwrap();

    buf.trim().to_string()
}

fn main() {
    loop {
        let input = take_input();
        if input == "q" {
            break;
        }

        process_expr(&input);
    }
}