lemon/dsl/
ops.rs

1//! Op vocabulary: DSL surface names ⇄ JSON `Expr` op tags + field layout.
2//! Field declaration order MUST match `spec.rs` so positional args line up.
3
4#[derive(Debug, Clone, Copy)]
5pub enum Field {
6    Expr(&'static str),
7    ExprOpt(&'static str),
8    ExprList(&'static str),
9    Num(&'static str),
10    NumOpt(&'static str),
11    BoolOpt(&'static str),
12    Str(&'static str),
13    StrOpt(&'static str),
14    StrListOpt(&'static str),
15}
16
17pub fn field_name(f: &Field) -> &'static str {
18    match f {
19        Field::Expr(n)
20        | Field::ExprOpt(n)
21        | Field::ExprList(n)
22        | Field::Num(n)
23        | Field::NumOpt(n)
24        | Field::BoolOpt(n)
25        | Field::Str(n)
26        | Field::StrOpt(n)
27        | Field::StrListOpt(n) => n,
28    }
29}
30
31/// JSON-schema kind for a field: `expr` (an `Expr` subtree), `number`, `string`,
32/// `bool`, or `list` (a list of `Expr`) / `list-string`.
33pub fn field_kind(f: &Field) -> &'static str {
34    match f {
35        Field::Expr(_) | Field::ExprOpt(_) => "expr",
36        Field::ExprList(_) => "list",
37        Field::Num(_) | Field::NumOpt(_) => "number",
38        Field::BoolOpt(_) => "bool",
39        Field::Str(_) | Field::StrOpt(_) => "string",
40        Field::StrListOpt(_) => "list-string",
41    }
42}
43
44/// Whether the parser requires this field (required scalar/Expr/list) or accepts
45/// it as optional (`*Opt` variants).
46pub fn field_required(f: &Field) -> bool {
47    matches!(
48        f,
49        Field::Expr(_) | Field::Num(_) | Field::Str(_) | Field::ExprList(_)
50    )
51}
52
53pub struct OpSig {
54    pub tag: &'static str,
55    pub fields: &'static [Field],
56}
57
58/// One row per function-style op. `names[0]` is the canonical DSL name (used by the
59/// printer); the rest are accepted aliases. Operator-style ops (Gt/And/Add/Neg/…)
60/// are NOT here — see `binop_tag`/`prefix_tag`.
61///
62/// `desc` is a one-line, machine-readable description consumed by the schema
63/// generator (`cargo run -p lemon-lang --example gen-schema`); keep it terse and
64/// in sync with the `Expr` doc-comment in `spec.rs`.
65struct Row {
66    names: &'static [&'static str],
67    sig: OpSig,
68    desc: &'static str,
69}
70
71use Field::*;
72
73static ROWS: &[Row] = &[
74    Row {
75        names: &["sma", "average"],
76        sig: OpSig {
77            tag: "Average",
78            fields: &[Expr("of"), Num("n")],
79        },
80        desc: "Simple moving average of `of` over `n` days.",
81    },
82    Row {
83        names: &["ema"],
84        sig: OpSig {
85            tag: "Ema",
86            fields: &[Expr("of"), Num("n")],
87        },
88        desc: "Exponential moving average of `of` over `n` days.",
89    },
90    Row {
91        names: &["std"],
92        sig: OpSig {
93            tag: "Std",
94            fields: &[Expr("of"), Num("n")],
95        },
96        desc: "Rolling standard deviation of `of` over `n` days.",
97    },
98    Row {
99        names: &["rsi"],
100        sig: OpSig {
101            tag: "Rsi",
102            fields: &[Expr("of"), Num("n")],
103        },
104        desc: "Relative Strength Index of `of` over `n` days.",
105    },
106    Row {
107        names: &["pct_change"],
108        sig: OpSig {
109            tag: "PctChange",
110            fields: &[Expr("of"), Num("n")],
111        },
112        desc: "Percentage change of `of` over `n` days.",
113    },
114    Row {
115        names: &["rise"],
116        sig: OpSig {
117            tag: "Rise",
118            fields: &[Expr("of"), Num("n")],
119        },
120        desc: "1 where `of` rose for `n` consecutive days, else 0.",
121    },
122    Row {
123        names: &["fall"],
124        sig: OpSig {
125            tag: "Fall",
126            fields: &[Expr("of"), Num("n")],
127        },
128        desc: "1 where `of` fell for `n` consecutive days, else 0.",
129    },
130    Row {
131        names: &["shift"],
132        sig: OpSig {
133            tag: "Shift",
134            fields: &[Expr("of"), Num("n")],
135        },
136        desc: "`of` lagged forward by `n` days.",
137    },
138    Row {
139        names: &["rolling_max"],
140        sig: OpSig {
141            tag: "RollingMax",
142            fields: &[Expr("of"), Num("n")],
143        },
144        desc: "Rolling maximum of `of` over `n` days.",
145    },
146    Row {
147        names: &["atr"],
148        sig: OpSig {
149            tag: "Atr",
150            fields: &[Expr("high"), Expr("low"), Expr("close"), Num("n")],
151        },
152        desc: "Average True Range over `n` days from high/low/close.",
153    },
154    Row {
155        names: &["natr"],
156        sig: OpSig {
157            tag: "Natr",
158            fields: &[Expr("high"), Expr("low"), Expr("close"), Num("n")],
159        },
160        desc: "Normalized ATR (percent) over `n` days.",
161    },
162    Row {
163        names: &["willr"],
164        sig: OpSig {
165            tag: "WillR",
166            fields: &[Expr("high"), Expr("low"), Expr("close"), Num("n")],
167        },
168        desc: "Williams %R over `n` days.",
169    },
170    Row {
171        names: &["cci"],
172        sig: OpSig {
173            tag: "Cci",
174            fields: &[Expr("high"), Expr("low"), Expr("close"), Num("n")],
175        },
176        desc: "Commodity Channel Index over `n` days.",
177    },
178    Row {
179        names: &["stoch_k"],
180        sig: OpSig {
181            tag: "StochK",
182            fields: &[Expr("high"), Expr("low"), Expr("close"), Num("n")],
183        },
184        desc: "Stochastic %K over `n` days.",
185    },
186    Row {
187        names: &["stoch_d"],
188        sig: OpSig {
189            tag: "StochD",
190            fields: &[
191                Expr("high"),
192                Expr("low"),
193                Expr("close"),
194                Num("n"),
195                NumOpt("d"),
196            ],
197        },
198        desc: "Stochastic %D: `d`-day average of %K over `n` days (d defaults to 3).",
199    },
200    Row {
201        names: &["aroon_up"],
202        sig: OpSig {
203            tag: "AroonUp",
204            fields: &[Expr("high"), Num("n")],
205        },
206        desc: "Aroon Up over `n` days from high.",
207    },
208    Row {
209        names: &["aroon_down"],
210        sig: OpSig {
211            tag: "AroonDown",
212            fields: &[Expr("low"), Num("n")],
213        },
214        desc: "Aroon Down over `n` days from low.",
215    },
216    Row {
217        names: &["adx"],
218        sig: OpSig {
219            tag: "Adx",
220            fields: &[Expr("high"), Expr("low"), Expr("close"), Num("n")],
221        },
222        desc: "Average Directional Index over `n` days.",
223    },
224    Row {
225        names: &["plus_di"],
226        sig: OpSig {
227            tag: "PlusDi",
228            fields: &[Expr("high"), Expr("low"), Expr("close"), Num("n")],
229        },
230        desc: "Plus Directional Indicator (+DI) over `n` days.",
231    },
232    Row {
233        names: &["minus_di"],
234        sig: OpSig {
235            tag: "MinusDi",
236            fields: &[Expr("high"), Expr("low"), Expr("close"), Num("n")],
237        },
238        desc: "Minus Directional Indicator (-DI) over `n` days.",
239    },
240    Row {
241        names: &["obv"],
242        sig: OpSig {
243            tag: "Obv",
244            fields: &[Expr("close"), Expr("volume")],
245        },
246        desc: "On-Balance Volume from close and volume.",
247    },
248    Row {
249        names: &["mfi"],
250        sig: OpSig {
251            tag: "Mfi",
252            fields: &[
253                Expr("high"),
254                Expr("low"),
255                Expr("close"),
256                Expr("volume"),
257                Num("n"),
258            ],
259        },
260        desc: "Money Flow Index over `n` days.",
261    },
262    Row {
263        names: &["vwap"],
264        sig: OpSig {
265            tag: "Vwap",
266            fields: &[
267                Expr("high"),
268                Expr("low"),
269                Expr("close"),
270                Expr("volume"),
271                Num("n"),
272            ],
273        },
274        desc: "Volume-Weighted Average Price over `n` days from high/low/close/volume.",
275    },
276    Row {
277        names: &["is_largest"],
278        sig: OpSig {
279            tag: "IsLargest",
280            fields: &[Expr("of"), Num("n")],
281        },
282        desc: "1 for the `n` highest values per row (cross-section), else 0.",
283    },
284    Row {
285        names: &["is_smallest"],
286        sig: OpSig {
287            tag: "IsSmallest",
288            fields: &[Expr("of"), Num("n")],
289        },
290        desc: "1 for the `n` lowest values per row (cross-section), else 0.",
291    },
292    Row {
293        names: &["sustain"],
294        sig: OpSig {
295            tag: "Sustain",
296            fields: &[Expr("of"), Num("nwindow"), NumOpt("nsatisfy")],
297        },
298        desc: "1 where `of` held true at least `nsatisfy` times within the last `nwindow` rows.",
299    },
300    Row {
301        names: &["is_entry"],
302        sig: OpSig {
303            tag: "IsEntry",
304            fields: &[Expr("of")],
305        },
306        desc: "1 on the row where `of` turns false->true (rising edge).",
307    },
308    Row {
309        names: &["is_exit"],
310        sig: OpSig {
311            tag: "IsExit",
312            fields: &[Expr("of")],
313        },
314        desc: "1 on the row where `of` turns true->false (falling edge).",
315    },
316    Row {
317        names: &["ceil"],
318        sig: OpSig {
319            tag: "Ceil",
320            fields: &[Expr("of")],
321        },
322        desc: "Ceiling of `of`.",
323    },
324    Row {
325        names: &["rank"],
326        sig: OpSig {
327            tag: "Rank",
328            fields: &[Expr("of"), BoolOpt("pct"), BoolOpt("ascending")],
329        },
330        desc: "Cross-sectional rank of `of` per row; `pct` for 0..1 percentile (default true), `ascending` sets direction (default true).",
331    },
332    Row {
333        names: &["mask"],
334        sig: OpSig {
335            tag: "Mask",
336            fields: &[Expr("of"), Expr("by")],
337        },
338        desc: "`of` kept only where `by` is true; elsewhere dropped.",
339    },
340    Row {
341        names: &["normalize_row"],
342        sig: OpSig {
343            tag: "NormalizeRow",
344            fields: &[Expr("of")],
345        },
346        desc: "Scale each row so gross weight (sum of |w|) is 1 — turns a raw signal into explicit portfolio weights. NaN preserved; zero rows unchanged.",
347    },
348    Row {
349        names: &["hold_until"],
350        sig: OpSig {
351            tag: "HoldUntil",
352            fields: &[
353                Expr("entry"),
354                Expr("exit"),
355                NumOpt("nstocks_limit"),
356                ExprOpt("rank"),
357                NumOpt("stop_loss"),
358                NumOpt("take_profit"),
359                NumOpt("trail_stop"),
360                NumOpt("trail_stop_activation"),
361            ],
362        },
363        desc: "Stateful rotation: enter on `entry`, exit on `exit`, hold up to `nstocks_limit` (prioritised by `rank`), with optional stop_loss/take_profit/trail_stop/trail_stop_activation.",
364    },
365    Row {
366        names: &["rebalance"],
367        sig: OpSig {
368            tag: "Rebalance",
369            fields: &[Expr("of"), StrOpt("freq"), ExprOpt("on")],
370        },
371        desc: "Hold `of`, refreshing on calendar `freq` (W/ME/QE) or on dates where `on` is true.",
372    },
373    Row {
374        names: &["neutralize"],
375        sig: OpSig {
376            tag: "Neutralize",
377            fields: &[Expr("of"), ExprList("by"), BoolOpt("add_const")],
378        },
379        desc: "Cross-sectionally regress `of` against the `by` factors, optionally adding a constant (default true).",
380    },
381    Row {
382        names: &["neutralize_industry"],
383        sig: OpSig {
384            tag: "NeutralizeIndustry",
385            fields: &[Expr("of"), BoolOpt("add_const")],
386        },
387        desc: "Neutralize `of` within each industry/sector (add_const defaults to true).",
388    },
389    Row {
390        names: &["industry_rank"],
391        sig: OpSig {
392            tag: "IndustryRank",
393            fields: &[Expr("of"), StrListOpt("categories")],
394        },
395        desc: "Rank `of` within each industry, optionally limited to `categories`.",
396    },
397    Row {
398        names: &["groupby_category"],
399        sig: OpSig {
400            tag: "GroupbyCategory",
401            fields: &[Expr("of"), Str("agg")],
402        },
403        desc: "Aggregate `of` within each industry using `agg` (e.g. mean).",
404    },
405];
406
407// ---------------------------------------------------------------------------
408// Public catalog API — the single source of truth consumed by the schema
409// generator (`cargo run -p lemon-lang --example gen-schema`). Everything below
410// is derived from `ROWS`, `BINOPS`, and `prefix_tag`, so the emitted JSON
411// artifacts can never drift from the parser.
412// ---------------------------------------------------------------------------
413
414/// A field's serde default, when it has one. `serde(default)` on an `Option`
415/// field means "absent" (no default value to emit); those return `None` here.
416/// Kept co-located with the field declarations so it stays in sync with
417/// `spec.rs`.
418pub fn field_default(tag: &str, field_name: &str) -> Option<serde_json::Value> {
419    use serde_json::json;
420    match (tag, field_name) {
421        ("StochD", "d") => Some(json!(3)),
422        ("Rank", "pct") => Some(json!(true)),
423        ("Rank", "ascending") => Some(json!(true)),
424        ("Neutralize", "add_const") => Some(json!(true)),
425        ("NeutralizeIndustry", "add_const") => Some(json!(true)),
426        _ => None,
427    }
428}
429
430/// One field of a callable op, in schema-friendly form.
431pub struct FieldInfo {
432    pub name: &'static str,
433    /// One of: `expr`, `number`, `string`, `bool`, `list`, `list-string`.
434    pub kind: &'static str,
435    pub required: bool,
436    pub default: Option<serde_json::Value>,
437}
438
439/// One callable op: its canonical name, aliases, op tag, ordered fields, and a
440/// one-line description.
441pub struct OpInfo {
442    pub name: &'static str,
443    pub aliases: &'static [&'static str],
444    pub tag: &'static str,
445    pub description: &'static str,
446    pub fields: Vec<FieldInfo>,
447}
448
449/// Every function-style op in `ROWS`, in declaration order.
450pub fn function_ops() -> Vec<OpInfo> {
451    ROWS.iter()
452        .map(|r| OpInfo {
453            name: r.names[0],
454            aliases: &r.names[1..],
455            tag: r.sig.tag,
456            description: r.desc,
457            fields: r
458                .sig
459                .fields
460                .iter()
461                .map(|f| {
462                    let name = field_name(f);
463                    FieldInfo {
464                        name,
465                        kind: field_kind(f),
466                        required: field_required(f),
467                        default: field_default(r.sig.tag, name),
468                    }
469                })
470                .collect(),
471        })
472        .collect()
473}
474
475/// The binary operator ops: `(symbol, op tag)`. Both operands are `l`/`r` exprs.
476pub fn binary_operators() -> &'static [(&'static str, &'static str)] {
477    BINOPS
478}
479
480pub fn op_by_name(name: &str) -> Option<&'static OpSig> {
481    ROWS.iter()
482        .find(|r| r.names.contains(&name))
483        .map(|r| &r.sig)
484}
485
486pub fn sig_by_tag(tag: &str) -> Option<&'static OpSig> {
487    ROWS.iter().find(|r| r.sig.tag == tag).map(|r| &r.sig)
488}
489
490/// Canonical DSL name for a function-style op tag (printer).
491pub fn dsl_name_for_tag(tag: &str) -> &'static str {
492    ROWS.iter()
493        .find(|r| r.sig.tag == tag)
494        .map(|r| r.names[0])
495        .unwrap_or_else(|| {
496            // Tag not in the function table — callers should not reach this for
497            // operator-style ops; return the tag itself via a leaked copy so the
498            // return type stays `&'static str`.
499            ALL_OP_TAGS
500                .iter()
501                .copied()
502                .find(|t| *t == tag)
503                .unwrap_or("")
504        })
505}
506
507static BINOPS: &[(&str, &str)] = &[
508    (">", "Gt"),
509    ("<", "Lt"),
510    (">=", "Ge"),
511    ("<=", "Le"),
512    ("and", "And"),
513    ("or", "Or"),
514    ("+", "Add"),
515    ("-", "Sub"),
516    ("*", "Mul"),
517    ("/", "Div"),
518];
519
520pub fn binop_tag(op: &str) -> Option<&'static str> {
521    BINOPS.iter().find(|(s, _)| *s == op).map(|(_, t)| *t)
522}
523
524pub fn binop_symbol_for_tag(tag: &str) -> Option<&'static str> {
525    BINOPS.iter().find(|(_, t)| *t == tag).map(|(s, _)| *s)
526}
527
528pub fn prefix_tag(op: &str) -> Option<&'static str> {
529    if op == "-" {
530        Some("Neg")
531    } else {
532        None
533    }
534}
535
536/// All op tags in `spec.rs` — the completeness gate (see ops::tests).
537pub static ALL_OP_TAGS: &[&str] = &[
538    "Data",
539    "Const",
540    "Average",
541    "Ema",
542    "Std",
543    "Rsi",
544    "PctChange",
545    "Rise",
546    "Shift",
547    "RollingMax",
548    "Atr",
549    "Natr",
550    "WillR",
551    "Cci",
552    "StochK",
553    "StochD",
554    "AroonUp",
555    "AroonDown",
556    "Adx",
557    "PlusDi",
558    "MinusDi",
559    "Obv",
560    "Mfi",
561    "Vwap",
562    "Fall",
563    "IsLargest",
564    "IsSmallest",
565    "Sustain",
566    "IsEntry",
567    "IsExit",
568    "Gt",
569    "Lt",
570    "Ge",
571    "Le",
572    "And",
573    "Or",
574    "Not",
575    "Add",
576    "Sub",
577    "Mul",
578    "Div",
579    "Neg",
580    "Ceil",
581    "Rank",
582    "Mask",
583    "NormalizeRow",
584    "HoldUntil",
585    "Rebalance",
586    "Neutralize",
587    "NeutralizeIndustry",
588    "IndustryRank",
589    "GroupbyCategory",
590];
591
592#[cfg(test)]
593mod tests {
594    use super::*;
595
596    #[test]
597    fn looks_up_alias_and_snake_case() {
598        assert_eq!(op_by_name("sma").unwrap().tag, "Average");
599        assert_eq!(op_by_name("average").unwrap().tag, "Average");
600        assert_eq!(op_by_name("rolling_max").unwrap().tag, "RollingMax");
601        assert!(op_by_name("not_an_op").is_none());
602    }
603
604    #[test]
605    fn maps_operators() {
606        assert_eq!(binop_tag(">"), Some("Gt"));
607        assert_eq!(binop_tag("and"), Some("And"));
608        assert_eq!(binop_tag("+"), Some("Add"));
609        assert_eq!(prefix_tag("-"), Some("Neg"));
610    }
611
612    #[test]
613    fn round_trip_names_for_printer() {
614        assert_eq!(dsl_name_for_tag("Average"), "sma");
615        assert_eq!(binop_symbol_for_tag("Gt"), Some(">"));
616        assert_eq!(binop_symbol_for_tag("Average"), None);
617    }
618
619    #[test]
620    fn function_ops_expose_every_row_with_schema_metadata() {
621        let ops = function_ops();
622        // One OpInfo per row in the table.
623        assert_eq!(ops.len(), ROWS.len());
624
625        // sma is the canonical name for Average and carries `average` as an alias.
626        let sma = ops.iter().find(|o| o.tag == "Average").unwrap();
627        assert_eq!(sma.name, "sma");
628        assert!(sma.aliases.contains(&"average"));
629        assert!(!sma.description.is_empty());
630
631        // Every field carries a name, a known kind, and a required flag; the field
632        // kinds collectively cover the full `field_kind` match.
633        let mut kinds = std::collections::BTreeSet::new();
634        for op in &ops {
635            for f in &op.fields {
636                assert!(!f.name.is_empty());
637                assert!(matches!(
638                    f.kind,
639                    "expr" | "number" | "string" | "bool" | "list" | "list-string"
640                ));
641                kinds.insert(f.kind);
642            }
643        }
644        // The table exercises expr and number kinds at minimum.
645        assert!(kinds.contains("expr"));
646        assert!(kinds.contains("number"));
647    }
648
649    #[test]
650    fn field_defaults_match_serde() {
651        // Documented serde defaults surface through field_default…
652        assert_eq!(field_default("StochD", "d"), Some(serde_json::json!(3)));
653        assert_eq!(field_default("Rank", "pct"), Some(serde_json::json!(true)));
654        assert_eq!(
655            field_default("Neutralize", "add_const"),
656            Some(serde_json::json!(true))
657        );
658        // …and everything else has no default.
659        assert_eq!(field_default("Average", "n"), None);
660        assert_eq!(field_default("Nope", "nope"), None);
661
662        // field_default is threaded through function_ops for the defaulted fields.
663        let stoch_d = function_ops()
664            .into_iter()
665            .find(|o| o.tag == "StochD")
666            .unwrap();
667        let d_field = stoch_d.fields.iter().find(|f| f.name == "d").unwrap();
668        assert_eq!(d_field.default, Some(serde_json::json!(3)));
669    }
670
671    #[test]
672    fn field_kind_and_required_cover_every_variant() {
673        use Field::*;
674        for (f, kind, required) in [
675            (Expr("x"), "expr", true),
676            (ExprOpt("x"), "expr", false),
677            (ExprList("x"), "list", true),
678            (Num("x"), "number", true),
679            (NumOpt("x"), "number", false),
680            (BoolOpt("x"), "bool", false),
681            (Str("x"), "string", true),
682            (StrOpt("x"), "string", false),
683            (StrListOpt("x"), "list-string", false),
684        ] {
685            assert_eq!(field_name(&f), "x");
686            assert_eq!(field_kind(&f), kind);
687            assert_eq!(field_required(&f), required);
688        }
689    }
690
691    #[test]
692    fn tag_and_operator_lookups() {
693        // sig_by_tag round-trips against op_by_name.
694        assert_eq!(sig_by_tag("Average").unwrap().tag, "Average");
695        assert!(sig_by_tag("NotATag").is_none());
696
697        // binary_operators exposes the BINOPS table.
698        let ops = binary_operators();
699        assert!(ops.contains(&(">", "Gt")));
700        assert!(ops.contains(&("/", "Div")));
701        assert_eq!(ops.len(), 10);
702
703        // Unknown / operator-style tags return "" from dsl_name_for_tag (fallback arm).
704        assert_eq!(dsl_name_for_tag("TotallyUnknownTag"), "");
705        assert_eq!(prefix_tag("+"), None);
706    }
707
708    #[test]
709    fn every_spec_op_has_a_signature_or_operator() {
710        // The 51 op tags from spec.rs. If a new op is added, add it here AND to the
711        // table/operator maps — this test is the completeness gate for this plan.
712        for tag in ALL_OP_TAGS {
713            let known = ROWS.iter().any(|r| r.sig.tag == *tag)
714                || binop_symbol_for_tag(tag).is_some()
715                || *tag == "Neg"
716                || *tag == "Not"
717                || *tag == "Const"
718                || *tag == "Data";
719            assert!(known, "op `{tag}` has no DSL handler");
720        }
721    }
722}
lemon/dsl/ops.rs

lemon/dsl/
ops.rs
