1use crate::Type;
8
9#[derive(Clone, Copy, Debug, Eq, PartialEq)]
11pub enum Operation {
12 Len,
14 Utf8,
16 Decode,
18 Append,
20 Get,
22 Shape,
24 Reshape,
26 Transpose,
28 Sum,
30 Matmul,
32 Convert,
34 Numeric(&'static str),
36}
37
38pub 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
66pub 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
93pub 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
190pub 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}