1use std::{future::Future, hash::Hash, ops::Deref};
2
3use anyhow::Result;
4use bincode::{Decode, Encode};
5use either::Either;
6use rustc_hash::FxHashMap;
7use smallvec::{SmallVec, smallvec};
8use turbo_tasks::{
9 FxIndexMap, NonLocalValue, ReadRef, ResolvedVc, TaskInput, TryFlatJoinIterExt, TryJoinIterExt,
10 Vc, trace::TraceRawVcs,
11};
12
13use crate::{
14 chunk::{ChunkItem, ChunkItemWithAsyncModuleInfo, ChunkType, ChunkableModule, ChunkingContext},
15 module_graph::{
16 ModuleGraph,
17 async_module_info::AsyncModulesInfo,
18 chunk_group_info::RoaringBitmapWrapper,
19 module_batch::{ChunkableModuleBatchGroup, ChunkableModuleOrBatch, ModuleBatch},
20 },
21};
22
23pub async fn attach_async_info_to_chunkable_module(
24 module: ResolvedVc<Box<dyn ChunkableModule>>,
25 async_module_info: &ReadRef<AsyncModulesInfo>,
26 module_graph: Vc<ModuleGraph>,
27 chunking_context: Vc<Box<dyn ChunkingContext>>,
28) -> Result<ChunkItemWithAsyncModuleInfo> {
29 let general_module = ResolvedVc::upcast(module);
30 let async_info = if async_module_info.contains(&general_module) {
31 Some(
32 module_graph
33 .referenced_async_modules(*general_module)
34 .to_resolved()
35 .await?,
36 )
37 } else {
38 None
39 };
40 let chunk_item = module
41 .as_chunk_item(module_graph, chunking_context)
42 .to_resolved()
43 .await?;
44 Ok(ChunkItemWithAsyncModuleInfo {
45 chunk_item,
46 module: Some(module),
47 async_info,
48 })
49}
50
51#[derive(
52 Debug, Clone, PartialEq, Eq, Hash, TraceRawVcs, NonLocalValue, TaskInput, Encode, Decode,
53)]
54pub enum ChunkItemOrBatchWithAsyncModuleInfo {
55 ChunkItem(ChunkItemWithAsyncModuleInfo),
56 Batch(ResolvedVc<ChunkItemBatchWithAsyncModuleInfo>),
57}
58
59type ChunkItemOrBatchWithAsyncModuleInfoByChunkType = Either<
60 ChunkItemBatchWithAsyncModuleInfoByChunkTypeData,
61 ReadRef<ChunkItemBatchWithAsyncModuleInfoByChunkType>,
62>;
63
64#[turbo_tasks::value(transparent)]
65pub struct ChunkItemOrBatchWithAsyncModuleInfos(Vec<ChunkItemOrBatchWithAsyncModuleInfo>);
66
67impl ChunkItemOrBatchWithAsyncModuleInfo {
68 pub async fn from_chunkable_module_or_batch(
69 chunkable_module_or_batch: ChunkableModuleOrBatch,
70 async_module_info: &ReadRef<AsyncModulesInfo>,
71 module_graph: Vc<ModuleGraph>,
72 chunking_context: Vc<Box<dyn ChunkingContext>>,
73 ) -> Result<Option<Self>> {
74 Ok(match chunkable_module_or_batch {
75 ChunkableModuleOrBatch::Module(module) => Some(Self::ChunkItem(
76 attach_async_info_to_chunkable_module(
77 module,
78 async_module_info,
79 module_graph,
80 chunking_context,
81 )
82 .await?,
83 )),
84 ChunkableModuleOrBatch::Batch(batch) => Some(Self::Batch(
85 ChunkItemBatchWithAsyncModuleInfo::from_module_batch(
86 *batch,
87 module_graph,
88 chunking_context,
89 )
90 .to_resolved()
91 .await?,
92 )),
93 ChunkableModuleOrBatch::None(_) => None,
94 })
95 }
96
97 pub async fn split_by_chunk_type(
98 &self,
99 ) -> Result<ChunkItemOrBatchWithAsyncModuleInfoByChunkType> {
100 Ok(match self {
101 Self::ChunkItem(item) => Either::Left(smallvec![(
102 item.chunk_item.ty().to_resolved().await?,
103 Self::ChunkItem(item.clone())
104 )]),
105 Self::Batch(batch) => Either::Right(batch.split_by_chunk_type().await?),
106 })
107 }
108}
109
110#[turbo_tasks::value]
111#[derive(Debug, Clone, Hash, TaskInput)]
112pub struct ChunkItemBatchWithAsyncModuleInfo {
113 pub chunk_items: Vec<ChunkItemWithAsyncModuleInfo>,
114 pub chunk_groups: Option<RoaringBitmapWrapper>,
115}
116
117#[turbo_tasks::value_impl]
118impl ChunkItemBatchWithAsyncModuleInfo {
119 #[turbo_tasks::function]
120 pub fn new(chunk_items: Vec<ChunkItemWithAsyncModuleInfo>) -> Vc<Self> {
121 Self {
122 chunk_items,
123 chunk_groups: None,
124 }
125 .cell()
126 }
127
128 #[turbo_tasks::function]
129 pub async fn from_module_batch(
130 batch: Vc<ModuleBatch>,
131 module_graph: Vc<ModuleGraph>,
132 chunking_context: Vc<Box<dyn ChunkingContext>>,
133 ) -> Result<Vc<Self>> {
134 let async_module_info = module_graph.async_module_info().await?;
135 let batch = batch.await?;
136 let chunk_items = batch
137 .modules
138 .iter()
139 .map(|module| {
140 attach_async_info_to_chunkable_module(
141 *module,
142 &async_module_info,
143 module_graph,
144 chunking_context,
145 )
146 })
147 .try_join()
148 .await?;
149 Ok(Self {
150 chunk_items,
151 chunk_groups: batch.chunk_groups.clone(),
152 }
153 .cell())
154 }
155
156 #[turbo_tasks::function]
157 pub async fn split_by_chunk_type(
158 self: Vc<Self>,
159 ) -> Result<Vc<ChunkItemBatchWithAsyncModuleInfoByChunkType>> {
160 let this = self.await?;
161 let mut iter = this.chunk_items.iter().enumerate();
162 let Some((_, first)) = iter.next() else {
163 return Ok(Vc::cell(SmallVec::new()));
164 };
165 let chunk_type = first.chunk_item.ty().to_resolved().await?;
166 while let Some((i, item)) = iter.next() {
167 let ty = item.chunk_item.ty().to_resolved().await?;
168 if ty != chunk_type {
169 let mut map = FxIndexMap::default();
170 map.insert(chunk_type, this.chunk_items[..i].to_vec());
171 map.insert(ty, vec![item.clone()]);
172 for (_, item) in iter {
173 map.entry(item.chunk_item.ty().to_resolved().await?)
174 .or_default()
175 .push(item.clone());
176 }
177 return Ok(Vc::cell(
178 map.into_iter()
179 .map(|(ty, chunk_items)| {
180 let item = if chunk_items.len() == 1 {
181 ChunkItemOrBatchWithAsyncModuleInfo::ChunkItem(
182 chunk_items.into_iter().next().unwrap(),
183 )
184 } else {
185 ChunkItemOrBatchWithAsyncModuleInfo::Batch(
186 Self {
187 chunk_items,
188 chunk_groups: this.chunk_groups.clone(),
189 }
190 .resolved_cell(),
191 )
192 };
193 (ty, item)
194 })
195 .collect(),
196 ));
197 }
198 }
199 Ok(Vc::cell(smallvec![(
200 chunk_type,
201 ChunkItemOrBatchWithAsyncModuleInfo::Batch(self.to_resolved().await?)
202 )]))
203 }
204}
205
206type ChunkItemBatchWithAsyncModuleInfoByChunkTypeData = SmallVec<
207 [(
208 ResolvedVc<Box<dyn ChunkType>>,
209 ChunkItemOrBatchWithAsyncModuleInfo,
210 ); 1],
211>;
212
213#[turbo_tasks::value(transparent)]
214pub struct ChunkItemBatchWithAsyncModuleInfoByChunkType(
215 ChunkItemBatchWithAsyncModuleInfoByChunkTypeData,
216);
217
218type ChunkItemBatchGroupByChunkTypeT = SmallVec<
219 [(
220 ResolvedVc<Box<dyn ChunkType>>,
221 ResolvedVc<ChunkItemBatchGroup>,
222 ); 1],
223>;
224
225#[turbo_tasks::value(transparent)]
226pub struct ChunkItemBatchGroupByChunkType(ChunkItemBatchGroupByChunkTypeT);
227
228#[turbo_tasks::value(transparent)]
229pub struct ChunkItemBatchGroups(Vec<ResolvedVc<ChunkItemBatchGroup>>);
230
231#[turbo_tasks::value]
232pub struct ChunkItemBatchGroup {
233 pub items: Vec<ChunkItemOrBatchWithAsyncModuleInfo>,
234 pub chunk_groups: RoaringBitmapWrapper,
235}
236
237#[turbo_tasks::value_impl]
238impl ChunkItemBatchGroup {
239 #[turbo_tasks::function]
240 pub async fn from_module_batch_group(
241 batch_group: Vc<ChunkableModuleBatchGroup>,
242 module_graph: Vc<ModuleGraph>,
243 chunking_context: Vc<Box<dyn ChunkingContext>>,
244 ) -> Result<Vc<Self>> {
245 let async_module_info = module_graph.async_module_info().await?;
246 let batch_group = batch_group.await?;
247 let items = batch_group
248 .items
249 .iter()
250 .map(|&batch| {
251 ChunkItemOrBatchWithAsyncModuleInfo::from_chunkable_module_or_batch(
252 batch,
253 &async_module_info,
254 module_graph,
255 chunking_context,
256 )
257 })
258 .try_flat_join()
259 .await?;
260 Ok(Self {
261 items,
262 chunk_groups: batch_group.chunk_groups.clone(),
263 }
264 .cell())
265 }
266
267 #[turbo_tasks::function]
268 pub async fn split_by_chunk_type(self: Vc<Self>) -> Result<Vc<ChunkItemBatchGroupByChunkType>> {
269 let this = self.await?;
270 let mut map: FxIndexMap<_, Vec<_>> = FxIndexMap::default();
272 for item in &this.items {
273 let split = item.split_by_chunk_type().await?;
274 for (ty, value) in split.iter() {
275 map.entry(*ty).or_default().push(value.clone());
276 }
277 }
278 let result = if map.len() == 1 {
279 let (ty, _) = map.into_iter().next().unwrap();
280 smallvec![(ty, self.to_resolved().await?)]
281 } else {
282 map.into_iter()
283 .map(|(ty, items)| {
284 (
285 ty,
286 ChunkItemBatchGroup {
287 items,
288 chunk_groups: this.chunk_groups.clone(),
289 }
290 .resolved_cell(),
291 )
292 })
293 .collect()
294 };
295 Ok(Vc::cell(result))
296 }
297}
298
299pub async fn batch_info<'a, BatchGroup, Item, Info, BatchGroupInfo, A, B>(
300 batch_groups: &[ResolvedVc<BatchGroup>],
301 items: &[Item],
302 get_batch_group_info: impl Fn(Vc<BatchGroup>) -> A + Send + 'a,
303 get_item_info: impl Fn(&Item) -> B + Send + 'a,
304) -> Result<Vec<Info>>
305where
306 A: Future<Output = Result<BatchGroupInfo>> + Send + 'a,
307 B: Future<Output = Result<Info>> + Send + 'a,
308 BatchGroup: Send,
309 Item: Send + Eq + Hash,
310 BatchGroupInfo: Deref<Target = FxHashMap<Item, Info>> + Send,
311 Info: Clone + Send,
312{
313 let batch_group_info: Vec<BatchGroupInfo> = batch_groups
314 .iter()
315 .map(|&batch_group| get_batch_group_info(*batch_group))
316 .try_join()
317 .await?;
318 let batch_group_info = batch_group_info
319 .iter()
320 .flat_map(|info| info.iter())
321 .collect::<FxHashMap<_, _>>();
322 items
323 .iter()
324 .map(async |item| {
325 Ok(if let Some(&info) = batch_group_info.get(item) {
326 info.clone()
327 } else {
328 get_item_info(item).await?
329 })
330 })
331 .try_join()
332 .await
333}