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