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.
6//!   `uuid::Uuid`, `std::collections::HashMap`) to a [`CodecExpr`] template.
7//! - [`WithWrapper`] maps a `#[rkyv(with = ...)]` *wrapper* path (e.g.
8//!   `rkyv::with::AsBox`) to a transformation of the underlying field codec.
9//!
10//! Both are keyed by fully-qualified path strings. Unknown-path lookups
11//! produce a did-you-mean suggestion when a registered key shares the last
12//! 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; type arguments are filled
22/// 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. The closure runs **once** with
41    /// `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. The closure runs **once** with
47    /// `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. The closure runs **once** with
53    /// `[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
77    /// arity — 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, unless
96    /// [`allow_trailing_args`](ExternalType::allow_trailing_args) was set, in
97    /// which case extra trailing arguments are ignored.
98    ///
99    /// The `rust_path` of a returned [`DiagnosticKind::GenericArity`] is left
100    /// empty; 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. The closure runs **once** with
143    /// `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
173    /// [`map`](WithWrapper::map) and [`identity`](WithWrapper::identity)
174    /// wrappers; `None` is returned for [`skip`](WithWrapper::skip).
175    pub(crate) fn apply(&self, underlying: Option<CodecExpr>) -> Option<CodecExpr> {
176        match &self.kind {
177            WithWrapperKind::Replace(expr) => Some(expr.clone()),
178            WithWrapperKind::Map(template) => {
179                let underlying = underlying.expect("map wrapper requires the underlying codec");
180                Some(template.substitute(&[underlying]))
181            }
182            WithWrapperKind::Identity => {
183                Some(underlying.expect("identity wrapper requires the underlying codec"))
184            }
185            WithWrapperKind::Skip => None,
186        }
187    }
188}
189
190/// The registries backing a [`CodeGenerator`](crate::CodeGenerator).
191#[derive(Debug, Clone)]
192pub(crate) struct Registry {
193    types: BTreeMap<String, ExternalType>,
194    wrappers: BTreeMap<String, WithWrapper>,
195}
196
197impl Registry {
198    /// An empty registry.
199    pub(crate) fn empty() -> Self {
200        Self {
201            types: BTreeMap::new(),
202            wrappers: BTreeMap::new(),
203        }
204    }
205
206    /// A registry pre-populated with the built-in rkyv mappings.
207    pub(crate) fn with_builtins() -> Self {
208        let mut registry = Self::empty();
209
210        registry.register_type(
211            "uuid::Uuid",
212            ExternalType::leaf(CodecExpr::import_from("rkyv-js/lib/uuid", "uuid")),
213        );
214        registry.register_type(
215            "bytes::Bytes",
216            ExternalType::leaf(CodecExpr::import_from("rkyv-js/lib/bytes", "bytes")),
217        );
218        registry.register_type("smol_str::SmolStr", ExternalType::leaf(codec::string()));
219
220        // Vec-shaped containers.
221        registry.register_type(
222            "std::collections::VecDeque",
223            ExternalType::generic1(codec::vec),
224        );
225        registry.register_type("thin_vec::ThinVec", ExternalType::generic1(codec::vec));
226        // `ArrayVec<T, N>`: the const-generic capacity is skipped during
227        // argument collection, but tolerate it anyway.
228        registry.register_type(
229            "arrayvec::ArrayVec",
230            ExternalType::generic1(codec::vec).allow_trailing_args(),
231        );
232        // `SmallVec<[T; N]>` / `TinyVec<[T; N]>`: the array argument is
233        // unwrapped to `T` during argument collection.
234        registry.register_type("smallvec::SmallVec", ExternalType::generic1(codec::vec));
235        registry.register_type("tinyvec::TinyVec", ExternalType::generic1(codec::vec));
236
237        // BTree collections.
238        registry.register_type(
239            "std::collections::BTreeMap",
240            ExternalType::generic2(|k, v| {
241                CodecExpr::call(CodecExpr::import_from("rkyv-js/lib/btreemap", "btreeMap"), [k, v])
242            }),
243        );
244        registry.register_type(
245            "std::collections::BTreeSet",
246            ExternalType::generic1(|t| {
247                CodecExpr::call(CodecExpr::import_from("rkyv-js/lib/btreemap", "btreeSet"), [t])
248            }),
249        );
250
251        // Hash collections (trailing hasher parameter allowed).
252        for path in ["std::collections::HashMap", "hashbrown::HashMap"] {
253            registry.register_type(
254                path,
255                ExternalType::generic2(|k, v| {
256                    CodecExpr::call(CodecExpr::import_from("rkyv-js/lib/hashmap", "hashMap"), [k, v])
257                })
258                .allow_trailing_args(),
259            );
260        }
261        for path in ["std::collections::HashSet", "hashbrown::HashSet"] {
262            registry.register_type(
263                path,
264                ExternalType::generic1(|t| {
265                    CodecExpr::call(CodecExpr::import_from("rkyv-js/lib/hashmap", "hashSet"), [t])
266                })
267                .allow_trailing_args(),
268            );
269        }
270
271        // Index collections (trailing hasher parameter allowed).
272        registry.register_type(
273            "indexmap::IndexMap",
274            ExternalType::generic2(|k, v| {
275                CodecExpr::call(CodecExpr::import_from("rkyv-js/lib/indexmap", "indexMap"), [k, v])
276            })
277            .allow_trailing_args(),
278        );
279        registry.register_type(
280            "indexmap::IndexSet",
281            ExternalType::generic1(|t| {
282                CodecExpr::call(CodecExpr::import_from("rkyv-js/lib/indexmap", "indexSet"), [t])
283            })
284            .allow_trailing_args(),
285        );
286
287        // Shared pointers.
288        for path in ["std::rc::Rc", "std::sync::Arc", "triomphe::Arc"] {
289            registry.register_type(path, ExternalType::generic1(codec::rc));
290        }
291        for path in ["std::rc::Weak", "std::sync::Weak"] {
292            registry.register_type(path, ExternalType::generic1(codec::weak));
293        }
294
295        // Built-in with-wrappers.
296        registry.register_wrapper("rkyv::with::AsBox", WithWrapper::map(codec::boxed));
297        registry.register_wrapper("rkyv::with::Inline", WithWrapper::identity());
298        registry.register_wrapper("rkyv::with::InlineAsBox", WithWrapper::map(codec::boxed));
299        registry.register_wrapper("rkyv::with::Skip", WithWrapper::skip());
300
301        registry
302    }
303
304    pub(crate) fn register_type(&mut self, path: impl Into<String>, external: ExternalType) {
305        self.types.insert(path.into(), external);
306    }
307
308    pub(crate) fn unregister_type(&mut self, path: &str) {
309        self.types.remove(path);
310    }
311
312    pub(crate) fn get_type(&self, path: &str) -> Option<&ExternalType> {
313        self.types.get(path)
314    }
315
316    pub(crate) fn register_wrapper(&mut self, path: impl Into<String>, wrapper: WithWrapper) {
317        self.wrappers.insert(path.into(), wrapper);
318    }
319
320    pub(crate) fn get_wrapper(&self, path: &str) -> Option<&WithWrapper> {
321        self.wrappers.get(path)
322    }
323
324    /// A registered type path sharing the last segment with `path`, if any.
325    pub(crate) fn suggest_type(&self, path: &str) -> Option<String> {
326        let last = path.rsplit("::").next()?;
327        self.types
328            .keys()
329            .find(|key| key.rsplit("::").next() == Some(last))
330            .cloned()
331    }
332}
333
334#[cfg(test)]
335mod tests {
336    use super::*;
337    use std::collections::BTreeMap as Map;
338
339    fn render(expr: &CodecExpr) -> String {
340        expr.render(&Map::new()).unwrap()
341    }
342
343    #[test]
344    fn leaf_instantiates_with_no_args() {
345        let uuid = ExternalType::leaf(CodecExpr::import_from("rkyv-js/lib/uuid", "uuid"));
346        let expr = uuid.instantiate(vec![]).unwrap();
347        assert_eq!(render(&expr), "uuid");
348    }
349
350    #[test]
351    fn generic1_builds_once_with_param0() {
352        let vec_like = ExternalType::generic1(codec::vec);
353        let expr = vec_like.instantiate(vec![codec::u32()]).unwrap();
354        assert_eq!(render(&expr), "r.vec(r.u32)");
355    }
356
357    #[test]
358    fn generic2_instantiates_in_order() {
359        let map_like = ExternalType::generic2(|k, v| {
360            CodecExpr::call(CodecExpr::import_from("m", "pair"), [k, v])
361        });
362        let expr = map_like
363            .instantiate(vec![codec::string(), codec::u32()])
364            .unwrap();
365        assert_eq!(render(&expr), "pair(r.string, r.u32)");
366    }
367
368    #[test]
369    fn too_few_args_is_an_arity_error() {
370        let map_like = ExternalType::generic2(|k, v| {
371            CodecExpr::call(CodecExpr::runtime("pair"), [k, v])
372        });
373        let err = map_like.instantiate(vec![codec::string()]).unwrap_err();
374        assert!(matches!(
375            err,
376            DiagnosticKind::GenericArity {
377                expected: 2,
378                found: 1,
379                ..
380            }
381        ));
382    }
383
384    #[test]
385    fn too_many_args_is_an_arity_error_without_trailing() {
386        let set_like = ExternalType::generic1(codec::vec);
387        let err = set_like
388            .instantiate(vec![codec::u8(), codec::u16()])
389            .unwrap_err();
390        assert!(matches!(
391            err,
392            DiagnosticKind::GenericArity {
393                expected: 1,
394                found: 2,
395                ..
396            }
397        ));
398    }
399
400    #[test]
401    fn trailing_args_are_ignored_when_allowed() {
402        let map_like = ExternalType::generic2(|k, v| {
403            CodecExpr::call(CodecExpr::import_from("m", "hashMap"), [k, v])
404        })
405        .allow_trailing_args();
406        let expr = map_like
407            .instantiate(vec![codec::string(), codec::u32(), codec::u8()])
408            .unwrap();
409        assert_eq!(render(&expr), "hashMap(r.string, r.u32)");
410        // Too few args still fail even with trailing allowed.
411        assert!(map_like.instantiate(vec![codec::string()]).is_err());
412    }
413
414    #[test]
415    #[should_panic(expected = "references Param(1)")]
416    fn generic_panics_on_out_of_range_param() {
417        let _ = ExternalType::generic(1, |_| {
418            CodecExpr::call(CodecExpr::runtime("x"), [CodecExpr::Param(1)])
419        });
420    }
421
422    #[test]
423    fn wrapper_replace_ignores_underlying() {
424        let wrapper = WithWrapper::replace(CodecExpr::import_from("./coord.ts", "Coord"));
425        assert!(!wrapper.needs_underlying());
426        let expr = wrapper.apply(None).unwrap();
427        assert_eq!(render(&expr), "Coord");
428    }
429
430    #[test]
431    fn wrapper_map_transforms_underlying() {
432        let wrapper = WithWrapper::map(codec::boxed);
433        assert!(wrapper.needs_underlying());
434        let expr = wrapper.apply(Some(codec::string())).unwrap();
435        assert_eq!(render(&expr), "r.box(r.string)");
436    }
437
438    #[test]
439    fn wrapper_identity_passes_through() {
440        let wrapper = WithWrapper::identity();
441        let expr = wrapper.apply(Some(codec::u32())).unwrap();
442        assert_eq!(render(&expr), "r.u32");
443    }
444
445    #[test]
446    fn wrapper_skip_omits() {
447        let wrapper = WithWrapper::skip();
448        assert!(!wrapper.needs_underlying());
449        assert!(wrapper.apply(None).is_none());
450    }
451
452    #[test]
453    fn builtins_are_registered() {
454        let registry = Registry::with_builtins();
455        for path in [
456            "uuid::Uuid",
457            "bytes::Bytes",
458            "smol_str::SmolStr",
459            "std::collections::VecDeque",
460            "thin_vec::ThinVec",
461            "arrayvec::ArrayVec",
462            "smallvec::SmallVec",
463            "tinyvec::TinyVec",
464            "std::collections::BTreeMap",
465            "std::collections::BTreeSet",
466            "std::collections::HashMap",
467            "std::collections::HashSet",
468            "hashbrown::HashMap",
469            "hashbrown::HashSet",
470            "indexmap::IndexMap",
471            "indexmap::IndexSet",
472            "std::rc::Rc",
473            "std::sync::Arc",
474            "triomphe::Arc",
475            "std::rc::Weak",
476            "std::sync::Weak",
477        ] {
478            assert!(registry.get_type(path).is_some(), "missing builtin {path}");
479        }
480        for path in [
481            "rkyv::with::AsBox",
482            "rkyv::with::Inline",
483            "rkyv::with::InlineAsBox",
484            "rkyv::with::Skip",
485        ] {
486            assert!(registry.get_wrapper(path).is_some(), "missing wrapper {path}");
487        }
488    }
489
490    #[test]
491    fn hashmap_accepts_trailing_hasher() {
492        let registry = Registry::with_builtins();
493        let map = registry.get_type("std::collections::HashMap").unwrap();
494        let expr = map
495            .instantiate(vec![codec::string(), codec::u32(), CodecExpr::raw("S")])
496            .unwrap();
497        assert_eq!(render(&expr), "hashMap(r.string, r.u32)");
498    }
499
500    #[test]
501    fn btreemap_rejects_trailing_args() {
502        let registry = Registry::with_builtins();
503        let map = registry.get_type("std::collections::BTreeMap").unwrap();
504        let err = map
505            .instantiate(vec![codec::string(), codec::u32(), CodecExpr::raw("S")])
506            .unwrap_err();
507        assert!(matches!(err, DiagnosticKind::GenericArity { .. }));
508    }
509
510    #[test]
511    fn suggestion_matches_last_segment() {
512        let registry = Registry::with_builtins();
513        assert_eq!(
514            registry.suggest_type("collections::HashMap"),
515            Some("hashbrown::HashMap".to_string()),
516        );
517        assert_eq!(
518            registry.suggest_type("other::Uuid"),
519            Some("uuid::Uuid".to_string()),
520        );
521        assert_eq!(registry.suggest_type("chrono::NaiveDate"), None);
522    }
523}
rkyv_js_codegen/registry.rs

rkyv_js_codegen/
registry.rs
