Skip to main content

konjure_lang/
methods.rs

1//! Checked signatures for native operations owned by their receiver types.
2//!
3//! These declarations deliberately exclude the implicit receiver.  The checker
4//! supplies that receiver when it resolves a method expression, while the
5//! runtime can dispatch [`operation`] using the same canonical name.
6
7use crate::Type;
8
9/// A native operation after method resolution has selected its receiver.
10#[derive(Clone, Copy, Debug, Eq, PartialEq)]
11pub enum Operation {
12    /// Count a receiver's elements, bytes, or scalar values.
13    Len,
14    /// Encode UTF-8 text.
15    Utf8,
16    /// Decode UTF-8 bytes.
17    Decode,
18    /// Append one value to a list.
19    Append,
20    /// Read one list, text, byte, or tensor element.
21    Get,
22    /// Return tensor dimensions.
23    Shape,
24    /// Reinterpret a tensor with a supplied shape.
25    Reshape,
26    /// Reorder tensor axes.
27    Transpose,
28    /// Sum tensor elements.
29    Sum,
30    /// Multiply two tensors.
31    Matmul,
32    /// Convert numeric scalar or tensor data to an explicit dtype.
33    Convert,
34    /// Compute the named numeric operation.
35    Numeric(&'static str),
36}
37
38/// Returns the runtime operation corresponding to a receiver-owned method.
39pub fn operation(name: &str) -> Option<Operation> {
40    Some(match name {
41        "len" => Operation::Len,
42        "utf8" => Operation::Utf8,
43        "decode" => Operation::Decode,
44        "append" => Operation::Append,
45        "get" => Operation::Get,
46        "shape" => Operation::Shape,
47        "reshape" => Operation::Reshape,
48        "transpose" => Operation::Transpose,
49        "sum" => Operation::Sum,
50        "matmul" => Operation::Matmul,
51        "convert" => Operation::Convert,
52        "sin" => Operation::Numeric("sin"),
53        "cos" => Operation::Numeric("cos"),
54        "sqrt" => Operation::Numeric("sqrt"),
55        "abs" => Operation::Numeric("abs"),
56        "floor" => Operation::Numeric("floor"),
57        "ceil" => Operation::Numeric("ceil"),
58        "min" => Operation::Numeric("min"),
59        "max" => Operation::Numeric("max"),
60        "clamp" => Operation::Numeric("clamp"),
61        "pow" => Operation::Numeric("pow"),
62        _ => return None,
63    })
64}
65
66/// Names which are no longer global functions because their receiver owns them.
67pub fn is_owned_global(name: &str) -> bool {
68    matches!(
69        name,
70        "len"
71            | "decode"
72            | "utf8"
73            | "shape"
74            | "reshape"
75            | "transpose"
76            | "sum"
77            | "convert"
78            | "matmul"
79            | "append"
80            | "sin"
81            | "cos"
82            | "sqrt"
83            | "abs"
84            | "floor"
85            | "ceil"
86            | "min"
87            | "max"
88            | "clamp"
89            | "pow"
90    )
91}
92
93/// Returns every applicable method signature, excluding the implicit receiver.
94///
95/// Multiple signatures are possible only for `Tensor[T].transpose`; callers
96/// select the one whose explicit parameter count matches the authored call.
97pub fn signatures(receiver: &Type, name: &str, type_arg: Option<&Type>) -> Vec<Type> {
98    let number = Type::named("f64");
99    let text = Type::named("Str");
100    let binary = Type::named("Bin");
101    let error = Type::named("DataError");
102    let list_number = Type::List(Box::new(number.clone()));
103    let result = |value| applied("Res", vec![value, error.clone()]);
104    let option = |value| applied("Opt", vec![value]);
105
106    match receiver {
107        Type::Named(value) if value == "Str" => match name {
108            "len" => vec![Type::function(vec![], number)],
109            "utf8" => vec![Type::function(vec![], binary)],
110            "get" => vec![Type::function(vec![Type::named("f64")], option(text))],
111            _ => vec![],
112        },
113        Type::Named(value) if value == "Bin" => match name {
114            "len" => vec![Type::function(vec![], number)],
115            "decode" => vec![Type::function(vec![], result(text))],
116            "get" => vec![Type::function(
117                vec![Type::named("f64")],
118                option(Type::named("u8")),
119            )],
120            _ => vec![],
121        },
122        Type::List(element) => match name {
123            "len" => vec![Type::function(vec![], number)],
124            "append" => vec![Type::function(
125                vec![element.as_ref().clone()],
126                Type::List(element.clone()),
127            )],
128            "get" => vec![Type::function(
129                vec![Type::named("f64")],
130                option(element.as_ref().clone()),
131            )],
132            _ => vec![],
133        },
134        Type::Tensor(element) => match name {
135            "len" => vec![Type::function(vec![], number)],
136            "shape" => vec![Type::function(vec![], list_number.clone())],
137            "reshape" => vec![Type::function(
138                vec![list_number.clone()],
139                result(Type::Tensor(element.clone())),
140            )],
141            "transpose" => vec![
142                Type::function(vec![], result(Type::Tensor(element.clone()))),
143                Type::function(
144                    vec![list_number.clone()],
145                    result(Type::Tensor(element.clone())),
146                ),
147            ],
148            "sum" => vec![Type::function(vec![], result(element.as_ref().clone()))],
149            "matmul" => vec![Type::function(
150                vec![Type::Tensor(element.clone())],
151                result(Type::Tensor(element.clone())),
152            )],
153            "get" => vec![Type::function(
154                vec![list_number],
155                option(element.as_ref().clone()),
156            )],
157            "convert" if is_dtype(type_arg) => vec![Type::function(
158                vec![],
159                result(Type::Tensor(Box::new(
160                    type_arg.expect("checked above").clone(),
161                ))),
162            )],
163            _ => vec![],
164        },
165        Type::Named(value) if value == "f64" => match name {
166            "sin" | "cos" | "sqrt" | "abs" | "floor" | "ceil" => {
167                vec![Type::function(vec![], result(number.clone()))]
168            }
169            "min" | "max" | "pow" => vec![Type::function(vec![number.clone()], result(number))],
170            "clamp" => vec![Type::function(
171                vec![number.clone(), number.clone()],
172                result(number),
173            )],
174            "convert" if is_dtype(type_arg) => vec![Type::function(
175                vec![],
176                result(type_arg.expect("checked above").clone()),
177            )],
178            _ => vec![],
179        },
180        Type::Named(_) if is_dtype(Some(receiver)) && name == "convert" && is_dtype(type_arg) => {
181            vec![Type::function(
182                vec![],
183                result(type_arg.expect("checked above").clone()),
184            )]
185        }
186        _ => vec![],
187    }
188}
189
190/// Returns the only applicable signature, excluding the implicit receiver.
191///
192/// Use [`signatures`] when choosing an overload such as `transpose` by arity.
193pub fn signature(receiver: &Type, name: &str, type_arg: Option<&Type>) -> Option<Type> {
194    let mut signatures = signatures(receiver, name, type_arg).into_iter();
195    let signature = signatures.next()?;
196    signatures.next().is_none().then_some(signature)
197}
198
199fn applied(name: &str, arguments: Vec<Type>) -> Type {
200    Type::Applied {
201        name: name.into(),
202        arguments,
203    }
204}
205
206fn is_dtype(ty: Option<&Type>) -> bool {
207    ty.is_some_and(|ty| crate::data::dtype(ty).is_some())
208}
209
210#[cfg(test)]
211mod tests {
212    use super::*;
213
214    #[test]
215    fn data_methods_have_receiver_specialized_results() {
216        let tensor = Type::Tensor(Box::new(Type::named("f32")));
217        assert_eq!(
218            signature(&Type::named("Bin"), "decode", None),
219            Some(Type::function(
220                vec![],
221                Type::Applied {
222                    name: "Res".into(),
223                    arguments: vec![Type::named("Str"), Type::named("DataError")],
224                },
225            ))
226        );
227        assert_eq!(
228            signature(&Type::List(Box::new(Type::named("u8"))), "append", None),
229            Some(Type::function(
230                vec![Type::named("u8")],
231                Type::List(Box::new(Type::named("u8"))),
232            ))
233        );
234        assert_eq!(
235            signature(&tensor, "convert", Some(&Type::named("u16"))),
236            Some(Type::function(
237                vec![],
238                Type::Applied {
239                    name: "Res".into(),
240                    arguments: vec![
241                        Type::Tensor(Box::new(Type::named("u16"))),
242                        Type::named("DataError"),
243                    ],
244                },
245            ))
246        );
247    }
248
249    #[test]
250    fn transpose_keeps_its_two_arity_overloads_explicit() {
251        let tensor = Type::Tensor(Box::new(Type::named("f64")));
252        assert_eq!(signature(&tensor, "transpose", None), None);
253        assert_eq!(signatures(&tensor, "transpose", None).len(), 2);
254    }
255
256    #[test]
257    fn owned_globals_are_not_left_in_the_function_namespace() {
258        for name in [
259            "len",
260            "decode",
261            "utf8",
262            "shape",
263            "reshape",
264            "transpose",
265            "sum",
266            "convert",
267            "matmul",
268            "append",
269        ] {
270            assert!(is_owned_global(name));
271            assert!(operation(name).is_some());
272        }
273        assert!(!is_owned_global("bytes"));
274        assert_eq!(operation("bytes"), None);
275    }
276}