Skip to main content

rspack_hook/
lib.rs

1use async_trait::async_trait;
2use rspack_error::Result;
3
4pub struct HookMetadata {
5  pub name: &'static str,
6}
7
8pub struct HookCommon {
9  metadata: HookMetadata,
10  tap_stages: Vec<i32>,
11  interceptor_count: usize,
12}
13
14impl HookCommon {
15  pub fn new(name: &'static str) -> Self {
16    Self {
17      metadata: HookMetadata { name },
18      tap_stages: Vec::new(),
19      interceptor_count: 0,
20    }
21  }
22
23  pub fn name(&self) -> &'static str {
24    self.metadata.name
25  }
26
27  pub fn tap_stages(&self) -> &[i32] {
28    &self.tap_stages
29  }
30
31  pub fn insert_tap_stage(&mut self, index: usize, stage: i32) {
32    self.tap_stages.insert(index, stage);
33  }
34
35  pub fn tap_insert_position(&self, stage: i32) -> usize {
36    self.tap_stages.partition_point(|&current| current <= stage)
37  }
38
39  pub fn increment_interceptor_count(&mut self) {
40    self.interceptor_count += 1;
41  }
42
43  pub fn interceptor_count(&self) -> usize {
44    self.interceptor_count
45  }
46
47  pub fn used_stages(&self) -> Vec<i32> {
48    let mut used_stages = self.tap_stages.clone();
49    // tap_stages is kept sorted by stage, so duplicate stages are adjacent.
50    used_stages.dedup();
51    used_stages
52  }
53
54  pub fn is_empty(&self) -> bool {
55    self.tap_stages.is_empty() && self.interceptor_count == 0
56  }
57}
58
59pub fn sort_indices_by_stage(stages: &[i32]) -> Vec<u16> {
60  debug_assert!(stages.len() <= HookTapIndex::INDEX_LIMIT);
61  let mut indices: Vec<_> = (0..stages.len()).map(|index| index as u16).collect();
62  indices.sort_by_key(|&index| {
63    let index = index as usize;
64    (stages[index], index)
65  });
66  debug_assert!(indices.windows(2).all(|indices| {
67    let prev = indices[0] as usize;
68    let next = indices[1] as usize;
69    (stages[prev], prev) <= (stages[next], next)
70  }));
71  indices
72}
73
74#[derive(Debug, Clone, Copy, PartialEq, Eq)]
75pub struct HookTapIndex(u16);
76
77impl HookTapIndex {
78  const INTERCEPT_FLAG: u16 = 1 << 15;
79  const INDEX_MASK: u16 = !Self::INTERCEPT_FLAG;
80  const INDEX_LIMIT: usize = Self::INDEX_MASK as usize + 1;
81
82  pub fn tap(index: u16) -> Self {
83    debug_assert!(index <= Self::INDEX_MASK);
84    Self(index)
85  }
86
87  pub fn intercept(index: u16) -> Self {
88    debug_assert!(index <= Self::INDEX_MASK);
89    Self(Self::INTERCEPT_FLAG | index)
90  }
91
92  pub fn is_tap(self) -> bool {
93    self.0 & Self::INTERCEPT_FLAG == 0
94  }
95
96  pub fn index(self) -> usize {
97    (self.0 & Self::INDEX_MASK) as usize
98  }
99}
100
101pub struct MergedTapIndicesByStage<'a> {
102  base_stages: &'a [i32],
103  additional_stages: &'a [i32],
104  additional_order: Vec<u16>,
105  base_index: u16,
106  additional_cursor: u16,
107}
108
109pub fn merged_tap_indices_by_stage<'a>(
110  base_stages: &'a [i32],
111  additional_stages: &'a [i32],
112) -> MergedTapIndicesByStage<'a> {
113  debug_assert!(base_stages.len() <= HookTapIndex::INDEX_LIMIT);
114  debug_assert!(additional_stages.len() <= HookTapIndex::INDEX_LIMIT);
115  debug_assert!(base_stages.windows(2).all(|stages| stages[0] <= stages[1]));
116  let additional_order = sort_indices_by_stage(additional_stages);
117  debug_assert_eq!(additional_order.len(), additional_stages.len());
118  debug_assert!(
119    additional_order
120      .iter()
121      .all(|&index| (index as usize) < additional_stages.len())
122  );
123  MergedTapIndicesByStage {
124    base_stages,
125    additional_stages,
126    additional_order,
127    base_index: 0,
128    additional_cursor: 0,
129  }
130}
131
132impl Iterator for MergedTapIndicesByStage<'_> {
133  type Item = HookTapIndex;
134
135  fn next(&mut self) -> Option<Self::Item> {
136    let base_index = self.base_index as usize;
137    let additional_cursor = self.additional_cursor as usize;
138    debug_assert!(base_index <= self.base_stages.len());
139    debug_assert!(additional_cursor <= self.additional_order.len());
140    if base_index == self.base_stages.len() && additional_cursor == self.additional_order.len() {
141      return None;
142    }
143
144    if additional_cursor == self.additional_order.len() {
145      let index = self.base_index;
146      debug_assert!(index <= HookTapIndex::INDEX_MASK);
147      self.base_index += 1;
148      return Some(HookTapIndex::tap(index));
149    }
150
151    if base_index == self.base_stages.len() {
152      let index = self.additional_order[additional_cursor];
153      debug_assert!(self.additional_cursor <= HookTapIndex::INDEX_MASK);
154      self.additional_cursor += 1;
155      return Some(HookTapIndex::intercept(index));
156    }
157
158    let additional_index = self.additional_order[additional_cursor];
159    if self.base_stages[base_index] <= self.additional_stages[additional_index as usize] {
160      let index = self.base_index;
161      debug_assert!(index <= HookTapIndex::INDEX_MASK);
162      self.base_index += 1;
163      Some(HookTapIndex::tap(index))
164    } else {
165      debug_assert!(self.additional_cursor <= HookTapIndex::INDEX_MASK);
166      self.additional_cursor += 1;
167      Some(HookTapIndex::intercept(additional_index))
168    }
169  }
170}
171
172#[async_trait]
173pub trait Interceptor<H: Hook> {
174  async fn call(&self, _hook: &H) -> Result<Vec<<H as Hook>::Tap>> {
175    unreachable!("Interceptor::call should only used in async hook")
176  }
177
178  fn call_blocking(&self, _hook: &H) -> Result<Vec<<H as Hook>::Tap>> {
179    unreachable!("Interceptor::call_blocking should only used in sync hook")
180  }
181}
182
183pub trait Hook {
184  type Tap;
185
186  fn used_stages(&self) -> Vec<i32>;
187
188  fn intercept(&mut self, interceptor: impl Interceptor<Self> + Send + Sync + 'static)
189  where
190    Self: Sized;
191}
192
193// pub trait Plugin<HookContainer> {
194//   fn apply(&self, hook_container: &mut HookContainer);
195// }
196
197#[doc(hidden)]
198pub mod __macro_helper {
199  pub use async_trait::async_trait;
200  pub use rspack_error::Result;
201  pub use tracing;
202}
203
204pub use rspack_macros::{define_hook, plugin, plugin_hook};