Skip to main content

fob_bundler/
module_collection_plugin.rs

1use rolldown_plugin::{
2    HookLoadArgs, HookLoadReturn, HookResolveIdArgs, HookResolveIdReturn, HookTransformArgs,
3    HookTransformReturn, HookUsage, Plugin, PluginContext, TransformPluginContext,
4};
5use std::sync::Arc;
6
7use crate::plugins::{FobPlugin, PluginPhase};
8use fob_graph::collection::{CollectedModule, CollectionState, parse_module_structure};
9
10/// Plugin that collects module information during the bundling process
11///
12/// Uses concurrent collections (DashMap/DashSet) internally for thread-safe access
13/// during parallel bundling, eliminating the need for Mutex locks.
14#[derive(Debug)]
15pub struct ModuleCollectionPlugin {
16    state: Arc<CollectionState>,
17}
18
19impl ModuleCollectionPlugin {
20    pub fn new() -> Self {
21        Self {
22            state: Arc::new(CollectionState::new()),
23        }
24    }
25
26    pub fn state(&self) -> Arc<CollectionState> {
27        Arc::clone(&self.state)
28    }
29
30    /// Extract the collected data.
31    ///
32    /// Creates a new CollectionState with cloned data from the concurrent collections.
33    /// The original state remains intact (since it's behind an Arc).
34    pub fn take_data(&self) -> CollectionState {
35        // Clone the data from DashMap/DashSet into a new CollectionState
36        // This preserves the concurrent collections for potential future use
37        let new_state = CollectionState::new();
38
39        // Clone modules
40        for entry in self.state.modules.iter() {
41            new_state
42                .modules
43                .insert(entry.key().clone(), entry.value().clone());
44        }
45
46        // Clone entry points
47        for entry in self.state.entry_points.iter() {
48            new_state.entry_points.insert(entry.key().clone());
49        }
50
51        // Clone resolved entry IDs
52        for entry in self.state.resolved_entry_ids.iter() {
53            new_state.resolved_entry_ids.insert(entry.key().clone());
54        }
55
56        new_state
57    }
58}
59
60impl Plugin for ModuleCollectionPlugin {
61    fn name(&self) -> std::borrow::Cow<'static, str> {
62        "module-collection-plugin".into()
63    }
64
65    fn register_hook_usage(&self) -> HookUsage {
66        HookUsage::ResolveId | HookUsage::Load | HookUsage::Transform
67    }
68
69    fn resolve_id(
70        &self,
71        _ctx: &PluginContext,
72        args: &HookResolveIdArgs,
73    ) -> impl std::future::Future<Output = HookResolveIdReturn> + Send {
74        let state = Arc::clone(&self.state);
75        let specifier = args.specifier.to_string();
76        let is_entry = args.importer.is_none();
77
78        async move {
79            // Track entry specifiers - we'll match resolved IDs in load hook
80            if is_entry {
81                state.mark_entry(specifier);
82            }
83
84            // Let Rolldown handle the actual resolution
85            Ok(None)
86        }
87    }
88
89    fn load(
90        &self,
91        _ctx: &PluginContext,
92        args: &HookLoadArgs<'_>,
93    ) -> impl std::future::Future<Output = HookLoadReturn> + Send {
94        let state = Arc::clone(&self.state);
95        let id = args.id.to_string();
96
97        async move {
98            // Check if this resolved ID corresponds to an entry point
99            let is_entry = state.entry_points.iter().any(|spec| {
100                id.ends_with(spec.key()) || id.contains(spec.key()) || spec.key() == &id
101            });
102
103            if is_entry {
104                state.resolved_entry_ids.insert(id);
105            }
106
107            // Don't modify load behavior
108            Ok(None)
109        }
110    }
111
112    fn transform(
113        &self,
114        _ctx: Arc<TransformPluginContext>,
115        args: &HookTransformArgs,
116    ) -> impl std::future::Future<Output = HookTransformReturn> + Send {
117        let state = Arc::clone(&self.state);
118        let code = args.code.to_string();
119        let id = args.id.to_string();
120
121        async move {
122            // Parse the module to extract imports/exports
123            // If parsing fails, treat as having side effects and no imports/exports
124            let (imports, exports, has_side_effects) =
125                parse_module_structure(&code).unwrap_or_else(|_| (vec![], vec![], true));
126
127            let is_entry = state.resolved_entry_ids.contains(&id);
128
129            let module = CollectedModule {
130                id: id.clone(),
131                code: Some(code),
132                is_entry,
133                is_external: false, // External modules won't go through transform
134                imports,
135                exports,
136                has_side_effects,
137            };
138
139            state.add_module(id, module);
140
141            // Don't modify the code
142            Ok(None)
143        }
144    }
145}
146
147impl FobPlugin for ModuleCollectionPlugin {
148    fn phase(&self) -> PluginPhase {
149        PluginPhase::PostProcess
150    }
151}
152
153impl Default for ModuleCollectionPlugin {
154    fn default() -> Self {
155        Self::new()
156    }
157}