presk_syntax/
ast.rs

1//! Abstract syntax of one Presk source file. Dimension expressions are
2//! ordinary [`Expr`]s; their interpretation as linear forms belongs to
3//! `presk-types`.
4
5use crate::source::{FileId, Span};
6
7#[derive(Clone, Debug)]
8pub struct Program {
9    pub file: FileId,
10    pub decls: Vec<Decl>,
11}
12
13#[derive(Clone, Debug)]
14pub struct Ident {
15    pub name: String,
16    pub span: Span,
17}
18
19/// `a.b.c`
20#[derive(Clone, Debug)]
21pub struct Path {
22    pub segs: Vec<Ident>,
23    pub span: Span,
24}
25
26impl Path {
27    pub fn single(id: Ident) -> Self {
28        Path { span: id.span, segs: vec![id] }
29    }
30
31    pub fn is_single(&self, name: &str) -> bool {
32        self.segs.len() == 1 && self.segs[0].name == name
33    }
34
35    pub fn dotted(&self) -> String {
36        self.segs.iter().map(|s| s.name.as_str()).collect::<Vec<_>>().join(".")
37    }
38}
39
40#[derive(Clone, Debug)]
41pub enum Decl {
42    /// `import "path" as alias`
43    Import { path: String, alias: Ident, span: Span },
44    /// `const NAME [: ty] = expr`
45    Const { name: Ident, ty: Option<Ty>, value: Expr, span: Span },
46    /// `static assert expr`
47    StaticAssert { cond: Expr, span: Span },
48    /// `dim NAME [<= expr]`
49    Dim { name: Ident, bound: Option<Expr>, span: Span },
50    /// `type NAME = ty`
51    TypeAlias { name: Ident, ty: Ty, span: Span },
52    Module(ModuleDecl),
53    Fn(FnDecl),
54    /// `optimizer NAME { items }`
55    Optimizer { name: Ident, items: Vec<ConfigItem>, span: Span },
56    /// `training run { items }`
57    Training { items: Vec<ConfigItem>, span: Span },
58    /// `inference serve { items }`
59    Inference { items: Vec<ConfigItem>, span: Span },
60    /// `effect NAME`
61    Effect { name: Ident, span: Span },
62    /// `model NAME = modexpr`
63    Model { name: Ident, body: ModExpr, span: Span },
64}
65
66impl Decl {
67    pub fn span(&self) -> Span {
68        match self {
69            Decl::Import { span, .. }
70            | Decl::Const { span, .. }
71            | Decl::StaticAssert { span, .. }
72            | Decl::Dim { span, .. }
73            | Decl::TypeAlias { span, .. }
74            | Decl::Optimizer { span, .. }
75            | Decl::Training { span, .. }
76            | Decl::Inference { span, .. }
77            | Decl::Effect { span, .. }
78            | Decl::Model { span, .. } => *span,
79            Decl::Module(m) => m.span,
80            Decl::Fn(f) => f.span,
81        }
82    }
83
84    pub fn name(&self) -> Option<&Ident> {
85        match self {
86            Decl::Import { alias, .. } => Some(alias),
87            Decl::Const { name, .. }
88            | Decl::Dim { name, .. }
89            | Decl::TypeAlias { name, .. }
90            | Decl::Optimizer { name, .. }
91            | Decl::Effect { name, .. }
92            | Decl::Model { name, .. } => Some(name),
93            Decl::Module(m) => Some(&m.name),
94            Decl::Fn(f) => Some(&f.name),
95            Decl::StaticAssert { .. } | Decl::Training { .. } | Decl::Inference { .. } => None,
96        }
97    }
98}
99
100#[derive(Clone, Debug)]
101pub struct ModuleDecl {
102    pub name: Ident,
103    pub typarams: Vec<TyParam>,
104    pub members: Vec<Member>,
105    pub span: Span,
106}
107
108/// `{In Out : dim, S : Shape}` — one entry per name.
109#[derive(Clone, Debug)]
110pub struct TyParam {
111    pub name: Ident,
112    /// `dim` or `Shape`
113    pub kind: Ident,
114}
115
116#[derive(Clone, Debug)]
117pub enum Member {
118    /// `param a, b : ty [init = expr]`
119    Param { names: Vec<Ident>, ty: Ty, init: Option<Expr>, span: Span },
120    /// `use a, b : modexpr`
121    Use { names: Vec<Ident>, module: ModExpr, span: Span },
122    Fn(FnDecl),
123}
124
125/// A module expression: `Linear{D, D}`, `ff.MLP{D, 4*D}`, `stack{L}(Block)`,
126/// `A ∘ B`.
127#[derive(Clone, Debug)]
128pub enum ModExpr {
129    App { path: Path, dim_args: Vec<Expr>, mod_args: Vec<ModExpr>, span: Span },
130    Compose { parts: Vec<ModExpr>, span: Span },
131}
132
133impl ModExpr {
134    pub fn span(&self) -> Span {
135        match self {
136            ModExpr::App { span, .. } | ModExpr::Compose { span, .. } => *span,
137        }
138    }
139}
140
141#[derive(Clone, Debug)]
142pub struct FnDecl {
143    pub name: Ident,
144    pub typarams: Vec<TyParam>,
145    pub params: Vec<Param>,
146    pub ret: Option<Ty>,
147    /// `uses Rand`
148    pub effects: Vec<Ident>,
149    pub body: Expr,
150    pub span: Span,
151}
152
153#[derive(Clone, Debug)]
154pub struct Param {
155    pub name: Ident,
156    pub ty: Option<Ty>,
157    pub default: Option<Expr>,
158    /// `axis i` parameter form
159    pub axis: bool,
160    pub span: Span,
161}
162
163#[derive(Clone, Debug)]
164pub enum ConfigItem {
165    /// `name = expr`, `name : ty = expr` or `name : expr`
166    Field { name: Ident, ty: Option<Ty>, value: Expr, span: Span },
167    /// `checkpoint every N steps to "path"`
168    Checkpoint { every: Expr, to: String, span: Span },
169}
170
171#[derive(Clone, Debug)]
172pub enum Ty {
173    /// `ℝ[axes]`; bare `ℝ` is `ℝ[]`
174    Real { axes: Vec<AxisSpec>, span: Span },
175    /// `idx[N][axes]`
176    Idx { bound: Box<Expr>, axes: Vec<AxisSpec>, span: Span },
177    /// `bool[axes]`: a tensor of truth values (plan 0.10.0 step 7b; `bool`
178    /// is a contextual word, and without axes the hyperparameter type, a
179    /// [`Ty::Named`]).
180    Bool { axes: Vec<AxisSpec>, span: Span },
181    /// `param ty`
182    Param { inner: Box<Ty>, span: Span },
183    /// `params path`
184    Params { module: Path, span: Span },
185    /// `nat`, `f64`, `Tokens`, `MLP{D, 4*D}`, `ff.Linear{D, D}`
186    Named { path: Path, args: Vec<Expr>, span: Span },
187    Tuple { elems: Vec<Ty>, span: Span },
188}
189
190impl Ty {
191    pub fn span(&self) -> Span {
192        match self {
193            Ty::Real { span, .. }
194            | Ty::Idx { span, .. }
195            | Ty::Bool { span, .. }
196            | Ty::Param { span, .. }
197            | Ty::Params { span, .. }
198            | Ty::Named { span, .. }
199            | Ty::Tuple { span, .. } => *span,
200        }
201    }
202}
203
204#[derive(Clone, Debug)]
205pub enum AxisSpec {
206    /// `[name:] dimexpr`
207    Axis { name: Option<Ident>, dim: Expr },
208    /// `...`
209    Ellipsis { span: Span },
210}
211
212#[derive(Clone, Debug)]
213pub enum Expr {
214    Lit { lit: Lit, span: Span },
215    Var { path: Path },
216    /// `let pat = init NEWLINE body`
217    Let { pat: Pat, init: Box<Expr>, body: Box<Expr>, span: Span },
218    /// `λ pat. body`
219    Lambda { pat: Pat, body: Box<Expr>, span: Span },
220    Call { callee: Box<Expr>, args: Vec<Arg>, span: Span },
221    /// `base.field` on a non-path base
222    Field { base: Box<Expr>, field: Ident, span: Span },
223    /// `base[sel, …]`
224    Index { base: Box<Expr>, sels: Vec<IndexSel>, span: Span },
225    /// `base{dim args}`
226    Inst { base: Box<Expr>, args: Vec<Expr>, span: Span },
227    /// `ein[axes](body)`
228    Ein { axes: Vec<Ident>, body: Box<Expr>, span: Span },
229    Binary { op: BinOp, lhs: Box<Expr>, rhs: Box<Expr>, span: Span },
230    Unary { op: UnOp, expr: Box<Expr>, span: Span },
231    /// `expr : ty`
232    Ascribe { expr: Box<Expr>, ty: Ty, span: Span },
233    Tuple { elems: Vec<Expr>, span: Span },
234}
235
236impl Expr {
237    pub fn span(&self) -> Span {
238        match self {
239            Expr::Lit { span, .. }
240            | Expr::Let { span, .. }
241            | Expr::Lambda { span, .. }
242            | Expr::Call { span, .. }
243            | Expr::Field { span, .. }
244            | Expr::Index { span, .. }
245            | Expr::Inst { span, .. }
246            | Expr::Ein { span, .. }
247            | Expr::Binary { span, .. }
248            | Expr::Unary { span, .. }
249            | Expr::Ascribe { span, .. }
250            | Expr::Tuple { span, .. } => *span,
251            Expr::Var { path } => path.span,
252        }
253    }
254}
255
256#[derive(Clone, Debug)]
257pub enum Pat {
258    Ident(Ident),
259    Tuple { elems: Vec<Pat>, span: Span },
260}
261
262impl Pat {
263    pub fn span(&self) -> Span {
264        match self {
265            Pat::Ident(i) => i.span,
266            Pat::Tuple { span, .. } => *span,
267        }
268    }
269}
270
271#[derive(Clone, Debug)]
272pub enum Arg {
273    Pos(Expr),
274    /// `name = expr`
275    Named { name: Ident, value: Expr },
276    /// `axis name`
277    Axis(Ident),
278    /// `name: expr` (an axis name with its length, e.g. `positions(seq: T)`)
279    NamedAxis { name: Ident, value: Expr },
280}
281
282#[derive(Clone, Debug)]
283pub enum IndexSel {
284    Pos(Expr),
285    /// `axis: expr`
286    Named { axis: Ident, value: Expr },
287    /// `axis: lo..hi`, or `lo..hi` on the first axis: a slice, which keeps
288    /// its axis at the length `hi - lo` (reference, rule E-Slice).
289    Slice { axis: Option<Ident>, lo: Expr, hi: Expr },
290}
291
292#[derive(Clone, Debug)]
293pub enum Lit {
294    Nat(String),
295    Float(String),
296    Quantity { num: String, unit: String },
297    Str(String),
298}
299
300#[derive(Clone, Copy, Debug, PartialEq, Eq)]
301pub enum BinOp {
302    /// `|>`
303    Pipe,
304    /// `|`
305    Alt,
306    /// `⊗`
307    Tensor,
308    /// `∘`
309    Compose,
310    Eq,
311    Le,
312    Ge,
313    Lt,
314    Gt,
315    Add,
316    Sub,
317    Mul,
318    Div,
319    Rem,
320    /// `·`
321    Dot,
322    Pow,
323}
324
325impl BinOp {
326    pub fn text(self) -> &'static str {
327        match self {
328            BinOp::Pipe => "|>",
329            BinOp::Alt => "|",
330            BinOp::Tensor => "⊗",
331            BinOp::Compose => "∘",
332            BinOp::Eq => "==",
333            BinOp::Le => "<=",
334            BinOp::Ge => ">=",
335            BinOp::Lt => "<",
336            BinOp::Gt => ">",
337            BinOp::Add => "+",
338            BinOp::Sub => "-",
339            BinOp::Mul => "*",
340            BinOp::Div => "/",
341            BinOp::Rem => "%",
342            BinOp::Dot => "·",
343            BinOp::Pow => "^",
344        }
345    }
346}
347
348#[derive(Clone, Copy, Debug, PartialEq, Eq)]
349pub enum UnOp {
350    Neg,
351}
presk_syntax/ast.rs

presk_syntax/
ast.rs
