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 still visited).
139    pub fn visit(&self, f: &mut impl FnMut(&CodecExpr)) {
140        f(self);
141        match self {
142            CodecExpr::Call(callee, args) => {
143                callee.visit(f);
144                for arg in args {
145                    arg.visit(f);
146                }
147            }
148            CodecExpr::Object(entries) => {
149                for (_, v) in entries {
150                    v.visit(f);
151                }
152            }
153            CodecExpr::Array(elements) => {
154                for element in elements {
155                    element.visit(f);
156                }
157            }
158            _ => {}
159        }
160    }
161
162    /// The highest [`CodecExpr::Param`] index in the tree, if any.
163    pub fn max_param(&self) -> Option<usize> {
164        let mut max: Option<usize> = None;
165        self.visit(&mut |node| {
166            if let CodecExpr::Param(i) = node {
167                max = Some(max.map_or(*i, |m| m.max(*i)));
168            }
169        });
170        max
171    }
172
173    /// Collect every [`Import`] referenced by the tree.
174    pub(crate) fn collect_imports(&self, into: &mut BTreeSet<Import>) {
175        self.visit(&mut |node| {
176            if let CodecExpr::Import(import) = node {
177                into.insert(import.clone());
178            }
179        });
180    }
181
182    /// Collect every [`CodecExpr::TypeRef`] name in the tree.
183    pub(crate) fn collect_type_refs(&self, into: &mut BTreeSet<String>) {
184        self.visit(&mut |node| {
185            if let CodecExpr::TypeRef(name) = node {
186                into.insert(name.clone());
187            }
188        });
189    }
190
191    /// Render the expression to TypeScript source.
192    ///
193    /// `archived_names` maps Rust type names to their exported archived names.
194    /// A [`CodecExpr::TypeRef`] missing from the map produces [`DiagnosticKind::UnresolvedTypeRef`].
195    pub fn render(
196        &self,
197        archived_names: &BTreeMap<String, String>,
198    ) -> Result<String, DiagnosticKind> {
199        match self {
200            CodecExpr::Runtime(name) => Ok(format!("r.{name}")),
201            CodecExpr::Import(import) => Ok(import.export.clone()),
202            CodecExpr::TypeRef(name) => archived_names.get(name).cloned().ok_or_else(|| {
203                DiagnosticKind::UnresolvedTypeRef { name: name.clone() }
204            }),
205            CodecExpr::Call(callee, args) => {
206                let callee = callee.render(archived_names)?;
207                let args = args
208                    .iter()
209                    .map(|a| a.render(archived_names))
210                    .collect::<Result<Vec<_>, _>>()?;
211                Ok(format!("{}({})", callee, args.join(", ")))
212            }
213            CodecExpr::Object(entries) => {
214                if entries.is_empty() {
215                    return Ok("{}".to_string());
216                }
217                let entries = entries
218                    .iter()
219                    .map(|(k, v)| Ok(format!("{}: {}", k, v.render(archived_names)?)))
220                    .collect::<Result<Vec<_>, DiagnosticKind>>()?;
221                Ok(format!("{{ {} }}", entries.join(", ")))
222            }
223            CodecExpr::Array(elements) => {
224                let elements = elements
225                    .iter()
226                    .map(|e| e.render(archived_names))
227                    .collect::<Result<Vec<_>, _>>()?;
228                Ok(format!("[{}]", elements.join(", ")))
229            }
230            CodecExpr::LitInt(n) => Ok(n.to_string()),
231            CodecExpr::Param(i) => panic!(
232                "CodecExpr::Param({i}) escaped template instantiation; registry templates \
233                 must be instantiated before rendering"
234            ),
235            CodecExpr::Raw(ts) => Ok(ts.clone()),
236        }
237    }
238}
239
240/// Builders mirroring the `rkyv-js` runtime combinators.
241///
242/// # Example
243///
244/// ```
245/// use rkyv_js_codegen::codec;
246/// use std::collections::BTreeMap;
247///
248/// let expr = codec::vec(codec::option(codec::u32()));
249/// assert_eq!(expr.render(&BTreeMap::new()).unwrap(), "r.vec(r.option(r.u32))");
250/// ```
251pub mod codec {
252    use super::CodecExpr;
253
254    /// `r.u8`
255    pub fn u8() -> CodecExpr {
256        CodecExpr::runtime("u8")
257    }
258    /// `r.i8`
259    pub fn i8() -> CodecExpr {
260        CodecExpr::runtime("i8")
261    }
262    /// `r.u16`
263    pub fn u16() -> CodecExpr {
264        CodecExpr::runtime("u16")
265    }
266    /// `r.i16`
267    pub fn i16() -> CodecExpr {
268        CodecExpr::runtime("i16")
269    }
270    /// `r.u32`
271    pub fn u32() -> CodecExpr {
272        CodecExpr::runtime("u32")
273    }
274    /// `r.i32`
275    pub fn i32() -> CodecExpr {
276        CodecExpr::runtime("i32")
277    }
278    /// `r.u64`
279    pub fn u64() -> CodecExpr {
280        CodecExpr::runtime("u64")
281    }
282    /// `r.i64`
283    pub fn i64() -> CodecExpr {
284        CodecExpr::runtime("i64")
285    }
286    /// `r.f32`
287    pub fn f32() -> CodecExpr {
288        CodecExpr::runtime("f32")
289    }
290    /// `r.f64`
291    pub fn f64() -> CodecExpr {
292        CodecExpr::runtime("f64")
293    }
294    /// `r.bool`
295    pub fn bool_() -> CodecExpr {
296        CodecExpr::runtime("bool")
297    }
298    /// `r.char`
299    pub fn char_() -> CodecExpr {
300        CodecExpr::runtime("char")
301    }
302    /// `r.unit`
303    pub fn unit() -> CodecExpr {
304        CodecExpr::runtime("unit")
305    }
306    /// `r.string`
307    pub fn string() -> CodecExpr {
308        CodecExpr::runtime("string")
309    }
310    /// `r.vec(inner)`
311    pub fn vec(inner: CodecExpr) -> CodecExpr {
312        CodecExpr::call(CodecExpr::runtime("vec"), [inner])
313    }
314    /// `r.option(inner)`
315    pub fn option(inner: CodecExpr) -> CodecExpr {
316        CodecExpr::call(CodecExpr::runtime("option"), [inner])
317    }
318    /// `r.box(inner)`
319    pub fn boxed(inner: CodecExpr) -> CodecExpr {
320        CodecExpr::call(CodecExpr::runtime("box"), [inner])
321    }
322    /// `r.rc(inner)`
323    pub fn rc(inner: CodecExpr) -> CodecExpr {
324        CodecExpr::call(CodecExpr::runtime("rc"), [inner])
325    }
326    /// `r.weak(inner)`
327    pub fn weak(inner: CodecExpr) -> CodecExpr {
328        CodecExpr::call(CodecExpr::runtime("weak"), [inner])
329    }
330    /// `r.array(inner, len)`
331    pub fn array(inner: CodecExpr, len: u64) -> CodecExpr {
332        CodecExpr::call(CodecExpr::runtime("array"), [inner, CodecExpr::LitInt(len)])
333    }
334    /// `r.tuple(e0, e1, ...)` - the empty tuple is `r.unit`.
335    pub fn tuple(elems: impl IntoIterator<Item = CodecExpr>) -> CodecExpr {
336        let elems: Vec<_> = elems.into_iter().collect();
337        if elems.is_empty() {
338            unit()
339        } else {
340            CodecExpr::call(CodecExpr::runtime("tuple"), elems)
341        }
342    }
343    /// A reference to a generated type by its Rust name;
344    /// resolved to the archived name at `generate()` time.
345    pub fn named(rust_name: impl Into<String>) -> CodecExpr {
346        CodecExpr::type_ref(rust_name)
347    }
348}
349
350/// Generate the import block for a set of expressions.
351///
352/// Always starts with `import * as r from 'rkyv-js';`, followed by named
353/// imports grouped by module and sorted by module specifier
354/// (exports sorted within each statement).
355///
356/// The same export name imported from two different modules is reported as
357/// [`DiagnosticKind::ImportConflict`].
358pub fn generate_import_block<'a>(
359    exprs: impl IntoIterator<Item = &'a CodecExpr>,
360) -> Result<String, Vec<DiagnosticKind>> {
361    let mut imports: BTreeSet<Import> = BTreeSet::new();
362    for expr in exprs {
363        expr.collect_imports(&mut imports);
364    }
365
366    // Detect the same export name pulled from different modules.
367    let mut export_modules: BTreeMap<&str, BTreeSet<&str>> = BTreeMap::new();
368    for import in &imports {
369        export_modules
370            .entry(&import.export)
371            .or_default()
372            .insert(&import.module);
373    }
374    let conflicts: Vec<DiagnosticKind> = export_modules
375        .iter()
376        .filter(|(_, modules)| modules.len() > 1)
377        .map(|(export, modules)| DiagnosticKind::ImportConflict {
378            export: export.to_string(),
379            modules: modules.iter().map(|m| m.to_string()).collect(),
380        })
381        .collect();
382    if !conflicts.is_empty() {
383        return Err(conflicts);
384    }
385
386    let mut by_module: BTreeMap<&str, Vec<&str>> = BTreeMap::new();
387    for import in &imports {
388        by_module
389            .entry(&import.module)
390            .or_default()
391            .push(&import.export);
392    }
393
394    let mut output = String::from("import * as r from 'rkyv-js';\n");
395    for (module, exports) in by_module {
396        output.push_str(&format!(
397            "import {{ {} }} from '{}';\n",
398            exports.join(", "),
399            module
400        ));
401    }
402    Ok(output)
403}
404
405#[cfg(test)]
406mod tests {
407    use super::*;
408
409    fn render(expr: &CodecExpr) -> String {
410        expr.render(&BTreeMap::new()).unwrap()
411    }
412
413    #[test]
414    fn renders_primitives() {
415        assert_eq!(render(&codec::u8()), "r.u8");
416        assert_eq!(render(&codec::u64()), "r.u64");
417        assert_eq!(render(&codec::bool_()), "r.bool");
418        assert_eq!(render(&codec::char_()), "r.char");
419        assert_eq!(render(&codec::unit()), "r.unit");
420        assert_eq!(render(&codec::string()), "r.string");
421    }
422
423    #[test]
424    fn renders_containers() {
425        assert_eq!(render(&codec::vec(codec::u32())), "r.vec(r.u32)");
426        assert_eq!(render(&codec::option(codec::string())), "r.option(r.string)");
427        assert_eq!(render(&codec::boxed(codec::u64())), "r.box(r.u64)");
428        assert_eq!(render(&codec::rc(codec::string())), "r.rc(r.string)");
429        assert_eq!(render(&codec::weak(codec::u32())), "r.weak(r.u32)");
430        assert_eq!(render(&codec::array(codec::u16(), 4)), "r.array(r.u16, 4)");
431        assert_eq!(
432            render(&codec::tuple([codec::u8(), codec::string()])),
433            "r.tuple(r.u8, r.string)"
434        );
435        assert_eq!(render(&codec::tuple([])), "r.unit");
436    }
437
438    #[test]
439    fn renders_nested() {
440        let expr = codec::vec(codec::option(codec::vec(codec::u16())));
441        assert_eq!(render(&expr), "r.vec(r.option(r.vec(r.u16)))");
442    }
443
444    #[test]
445    fn renders_object() {
446        let expr = CodecExpr::object([("a", codec::u8()), ("b", codec::u32())]);
447        assert_eq!(render(&expr), "{ a: r.u8, b: r.u32 }");
448        assert_eq!(render(&CodecExpr::Object(Vec::new())), "{}");
449    }
450
451    #[test]
452    fn renders_array() {
453        let expr = CodecExpr::array([codec::u8(), codec::string()]);
454        assert_eq!(render(&expr), "[r.u8, r.string]");
455        assert_eq!(render(&CodecExpr::Array(Vec::new())), "[]");
456        let template = CodecExpr::array([CodecExpr::Param(0)]);
457        assert_eq!(render(&template.substitute(&[codec::u32()])), "[r.u32]");
458    }
459
460    #[test]
461    fn renders_import_and_raw() {
462        let expr = CodecExpr::call(
463            CodecExpr::import_from("rkyv-js/lib/hashmap", "hashMap"),
464            [codec::string(), codec::u32()],
465        );
466        assert_eq!(render(&expr), "hashMap(r.string, r.u32)");
467        assert_eq!(render(&CodecExpr::raw("myCustom(r.u8)")), "myCustom(r.u8)");
468    }
469
470    #[test]
471    fn renders_type_ref_via_map() {
472        let mut names = BTreeMap::new();
473        names.insert("Point".to_string(), "ArchivedPoint".to_string());
474        let expr = codec::vec(codec::named("Point"));
475        assert_eq!(expr.render(&names).unwrap(), "r.vec(ArchivedPoint)");
476    }
477
478    #[test]
479    fn missing_type_ref_is_an_error() {
480        let expr = codec::named("Missing");
481        let err = expr.render(&BTreeMap::new()).unwrap_err();
482        assert!(matches!(
483            err,
484            DiagnosticKind::UnresolvedTypeRef { ref name } if name == "Missing"
485        ));
486    }
487
488    #[test]
489    fn substitute_replaces_params() {
490        let template = codec::vec(CodecExpr::Param(0));
491        let out = template.substitute(&[codec::u32()]);
492        assert_eq!(render(&out), "r.vec(r.u32)");
493
494        let template = CodecExpr::call(
495            CodecExpr::import_from("m", "pair"),
496            [CodecExpr::Param(0), CodecExpr::Param(1)],
497        );
498        let out = template.substitute(&[codec::string(), codec::u8()]);
499        assert_eq!(render(&out), "pair(r.string, r.u8)");
500    }
501
502    #[test]
503    fn substitute_inside_objects() {
504        let template = CodecExpr::object([("inner", CodecExpr::Param(0))]);
505        let out = template.substitute(&[codec::u8()]);
506        assert_eq!(render(&out), "{ inner: r.u8 }");
507    }
508
509    #[test]
510    fn max_param_walks_the_tree() {
511        assert_eq!(codec::u8().max_param(), None);
512        let expr = CodecExpr::call(
513            CodecExpr::runtime("x"),
514            [CodecExpr::Param(0), codec::vec(CodecExpr::Param(3))],
515        );
516        assert_eq!(expr.max_param(), Some(3));
517    }
518
519    #[test]
520    fn raw_is_never_inspected() {
521        let expr = CodecExpr::raw("hashMap(r.u8)");
522        let mut imports = BTreeSet::new();
523        expr.collect_imports(&mut imports);
524        assert!(imports.is_empty());
525        let mut refs = BTreeSet::new();
526        expr.collect_type_refs(&mut refs);
527        assert!(refs.is_empty());
528        assert_eq!(expr.max_param(), None);
529    }
530
531    #[test]
532    fn import_block_groups_and_sorts() {
533        let exprs = [
534            CodecExpr::import_from("rkyv-js/lib/indexmap", "indexSet"),
535            CodecExpr::import_from("rkyv-js/lib/indexmap", "indexMap"),
536            CodecExpr::import_from("rkyv-js/lib/bytes", "bytes"),
537            codec::u8(),
538        ];
539        let block = generate_import_block(exprs.iter()).unwrap();
540        assert_eq!(
541            block,
542            "import * as r from 'rkyv-js';\n\
543             import { bytes } from 'rkyv-js/lib/bytes';\n\
544             import { indexMap, indexSet } from 'rkyv-js/lib/indexmap';\n"
545        );
546    }
547
548    #[test]
549    fn import_block_dedups() {
550        let exprs = [
551            CodecExpr::import_from("rkyv-js/lib/uuid", "uuid"),
552            CodecExpr::import_from("rkyv-js/lib/uuid", "uuid"),
553        ];
554        let block = generate_import_block(exprs.iter()).unwrap();
555        assert_eq!(
556            block,
557            "import * as r from 'rkyv-js';\nimport { uuid } from 'rkyv-js/lib/uuid';\n"
558        );
559    }
560
561    #[test]
562    fn import_block_detects_conflicts() {
563        let exprs = [
564            CodecExpr::import_from("pkg-a", "codec"),
565            CodecExpr::import_from("pkg-b", "codec"),
566        ];
567        let errs = generate_import_block(exprs.iter()).unwrap_err();
568        assert_eq!(errs.len(), 1);
569        assert!(matches!(
570            &errs[0],
571            DiagnosticKind::ImportConflict { export, modules }
572                if export == "codec" && modules == &vec!["pkg-a".to_string(), "pkg-b".to_string()]
573        ));
574    }
575
576    #[test]
577    #[should_panic(expected = "escaped template instantiation")]
578    fn rendering_a_param_panics() {
579        let _ = CodecExpr::Param(0).render(&BTreeMap::new());
580    }
581}
rkyv_js_codegen/expr.rs

rkyv_js_codegen/
expr.rs
