1use std::collections::{HashMap, HashSet};
5use std::path::{Path, PathBuf};
6
7#[derive(Debug, Default)]
8pub struct ModuleNode {
9 pub importers: HashSet<PathBuf>,
10 pub imports: HashSet<PathBuf>,
11 pub is_self_accepting: bool,
12}
13
14#[derive(Debug, Default)]
15pub struct ModuleGraph {
16 modules: HashMap<PathBuf, ModuleNode>,
17}
18
19#[derive(Debug, PartialEq, Eq)]
20pub struct UpdatePlan {
21 pub dirty: Vec<PathBuf>,
22 pub boundaries: Vec<PathBuf>,
23}
24
25#[derive(Debug, PartialEq, Eq)]
26pub enum HmrDecision {
27 Update { boundaries: Vec<PathBuf> },
28 FullReload { reason: String },
29}
30
31impl ModuleGraph {
32 pub fn new() -> Self {
33 Self::default()
34 }
35
36 pub fn ensure_module(&mut self, path: &Path) -> &mut ModuleNode {
37 self.modules.entry(path.to_path_buf()).or_default()
38 }
39
40 pub fn add_import(&mut self, importer: &Path, imported: &Path) {
41 self.ensure_module(importer)
42 .imports
43 .insert(imported.to_path_buf());
44 self.ensure_module(imported)
45 .importers
46 .insert(importer.to_path_buf());
47 }
48
49 pub fn set_self_accepting(&mut self, path: &Path, accepting: bool) {
50 self.ensure_module(path).is_self_accepting = accepting;
51 }
52
53 pub fn set_imports(&mut self, importer: &Path, imports: &[PathBuf]) {
54 let stale: Vec<PathBuf> = self
55 .ensure_module(importer)
56 .imports
57 .iter()
58 .filter(|old| !imports.contains(old))
59 .cloned()
60 .collect();
61 for old in stale {
62 if let Some(node) = self.modules.get_mut(&old) {
63 node.importers.remove(importer);
64 }
65 self.ensure_module(importer).imports.remove(&old);
66 }
67 for import in imports {
68 self.add_import(importer, import);
69 }
70 }
71
72 pub fn contains(&self, path: &Path) -> bool {
73 self.modules.contains_key(path)
74 }
75
76 pub fn module_paths(&self) -> Vec<PathBuf> {
77 let mut paths: Vec<PathBuf> = self.modules.keys().cloned().collect();
78 paths.sort();
79 paths
80 }
81
82 pub fn propagate_update(&self, changed: &Path) -> HmrDecision {
83 if !self.modules.contains_key(changed) {
84 return HmrDecision::FullReload {
85 reason: format!("{} is not in the module graph", changed.display()),
86 };
87 }
88 match self.collect_boundaries(&[changed], &[]) {
89 Ok(mut boundaries) => {
90 boundaries.sort();
91 boundaries.dedup();
92 HmrDecision::Update { boundaries }
93 }
94 Err(reason) => HmrDecision::FullReload { reason },
95 }
96 }
97
98 pub fn update_plan(&self, changed: &Path) -> Result<UpdatePlan, String> {
99 match self.propagate_update(changed) {
100 HmrDecision::FullReload { reason } => Err(reason),
101 HmrDecision::Update { boundaries } => Ok(UpdatePlan {
102 dirty: self.dirty_closure(changed),
103 boundaries,
104 }),
105 }
106 }
107
108 pub fn update_plan_from_importers(&self, changed: &Path) -> Result<UpdatePlan, String> {
109 match self.propagate_update_from_importers(changed) {
110 HmrDecision::FullReload { reason } => Err(reason),
111 HmrDecision::Update { boundaries } => Ok(UpdatePlan {
112 dirty: self.dirty_closure(changed),
113 boundaries,
114 }),
115 }
116 }
117
118 fn dirty_closure(&self, changed: &Path) -> Vec<PathBuf> {
119 let mut dirty: Vec<PathBuf> = vec![changed.to_path_buf()];
120 let mut queue = vec![changed.to_path_buf()];
121 let mut seen: HashSet<PathBuf> = queue.iter().cloned().collect();
122 while let Some(current) = queue.pop() {
123 let Some(node) = self.modules.get(¤t) else {
124 continue;
125 };
126 if node.is_self_accepting && current != changed {
127 continue;
128 }
129 for importer in &node.importers {
130 if seen.insert(importer.clone()) {
131 dirty.push(importer.clone());
132 queue.push(importer.clone());
133 }
134 }
135 }
136 dirty.sort();
137 dirty
138 }
139
140 pub fn propagate_update_from_importers(&self, changed: &Path) -> HmrDecision {
141 let Some(node) = self.modules.get(changed) else {
142 return HmrDecision::FullReload {
143 reason: format!("{} is not in the module graph", changed.display()),
144 };
145 };
146 if node.importers.is_empty() {
147 return HmrDecision::FullReload {
148 reason: format!("{} invalidated at an entry", changed.display()),
149 };
150 }
151
152 let seeds: Vec<&Path> = node.importers.iter().map(PathBuf::as_path).collect();
153 match self.collect_boundaries(&seeds, &[changed]) {
154 Ok(mut boundaries) => {
155 boundaries.sort();
156 boundaries.dedup();
157 HmrDecision::Update { boundaries }
158 }
159 Err(reason) => HmrDecision::FullReload { reason },
160 }
161 }
162
163 fn collect_boundaries<'a>(
164 &'a self,
165 seeds: &[&'a Path],
166 pre_stack: &[&'a Path],
167 ) -> Result<Vec<PathBuf>, String> {
168 let mut colors: HashMap<&'a Path, Color> = HashMap::new();
169 for module in pre_stack {
170 colors.insert(module, Color::Gray);
171 }
172 let mut boundaries: Vec<PathBuf> = Vec::new();
173 for seed in seeds {
174 self.climb(seed, &mut colors, &mut boundaries)?;
175 }
176 Ok(boundaries)
177 }
178
179 fn climb<'a>(
180 &'a self,
181 seed: &'a Path,
182 colors: &mut HashMap<&'a Path, Color>,
183 boundaries: &mut Vec<PathBuf>,
184 ) -> Result<(), String> {
185 enum Step<'s> {
186 Enter(&'s Path),
187 Blacken(&'s Path),
188 }
189
190 let mut stack = vec![Step::Enter(seed)];
191 while let Some(step) = stack.pop() {
192 let current = match step {
193 Step::Blacken(path) => {
194 colors.insert(path, Color::Black);
195 continue;
196 }
197 Step::Enter(path) => path,
198 };
199 match colors.get(current) {
200 Some(Color::Black) => continue,
201 Some(Color::Gray) => {
202 return Err(format!("circular import involving {}", current.display()));
203 }
204 None => {}
205 }
206 let Some(node) = self.modules.get(current) else {
207 return Err(format!("{} is not in the module graph", current.display()));
208 };
209 if node.is_self_accepting {
210 boundaries.push(current.to_path_buf());
211 colors.insert(current, Color::Black);
212 continue;
213 }
214 if node.importers.is_empty() {
215 return Err(format!(
216 "update reached entry {} with no accepting boundary",
217 current.display()
218 ));
219 }
220 colors.insert(current, Color::Gray);
221 stack.push(Step::Blacken(current));
222 let mut importers: Vec<&Path> = node.importers.iter().map(PathBuf::as_path).collect();
223 importers.sort_unstable();
224 for importer in importers.into_iter().rev() {
225 stack.push(Step::Enter(importer));
226 }
227 }
228 Ok(())
229 }
230}
231
232#[derive(Clone, Copy, PartialEq)]
233enum Color {
234 Gray,
235 Black,
236}
237
238#[cfg(test)]
239mod tests {
240 use super::*;
241
242 fn p(s: &str) -> PathBuf {
243 PathBuf::from(s)
244 }
245
246 fn graph() -> ModuleGraph {
247 let mut g = ModuleGraph::new();
248 g.add_import(&p("main.tsx"), &p("App.tsx"));
249 g.add_import(&p("App.tsx"), &p("Button.tsx"));
250 g.set_self_accepting(&p("App.tsx"), true);
251 g.set_self_accepting(&p("Button.tsx"), true);
252 g
253 }
254
255 #[test]
256 fn change_to_self_accepting_leaf_updates_in_place() {
257 let decision = graph().propagate_update(&p("Button.tsx"));
258 assert_eq!(
259 decision,
260 HmrDecision::Update {
261 boundaries: vec![p("Button.tsx")]
262 }
263 );
264 }
265
266 #[test]
267 fn change_propagates_up_to_nearest_boundary() {
268 let mut g = graph();
269 g.add_import(&p("Button.tsx"), &p("utils.ts"));
270 let decision = g.propagate_update(&p("utils.ts"));
271 assert_eq!(
272 decision,
273 HmrDecision::Update {
274 boundaries: vec![p("Button.tsx")]
275 }
276 );
277 }
278
279 #[test]
280 fn dead_end_at_entry_forces_full_reload() {
281 let mut g = graph();
282 g.set_self_accepting(&p("App.tsx"), false);
283 g.add_import(&p("App.tsx"), &p("config.ts"));
284 assert!(matches!(
285 g.propagate_update(&p("config.ts")),
286 HmrDecision::FullReload { .. }
287 ));
288 }
289
290 #[test]
291 fn unknown_module_forces_full_reload() {
292 assert!(matches!(
293 graph().propagate_update(&p("never-imported.ts")),
294 HmrDecision::FullReload { .. }
295 ));
296 }
297
298 #[test]
299 fn update_plan_collects_dirty_chain_through_non_boundaries() {
300 let mut g = graph();
301 g.add_import(&p("Button.tsx"), &p("utils.ts"));
302 let plan = g.update_plan(&p("utils.ts")).unwrap();
303 assert_eq!(plan.boundaries, vec![p("Button.tsx")]);
304 assert_eq!(plan.dirty, vec![p("Button.tsx"), p("utils.ts")]);
305 assert!(!plan.dirty.contains(&p("App.tsx")));
306 assert!(g.update_plan(&p("main.tsx")).is_err());
307 }
308
309 #[test]
310 fn invalidate_skips_own_acceptance_and_climbs_to_importer() {
311 let decision = graph().propagate_update_from_importers(&p("Button.tsx"));
312 assert_eq!(
313 decision,
314 HmrDecision::Update {
315 boundaries: vec![p("App.tsx")]
316 }
317 );
318 assert!(matches!(
319 graph().propagate_update_from_importers(&p("App.tsx")),
320 HmrDecision::FullReload { .. }
321 ));
322 }
323
324 #[test]
325 fn set_imports_drops_stale_reverse_edges() {
326 let mut g = graph();
327 g.add_import(&p("App.tsx"), &p("old-dep.ts"));
328 g.set_imports(&p("App.tsx"), &[p("Button.tsx")]);
329 assert!(matches!(
330 g.propagate_update(&p("old-dep.ts")),
331 HmrDecision::FullReload { .. }
332 ));
333 assert_eq!(
334 g.propagate_update(&p("Button.tsx")),
335 HmrDecision::Update {
336 boundaries: vec![p("Button.tsx")]
337 }
338 );
339 }
340
341 #[test]
342 fn diamond_graphs_dedupe_boundaries() {
343 let mut g = ModuleGraph::new();
344 for parent in ["A.tsx", "B.tsx"] {
345 g.add_import(&p("main.tsx"), &p(parent));
346 g.add_import(&p(parent), &p("shared.ts"));
347 g.set_self_accepting(&p(parent), true);
348 }
349 let decision = g.propagate_update(&p("shared.ts"));
350 assert_eq!(
351 decision,
352 HmrDecision::Update {
353 boundaries: vec![p("A.tsx"), p("B.tsx")]
354 }
355 );
356 let plan = g.update_plan(&p("shared.ts")).unwrap();
357 assert_eq!(plan.dirty, vec![p("A.tsx"), p("B.tsx"), p("shared.ts")]);
358 }
359
360 #[test]
361 fn invalidate_plan_escalates_to_importer_boundaries() {
362 let plan = graph()
363 .update_plan_from_importers(&p("Button.tsx"))
364 .unwrap();
365 assert_eq!(plan.boundaries, vec![p("App.tsx")]);
366 assert!(
367 plan.dirty.contains(&p("Button.tsx")),
368 "changed module is dirty"
369 );
370 assert!(
371 plan.dirty.contains(&p("App.tsx")),
372 "importer boundary is dirty"
373 );
374 assert!(graph().update_plan_from_importers(&p("App.tsx")).is_err());
375 }
376
377 #[test]
378 fn diamond_ladder_is_linear_not_exponential() {
379 let n = 60;
380 let mut g = ModuleGraph::new();
381 let m = |i: usize| p(&format!("M{i}.ts"));
382 for i in 0..n {
383 for side in ["A", "B"] {
384 let mid = p(&format!("{side}{i}.ts"));
385 g.add_import(&mid, &m(i));
386 g.add_import(&m(i + 1), &mid);
387 }
388 }
389 g.add_import(&p("Root.tsx"), &m(n));
390 g.set_self_accepting(&p("Root.tsx"), true);
391
392 assert_eq!(
393 g.propagate_update(&m(0)),
394 HmrDecision::Update {
395 boundaries: vec![p("Root.tsx")]
396 }
397 );
398 }
399
400 #[test]
401 fn import_cycles_do_not_hang_propagation() {
402 let mut g = ModuleGraph::new();
403 g.add_import(&p("a.ts"), &p("b.ts"));
404 g.add_import(&p("b.ts"), &p("a.ts"));
405 assert!(matches!(
406 g.propagate_update(&p("a.ts")),
407 HmrDecision::FullReload { .. }
408 ));
409 }
410}