rkyv_js_codegen/
registry.rs

1//! Registries mapping Rust paths to codec templates.
2//!
3//! Two registries drive source extraction:
4//!
5//! - [`ExternalType`] maps a fully-qualified Rust *type* path (e.g. `uuid::Uuid`, `std::collections::HashMap`) 
6//!   to a [`CodecExpr`] template.
7//! - [`WithWrapper`] maps a `#[rkyv(with = ...)]` *wrapper* path (e.g. `rkyv::with::AsBox`)
8//!   to a transformation of the underlying field codec.
9//!
10//! Both are keyed by fully-qualified path strings. 
11//!
12//! Unknown-path lookups produce a did-you-mean suggestion when a registered key shares the last path segment.
13
14use std::collections::BTreeMap;
15
16use crate::error::DiagnosticKind;
17use crate::expr::{CodecExpr, codec};
18
19/// A codec template for an external Rust type.
20///
21/// Templates are built once at registration time;
22/// type arguments are filled in per use site via [`CodecExpr::Param`] placeholders.
23#[derive(Debug, Clone)]
24pub struct ExternalType {
25    arity: usize,
26    allow_trailing: bool,
27    template: CodecExpr,
28}
29
30impl ExternalType {
31    /// A type with no type parameters (e.g. `uuid::Uuid`).
32    pub fn leaf(expr: CodecExpr) -> Self {
33        Self {
34            arity: 0,
35            allow_trailing: false,
36            template: expr,
37        }
38    }
39
40    /// A type with one type parameter.
41    /// The closure runs **once** with `Param(0)` to build the template.
42    pub fn generic1(build: impl FnOnce(CodecExpr) -> CodecExpr) -> Self {
43        Self::generic(1, |params| build(params[0].clone()))
44    }
45
46    /// A type with two type parameters.
47    /// The closure runs **once** with `Param(0)` and `Param(1)`.
48    pub fn generic2(build: impl FnOnce(CodecExpr, CodecExpr) -> CodecExpr) -> Self {
49        Self::generic(2, |params| build(params[0].clone(), params[1].clone()))
50    }
51
52    /// A type with `arity` type parameters.
53    /// The closure runs **once** with `[Param(0), ..., Param(arity - 1)]`.
54    ///
55    /// # Panics
56    ///
57    /// Panics if the produced template references a `Param` index `>= arity`.
58    pub fn generic(arity: usize, build: impl FnOnce(&[CodecExpr]) -> CodecExpr) -> Self {
59        let params: Vec<CodecExpr> = (0..arity).map(CodecExpr::Param).collect();
60        let template = build(&params);
61        if let Some(max) = template.max_param()
62            && max >= arity
63        {
64            panic!(
65                "ExternalType::generic template references Param({max}), but the declared \
66                 arity is {arity}"
67            );
68        }
69        Self {
70            arity,
71            allow_trailing: false,
72            template,
73        }
74    }
75
76    /// Accept (and ignore) extra trailing type arguments beyond the declared arity.
77    /// e.g. the hasher parameter of `HashMap<K, V, S>`.
78    pub fn allow_trailing_args(mut self) -> Self {
79        self.allow_trailing = true;
80        self
81    }
82
83    /// The number of type parameters the template consumes.
84    pub(crate) fn arity(&self) -> usize {
85        self.arity
86    }
87
88    /// Whether extra trailing type arguments are tolerated.
89    pub(crate) fn allows_trailing(&self) -> bool {
90        self.allow_trailing
91    }
92
93    /// Fill in the template with concrete type arguments.
94    ///
95    /// The argument count must match the declared arity exactly,
96    /// unless [`allow_trailing_args`](ExternalType::allow_trailing_args) was set,
97    /// in which case extra trailing arguments are ignored.
98    ///
99    /// The `rust_path` of a returned [`DiagnosticKind::GenericArity`] is left empty;
100    /// the caller fills it in with the use-site path.
101    pub(crate) fn instantiate(&self, args: Vec<CodecExpr>) -> Result<CodecExpr, DiagnosticKind> {
102        let acceptable = args.len() == self.arity || (self.allow_trailing && args.len() > self.arity);
103        if !acceptable {
104            return Err(DiagnosticKind::GenericArity {
105                rust_path: String::new(),
106                expected: self.arity,
107                found: args.len(),
108            });
109        }
110        Ok(self.template.substitute(&args[..self.arity]))
111    }
112}
113
114/// The behavior of a with-wrapper.
115#[derive(Debug, Clone)]
116enum WithWrapperKind {
117    /// Emit a fixed expression, ignoring the underlying field type.
118    Replace(CodecExpr),
119    /// Transform the underlying field codec (template with `Param(0)`).
120    Map(CodecExpr),
121    /// Use the underlying field codec unchanged.
122    Identity,
123    /// Omit the field from the generated bindings entirely.
124    Skip,
125}
126
127/// A handler for a `#[rkyv(with = W)]` field wrapper.
128#[derive(Debug, Clone)]
129pub struct WithWrapper {
130    kind: WithWrapperKind,
131}
132
133impl WithWrapper {
134    /// Emit `expr` for the field, ignoring the underlying Rust type
135    /// (e.g. an `AsJson` wrapper backed by a custom codec).
136    pub fn replace(expr: CodecExpr) -> Self {
137        Self {
138            kind: WithWrapperKind::Replace(expr),
139        }
140    }
141
142    /// Transform the underlying field codec.
143    /// The closure runs **once** with `Param(0)` standing in for the underlying codec expression.
144    pub fn map(build: impl FnOnce(CodecExpr) -> CodecExpr) -> Self {
145        Self {
146            kind: WithWrapperKind::Map(build(CodecExpr::Param(0))),
147        }
148    }
149
150    /// Use the underlying field codec unchanged (e.g. `rkyv::with::Inline`).
151    pub fn identity() -> Self {
152        Self {
153            kind: WithWrapperKind::Identity,
154        }
155    }
156
157    /// Omit the field entirely (`rkyv::with::Skip`).
158    pub fn skip() -> Self {
159        Self {
160            kind: WithWrapperKind::Skip,
161        }
162    }
163
164    /// Whether this wrapper needs the underlying field type resolved.
165    pub(crate) fn needs_underlying(&self) -> bool {
166        matches!(
167            self.kind,
168            WithWrapperKind::Map(_) | WithWrapperKind::Identity
169        )
170    }
171
172    /// Apply the wrapper. `underlying` is only consulted for [`map`](WithWrapper::map) and [`identity`](WithWrapper::identity) wrappers;
173    /// `None` is returned for [`skip`](WithWrapper::skip).
174    pub(crate) fn apply(&self, underlying: Option<CodecExpr>) -> Option<CodecExpr> {
175        match &self.kind {
176            WithWrapperKind::Replace(expr) => Some(expr.clone()),
177            WithWrapperKind::Map(template) => {
178                let underlying = underlying.expect("map wrapper requires the underlying codec");
179                Some(template.substitute(&[underlying]))
180            }
181            WithWrapperKind::Identity => {
182                Some(underlying.expect("identity wrapper requires the underlying codec"))
183            }
184            WithWrapperKind::Skip => None,
185        }
186    }
187}
188
189/// The registries backing a [`CodeGenerator`](crate::CodeGenerator).
190#[derive(Debug, Clone)]
191pub(crate) struct Registry {
192    types: BTreeMap<String, ExternalType>,
193    wrappers: BTreeMap<String, WithWrapper>,
194}
195
196impl Registry {
197    /// An empty registry.
198    pub(crate) fn empty() -> Self {
199        Self {
200            types: BTreeMap::new(),
201            wrappers: BTreeMap::new(),
202        }
203    }
204
205    /// A registry pre-populated with the built-in rkyv mappings.
206    pub(crate) fn with_builtins() -> Self {
207        let mut registry = Self::empty();
208
209        registry.register_type(
210            "uuid::Uuid",
211            ExternalType::leaf(CodecExpr::import_from("rkyv-js/lib/uuid", "uuid")),
212        );
213        registry.register_type(
214            "bytes::Bytes",
215            ExternalType::leaf(CodecExpr::import_from("rkyv-js/lib/bytes", "bytes")),
216        );
217        registry.register_type("smol_str::SmolStr", ExternalType::leaf(codec::string()));
218
219        // Vec-shaped containers.
220        registry.register_type(
221            "std::collections::VecDeque",
222            ExternalType::generic1(codec::vec),
223        );
224        registry.register_type("thin_vec::ThinVec", ExternalType::generic1(codec::vec));
225        // `ArrayVec<T, N>`: the const-generic capacity is skipped during argument collection, but tolerate it anyway.
226        registry.register_type(
227            "arrayvec::ArrayVec",
228            ExternalType::generic1(codec::vec).allow_trailing_args(),
229        );
230        // `SmallVec<[T; N]>` / `TinyVec<[T; N]>`: the array argument is unwrapped to `T` during argument collection.
231        registry.register_type("smallvec::SmallVec", ExternalType::generic1(codec::vec));
232        registry.register_type("tinyvec::TinyVec", ExternalType::generic1(codec::vec));
233
234        // BTree collections.
235        registry.register_type(
236            "std::collections::BTreeMap",
237            ExternalType::generic2(|k, v| {
238                CodecExpr::call(CodecExpr::import_from("rkyv-js/lib/btreemap", "btreeMap"), [k, v])
239            }),
240        );
241        registry.register_type(
242            "std::collections::BTreeSet",
243            ExternalType::generic1(|t| {
244                CodecExpr::call(CodecExpr::import_from("rkyv-js/lib/btreemap", "btreeSet"), [t])
245            }),
246        );
247
248        // Hash collections (trailing hasher parameter allowed).
249        for path in ["std::collections::HashMap", "hashbrown::HashMap"] {
250            registry.register_type(
251                path,
252                ExternalType::generic2(|k, v| {
253                    CodecExpr::call(CodecExpr::import_from("rkyv-js/lib/hashmap", "hashMap"), [k, v])
254                })
255                .allow_trailing_args(),
256            );
257        }
258        for path in ["std::collections::HashSet", "hashbrown::HashSet"] {
259            registry.register_type(
260                path,
261                ExternalType::generic1(|t| {
262                    CodecExpr::call(CodecExpr::import_from("rkyv-js/lib/hashmap", "hashSet"), [t])
263                })
264                .allow_trailing_args(),
265            );
266        }
267
268        // Index collections (trailing hasher parameter allowed).
269        registry.register_type(
270            "indexmap::IndexMap",
271            ExternalType::generic2(|k, v| {
272                CodecExpr::call(CodecExpr::import_from("rkyv-js/lib/indexmap", "indexMap"), [k, v])
273            })
274            .allow_trailing_args(),
275        );
276        registry.register_type(
277            "indexmap::IndexSet",
278            ExternalType::generic1(|t| {
279                CodecExpr::call(CodecExpr::import_from("rkyv-js/lib/indexmap", "indexSet"), [t])
280            })
281            .allow_trailing_args(),
282        );
283
284        // Shared pointers.
285        for path in ["std::rc::Rc", "std::sync::Arc", "triomphe::Arc"] {
286            registry.register_type(path, ExternalType::generic1(codec::rc));
287        }
288        for path in ["std::rc::Weak", "std::sync::Weak"] {
289            registry.register_type(path, ExternalType::generic1(codec::weak));
290        }
291
292        // Built-in with-wrappers.
293        registry.register_wrapper("rkyv::with::AsBox", WithWrapper::map(codec::boxed));
294        registry.register_wrapper("rkyv::with::Inline", WithWrapper::identity());
295        registry.register_wrapper("rkyv::with::InlineAsBox", WithWrapper::map(codec::boxed));
296        registry.register_wrapper("rkyv::with::Skip", WithWrapper::skip());
297
298        registry
299    }
300
301    pub(crate) fn register_type(&mut self, path: impl Into<String>, external: ExternalType) {
302        self.types.insert(path.into(), external);
303    }
304
305    pub(crate) fn unregister_type(&mut self, path: &str) {
306        self.types.remove(path);
307    }
308
309    pub(crate) fn get_type(&self, path: &str) -> Option<&ExternalType> {
310        self.types.get(path)
311    }
312
313    pub(crate) fn register_wrapper(&mut self, path: impl Into<String>, wrapper: WithWrapper) {
314        self.wrappers.insert(path.into(), wrapper);
315    }
316
317    pub(crate) fn get_wrapper(&self, path: &str) -> Option<&WithWrapper> {
318        self.wrappers.get(path)
319    }
320
321    /// A registered type path sharing the last segment with `path`, if any.
322    pub(crate) fn suggest_type(&self, path: &str) -> Option<String> {
323        let last = path.rsplit("::").next()?;
324        self.types
325            .keys()
326            .find(|key| key.rsplit("::").next() == Some(last))
327            .cloned()
328    }
329}
330
331#[cfg(test)]
332mod tests {
333    use super::*;
334    use std::collections::BTreeMap as Map;
335
336    fn render(expr: &CodecExpr) -> String {
337        expr.render(&Map::new()).unwrap()
338    }
339
340    #[test]
341    fn leaf_instantiates_with_no_args() {
342        let uuid = ExternalType::leaf(CodecExpr::import_from("rkyv-js/lib/uuid", "uuid"));
343        let expr = uuid.instantiate(vec![]).unwrap();
344        assert_eq!(render(&expr), "uuid");
345    }
346
347    #[test]
348    fn generic1_builds_once_with_param0() {
349        let vec_like = ExternalType::generic1(codec::vec);
350        let expr = vec_like.instantiate(vec![codec::u32()]).unwrap();
351        assert_eq!(render(&expr), "r.vec(r.u32)");
352    }
353
354    #[test]
355    fn generic2_instantiates_in_order() {
356        let map_like = ExternalType::generic2(|k, v| {
357            CodecExpr::call(CodecExpr::import_from("m", "pair"), [k, v])
358        });
359        let expr = map_like
360            .instantiate(vec![codec::string(), codec::u32()])
361            .unwrap();
362        assert_eq!(render(&expr), "pair(r.string, r.u32)");
363    }
364
365    #[test]
366    fn too_few_args_is_an_arity_error() {
367        let map_like = ExternalType::generic2(|k, v| {
368            CodecExpr::call(CodecExpr::runtime("pair"), [k, v])
369        });
370        let err = map_like.instantiate(vec![codec::string()]).unwrap_err();
371        assert!(matches!(
372            err,
373            DiagnosticKind::GenericArity {
374                expected: 2,
375                found: 1,
376                ..
377            }
378        ));
379    }
380
381    #[test]
382    fn too_many_args_is_an_arity_error_without_trailing() {
383        let set_like = ExternalType::generic1(codec::vec);
384        let err = set_like
385            .instantiate(vec![codec::u8(), codec::u16()])
386            .unwrap_err();
387        assert!(matches!(
388            err,
389            DiagnosticKind::GenericArity {
390                expected: 1,
391                found: 2,
392                ..
393            }
394        ));
395    }
396
397    #[test]
398    fn trailing_args_are_ignored_when_allowed() {
399        let map_like = ExternalType::generic2(|k, v| {
400            CodecExpr::call(CodecExpr::import_from("m", "hashMap"), [k, v])
401        })
402        .allow_trailing_args();
403        let expr = map_like
404            .instantiate(vec![codec::string(), codec::u32(), codec::u8()])
405            .unwrap();
406        assert_eq!(render(&expr), "hashMap(r.string, r.u32)");
407        // Too few args still fail even with trailing allowed.
408        assert!(map_like.instantiate(vec![codec::string()]).is_err());
409    }
410
411    #[test]
412    #[should_panic(expected = "references Param(1)")]
413    fn generic_panics_on_out_of_range_param() {
414        let _ = ExternalType::generic(1, |_| {
415            CodecExpr::call(CodecExpr::runtime("x"), [CodecExpr::Param(1)])
416        });
417    }
418
419    #[test]
420    fn wrapper_replace_ignores_underlying() {
421        let wrapper = WithWrapper::replace(CodecExpr::import_from("./coord.ts", "Coord"));
422        assert!(!wrapper.needs_underlying());
423        let expr = wrapper.apply(None).unwrap();
424        assert_eq!(render(&expr), "Coord");
425    }
426
427    #[test]
428    fn wrapper_map_transforms_underlying() {
429        let wrapper = WithWrapper::map(codec::boxed);
430        assert!(wrapper.needs_underlying());
431        let expr = wrapper.apply(Some(codec::string())).unwrap();
432        assert_eq!(render(&expr), "r.box(r.string)");
433    }
434
435    #[test]
436    fn wrapper_identity_passes_through() {
437        let wrapper = WithWrapper::identity();
438        let expr = wrapper.apply(Some(codec::u32())).unwrap();
439        assert_eq!(render(&expr), "r.u32");
440    }
441
442    #[test]
443    fn wrapper_skip_omits() {
444        let wrapper = WithWrapper::skip();
445        assert!(!wrapper.needs_underlying());
446        assert!(wrapper.apply(None).is_none());
447    }
448
449    #[test]
450    fn builtins_are_registered() {
451        let registry = Registry::with_builtins();
452        for path in [
453            "uuid::Uuid",
454            "bytes::Bytes",
455            "smol_str::SmolStr",
456            "std::collections::VecDeque",
457            "thin_vec::ThinVec",
458            "arrayvec::ArrayVec",
459            "smallvec::SmallVec",
460            "tinyvec::TinyVec",
461            "std::collections::BTreeMap",
462            "std::collections::BTreeSet",
463            "std::collections::HashMap",
464            "std::collections::HashSet",
465            "hashbrown::HashMap",
466            "hashbrown::HashSet",
467            "indexmap::IndexMap",
468            "indexmap::IndexSet",
469            "std::rc::Rc",
470            "std::sync::Arc",
471            "triomphe::Arc",
472            "std::rc::Weak",
473            "std::sync::Weak",
474        ] {
475            assert!(registry.get_type(path).is_some(), "missing builtin {path}");
476        }
477        for path in [
478            "rkyv::with::AsBox",
479            "rkyv::with::Inline",
480            "rkyv::with::InlineAsBox",
481            "rkyv::with::Skip",
482        ] {
483            assert!(registry.get_wrapper(path).is_some(), "missing wrapper {path}");
484        }
485    }
486
487    #[test]
488    fn hashmap_accepts_trailing_hasher() {
489        let registry = Registry::with_builtins();
490        let map = registry.get_type("std::collections::HashMap").unwrap();
491        let expr = map
492            .instantiate(vec![codec::string(), codec::u32(), CodecExpr::raw("S")])
493            .unwrap();
494        assert_eq!(render(&expr), "hashMap(r.string, r.u32)");
495    }
496
497    #[test]
498    fn btreemap_rejects_trailing_args() {
499        let registry = Registry::with_builtins();
500        let map = registry.get_type("std::collections::BTreeMap").unwrap();
501        let err = map
502            .instantiate(vec![codec::string(), codec::u32(), CodecExpr::raw("S")])
503            .unwrap_err();
504        assert!(matches!(err, DiagnosticKind::GenericArity { .. }));
505    }
506
507    #[test]
508    fn suggestion_matches_last_segment() {
509        let registry = Registry::with_builtins();
510        assert_eq!(
511            registry.suggest_type("collections::HashMap"),
512            Some("hashbrown::HashMap".to_string()),
513        );
514        assert_eq!(
515            registry.suggest_type("other::Uuid"),
516            Some("uuid::Uuid".to_string()),
517        );
518        assert_eq!(registry.suggest_type("chrono::NaiveDate"), None);
519    }
520}
rkyv_js_codegen/registry.rs

rkyv_js_codegen/
registry.rs
