Skip to main content

alphawinnow/
operators.rs

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    /// Parse and validate a portable public catalog.
108    ///
109    /// # Errors
110    /// Returns a JSON or catalog-invariant diagnostic.
111    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    /// Validate names, kinds, domains, weights, and uniqueness.
118    ///
119    /// # Errors
120    /// Returns the first violated catalog invariant.
121    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    /// SHA-256 of normalized catalog data consumed by the tool.
176    #[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/// The embedded, synthetic, license-safe public catalog.
266///
267/// # Panics
268/// Panics only when the catalog embedded at compile time is invalid, which is
269/// guarded by tests and prevents shipping an inconsistent binary.
270#[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
284/// Construct a typed call using the embedded public catalog.
285///
286/// # Errors
287/// Returns an error for unknown operators, invalid arity, kinds, or parameters.
288pub 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
296/// Construct a typed call using an explicitly supplied catalog.
297///
298/// # Errors
299/// Returns an error for unknown operators, invalid arity, kinds, or parameters.
300///
301/// # Panics
302/// A programmatically constructed catalog must pass [`Catalog::validate`]
303/// before use; violating that contract can leave missing signature slots.
304pub 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}