use std::collections::HashMap;
use logicaffeine_base::quantity::units;
use logicaffeine_base::Dimension;
use crate::ast::stmt::{BinaryOpKind, Expr, Literal, Stmt, TypeExpr};
use crate::intern::{Interner, Symbol};
use crate::token::Span;
#[derive(Clone, Copy, PartialEq)]
enum QDim {
Known(Dimension),
Unknown,
NotQuantity,
}
impl QDim {
fn is_quantity(self) -> bool {
!matches!(self, QDim::NotQuantity)
}
}
#[derive(Clone, PartialEq, Eq)]
enum CurInfo {
Known(String),
Unknown,
NotMoney,
}
impl CurInfo {
fn is_money(&self) -> bool {
!matches!(self, CurInfo::NotMoney)
}
}
pub struct DimensionError {
pub message: String,
pub span: Span,
}
impl std::fmt::Display for DimensionError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.message)
}
}
pub struct DimensionChecker<'a> {
vars: HashMap<Symbol, QDim>,
fn_returns: HashMap<Symbol, QDim>,
cur_vars: HashMap<Symbol, CurInfo>,
interner: &'a Interner,
}
impl<'a> DimensionChecker<'a> {
pub fn new(interner: &'a Interner) -> Self {
Self {
vars: HashMap::new(),
fn_returns: HashMap::new(),
cur_vars: HashMap::new(),
interner,
}
}
pub fn check_program(&mut self, stmts: &[Stmt<'_>]) -> Result<(), DimensionError> {
self.collect_fn_returns(stmts);
self.check_block(stmts)
}
fn collect_fn_returns(&mut self, stmts: &[Stmt<'_>]) {
for stmt in stmts {
if let Stmt::FunctionDef { name, return_type, .. } = stmt {
let qd = return_type.map(|t| self.dim_from_type(t)).unwrap_or(QDim::NotQuantity);
self.fn_returns.insert(*name, qd);
}
}
}
fn dim_from_type(&self, ty: &TypeExpr<'_>) -> QDim {
match ty {
TypeExpr::Generic { base, params }
if self.interner.resolve(*base) == "Quantity" && params.len() == 1 =>
{
match ¶ms[0] {
TypeExpr::Primitive(s) | TypeExpr::Named(s) => {
match Dimension::by_name(self.interner.resolve(*s)) {
Some(d) => QDim::Known(d),
None => QDim::Unknown,
}
}
_ => QDim::Unknown,
}
}
TypeExpr::Primitive(s) | TypeExpr::Named(s)
if self.interner.resolve(*s) == "Quantity" =>
{
QDim::Unknown
}
_ => QDim::NotQuantity,
}
}
fn check_block(&mut self, stmts: &[Stmt<'_>]) -> Result<(), DimensionError> {
for stmt in stmts {
self.check_stmt(stmt)?;
}
Ok(())
}
fn check_stmt(&mut self, stmt: &Stmt<'_>) -> Result<(), DimensionError> {
match stmt {
Stmt::Let { var, value, ty, .. } => {
let inferred = self.infer(value)?;
let declared = ty.map(|t| self.dim_from_type(t));
let d = match declared {
Some(QDim::Known(k)) => QDim::Known(k),
Some(QDim::Unknown) if !inferred.is_quantity() => QDim::Unknown,
_ => inferred,
};
self.vars.insert(*var, d);
let c = self.currency_of(value);
self.cur_vars.insert(*var, c);
}
Stmt::Set { value, .. } => {
self.infer(value)?;
}
Stmt::SetField { object, value, .. } => {
self.infer(object)?;
self.infer(value)?;
}
Stmt::If { cond, then_block, else_block } => {
self.infer(cond)?;
self.check_block(then_block)?;
if let Some(e) = else_block {
self.check_block(e)?;
}
}
Stmt::While { cond, body, .. } => {
self.infer(cond)?;
self.check_block(body)?;
}
Stmt::Repeat { iterable, body, .. } => {
self.infer(iterable)?;
self.check_block(body)?;
}
Stmt::Return { value: Some(e) } => {
self.infer(e)?;
}
Stmt::Show { object, .. } => {
self.infer(object)?;
}
Stmt::RuntimeAssert { condition, .. } => {
self.infer(condition)?;
}
Stmt::Call { args, .. } => {
for a in args {
self.infer(a)?;
}
}
Stmt::FunctionDef { params, body, .. } => {
let saved = self.vars.clone();
let saved_cur = self.cur_vars.clone();
for (name, ty) in params {
let qd = self.dim_from_type(ty);
if qd.is_quantity() {
self.vars.insert(*name, qd);
}
}
self.check_block(body)?;
self.vars = saved;
self.cur_vars = saved_cur;
}
_ => {}
}
Ok(())
}
fn unit_dim(&self, expr: &Expr<'_>) -> QDim {
if let Expr::Literal(Literal::Text(sym)) = expr {
if let Some(unit) = units::by_name(self.interner.resolve(*sym)) {
return QDim::Known(unit.dimension);
}
}
QDim::Unknown
}
fn currency_of(&self, expr: &Expr<'_>) -> CurInfo {
match expr {
Expr::Identifier(s) => self.cur_vars.get(s).cloned().unwrap_or(CurInfo::NotMoney),
Expr::Call { function, args } => {
if self.interner.resolve(*function) == "money" && args.len() == 2 {
if let Expr::Literal(Literal::Text(code)) = args[1] {
return CurInfo::Known(self.interner.resolve(*code).to_ascii_uppercase());
}
return CurInfo::Unknown;
}
CurInfo::NotMoney
}
Expr::BinaryOp { op, left, right } => match op {
BinaryOpKind::Add | BinaryOpKind::Subtract => {
let (l, r) = (self.currency_of(left), self.currency_of(right));
if let CurInfo::Known(_) = l {
l
} else {
r
}
}
BinaryOpKind::Multiply => {
let (l, r) = (self.currency_of(left), self.currency_of(right));
if matches!(l, CurInfo::Known(_) | CurInfo::Unknown) {
l
} else {
r
}
}
BinaryOpKind::Divide => {
let (l, r) = (self.currency_of(left), self.currency_of(right));
if r.is_money() {
CurInfo::NotMoney
} else {
l
}
}
_ => CurInfo::NotMoney,
},
Expr::Copy { expr } | Expr::Give { value: expr } => self.currency_of(expr),
_ => CurInfo::NotMoney,
}
}
fn check_currency_match(
&self,
left: &Expr<'_>,
right: &Expr<'_>,
verb: &str,
) -> Result<(), DimensionError> {
if let (CurInfo::Known(a), CurInfo::Known(b)) = (self.currency_of(left), self.currency_of(right))
{
if a != b {
return Err(DimensionError {
message: format!("cannot {verb} money of different currencies ({a} vs {b})"),
span: Span::default(),
});
}
}
Ok(())
}
fn infer(&self, expr: &Expr<'_>) -> Result<QDim, DimensionError> {
match expr {
Expr::Literal(_) => Ok(QDim::NotQuantity),
Expr::Identifier(s) => Ok(self.vars.get(s).copied().unwrap_or(QDim::NotQuantity)),
Expr::Call { function, args } => {
for a in args {
self.infer(a)?;
}
match self.interner.resolve(*function) {
"quantity" | "convert" if args.len() == 2 => Ok(self.unit_dim(args[1])),
_ => Ok(self.fn_returns.get(function).copied().unwrap_or(QDim::Unknown)),
}
}
Expr::BinaryOp { op, left, right } => {
let l = self.infer(left)?;
let r = self.infer(right)?;
match op {
BinaryOpKind::Add | BinaryOpKind::Subtract => {
self.check_currency_match(
left,
right,
if matches!(op, BinaryOpKind::Add) { "add" } else { "subtract" },
)?;
if let (QDim::Known(a), QDim::Known(b)) = (l, r) {
if a != b {
return Err(self.mismatch_err(
if matches!(op, BinaryOpKind::Add) { "add" } else { "subtract" },
a,
b,
));
}
return Ok(QDim::Known(a));
}
Ok(self.propagate(l, r))
}
BinaryOpKind::Lt | BinaryOpKind::Gt | BinaryOpKind::LtEq | BinaryOpKind::GtEq => {
self.check_currency_match(left, right, "compare")?;
if let (QDim::Known(a), QDim::Known(b)) = (l, r) {
if a != b {
return Err(self.mismatch_err("compare", a, b));
}
}
Ok(QDim::NotQuantity)
}
BinaryOpKind::Multiply => Ok(match (l, r) {
(QDim::Known(a), QDim::Known(b)) => QDim::Known(a.mul(b)),
(QDim::Known(a), QDim::NotQuantity) | (QDim::NotQuantity, QDim::Known(a)) => {
QDim::Known(a)
}
_ if l.is_quantity() || r.is_quantity() => QDim::Unknown,
_ => QDim::NotQuantity,
}),
BinaryOpKind::Divide => Ok(match (l, r) {
(QDim::Known(a), QDim::Known(b)) => QDim::Known(a.div(b)),
(QDim::Known(a), QDim::NotQuantity) => QDim::Known(a),
_ if l.is_quantity() || r.is_quantity() => QDim::Unknown,
_ => QDim::NotQuantity,
}),
_ => Ok(QDim::NotQuantity),
}
}
Expr::Not { operand } => {
self.infer(operand)?;
Ok(QDim::NotQuantity)
}
Expr::Index { collection, index } => {
self.infer(collection)?;
self.infer(index)?;
Ok(QDim::Unknown)
}
Expr::Slice { collection, start, end } => {
self.infer(collection)?;
self.infer(start)?;
self.infer(end)?;
Ok(QDim::Unknown)
}
Expr::Copy { expr } | Expr::Give { value: expr } => self.infer(expr),
Expr::Length { collection } => {
self.infer(collection)?;
Ok(QDim::NotQuantity)
}
Expr::Contains { collection, value } => {
self.infer(collection)?;
self.infer(value)?;
Ok(QDim::NotQuantity)
}
_ => Ok(QDim::NotQuantity),
}
}
fn propagate(&self, l: QDim, r: QDim) -> QDim {
match (l, r) {
(QDim::Known(d), _) | (_, QDim::Known(d)) => QDim::Known(d),
_ if l.is_quantity() || r.is_quantity() => QDim::Unknown,
_ => QDim::NotQuantity,
}
}
fn mismatch_err(&self, verb: &str, a: Dimension, b: Dimension) -> DimensionError {
DimensionError {
message: format!("cannot {verb} quantities of different dimensions ({a} vs {b})"),
span: Span::default(),
}
}
}