asklyphe/asklyphe-frontend/src/math.rs

use tracing::{debug, error};
use once_cell::sync::Lazy;

#[derive(Debug)]
pub struct Calculation {
	pub equation: String,
	pub result: String,
}

pub fn calculate(query: &str) -> Option<Calculation> {
	debug!("Got query {}", query);
	let mut parser = Parser::new(Lexer::new(query));
	let tree = parser.parse()?;
	let res = tree.eval();

	debug!("Calculation: {}", query);
	debug!("Tree: {:?}", tree);
	debug!("Result: {}", res);
	Some(Calculation {equation: query.to_string(), result: res.to_string()})
}

// TODO: put into own crate with dependency astro-float = "0.9.2" so I can use more than f64
#[derive(Debug, Copy, Clone)]
enum Token {
	Op(Op),
	Atom(Atom),
/*	Number(f64),
	Func(Func),*/
}

#[derive(Debug, Copy, Clone)]
enum Op {
	BinOp(BinOp),
	Func(Func), // A function is an Op that takes whatever the next thing is and binds it, either the next number or whatever is in parens
}

impl Op {
	fn get_lbp(&self) -> f64 {
		match self {
			Op::BinOp(op) => match op {
				BinOp::LParen => 0.0,
				BinOp::RParen => 0.0,
				BinOp::Add => 1.0,
				BinOp::Subtract => 1.0,
				BinOp::Multiply => 2.0,
				BinOp::Divide => 2.0,
				BinOp::Exponent => 3.1,
			},
			Op::Func(_) => 0.0, // TODO: decide if this is a good LBP
		}
	}

	fn get_rbp(&self) -> f64 {
		match self {
			Op::BinOp(op) => match op {
				BinOp::LParen => 0.0,
				BinOp::RParen => 0.0,
				BinOp::Add => 1.1,
				BinOp::Subtract => 1.1,
				BinOp::Multiply => 2.1,
				BinOp::Divide => 2.1,
				BinOp::Exponent => 3.0,
			},
			Op::Func(_) => 4.0, // TODO: decide if this is a good RBP
		}

	}

	fn apply_to(&self, args: &Vec<Expr>) -> f64 {
		match (self) {
			Op::BinOp(op) => match op {
				BinOp::LParen => args[0].eval(),
				BinOp::RParen => args[0].eval(),
				BinOp::Add => args[0].eval() + args[1].eval(),
				BinOp::Subtract => args[0].eval() - args[1].eval(),
				BinOp::Multiply => args[0].eval() * args[1].eval(),
				BinOp::Divide => args[0].eval() / args[1].eval(),
				BinOp::Exponent => args[0].eval().powf(args[1].eval()),
			},
			Op::Func(f) => match f {
				Func::Sine => args[0].eval().sin(),
				Func::ArcSine => args[0].eval().asin(),
				Func::SquareRoot => args[0].eval().sqrt(),
				_ => todo!("{:?}", f),
			},
		}
	}
}

#[derive(Debug, Copy, Clone)]
enum BinOp {
	Add,
	Subtract,
	Multiply,
	Divide,
	Exponent,
	LParen,
	RParen,
}

#[derive(Debug, Copy, Clone)]
enum Atom {
	Number(f64), // TODO: use the high precision floats library instead
	Const(Const),
}

impl Atom {
	fn get_val(&self) -> f64 {
		match self {
			Atom::Number(val) => *val,
			Atom::Const(c) => match c {
				Const::Pi => 3.141592653589793238462643383279,
				Const::E => 2.718281828459045235360287471352,
			}
		}
	}
}

#[derive(Debug, Copy, Clone)]
enum Func {
	Sine,
	Cosine,
	Tangent,
	// sin-1, cos-1, tan-1
	ArcSine,
	ArcCosine,
	ArcTangent,
	Log2,
	Log10,
	LogN,
	Square,
	SquareRoot,
}

#[derive(Debug, Copy, Clone)]
enum Const {
	Pi,
	E,
}

