tancore/math/
complex.rs

1// #todo implement with Tan
2// #todo consider putting in Prelude.
3
4// (let z (Complex 1.0 0.3))
5// (let r (* z z))
6
7// #todo Should definitely stay out of core.
8
9use tan::{
10    context::Context,
11    error::Error,
12    expr::{annotate_type, has_type_annotation, Expr},
13    util::{expect_lock_read, module_util::require_module},
14};
15
16// #todo how to extend to quaternions and octonions.
17
18fn try_complex(expr: &Expr) -> Option<(f64, f64)> {
19    if !has_type_annotation(expr, "Complex") {
20        return None;
21    }
22
23    let Expr::Array(v) = expr.unpack() else {
24        return None;
25    };
26
27    let v = expect_lock_read(v);
28
29    // #todo use ForeignStruct.
30    let re = v[0].as_float().unwrap();
31    let im = v[1].as_float().unwrap();
32
33    Some((re, im))
34}
35
36#[inline(always)]
37fn make_complex(re: impl Into<Expr>, im: impl Into<Expr>) -> Expr {
38    let expr = Expr::array(vec![re.into(), im.into()]);
39    annotate_type(expr, "Complex")
40}
41
42// (Complex)
43// (Complex re)
44// (Complex re im)
45pub fn complex_new(args: &[Expr]) -> Result<Expr, Error> {
46    // #todo consider a ForeignStruct.
47
48    let re = args.first().unwrap_or_else(|| &Expr::Float(0.0));
49    let im = args.get(1).unwrap_or_else(|| &Expr::Float(0.0));
50
51    // #todo optimize the clones.
52    Ok(make_complex(re.clone(), im.clone()))
53}
54
55pub fn complex_add(args: &[Expr]) -> Result<Expr, Error> {
56    // #todo initialize with first arg, to save one op.
57
58    let mut re_acc = 0.0;
59    let mut im_acc = 0.0;
60
61    for arg in args {
62        let Some((re, im)) = try_complex(arg) else {
63            return Err(Error::invalid_arguments(
64                &format!("{arg} is not a Complex"),
65                arg.range(),
66            ));
67        };
68        re_acc += re;
69        im_acc += im;
70    }
71
72    Ok(make_complex(re_acc, im_acc))
73}
74
75// (ac - bd) + (ad + bc)i.
76// #todo only supports 2 arguments for the moment.
77// #todo extract multiplication helper and support multiple arguments.
78pub fn complex_mul(args: &[Expr]) -> Result<Expr, Error> {
79    let [c, z] = args else {
80        return Err(Error::invalid_arguments("requires two arguments", None));
81    };
82
83    let Some((a, b)) = try_complex(c) else {
84        return Err(Error::invalid_arguments(
85            &format!("{c} is not a Complex"),
86            c.range(),
87        ));
88    };
89
90    let Some((c, d)) = try_complex(z) else {
91        return Err(Error::invalid_arguments(
92            &format!("{z} is not a Complex"),
93            c.range(),
94        ));
95    };
96
97    // complex multiplication formula: (ac - bd) + (ad + bc)i.
98
99    let re = (a * c) - (b * d);
100    let im = (a * d) + (b * c);
101
102    Ok(make_complex(re, im))
103}
104
105// |z| = √(a² + b²)
106pub fn complex_abs(args: &[Expr]) -> Result<Expr, Error> {
107    let [c] = args else {
108        return Err(Error::invalid_arguments("requires `self` argument", None));
109    };
110
111    let Some((a, b)) = try_complex(c) else {
112        return Err(Error::invalid_arguments(
113            &format!("{c} is not a Complex"),
114            c.range(),
115        ));
116    };
117
118    // complex abs formula: |z| = √(a² + b²)
119    // #insight the complex abs is the 'magnitude' of the complex number.
120
121    // #todo detect and optimize pure real (a + 0i) and pure imaginary (0 + bi) cases.
122
123    let magnitude = ((a * a) + (b * b)).sqrt();
124
125    Ok(magnitude.into())
126}
127
128pub fn setup_lib_math_complex(context: &mut Context) {
129    // #todo skip the `math/` prefix?
130    let module = require_module("math/complex", context);
131
132    // #todo make type-paremetric.
133    // #todo better name?
134    // (let z (Complex 1.0 0.3))
135    module.insert_invocable("Complex", Expr::foreign_func(&complex_new));
136
137    module.insert_invocable("+$$Complex$$Complex", Expr::foreign_func(&complex_add));
138    module.insert_invocable(
139        "+$$Complex$$Complex$$Complex",
140        Expr::foreign_func(&complex_add),
141    );
142
143    module.insert_invocable("*$$Complex$$Complex", Expr::foreign_func(&complex_mul));
144
145    // #todo move this to arithmetic or something similar.
146    module.insert_invocable("abs", Expr::foreign_func(&complex_abs));
147    module.insert_invocable("abs$$Complex", Expr::foreign_func(&complex_abs));
148
149    // #todo also consider Complex:one, Complex:zero ~~ (Complex :zero) -> Complex:zero
150    // #todo `Complex/one`
151    // #todo `Complex/zero`
152    // #todo `Complex/re`
153    // #todo `Complex/im`
154    // #todo `(re c)`, `(re-of c)`, `(get-re c)`
155    // #todo `(im c)`, `(im-of c)`, `(get-im c)`
156}
tancore/math/complex.rs

tancore/math/
complex.rs
