Skip to main content

tla_syntax/
ast.rs

1use crate::token::Op;
2
3#[derive(Debug, Clone, PartialEq, Eq)]
4pub struct Module {
5    pub name: String,
6    pub extends: Vec<String>,
7    pub units: Vec<Unit>,
8}
9
10impl Module {
11    pub fn definition(&self, name: &str) -> Option<&Def> {
12        self.units.iter().find_map(|u| match u {
13            Unit::Def(d) if d.name == name => Some(d),
14            _ => None,
15        })
16    }
17
18    pub fn constants(&self) -> impl Iterator<Item = &Decl> {
19        self.units
20            .iter()
21            .filter_map(|u| match u {
22                Unit::Constants(ds) => Some(ds),
23                _ => None,
24            })
25            .flatten()
26    }
27
28    pub fn variables(&self) -> impl Iterator<Item = &String> {
29        self.units
30            .iter()
31            .filter_map(|u| match u {
32                Unit::Variables(vs) => Some(vs),
33                _ => None,
34            })
35            .flatten()
36    }
37}
38
39#[derive(Debug, Clone, PartialEq, Eq)]
40pub enum Unit {
41    Constants(Vec<Decl>),
42    Variables(Vec<String>),
43    Recursive(Vec<Decl>),
44    Def(Def),
45    /// `S == INSTANCE M WITH x <- e` — named when it introduces a prefix,
46    /// anonymous when the module's definitions are pulled in directly.
47    Instance {
48        name: Option<String>,
49        module: String,
50        subs: Vec<(String, Expr)>,
51    },
52    Assume(Expr),
53    Theorem(Expr),
54    /// A module declared inside another one.
55    Inner(Box<Module>),
56    /// A unit the evaluator has no use for and the parser did not keep: a
57    /// TLAPS proof, or a theorem stated in the `ASSUME ... PROVE` form.
58    Opaque,
59}
60
61/// A declared name together with its arity; zero for a plain constant.
62#[derive(Debug, Clone, PartialEq, Eq)]
63pub struct Decl {
64    pub name: String,
65    pub arity: usize,
66}
67
68#[derive(Debug, Clone, PartialEq, Eq)]
69pub struct Def {
70    pub name: String,
71    pub params: Vec<Param>,
72    pub body: Expr,
73    pub local: bool,
74}
75
76/// A formal parameter. Arity is zero for an ordinary one and positive for an
77/// operator parameter, which is declared as `f(_)` and must be applied.
78#[derive(Debug, Clone, PartialEq, Eq)]
79pub struct Param {
80    pub name: String,
81    pub arity: usize,
82}
83
84impl Param {
85    pub fn value(name: impl Into<String>) -> Self {
86        Self {
87            name: name.into(),
88            arity: 0,
89        }
90    }
91}
92
93/// One `x, y \in S` group of a quantifier, function or set constructor.
94#[derive(Debug, Clone, PartialEq, Eq)]
95pub struct Bound {
96    pub names: Vec<String>,
97    /// Absent for the unbounded forms, as in `\E x : P`.
98    pub domain: Option<Expr>,
99    /// True for `<<x, y>> \in S`, which destructures each element.
100    pub destructure: bool,
101}
102
103impl Bound {
104    pub fn mentions_next_state(&self) -> bool {
105        self.domain.as_ref().is_some_and(Expr::mentions_next_state)
106    }
107}
108
109/// An `INSTANCE` introduced by a `LET`, in scope only for its body.
110#[derive(Debug, Clone, PartialEq, Eq)]
111pub struct LetInstance {
112    pub name: Option<String>,
113    pub module: String,
114    pub subs: Vec<(String, Expr)>,
115}
116
117#[derive(Debug, Clone, Copy, PartialEq, Eq)]
118pub enum QuantKind {
119    Forall,
120    Exists,
121    /// `\AA` and `\EE`, which quantify over a hidden *variable* rather than a
122    /// value, and so describe behaviours rather than states.
123    TemporalForall,
124    TemporalExists,
125}
126
127impl QuantKind {
128    pub fn is_temporal(self) -> bool {
129        matches!(self, Self::TemporalForall | Self::TemporalExists)
130    }
131}
132
133#[derive(Debug, Clone, PartialEq, Eq)]
134pub enum ExceptPath {
135    Index(Expr),
136    Field(String),
137}
138
139#[derive(Debug, Clone, PartialEq, Eq)]
140pub enum Expr {
141    Num(i64),
142    /// `123.456`, as written. TLA+ decimals are exact rationals.
143    Decimal(String),
144    Str(String),
145    Bool(bool),
146    Ident(String),
147    /// `x'`
148    Prime(Box<Expr>),
149    /// `@`, legal only inside an `EXCEPT` update.
150    At,
151    /// `Op(a, b)` — application of a defined operator.
152    Apply(Box<Expr>, Vec<Expr>),
153    /// `f[a]` — function application.
154    FnApply(Box<Expr>, Vec<Expr>),
155    /// `r.field`
156    Field(Box<Expr>, String),
157    /// `Inst!Name(args)`
158    Qualified {
159        instance: String,
160        name: String,
161        args: Vec<Expr>,
162    },
163    Unary(Op, Box<Expr>),
164    Binary(Op, Box<Expr>, Box<Expr>),
165    Tuple(Vec<Expr>),
166    SetEnum(Vec<Expr>),
167    /// `{x \in S : P}`
168    SetFilter {
169        bound: Box<Bound>,
170        pred: Box<Expr>,
171    },
172    /// `{e : x \in S, y \in T}`
173    SetMap {
174        expr: Box<Expr>,
175        bounds: Vec<Bound>,
176    },
177    /// `[a |-> 1, b |-> 2]`
178    Record(Vec<(String, Expr)>),
179    /// `[a : S, b : T]`
180    RecordSet(Vec<(String, Expr)>),
181    /// `[x \in S |-> e]`
182    FnDef {
183        bounds: Vec<Bound>,
184        body: Box<Expr>,
185    },
186    /// `[S -> T]`
187    FnSet {
188        domain: Box<Expr>,
189        range: Box<Expr>,
190    },
191    /// `[f EXCEPT ![a] = e, !.g = e2]`
192    Except {
193        base: Box<Expr>,
194        updates: Vec<(Vec<ExceptPath>, Expr)>,
195    },
196    Quant {
197        kind: QuantKind,
198        bounds: Vec<Bound>,
199        body: Box<Expr>,
200    },
201    Choose {
202        bound: Box<Bound>,
203        body: Box<Expr>,
204    },
205    Let {
206        defs: Vec<Def>,
207        /// `LET I == INSTANCE M IN ...` — an instance whose scope is the body.
208        instances: Vec<LetInstance>,
209        body: Box<Expr>,
210    },
211    If {
212        cond: Box<Expr>,
213        then: Box<Expr>,
214        otherwise: Box<Expr>,
215    },
216    Case {
217        arms: Vec<(Expr, Expr)>,
218        other: Option<Box<Expr>>,
219    },
220    /// `LAMBDA x, y : e` — an operator written where one is expected.
221    Lambda {
222        params: Vec<Param>,
223        body: Box<Expr>,
224    },
225    /// `[A]_vars`
226    ActionBox {
227        action: Box<Expr>,
228        subscript: Box<Expr>,
229    },
230    /// `<<A>>_vars`
231    ActionAngle {
232        action: Box<Expr>,
233        subscript: Box<Expr>,
234    },
235    /// `WF_vars(A)` / `SF_vars(A)`
236    Fairness {
237        strong: bool,
238        subscript: Box<Expr>,
239        action: Box<Expr>,
240    },
241}
242
243impl Expr {
244    /// Does this constrain the successor state? Distinguishes an action's
245    /// guard, which says when it may happen, from its effect, which says what
246    /// it does — the two call for different advice when one of them fails.
247    pub fn mentions_next_state(&self) -> bool {
248        match self {
249            Expr::Prime(_)
250            | Expr::ActionBox { .. }
251            | Expr::ActionAngle { .. }
252            | Expr::Fairness { .. } => true,
253            Expr::Unary(op, inner) => {
254                *op == Op::Unchanged || *op == Op::Enabled || inner.mentions_next_state()
255            }
256            Expr::Binary(_, lhs, rhs) => lhs.mentions_next_state() || rhs.mentions_next_state(),
257            Expr::Apply(head, args) | Expr::FnApply(head, args) => {
258                head.mentions_next_state() || args.iter().any(Expr::mentions_next_state)
259            }
260            Expr::Field(inner, _) => inner.mentions_next_state(),
261            Expr::Qualified { args, .. } => args.iter().any(Expr::mentions_next_state),
262            Expr::Tuple(items) | Expr::SetEnum(items) => {
263                items.iter().any(Expr::mentions_next_state)
264            }
265            Expr::SetFilter { bound, pred } => {
266                bound.mentions_next_state() || pred.mentions_next_state()
267            }
268            Expr::SetMap { expr, bounds } => {
269                expr.mentions_next_state() || bounds.iter().any(Bound::mentions_next_state)
270            }
271            Expr::Record(fields) | Expr::RecordSet(fields) => {
272                fields.iter().any(|(_, v)| v.mentions_next_state())
273            }
274            Expr::FnSet { domain, range } => {
275                domain.mentions_next_state() || range.mentions_next_state()
276            }
277            Expr::Except { base, updates } => {
278                base.mentions_next_state()
279                    || updates.iter().any(|(path, value)| {
280                        value.mentions_next_state()
281                            || path.iter().any(|step| match step {
282                                ExceptPath::Index(e) => e.mentions_next_state(),
283                                ExceptPath::Field(_) => false,
284                            })
285                    })
286            }
287            Expr::FnDef { bounds, body } | Expr::Quant { bounds, body, .. } => {
288                body.mentions_next_state() || bounds.iter().any(Bound::mentions_next_state)
289            }
290            Expr::Choose { bound, body } => {
291                bound.mentions_next_state() || body.mentions_next_state()
292            }
293            Expr::Lambda { body, .. } => body.mentions_next_state(),
294            Expr::Let {
295                defs,
296                instances,
297                body,
298            } => {
299                body.mentions_next_state()
300                    || defs.iter().any(|d| d.body.mentions_next_state())
301                    || instances
302                        .iter()
303                        .any(|i| i.subs.iter().any(|(_, e)| e.mentions_next_state()))
304            }
305            Expr::If {
306                cond,
307                then,
308                otherwise,
309            } => {
310                cond.mentions_next_state()
311                    || then.mentions_next_state()
312                    || otherwise.mentions_next_state()
313            }
314            Expr::Case { arms, other } => {
315                arms.iter()
316                    .any(|(g, r)| g.mentions_next_state() || r.mentions_next_state())
317                    || other.as_ref().is_some_and(|o| o.mentions_next_state())
318            }
319            Expr::Num(_)
320            | Expr::Decimal(_)
321            | Expr::Str(_)
322            | Expr::Bool(_)
323            | Expr::Ident(_)
324            | Expr::At => false,
325        }
326    }
327
328    pub fn conjunction(items: Vec<Expr>) -> Expr {
329        Self::fold(items, Op::And)
330    }
331
332    pub fn disjunction(items: Vec<Expr>) -> Expr {
333        Self::fold(items, Op::Or)
334    }
335
336    fn fold(items: Vec<Expr>, op: Op) -> Expr {
337        let mut iter = items.into_iter();
338        let first = iter.next().expect("junction list has at least one item");
339        iter.fold(first, |acc, e| Expr::Binary(op, Box::new(acc), Box::new(e)))
340    }
341}