#[derive(Debug, Copy, Clone)]
enum ParseErr {
	Eof,
	Invalid,
}

// this can probably be swapped out with a lexer generator like Logos if needed

struct Lexer<'a> {
	data: &'a str,
	data_ptr: &'a str,
	idx: usize,
	next_tok: Result<Token, ParseErr>,
}

// TODO: refactor with iterator that returns Option(Token) where one token option is Eof (or a enum of Token(Token) and Eof, or just Option(Option(Token)))
impl Lexer<'_> {

	fn new(data: &str) -> Lexer<'_> {
		let mut n: Lexer = Lexer {data, data_ptr: data, idx: 0, next_tok: Err(ParseErr::Eof)};
		n.next();
		debug!("New finished!");
		n
	}
	
	fn _next(&mut self) -> Result<Token, ParseErr> {
		match self.data.chars().nth(self.idx) {
			Some(val) => {
				debug!("lexing char '{}' at idx {}", val, self.idx);
				// debug!("current char '{}'", self.data.chars().nth(0).unwrap());
				self.idx += 1;
				// TODO: make more efficient
				self.data_ptr = &self.data[self.idx..];
				match val {
					'+' => Ok(Token::Op(Op::BinOp(BinOp::Add))),
					'-' => Ok(Token::Op(Op::BinOp(BinOp::Subtract))),
					'×' | '*' => Ok(Token::Op(Op::BinOp(BinOp::Multiply))),
					'÷' | '/' => Ok(Token::Op(Op::BinOp(BinOp::Divide))),
					'^' => Ok(Token::Op(Op::BinOp(BinOp::Exponent))),
					'(' => Ok(Token::Op(Op::BinOp(BinOp::LParen))),
					')' => Ok(Token::Op(Op::BinOp(BinOp::RParen))),
					_ if val.is_whitespace() => self._next(),
					// TODO: maybe parse '-' as part of number so I can do '1 + -1' and similar
					_ if val.is_digit(10) => {
						let start = self.idx - 1;

						self.data_ptr.chars().take_while(|c| c.is_digit(10)).for_each(|_| self.idx += 1);//.next().unwrap_or(' ').is_digit(10) {self.idx += 1;}

						match self.data[start..self.idx].parse() {
							Ok(val) => Ok(Token::Atom(Atom::Number(val))),
							Err(e) => Err(ParseErr::Invalid),
						}
					},
					_ => {
						let mut l: usize;
						l = matches(&self.data[self.idx - 1..], "sin");
						if l != 0 {
							self.idx += l;
							return Ok(Token::Op(Op::Func(Func::Sine)));
						}
						l = matches(&self.data[self.idx - 1..], "cos");
						if l != 0 {
							self.idx += l;
							return Ok(Token::Op(Op::Func(Func::Cosine)));
						}
						l = matches(&self.data[self.idx - 1..], "tan");
						if l != 0 {
							self.idx += l;
							return Ok(Token::Op(Op::Func(Func::Tangent)));
						}
						l = matches(&self.data[self.idx - 1..], "arcsin");
						if l != 0 {
							self.idx += l;
							return Ok(Token::Op(Op::Func(Func::ArcSine)));
						}
						l = matches(&self.data[self.idx - 1..], "arccos");
						if l != 0 {
							self.idx += l;
							return Ok(Token::Op(Op::Func(Func::ArcCosine)));
						}
						l = matches(&self.data[self.idx - 1..], "arctan");
						if l != 0 {
							self.idx += l;
							return Ok(Token::Op(Op::Func(Func::ArcTangent)));
						}
						l = matches(&self.data[self.idx - 1..], "sqrt");
						if l != 0 {
							self.idx += l;
							return Ok(Token::Op(Op::Func(Func::SquareRoot)));
						}

						debug!("got invalid char '{}'", val);
						Err(ParseErr::Invalid)
					}
				}
			}
			None => Err(ParseErr::Eof),
		}
	}

	fn next(&mut self) -> Result<Token, ParseErr> {
		let val = self.next_tok;
		self.next_tok = self._next();
		val
	}

	fn peek(&mut self) -> Result<Token, ParseErr> {
		self.next_tok
	}

	// TODO: replace with iterator so I can do parser.parse(lexer.iter()) and parse does lex_iter.next() & such
	fn lex_all(&mut self) -> Option<Vec<Token>> {
		let mut tokens: Vec<Token> = vec![];
		loop {
			match self.next() {
				Err(ParseErr::Eof) => return Some(tokens),
				Err(ParseErr::Invalid) => return None,
				Ok(tok) => tokens.push(tok),
			}
			// debug!("tokens: {:?}", tokens);
		}
	}
}

