specta_typescript/
references.rs

1use std::{cell::RefCell, collections::HashSet};
2
3use specta::datatype::NamedReference;
4
5thread_local! {
6    static REFERENCED_TYPES: RefCell<Option<Vec<HashSet<NamedReference>>>> = const { RefCell::new(None) };
7    static MODULE_PATH_CONTEXT: RefCell<Vec<String>> = const { RefCell::new(Vec::new()) };
8}
9
10pub(crate) fn with_module_path<R>(module_path: &str, func: impl FnOnce() -> R) -> R {
11    struct Guard;
12    impl Drop for Guard {
13        fn drop(&mut self) {
14            MODULE_PATH_CONTEXT.with_borrow_mut(|ctx| {
15                ctx.pop();
16            });
17        }
18    }
19
20    MODULE_PATH_CONTEXT.with_borrow_mut(|ctx| {
21        ctx.push(module_path.to_string());
22    });
23
24    let guard = Guard;
25    let result = func();
26    std::mem::forget(guard);
27    MODULE_PATH_CONTEXT.with_borrow_mut(|ctx| {
28        ctx.pop();
29    });
30
31    result
32}
33
34pub(crate) fn current_module_path() -> Option<String> {
35    MODULE_PATH_CONTEXT.with_borrow(|ctx| ctx.last().cloned())
36}
37
38/// This function collects all Typescript references which are created within the given closure.
39///
40/// This can be used for determining the imports required in a particular file.
41pub fn collect_references<R>(func: impl FnOnce() -> R) -> (R, HashSet<NamedReference>) {
42    struct Guard;
43    impl Drop for Guard {
44        fn drop(&mut self) {
45            REFERENCED_TYPES.with_borrow_mut(|types| {
46                if let Some(v) = types {
47                    // Last collection means we can drop all memory
48                    if v.len() == 1 {
49                        *types = None;
50                    } else {
51                        // Otherwise just remove the current collection.
52                        v.pop();
53                    }
54                }
55            })
56        }
57    }
58
59    // If we have no collection, register one
60    // If we already have one create a new context.
61    REFERENCED_TYPES.with_borrow_mut(|v| {
62        if let Some(v) = v {
63            v.push(Default::default());
64        } else {
65            *v = Some(vec![Default::default()]);
66        }
67    });
68
69    let guard = Guard;
70    let result = func();
71    // We only use the guard when unwinding
72    std::mem::forget(guard);
73
74    (
75        result,
76        REFERENCED_TYPES.with_borrow_mut(|types| {
77            types
78                .as_mut()
79                .expect("COLLECTED_TYPES is unset but it should be set")
80                .pop()
81                .expect("COLLECTED_TYPES is missing a valid collection context")
82        }),
83    )
84}
85
86/// Used internally to track a named references.
87pub(crate) fn track_nr(r: &NamedReference) {
88    REFERENCED_TYPES.with_borrow_mut(|ctxs| {
89        if let Some(ctxs) = ctxs {
90            for ctx in ctxs {
91                ctx.insert(r.clone());
92            }
93        }
94    });
95}
specta_typescript/references.rs

specta_typescript/
references.rs
