1use std::{collections::BTreeMap, sync::OnceLock};
2
3use serde::{Deserialize, Serialize};
4
5use crate::{
6 ast::{Expr, ExprKind},
7 canonical::digest,
8};
9
10pub const CATALOG_SCHEMA: u32 = 1;
11const BUILTIN_CATALOG_JSON: &str = include_str!("../catalog/public-v1.json");
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
14#[serde(tag = "kind", rename_all = "snake_case")]
15pub enum Arity {
16 Unary,
17 Binary,
18 Exact { count: usize },
19 Variadic { min: usize, max: usize },
20}
21
22#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
23#[serde(tag = "kind", rename_all = "snake_case")]
24pub enum ValueDomain {
25 Any,
26 Window { min: u16, max: u16 },
27 WindowSet { values: Vec<u16> },
28 NonZeroScalar,
29}
30
31#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
32pub struct ArgumentSpec {
33 pub kinds: Vec<ExprKind>,
34 #[serde(default = "any_domain")]
35 pub domain: ValueDomain,
36}
37
38#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
39pub struct KeywordSpec {
40 pub name: String,
41 pub min: f64,
42 pub max: f64,
43 pub step: f64,
44 #[serde(default = "required_keyword")]
45 pub required: bool,
46}
47
48#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
49#[serde(tag = "kind", rename_all = "snake_case")]
50pub enum KeywordConstraint {
51 LessThan { left: String, right: String },
52}
53
54#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
55pub struct OperatorSpec {
56 pub name: String,
57 pub arity: Arity,
58 pub inputs: Vec<ArgumentSpec>,
59 #[serde(default)]
60 pub required_input_kinds: Vec<ExprKind>,
61 pub output: ExprKind,
62 #[serde(default)]
63 pub commutative: bool,
64 #[serde(default)]
65 pub associative: bool,
66 #[serde(default)]
67 pub keywords: Vec<KeywordSpec>,
68 #[serde(default)]
69 pub keyword_constraints: Vec<KeywordConstraint>,
70 pub generation_weight: u16,
71 #[serde(default)]
72 pub parse_only_reason: Option<String>,
73}
74
75#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
76pub struct FieldSpec {
77 pub name: String,
78 pub kind: ExprKind,
79 pub family: String,
80 pub allowed_roles: Vec<String>,
81}
82
83#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
84pub struct GroupSpec {
85 pub name: String,
86 pub kind: ExprKind,
87}
88
89#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
90pub struct ScalarDomain {
91 pub min: f64,
92 pub max: f64,
93 pub step: f64,
94}
95
96#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
97pub struct Catalog {
98 pub schema: u32,
99 pub name: String,
100 pub operators: Vec<OperatorSpec>,
101 pub fields: Vec<FieldSpec>,
102 pub groups: Vec<GroupSpec>,
103 pub scalar_domain: ScalarDomain,
104}
105
106impl Catalog {
107 pub fn from_json(input: &str) -> Result<Self, String> {
112 let catalog: Self = serde_json::from_str(input).map_err(|error| error.to_string())?;
113 catalog.validate()?;
114 Ok(catalog)
115 }
116
117 pub fn validate(&self) -> Result<(), String> {
122 if self.schema != CATALOG_SCHEMA {
123 return Err(format!("unsupported catalog schema {}", self.schema));
124 }
125 if self.operators.is_empty() || self.fields.is_empty() || self.groups.is_empty() {
126 return Err("catalog operators, fields, and groups must be non-empty".to_owned());
127 }
128 if !valid_domain(
129 self.scalar_domain.min,
130 self.scalar_domain.max,
131 self.scalar_domain.step,
132 ) {
133 return Err("invalid scalar domain".to_owned());
134 }
135 unique_names(
136 self.operators.iter().map(|item| item.name.as_str()),
137 "operator",
138 )?;
139 unique_names(self.fields.iter().map(|item| item.name.as_str()), "field")?;
140 unique_names(self.groups.iter().map(|item| item.name.as_str()), "group")?;
141 for field in &self.fields {
142 if field.kind != ExprKind::Signal
143 || field.family.is_empty()
144 || field.allowed_roles.is_empty()
145 {
146 return Err(format!("invalid public field `{}`", field.name));
147 }
148 }
149 for group in &self.groups {
150 if group.kind != ExprKind::Group {
151 return Err(format!("group `{}` must have group kind", group.name));
152 }
153 }
154 for operator in &self.operators {
155 validate_operator_spec(operator)?;
156 }
157 Ok(())
158 }
159
160 #[must_use]
161 pub fn lookup(&self, name: &str) -> Option<&OperatorSpec> {
162 self.operators.iter().find(|spec| spec.name == name)
163 }
164
165 #[must_use]
166 pub fn field(&self, name: &str) -> Option<&FieldSpec> {
167 self.fields.iter().find(|field| field.name == name)
168 }
169
170 #[must_use]
171 pub fn group(&self, name: &str) -> Option<&GroupSpec> {
172 self.groups.iter().find(|group| group.name == name)
173 }
174
175 #[must_use]
177 pub fn checksum(&self) -> String {
178 digest(&serde_json::to_string(self).unwrap_or_default())
179 }
180}
181
182const fn any_domain() -> ValueDomain {
183 ValueDomain::Any
184}
185
186const fn required_keyword() -> bool {
187 true
188}
189
190fn valid_name(name: &str) -> bool {
191 !name.is_empty()
192 && name.chars().enumerate().all(|(index, ch)| {
193 ch == '_' || ch.is_ascii_alphanumeric() && (index > 0 || !ch.is_ascii_digit())
194 })
195}
196
197fn unique_names<'a>(names: impl Iterator<Item = &'a str>, label: &str) -> Result<(), String> {
198 let mut seen = std::collections::BTreeSet::new();
199 for name in names {
200 if !valid_name(name) || !seen.insert(name) {
201 return Err(format!("invalid or duplicate {label} name `{name}`"));
202 }
203 }
204 Ok(())
205}
206
207fn valid_domain(min: f64, max: f64, step: f64) -> bool {
208 min.is_finite() && max.is_finite() && step.is_finite() && min < max && step > 0.0
209}
210
211fn validate_operator_spec(spec: &OperatorSpec) -> Result<(), String> {
212 if !valid_name(&spec.name) {
213 return Err(format!("invalid operator name `{}`", spec.name));
214 }
215 let valid_input_count = match spec.arity {
216 Arity::Unary => spec.inputs.len() == 1,
217 Arity::Binary => spec.inputs.len() == 2,
218 Arity::Exact { count } => count > 0 && spec.inputs.len() == count,
219 Arity::Variadic { min, max } => min > 0 && min <= max && spec.inputs.len() == 1,
220 };
221 if !valid_input_count || spec.inputs.iter().any(|input| input.kinds.is_empty()) {
222 return Err(format!("invalid input signature for `{}`", spec.name));
223 }
224 for input in &spec.inputs {
225 match &input.domain {
226 ValueDomain::Window { min, max } if *min == 0 || min > max => {
227 return Err(format!("invalid window domain for `{}`", spec.name));
228 }
229 ValueDomain::WindowSet { values }
230 if values.is_empty()
231 || values.contains(&0)
232 || values.windows(2).any(|pair| pair[0] >= pair[1]) =>
233 {
234 return Err(format!("invalid window set for `{}`", spec.name));
235 }
236 _ => {}
237 }
238 }
239 if spec.generation_weight == 0 && spec.parse_only_reason.is_none() {
240 return Err(format!(
241 "parse-only operator `{}` requires a documented reason",
242 spec.name
243 ));
244 }
245 for keyword in &spec.keywords {
246 if !valid_name(&keyword.name) || !valid_domain(keyword.min, keyword.max, keyword.step) {
247 return Err(format!("invalid keyword domain for `{}`", spec.name));
248 }
249 }
250 unique_names(
251 spec.keywords.iter().map(|item| item.name.as_str()),
252 "keyword",
253 )?;
254 for constraint in &spec.keyword_constraints {
255 let KeywordConstraint::LessThan { left, right } = constraint;
256 if !spec.keywords.iter().any(|item| item.name == *left)
257 || !spec.keywords.iter().any(|item| item.name == *right)
258 {
259 return Err(format!("invalid keyword constraint for `{}`", spec.name));
260 }
261 }
262 Ok(())
263}
264
265#[must_use]
271pub fn builtin_catalog() -> &'static Catalog {
272 static CATALOG: OnceLock<Catalog> = OnceLock::new();
273 CATALOG.get_or_init(|| {
274 Catalog::from_json(BUILTIN_CATALOG_JSON)
275 .expect("the embedded public catalog must remain valid")
276 })
277}
278
279#[must_use]
280pub fn lookup(name: &str) -> Option<&'static OperatorSpec> {
281 builtin_catalog().lookup(name)
282}
283
284pub fn build_call(
289 name: &str,
290 args: Vec<Expr>,
291 kwargs: BTreeMap<String, Expr>,
292) -> Result<Expr, String> {
293 build_call_with_catalog(builtin_catalog(), name, args, kwargs)
294}
295
296pub fn build_call_with_catalog(
305 catalog: &Catalog,
306 name: &str,
307 args: Vec<Expr>,
308 kwargs: BTreeMap<String, Expr>,
309) -> Result<Expr, String> {
310 let spec = catalog
311 .lookup(name)
312 .ok_or_else(|| format!("unknown operator `{name}`"))?;
313 validate_arity(spec, args.len())?;
314 validate_args(spec, &args)?;
315 validate_kwargs(spec, &kwargs)?;
316
317 Ok(match spec.arity {
318 Arity::Unary => {
319 let mut args = args.into_iter();
320 let Some(arg) = args.next() else {
321 return Err(format!("operator `{name}` requires one argument"));
322 };
323 Expr::UnaryCall {
324 op: name.to_owned(),
325 arg: Box::new(arg),
326 kwargs,
327 kind: spec.output,
328 }
329 }
330 Arity::Binary => {
331 let mut args = args.into_iter();
332 let Some(left) = args.next() else {
333 return Err(format!("operator `{name}` requires two arguments"));
334 };
335 let Some(right) = args.next() else {
336 return Err(format!("operator `{name}` requires two arguments"));
337 };
338 Expr::BinaryCall {
339 op: name.to_owned(),
340 left: Box::new(left),
341 right: Box::new(right),
342 kwargs,
343 kind: spec.output,
344 }
345 }
346 Arity::Exact { .. } | Arity::Variadic { .. } => Expr::VariadicCall {
347 op: name.to_owned(),
348 args,
349 kwargs,
350 kind: spec.output,
351 },
352 })
353}
354
355fn validate_arity(spec: &OperatorSpec, actual: usize) -> Result<(), String> {
356 let valid = match spec.arity {
357 Arity::Unary => actual == 1,
358 Arity::Binary => actual == 2,
359 Arity::Exact { count } => actual == count,
360 Arity::Variadic { min, max } => (min..=max).contains(&actual),
361 };
362 valid
363 .then_some(())
364 .ok_or_else(|| format!("operator `{}` received invalid arity {actual}", spec.name))
365}
366
367fn validate_args(spec: &OperatorSpec, args: &[Expr]) -> Result<(), String> {
368 for (index, arg) in args.iter().enumerate() {
369 let input = match spec.arity {
370 Arity::Variadic { .. } => &spec.inputs[0],
371 _ => &spec.inputs[index],
372 };
373 if !input.kinds.contains(&arg.kind()) {
374 return Err(format!(
375 "operator `{}` argument {} does not accept {:?}",
376 spec.name,
377 index + 1,
378 arg.kind()
379 ));
380 }
381 validate_value_domain(&spec.name, index, arg, &input.domain)?;
382 }
383 for required in &spec.required_input_kinds {
384 if !args.iter().any(|arg| arg.kind() == *required) {
385 return Err(format!(
386 "operator `{}` requires at least one {required:?} argument",
387 spec.name
388 ));
389 }
390 }
391 Ok(())
392}
393
394fn validate_value_domain(
395 operator: &str,
396 index: usize,
397 value: &Expr,
398 domain: &ValueDomain,
399) -> Result<(), String> {
400 match domain {
401 ValueDomain::Any => Ok(()),
402 ValueDomain::Window { min, max } => match value {
403 Expr::Scalar { value }
404 if value.fract() == 0.0 && (f64::from(*min)..=f64::from(*max)).contains(value) =>
405 {
406 Ok(())
407 }
408 _ => Err(format!(
409 "operator `{operator}` argument {} must be an integer window in {min}..={max}",
410 index + 1
411 )),
412 },
413 ValueDomain::WindowSet { values } => match value {
414 Expr::Scalar { value }
415 if value.fract() == 0.0
416 && values
417 .iter()
418 .any(|window| f64::from(*window).to_bits() == value.to_bits()) =>
419 {
420 Ok(())
421 }
422 _ => Err(format!(
423 "operator `{operator}` argument {} must be an allowed window",
424 index + 1
425 )),
426 },
427 ValueDomain::NonZeroScalar => match value {
428 Expr::Scalar { value } if *value != 0.0 => Ok(()),
429 _ => Err(format!(
430 "operator `{operator}` argument {} must be a non-zero scalar",
431 index + 1
432 )),
433 },
434 }
435}
436
437fn validate_kwargs(spec: &OperatorSpec, kwargs: &BTreeMap<String, Expr>) -> Result<(), String> {
438 for keyword in &spec.keywords {
439 if keyword.required && !kwargs.contains_key(&keyword.name) {
440 return Err(format!(
441 "operator `{}` requires keyword `{}`",
442 spec.name, keyword.name
443 ));
444 }
445 }
446 if kwargs.len() > spec.keywords.len()
447 || kwargs
448 .keys()
449 .any(|name| !spec.keywords.iter().any(|item| item.name == *name))
450 {
451 return Err(format!(
452 "operator `{}` received unsupported keyword arguments",
453 spec.name
454 ));
455 }
456 for (name, value) in kwargs {
457 let keyword = spec
458 .keywords
459 .iter()
460 .find(|item| item.name == *name)
461 .expect("keyword membership checked");
462 let Expr::Scalar { value } = value else {
463 return Err(format!("keyword `{name}` must be a scalar literal"));
464 };
465 if !(keyword.min..=keyword.max).contains(value) {
466 return Err(format!(
467 "keyword `{name}` must be in {}..={}",
468 keyword.min, keyword.max
469 ));
470 }
471 }
472 for constraint in &spec.keyword_constraints {
473 match constraint {
474 KeywordConstraint::LessThan { left, right } => {
475 let (Expr::Scalar { value: left_value }, Expr::Scalar { value: right_value }) =
476 (&kwargs[left], &kwargs[right])
477 else {
478 unreachable!("keyword literals validated")
479 };
480 if left_value >= right_value {
481 return Err(format!(
482 "operator `{}` requires keyword constraint {left} < {right}",
483 spec.name
484 ));
485 }
486 }
487 }
488 }
489 Ok(())
490}
491
492#[cfg(test)]
493mod tests {
494 use super::*;
495
496 #[test]
497 fn embedded_catalog_is_valid_and_stably_checksummed() {
498 let catalog = builtin_catalog();
499 catalog.validate().unwrap();
500 assert_eq!(catalog.checksum(), catalog.checksum());
501 assert_eq!(catalog.schema, CATALOG_SCHEMA);
502 }
503
504 #[test]
505 fn synthetic_field_is_added_as_data() {
506 let mut catalog = builtin_catalog().clone();
507 catalog.fields.push(FieldSpec {
508 name: "synthetic_quality".to_owned(),
509 kind: ExprKind::Signal,
510 family: "synthetic".to_owned(),
511 allowed_roles: vec!["signal_input".to_owned()],
512 });
513 catalog.validate().unwrap();
514 assert!(catalog.field("synthetic_quality").is_some());
515 crate::parse_expression_with_catalog("rank(synthetic_quality)", &catalog).unwrap();
516 }
517
518 #[test]
519 fn alpha158_like_rolling_operators_are_typed_and_parseable() {
520 for expression in [
521 "abs(ts_delay(close, 20))",
522 "sign(ts_sum(volume, 5))",
523 "ts_var(close, 10)",
524 "ts_skew(close, 10)",
525 "ts_kurt(close, 10)",
526 "ts_max(high, 20)",
527 "ts_min(low, 20)",
528 "ts_median(open, 10)",
529 "ts_mad(close, 10)",
530 "ts_wma(close, 20)",
531 "ts_ema(close, 20)",
532 "ts_cov(close, volume, 20)",
533 "ts_corr(close, volume, 20)",
534 ] {
535 let parsed = crate::parse_expression(expression).unwrap();
536 assert_eq!(parsed.kind(), ExprKind::Signal, "{expression}");
537 }
538 }
539}