fn matches(s: &str, check: &str) -> usize {
	// debug!("s: \"{}\", check: \"{}\"c_len: {}, s_len: {}, s[c_len]: {:?}, s[c_len + 1]: {:?}", s, check, check.chars().count(), s.chars().count(), s.chars().nth(check.chars().count()), s.chars().nth(check.chars().count() + 1));
	match (s.chars().count(), check.chars().count()) {
		(s_len, c_len) if s_len < c_len => 0,
		(s_len, c_len) if s_len == c_len && s == check => c_len - 1,
		(s_len, c_len) if s_len > c_len && s.starts_with(check) && s.chars().nth(c_len).unwrap().is_whitespace() => c_len,
		(_, _) => 0,
	}
}

struct Parser<'a> {
	lex: Lexer<'a>,
}

impl Parser<'_> {
	fn new(lex: Lexer) -> Parser { Parser {lex} }

	fn parse(&mut self) -> Option<Expr> {
		self.parse_expr(0.0).ok()
	}

	fn parse_expr(&mut self, min_bp: f64) -> Result<Expr, ParseErr> {
		/*while let Ok(val) = self.lex.next() {debug!("token: {:?}", val)}
		match self.lex.next().err() {
			
			_ => return Err(ParseErr::Invalid),
		}*/
		let mut lhs: Expr = match self.lex.next() {
			Ok(val) => match val {
				Token::Atom(val) => Ok(Expr::Atom(val)),
				Token::Op(op) => match op {
					Op::BinOp(BinOp::LParen) => self.parse_expr(op.get_lbp()),
					Op::Func(f) => Ok(Expr::Node(Op::Func(f), vec![self.parse_expr(op.get_lbp())?])),
					_ => Err(ParseErr::Invalid),
				},
			},
			Err(err) => Err(err),
		}.map_err(|err| { debug!("Unexpected error at start of expr: {:?}", err); err })?;
		debug!("lhs of expression is {:?}", lhs);
		loop {
			let op: Op = match self.lex.peek() {
				Err(ParseErr::Eof) => break,
				Err(e) => { debug!("In expr got err {:?}", e); Err(e) },
				Ok(tok) => match tok {
					Token::Op(op) => match op {
						Op::BinOp(op) => match op {
							BinOp::RParen => break,
							_ => Ok(Op::BinOp(op)),
						},
						Op::Func(f) => {
							lhs = Expr::Node(Op::Func(f), vec![self.parse_expr(Op::Func(f).get_lbp())?]);
							continue;
						},
					}
					v => { debug!("Got unexpected token {:?}", v); Err(ParseErr::Invalid) },
				}
			}.map_err(|err| { debug!("Unexpected error inside expr at {:?}", err); err })?;
			if (op.get_lbp() < min_bp) { break; }
			self.lex.next();
			let rhs: Expr = self.parse_expr(op.get_rbp())?;
			lhs = Expr::Node(op, vec![lhs, rhs]);
		}
		Ok(lhs)
	}
}

#[derive(Debug)]
enum Expr {
	Atom(Atom),
	Node(Op, Vec<Expr>),
}

impl Expr {
	fn eval(&self) -> f64 {
		let res = match (self) {
			Expr::Atom(at) => at.get_val(),
			Expr::Node(op, exprs) => op.apply_to(exprs),
		};
		// debug!("{:?} evaluated to {}", self, res);
		res
	}
}