talc/talc-lang/src/optimize.rs

use core::panic;

use crate::{
	parser::ast::{CatchBlock, Expr, ExprKind, LValue, LValueKind, ListItem, TableItem},
	value::Value,
	vm,
};

impl Expr {
	fn value(&self) -> Option<&Value> {
		if let ExprKind::Literal(v) = &self.kind {
			Some(v)
		} else {
			None
		}
	}
}

fn expr_take(e: &mut Expr) -> Expr {
	std::mem::replace(e, ExprKind::Literal(Value::Nil).span(e.span))
}

#[derive(Clone, Copy, Debug)]
struct OptState {
	ret: bool,
}

impl OptState {
	#[inline]
	const fn empty() -> Self {
		Self { ret: false }
	}

	#[inline]
	const fn ret(ret: bool) -> Self {
		Self { ret }
	}
}

pub fn optimize(body: &mut Expr) {
	let state = OptState::ret(true);
	optimize_ex_with(body, state);
}

#[inline]
fn optimize_ex(expr: &mut Expr) {
	optimize_ex_with(expr, OptState::empty());
}

fn optimize_ex_with(expr: &mut Expr, state: OptState) {
	let span = expr.span;
	match &mut expr.kind {
		ExprKind::Literal(_)
		| ExprKind::Ident(_)
		| ExprKind::Var(_)
		| ExprKind::Global(_) => (),
		ExprKind::UnaryOp(o, e) => {
			optimize_ex(e);
			if let Some(a) = e.value() {
				if let Ok(v) = vm::unary_op(*o, a.clone()) {
					*expr = ExprKind::Literal(v).span(span);
				}
			}
		}
		ExprKind::BinaryOp(o, l, r) => {
			optimize_ex(l);
			optimize_ex(r);
			if let (Some(a), Some(b)) = (l.value(), r.value()) {
				if let Ok(v) = vm::binary_op(*o, a.clone(), b.clone()) {
					*expr = ExprKind::Literal(v).span(span);
				}
			}
		}
		ExprKind::Assign(_, l, r) => {
			optimize_lv(l);
			optimize_ex(r);
		}
		ExprKind::FnDef(_, _, e) => {
			optimize_ex_with(e, OptState::ret(true));
		}
		ExprKind::Index(l, r) => {
			optimize_ex(l);
			optimize_ex(r);
		}
		ExprKind::FnCall(f, xs) => {
			optimize_ex(f);
			for e in &mut *xs {
				optimize_ex(e);
			}
			if state.ret {
				*expr = ExprKind::TailCall(Box::new(expr_take(f)), std::mem::take(xs))
					.span(span);
			}
		}
		ExprKind::AssocFnCall(f, a, xs) => {
			optimize_ex(f);
			for e in &mut *xs {
				optimize_ex(e);
			}
			if state.ret {
				*expr =
					ExprKind::AssocTailCall(Box::new(expr_take(f)), *a, std::mem::take(xs))
						.span(span);
			}
		}
		ExprKind::Pipe(l, r) => {
			optimize_ex(l);
			optimize_ex(r);
		}
		ExprKind::Block(xs) => {
			let len = xs.len();
			if len > 1 {
				for e in &mut xs[..len - 1] {
					optimize_ex(e);
				}
			}
			if let Some(e) = xs.last_mut() {
				optimize_ex_with(e, state);
			}
		}
		ExprKind::List(xs) => {
			for e in xs {
				match e {
					ListItem::Item(e) => optimize_ex(e),
					ListItem::Interpolate(e) => optimize_ex(e),
				}
			}
		}
		ExprKind::Table(t) => {
			for e in t {
				match e {
					TableItem::Pair(k, v) => {
						optimize_ex(k);
						optimize_ex(v);
					}
					TableItem::Interpolate(ext) => {
						optimize_ex(ext);
					}
				}
			}
		}
		ExprKind::Return(e) => {
			if let Some(e) = e {
				optimize_ex_with(e, OptState::ret(true));
			}
		}
		ExprKind::Break(e) => {
			if let Some(e) = e {
				optimize_ex(e);
			}
		}
		ExprKind::Continue => (),
		ExprKind::And(l, r) => {
			optimize_ex(l);
			optimize_ex(r);
			if let Some(v) = l.value() {
				if v.truthy() {
					*expr = expr_take(r);
				} else {
					*expr = expr_take(l);
				}
			}
		}
		ExprKind::Or(l, r) => {
			optimize_ex(l);
			optimize_ex(r);
			if let Some(v) = l.value() {
				if v.truthy() {
					*expr = expr_take(l);
				} else {
					*expr = expr_take(r);
				}
			}
		}
		ExprKind::If(c, t, e) => {
			optimize_ex(c);
			optimize_ex_with(t, state);
			if let Some(e) = e {
				optimize_ex_with(e, state);
			}
			if let Some(v) = c.value() {
				if v.truthy() {
					*expr = expr_take(t);
				} else if let Some(e) = e {
					*expr = expr_take(e);
				} else {
					*expr = ExprKind::Literal(Value::Nil).span(span);
				}
			}
		}
		ExprKind::While(c, b) => {
			optimize_ex(c);
			optimize_ex(b);
			if let Some(v) = c.value() {
				if !v.truthy() {
					*expr = ExprKind::Literal(Value::Nil).span(span);
				}
			}
		}
		ExprKind::For(_, i, b) => {
			optimize_ex(i);
			optimize_ex(b);
		}
		ExprKind::Lambda(_, b) => {
			optimize_ex_with(b, OptState::ret(true));
		}
		ExprKind::Try(b, cx) => {
			optimize_ex_with(b, state);
			for c in cx {
				optimize_cb(c);
			}
		}
		ExprKind::TailCall(..) | ExprKind::AssocTailCall(..) => {
			panic!("cannot optimize expression generated by optimizer")
		}
	}
}

fn optimize_lv(e: &mut LValue) {
	match &mut e.kind {
		LValueKind::Ident(_) | LValueKind::Var(_) | LValueKind::Global(_) => (),
		LValueKind::Index(l, r) => {
			optimize_ex(l);
			optimize_ex(r);
		}
		LValueKind::List(ls, _) => {
			for l in ls {
				optimize_lv(l);
			}
		}
	}
}

fn optimize_cb(e: &mut CatchBlock) {
	optimize_ex(&mut e.body);
}