geo_aid_script/unroll/library/
complex.rs

1//! All functions for complex number manipulation
2
3use num_rational::Ratio;
4
5use crate::{
6    parser::Type,
7    token::number::ProcNum,
8    unroll::{AnyExpr, Number},
9};
10
11use super::{prelude::*, Overload};
12
13struct Real;
14
15impl Overload for Real {
16    fn get_returned_type(&self, params: &[AnyExpr]) -> Option<Type> {
17        if params.len() != 1 {
18            return None;
19        }
20
21        params[0].can_convert_to_scalar(None).map(Type::Number)
22    }
23
24    fn unroll(
25        &self,
26        params: Vec<AnyExpr>,
27        context: &mut CompileContext,
28        props: Properties,
29    ) -> AnyExpr {
30        let ty = self.get_returned_type(&params).unwrap();
31        context
32            .real_display(
33                params[0]
34                    .clone_without_node()
35                    .convert_to(ty, context)
36                    .to_scalar()
37                    .unwrap(),
38                props,
39            )
40            .into()
41    }
42}
43
44struct Imaginary;
45
46impl Overload for Imaginary {
47    fn get_returned_type(&self, params: &[AnyExpr]) -> Option<Type> {
48        if params.len() != 1 {
49            return None;
50        }
51
52        params[0].can_convert_to_scalar(None).map(Type::Number)
53    }
54
55    fn unroll(
56        &self,
57        params: Vec<AnyExpr>,
58        context: &mut CompileContext,
59        props: Properties,
60    ) -> AnyExpr {
61        let ty = self.get_returned_type(&params).unwrap();
62        context
63            .imaginary_display(
64                params[0]
65                    .clone_without_node()
66                    .convert_to(ty, context)
67                    .to_scalar()
68                    .unwrap(),
69                props,
70            )
71            .into()
72    }
73}
74
75struct Conjugate;
76
77impl Overload for Conjugate {
78    fn get_returned_type(&self, params: &[AnyExpr]) -> Option<Type> {
79        if params.len() != 1 {
80            return None;
81        }
82
83        params[0].can_convert_to_scalar(None).map(Type::Number)
84    }
85
86    fn unroll(
87        &self,
88        params: Vec<AnyExpr>,
89        context: &mut CompileContext,
90        props: Properties,
91    ) -> AnyExpr {
92        let ty = self.get_returned_type(&params).unwrap();
93        let v = params[0]
94            .clone_without_node()
95            .convert_to(ty, context)
96            .to_scalar()
97            .unwrap();
98
99        let real = context.real(v.clone_without_node());
100        let imaginary = context.imaginary(v.clone_without_node());
101        let i_imaginary = context.mult(
102            imaginary,
103            Expr::new_spanless(Number {
104                unit: Some(unit::SCALAR),
105                data: NumberData::Number(ProcNum::i()),
106            }),
107        );
108
109        context.sub_display(real, i_imaginary, props).into()
110    }
111}
112
113fn vector(
114    origin: Expr<Point>,
115    target: Expr<Point>,
116    context: &CompileContext,
117    props: Properties,
118) -> Distance {
119    context
120        .sub_display(
121            context.to_complex(target),
122            context.to_complex(origin),
123            props,
124        )
125        .into()
126}
127
128/// Absolute value of a number, real or complex.
129#[must_use]
130pub fn abs(num: Expr<Number>, context: &CompileContext, props: Properties) -> Expr<Number> {
131    let re2 = context.mult(
132        context.real(num.clone_without_node()),
133        context.real(num.clone_without_node()),
134    );
135    let im2 = context.mult(
136        context.imaginary(num.clone_without_node()),
137        context.imaginary(num.clone_without_node()),
138    );
139    let norm = context.add(re2, im2);
140    context.pow_display(norm, Ratio::new(1, 2), props)
141}
142
143#[derive(Debug)]
144struct Abs;
145
146impl Overload for Abs {
147    fn get_returned_type(&self, params: &[AnyExpr]) -> Option<Type> {
148        if params.len() != 1 {
149            return None;
150        }
151
152        params[0].can_convert_to_scalar(None).map(Type::Number)
153    }
154
155    fn unroll(
156        &self,
157        mut params: Vec<AnyExpr>,
158        context: &mut CompileContext,
159        props: Properties,
160    ) -> AnyExpr {
161        let num = params
162            .swap_remove(0)
163            .convert_to(Type::Number(None), context)
164            .to_scalar()
165            .unwrap();
166
167        abs(num, context, props).into()
168    }
169}
170
171/// Register all the functions that need to be registered.
172pub fn register(library: &mut Library) {
173    library
174        .add(
175            Function::new("real")
176                .alias("re")
177                .alias_method(ty::SCALAR_UNKNOWN, "real")
178                .alias_method(ty::SCALAR_UNKNOWN, "re")
179                .overload(Real)
180                .overload(|context: &CompileContext, props| {
181                    Unitless::from(context.free_scalar_display(props))
182                }),
183        )
184        .add(
185            Function::new("imaginary")
186                .alias("im")
187                .alias_method(ty::SCALAR_UNKNOWN, "imaginary")
188                .alias_method(ty::SCALAR_UNKNOWN, "im")
189                .overload(Imaginary),
190        )
191        .add(
192            Function::new("tocomplex")
193                .alias_method(ty::POINT, "tocomplex")
194                .overload(|point: Expr<Point>, context: &CompileContext, props| {
195                    Distance::from(context.to_complex_display(point, props))
196                }),
197        )
198        .add(
199            Function::new("conjugate")
200                .alias_method(ty::SCALAR_UNKNOWN, "conjugate")
201                .overload(Conjugate),
202        )
203        .add(
204            Function::new("vector")
205                .alias("vec")
206                .alias_method(ty::collection(2), "vector")
207                .alias_method(ty::collection(2), "vec")
208                .overload(vector)
209                .overload(|mut col: Pc<2>, context: &CompileContext, props| {
210                    vector(
211                        col.index_with_node(0),
212                        col.index_without_node(1),
213                        context,
214                        props,
215                    )
216                }),
217        )
218        .add(
219            Function::new("abs")
220                .alias("module")
221                .alias_method(Type::Number(None), "abs")
222                .alias_method(Type::Number(None), "module")
223                .overload(Abs),
224        );
225}
geo_aid_script/unroll/library/complex.rs

geo_aid_script/unroll/library/
complex.rs
