use tan::{
context::Context,
error::Error,
expr::{annotate_type, has_type_annotation, Expr},
util::{expect_lock_read, module_util::require_module},
};
fn try_complex(expr: &Expr) -> Option<(f64, f64)> {
if !has_type_annotation(expr, "Complex") {
return None;
}
let Expr::Array(v) = expr.unpack() else {
return None;
};
let v = expect_lock_read(v);
let re = v[0].as_float().unwrap();
let im = v[1].as_float().unwrap();
Some((re, im))
}
#[inline(always)]
fn make_complex(re: impl Into<Expr>, im: impl Into<Expr>) -> Expr {
let expr = Expr::array(vec![re.into(), im.into()]);
annotate_type(expr, "Complex")
}
pub fn complex_new(args: &[Expr]) -> Result<Expr, Error> {
let re = args.first().unwrap_or_else(|| &Expr::Float(0.0));
let im = args.get(1).unwrap_or_else(|| &Expr::Float(0.0));
Ok(make_complex(re.clone(), im.clone()))
}
pub fn complex_add(args: &[Expr]) -> Result<Expr, Error> {
let mut re_acc = 0.0;
let mut im_acc = 0.0;
for arg in args {
let Some((re, im)) = try_complex(arg) else {
return Err(Error::invalid_arguments(
&format!("{arg} is not a Complex"),
arg.range(),
));
};
re_acc += re;
im_acc += im;
}
Ok(make_complex(re_acc, im_acc))
}
pub fn complex_mul(args: &[Expr]) -> Result<Expr, Error> {
let [c, z] = args else {
return Err(Error::invalid_arguments("requires two arguments", None));
};
let Some((a, b)) = try_complex(c) else {
return Err(Error::invalid_arguments(
&format!("{c} is not a Complex"),
c.range(),
));
};
let Some((c, d)) = try_complex(z) else {
return Err(Error::invalid_arguments(
&format!("{z} is not a Complex"),
c.range(),
));
};
let re = (a * c) - (b * d);
let im = (a * d) + (b * c);
Ok(make_complex(re, im))
}
pub fn complex_abs(args: &[Expr]) -> Result<Expr, Error> {
let [c] = args else {
return Err(Error::invalid_arguments("requires `self` argument", None));
};
let Some((a, b)) = try_complex(c) else {
return Err(Error::invalid_arguments(
&format!("{c} is not a Complex"),
c.range(),
));
};
let magnitude = ((a * a) + (b * b)).sqrt();
Ok(magnitude.into())
}
pub fn setup_lib_math_complex(context: &mut Context) {
let module = require_module("math/complex", context);
module.insert_invocable("Complex", Expr::foreign_func(&complex_new));
module.insert_invocable("+$$Complex$$Complex", Expr::foreign_func(&complex_add));
module.insert_invocable(
"+$$Complex$$Complex$$Complex",
Expr::foreign_func(&complex_add),
);
module.insert_invocable("*$$Complex$$Complex", Expr::foreign_func(&complex_mul));
module.insert_invocable("abs", Expr::foreign_func(&complex_abs));
module.insert_invocable("abs$$Complex", Expr::foreign_func(&complex_abs));
}