rkyv_js_codegen/
expr.rs

1//! Typed TypeScript expression tree for codec bindings.
2//!
3//! Every codec expression the generator emits is built from [`CodecExpr`]
4//! nodes instead of format strings. The tree knows which named imports it
5//! needs, which generated types it references, and how to render itself to
6//! TypeScript source.
7
8use std::collections::{BTreeMap, BTreeSet};
9
10use crate::error::DiagnosticKind;
11
12/// A named import contributed by a [`CodecExpr::Import`] node.
13///
14/// # Example
15///
16/// ```
17/// use rkyv_js_codegen::Import;
18///
19/// let import = Import::new("rkyv-js/lib/uuid", "uuid");
20/// assert_eq!(import.module, "rkyv-js/lib/uuid");
21/// assert_eq!(import.export, "uuid");
22/// ```
23#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
24pub struct Import {
25    /// The module specifier to import from (e.g. `"rkyv-js/lib/uuid"`).
26    pub module: String,
27    /// The named export to import (e.g. `"uuid"`).
28    pub export: String,
29}
30
31impl Import {
32    /// Create a new named import.
33    pub fn new(module: impl Into<String>, export: impl Into<String>) -> Self {
34        Self {
35            module: module.into(),
36            export: export.into(),
37        }
38    }
39}
40
41/// A TypeScript codec expression.
42///
43/// Expressions are composed structurally; rendering happens once at
44/// [`generate`](crate::CodeGenerator::generate) time, when all generated type
45/// names are known.
46#[derive(Debug, Clone, PartialEq, Eq, Hash)]
47pub enum CodecExpr {
48    /// A member of the core `rkyv-js` namespace import: `Runtime("u32")`
49    /// renders as `r.u32`.
50    Runtime(&'static str),
51    /// A named import; renders as the bare export name and contributes an
52    /// entry to the generated import block.
53    Import(Import),
54    /// A reference to a generated type by its *Rust* name. Resolved to the
55    /// archived (exported) name at `generate()` time.
56    TypeRef(String),
57    /// A call expression: `callee(args...)`.
58    Call(Box<CodecExpr>, Vec<CodecExpr>),
59    /// An object literal `{ k: v, ... }` — used for enum struct-variant
60    /// records.
61    Object(Vec<(String, CodecExpr)>),
62    /// An array literal `[a, b, ...]` — used for enum tuple variants.
63    Array(Vec<CodecExpr>),
64    /// An integer literal (array lengths).
65    LitInt(u64),
66    /// A placeholder inside registry templates, replaced by
67    /// [`substitute`](CodecExpr::substitute).
68    Param(usize),
69    /// Escape hatch: verbatim TypeScript. Never inspected for imports or
70    /// type references.
71    Raw(String),
72}
73
74impl CodecExpr {
75    /// A member of the core namespace import (`r.{name}`).
76    pub fn runtime(name: &'static str) -> Self {
77        CodecExpr::Runtime(name)
78    }
79
80    /// A named import from an arbitrary module.
81    pub fn import_from(module: impl Into<String>, export: impl Into<String>) -> Self {
82        CodecExpr::Import(Import::new(module, export))
83    }
84
85    /// A reference to a generated type by its Rust name.
86    pub fn type_ref(name: impl Into<String>) -> Self {
87        CodecExpr::TypeRef(name.into())
88    }
89
90    /// A call expression `callee(args...)`.
91    pub fn call(callee: CodecExpr, args: impl IntoIterator<Item = CodecExpr>) -> Self {
92        CodecExpr::Call(Box::new(callee), args.into_iter().collect())
93    }
94
95    /// An object literal `{ k: v, ... }`.
96    pub fn object(
97        entries: impl IntoIterator<Item = (impl Into<String>, CodecExpr)>,
98    ) -> Self {
99        CodecExpr::Object(entries.into_iter().map(|(k, v)| (k.into(), v)).collect())
100    }
101
102    /// An array literal `[a, b, ...]`.
103    pub fn array(elements: impl IntoIterator<Item = CodecExpr>) -> Self {
104        CodecExpr::Array(elements.into_iter().collect())
105    }
106
107    /// Verbatim TypeScript. The generator never inspects the contents.
108    pub fn raw(ts: impl Into<String>) -> Self {
109        CodecExpr::Raw(ts.into())
110    }
111
112    /// Replace every [`CodecExpr::Param`] `i` with `args[i]`.
113    ///
114    /// Parameters without a matching argument are left in place; the registry
115    /// checks arity before substituting.
116    pub fn substitute(&self, args: &[CodecExpr]) -> CodecExpr {
117        match self {
118            CodecExpr::Param(i) => args.get(*i).cloned().unwrap_or(CodecExpr::Param(*i)),
119            CodecExpr::Call(callee, call_args) => CodecExpr::Call(
120                Box::new(callee.substitute(args)),
121                call_args.iter().map(|a| a.substitute(args)).collect(),
122            ),
123            CodecExpr::Object(entries) => CodecExpr::Object(
124                entries
125                    .iter()
126                    .map(|(k, v)| (k.clone(), v.substitute(args)))
127                    .collect(),
128            ),
129            CodecExpr::Array(elements) => {
130                CodecExpr::Array(elements.iter().map(|e| e.substitute(args)).collect())
131            }
132            other => other.clone(),
133        }
134    }
135
136    /// Walk the expression tree in pre-order, calling `f` on every node.
137    ///
138    /// [`CodecExpr::Raw`] contents are never inspected (the node itself is
139    /// still visited).
140    pub fn visit(&self, f: &mut impl FnMut(&CodecExpr)) {
141        f(self);
142        match self {
143            CodecExpr::Call(callee, args) => {
144                callee.visit(f);
145                for arg in args {
146                    arg.visit(f);
147                }
148            }
149            CodecExpr::Object(entries) => {
150                for (_, v) in entries {
151                    v.visit(f);
152                }
153            }
154            CodecExpr::Array(elements) => {
155                for element in elements {
156                    element.visit(f);
157                }
158            }
159            _ => {}
160        }
161    }
162
163    /// The highest [`CodecExpr::Param`] index in the tree, if any.
164    pub fn max_param(&self) -> Option<usize> {
165        let mut max: Option<usize> = None;
166        self.visit(&mut |node| {
167            if let CodecExpr::Param(i) = node {
168                max = Some(max.map_or(*i, |m| m.max(*i)));
169            }
170        });
171        max
172    }
173
174    /// Collect every [`Import`] referenced by the tree.
175    pub(crate) fn collect_imports(&self, into: &mut BTreeSet<Import>) {
176        self.visit(&mut |node| {
177            if let CodecExpr::Import(import) = node {
178                into.insert(import.clone());
179            }
180        });
181    }
182
183    /// Collect every [`CodecExpr::TypeRef`] name in the tree.
184    pub(crate) fn collect_type_refs(&self, into: &mut BTreeSet<String>) {
185        self.visit(&mut |node| {
186            if let CodecExpr::TypeRef(name) = node {
187                into.insert(name.clone());
188            }
189        });
190    }
191
192    /// Render the expression to TypeScript source.
193    ///
194    /// `archived_names` maps Rust type names to their exported archived
195    /// names. A [`CodecExpr::TypeRef`] missing from the map produces
196    /// [`DiagnosticKind::UnresolvedTypeRef`].
197    pub fn render(
198        &self,
199        archived_names: &BTreeMap<String, String>,
200    ) -> Result<String, DiagnosticKind> {
201        match self {
202            CodecExpr::Runtime(name) => Ok(format!("r.{name}")),
203            CodecExpr::Import(import) => Ok(import.export.clone()),
204            CodecExpr::TypeRef(name) => archived_names.get(name).cloned().ok_or_else(|| {
205                DiagnosticKind::UnresolvedTypeRef { name: name.clone() }
206            }),
207            CodecExpr::Call(callee, args) => {
208                let callee = callee.render(archived_names)?;
209                let args = args
210                    .iter()
211                    .map(|a| a.render(archived_names))
212                    .collect::<Result<Vec<_>, _>>()?;
213                Ok(format!("{}({})", callee, args.join(", ")))
214            }
215            CodecExpr::Object(entries) => {
216                if entries.is_empty() {
217                    return Ok("{}".to_string());
218                }
219                let entries = entries
220                    .iter()
221                    .map(|(k, v)| Ok(format!("{}: {}", k, v.render(archived_names)?)))
222                    .collect::<Result<Vec<_>, DiagnosticKind>>()?;
223                Ok(format!("{{ {} }}", entries.join(", ")))
224            }
225            CodecExpr::Array(elements) => {
226                let elements = elements
227                    .iter()
228                    .map(|e| e.render(archived_names))
229                    .collect::<Result<Vec<_>, _>>()?;
230                Ok(format!("[{}]", elements.join(", ")))
231            }
232            CodecExpr::LitInt(n) => Ok(n.to_string()),
233            CodecExpr::Param(i) => panic!(
234                "CodecExpr::Param({i}) escaped template instantiation; registry templates \
235                 must be instantiated before rendering"
236            ),
237            CodecExpr::Raw(ts) => Ok(ts.clone()),
238        }
239    }
240}
241
242/// Builders mirroring the `rkyv-js` runtime combinators.
243///
244/// # Example
245///
246/// ```
247/// use rkyv_js_codegen::codec;
248/// use std::collections::BTreeMap;
249///
250/// let expr = codec::vec(codec::option(codec::u32()));
251/// assert_eq!(expr.render(&BTreeMap::new()).unwrap(), "r.vec(r.option(r.u32))");
252/// ```
253pub mod codec {
254    use super::CodecExpr;
255
256    /// `r.u8`
257    pub fn u8() -> CodecExpr {
258        CodecExpr::runtime("u8")
259    }
260    /// `r.i8`
261    pub fn i8() -> CodecExpr {
262        CodecExpr::runtime("i8")
263    }
264    /// `r.u16`
265    pub fn u16() -> CodecExpr {
266        CodecExpr::runtime("u16")
267    }
268    /// `r.i16`
269    pub fn i16() -> CodecExpr {
270        CodecExpr::runtime("i16")
271    }
272    /// `r.u32`
273    pub fn u32() -> CodecExpr {
274        CodecExpr::runtime("u32")
275    }
276    /// `r.i32`
277    pub fn i32() -> CodecExpr {
278        CodecExpr::runtime("i32")
279    }
280    /// `r.u64`
281    pub fn u64() -> CodecExpr {
282        CodecExpr::runtime("u64")
283    }
284    /// `r.i64`
285    pub fn i64() -> CodecExpr {
286        CodecExpr::runtime("i64")
287    }
288    /// `r.f32`
289    pub fn f32() -> CodecExpr {
290        CodecExpr::runtime("f32")
291    }
292    /// `r.f64`
293    pub fn f64() -> CodecExpr {
294        CodecExpr::runtime("f64")
295    }
296    /// `r.bool`
297    pub fn bool_() -> CodecExpr {
298        CodecExpr::runtime("bool")
299    }
300    /// `r.char`
301    pub fn char_() -> CodecExpr {
302        CodecExpr::runtime("char")
303    }
304    /// `r.unit`
305    pub fn unit() -> CodecExpr {
306        CodecExpr::runtime("unit")
307    }
308    /// `r.string`
309    pub fn string() -> CodecExpr {
310        CodecExpr::runtime("string")
311    }
312    /// `r.vec(inner)`
313    pub fn vec(inner: CodecExpr) -> CodecExpr {
314        CodecExpr::call(CodecExpr::runtime("vec"), [inner])
315    }
316    /// `r.option(inner)`
317    pub fn option(inner: CodecExpr) -> CodecExpr {
318        CodecExpr::call(CodecExpr::runtime("option"), [inner])
319    }
320    /// `r.box(inner)`
321    pub fn boxed(inner: CodecExpr) -> CodecExpr {
322        CodecExpr::call(CodecExpr::runtime("box"), [inner])
323    }
324    /// `r.rc(inner)`
325    pub fn rc(inner: CodecExpr) -> CodecExpr {
326        CodecExpr::call(CodecExpr::runtime("rc"), [inner])
327    }
328    /// `r.weak(inner)`
329    pub fn weak(inner: CodecExpr) -> CodecExpr {
330        CodecExpr::call(CodecExpr::runtime("weak"), [inner])
331    }
332    /// `r.array(inner, len)`
333    pub fn array(inner: CodecExpr, len: u64) -> CodecExpr {
334        CodecExpr::call(CodecExpr::runtime("array"), [inner, CodecExpr::LitInt(len)])
335    }
336    /// `r.tuple(e0, e1, ...)` — the empty tuple is `r.unit`.
337    pub fn tuple(elems: impl IntoIterator<Item = CodecExpr>) -> CodecExpr {
338        let elems: Vec<_> = elems.into_iter().collect();
339        if elems.is_empty() {
340            unit()
341        } else {
342            CodecExpr::call(CodecExpr::runtime("tuple"), elems)
343        }
344    }
345    /// A reference to a generated type by its Rust name; resolved to the
346    /// archived name at `generate()` time.
347    pub fn named(rust_name: impl Into<String>) -> CodecExpr {
348        CodecExpr::type_ref(rust_name)
349    }
350}
351
352/// Generate the import block for a set of expressions.
353///
354/// Always starts with `import * as r from 'rkyv-js';`, followed by named
355/// imports grouped by module and sorted by module specifier (exports sorted
356/// within each statement).
357///
358/// The same export name imported from two different modules is reported as
359/// [`DiagnosticKind::ImportConflict`].
360pub fn generate_import_block<'a>(
361    exprs: impl IntoIterator<Item = &'a CodecExpr>,
362) -> Result<String, Vec<DiagnosticKind>> {
363    let mut imports: BTreeSet<Import> = BTreeSet::new();
364    for expr in exprs {
365        expr.collect_imports(&mut imports);
366    }
367
368    // Detect the same export name pulled from different modules.
369    let mut export_modules: BTreeMap<&str, BTreeSet<&str>> = BTreeMap::new();
370    for import in &imports {
371        export_modules
372            .entry(&import.export)
373            .or_default()
374            .insert(&import.module);
375    }
376    let conflicts: Vec<DiagnosticKind> = export_modules
377        .iter()
378        .filter(|(_, modules)| modules.len() > 1)
379        .map(|(export, modules)| DiagnosticKind::ImportConflict {
380            export: export.to_string(),
381            modules: modules.iter().map(|m| m.to_string()).collect(),
382        })
383        .collect();
384    if !conflicts.is_empty() {
385        return Err(conflicts);
386    }
387
388    let mut by_module: BTreeMap<&str, Vec<&str>> = BTreeMap::new();
389    for import in &imports {
390        by_module
391            .entry(&import.module)
392            .or_default()
393            .push(&import.export);
394    }
395
396    let mut output = String::from("import * as r from 'rkyv-js';\n");
397    for (module, exports) in by_module {
398        output.push_str(&format!(
399            "import {{ {} }} from '{}';\n",
400            exports.join(", "),
401            module
402        ));
403    }
404    Ok(output)
405}
406
407#[cfg(test)]
408mod tests {
409    use super::*;
410
411    fn render(expr: &CodecExpr) -> String {
412        expr.render(&BTreeMap::new()).unwrap()
413    }
414
415    #[test]
416    fn renders_primitives() {
417        assert_eq!(render(&codec::u8()), "r.u8");
418        assert_eq!(render(&codec::u64()), "r.u64");
419        assert_eq!(render(&codec::bool_()), "r.bool");
420        assert_eq!(render(&codec::char_()), "r.char");
421        assert_eq!(render(&codec::unit()), "r.unit");
422        assert_eq!(render(&codec::string()), "r.string");
423    }
424
425    #[test]
426    fn renders_containers() {
427        assert_eq!(render(&codec::vec(codec::u32())), "r.vec(r.u32)");
428        assert_eq!(render(&codec::option(codec::string())), "r.option(r.string)");
429        assert_eq!(render(&codec::boxed(codec::u64())), "r.box(r.u64)");
430        assert_eq!(render(&codec::rc(codec::string())), "r.rc(r.string)");
431        assert_eq!(render(&codec::weak(codec::u32())), "r.weak(r.u32)");
432        assert_eq!(render(&codec::array(codec::u16(), 4)), "r.array(r.u16, 4)");
433        assert_eq!(
434            render(&codec::tuple([codec::u8(), codec::string()])),
435            "r.tuple(r.u8, r.string)"
436        );
437        assert_eq!(render(&codec::tuple([])), "r.unit");
438    }
439
440    #[test]
441    fn renders_nested() {
442        let expr = codec::vec(codec::option(codec::vec(codec::u16())));
443        assert_eq!(render(&expr), "r.vec(r.option(r.vec(r.u16)))");
444    }
445
446    #[test]
447    fn renders_object() {
448        let expr = CodecExpr::object([("a", codec::u8()), ("b", codec::u32())]);
449        assert_eq!(render(&expr), "{ a: r.u8, b: r.u32 }");
450        assert_eq!(render(&CodecExpr::Object(Vec::new())), "{}");
451    }
452
453    #[test]
454    fn renders_array() {
455        let expr = CodecExpr::array([codec::u8(), codec::string()]);
456        assert_eq!(render(&expr), "[r.u8, r.string]");
457        assert_eq!(render(&CodecExpr::Array(Vec::new())), "[]");
458        let template = CodecExpr::array([CodecExpr::Param(0)]);
459        assert_eq!(render(&template.substitute(&[codec::u32()])), "[r.u32]");
460    }
461
462    #[test]
463    fn renders_import_and_raw() {
464        let expr = CodecExpr::call(
465            CodecExpr::import_from("rkyv-js/lib/hashmap", "hashMap"),
466            [codec::string(), codec::u32()],
467        );
468        assert_eq!(render(&expr), "hashMap(r.string, r.u32)");
469        assert_eq!(render(&CodecExpr::raw("myCustom(r.u8)")), "myCustom(r.u8)");
470    }
471
472    #[test]
473    fn renders_type_ref_via_map() {
474        let mut names = BTreeMap::new();
475        names.insert("Point".to_string(), "ArchivedPoint".to_string());
476        let expr = codec::vec(codec::named("Point"));
477        assert_eq!(expr.render(&names).unwrap(), "r.vec(ArchivedPoint)");
478    }
479
480    #[test]
481    fn missing_type_ref_is_an_error() {
482        let expr = codec::named("Missing");
483        let err = expr.render(&BTreeMap::new()).unwrap_err();
484        assert!(matches!(
485            err,
486            DiagnosticKind::UnresolvedTypeRef { ref name } if name == "Missing"
487        ));
488    }
489
490    #[test]
491    fn substitute_replaces_params() {
492        let template = codec::vec(CodecExpr::Param(0));
493        let out = template.substitute(&[codec::u32()]);
494        assert_eq!(render(&out), "r.vec(r.u32)");
495
496        let template = CodecExpr::call(
497            CodecExpr::import_from("m", "pair"),
498            [CodecExpr::Param(0), CodecExpr::Param(1)],
499        );
500        let out = template.substitute(&[codec::string(), codec::u8()]);
501        assert_eq!(render(&out), "pair(r.string, r.u8)");
502    }
503
504    #[test]
505    fn substitute_inside_objects() {
506        let template = CodecExpr::object([("inner", CodecExpr::Param(0))]);
507        let out = template.substitute(&[codec::u8()]);
508        assert_eq!(render(&out), "{ inner: r.u8 }");
509    }
510
511    #[test]
512    fn max_param_walks_the_tree() {
513        assert_eq!(codec::u8().max_param(), None);
514        let expr = CodecExpr::call(
515            CodecExpr::runtime("x"),
516            [CodecExpr::Param(0), codec::vec(CodecExpr::Param(3))],
517        );
518        assert_eq!(expr.max_param(), Some(3));
519    }
520
521    #[test]
522    fn raw_is_never_inspected() {
523        let expr = CodecExpr::raw("hashMap(r.u8)");
524        let mut imports = BTreeSet::new();
525        expr.collect_imports(&mut imports);
526        assert!(imports.is_empty());
527        let mut refs = BTreeSet::new();
528        expr.collect_type_refs(&mut refs);
529        assert!(refs.is_empty());
530        assert_eq!(expr.max_param(), None);
531    }
532
533    #[test]
534    fn import_block_groups_and_sorts() {
535        let exprs = [
536            CodecExpr::import_from("rkyv-js/lib/indexmap", "indexSet"),
537            CodecExpr::import_from("rkyv-js/lib/indexmap", "indexMap"),
538            CodecExpr::import_from("rkyv-js/lib/bytes", "bytes"),
539            codec::u8(),
540        ];
541        let block = generate_import_block(exprs.iter()).unwrap();
542        assert_eq!(
543            block,
544            "import * as r from 'rkyv-js';\n\
545             import { bytes } from 'rkyv-js/lib/bytes';\n\
546             import { indexMap, indexSet } from 'rkyv-js/lib/indexmap';\n"
547        );
548    }
549
550    #[test]
551    fn import_block_dedups() {
552        let exprs = [
553            CodecExpr::import_from("rkyv-js/lib/uuid", "uuid"),
554            CodecExpr::import_from("rkyv-js/lib/uuid", "uuid"),
555        ];
556        let block = generate_import_block(exprs.iter()).unwrap();
557        assert_eq!(
558            block,
559            "import * as r from 'rkyv-js';\nimport { uuid } from 'rkyv-js/lib/uuid';\n"
560        );
561    }
562
563    #[test]
564    fn import_block_detects_conflicts() {
565        let exprs = [
566            CodecExpr::import_from("pkg-a", "codec"),
567            CodecExpr::import_from("pkg-b", "codec"),
568        ];
569        let errs = generate_import_block(exprs.iter()).unwrap_err();
570        assert_eq!(errs.len(), 1);
571        assert!(matches!(
572            &errs[0],
573            DiagnosticKind::ImportConflict { export, modules }
574                if export == "codec" && modules == &vec!["pkg-a".to_string(), "pkg-b".to_string()]
575        ));
576    }
577
578    #[test]
579    #[should_panic(expected = "escaped template instantiation")]
580    fn rendering_a_param_panics() {
581        let _ = CodecExpr::Param(0).render(&BTreeMap::new());
582    }
583}
rkyv_js_codegen/expr.rs

rkyv_js_codegen/
expr.rs
