steel_core/chunk/
chunk_generation_task.rs

1//! `ChunkGenerationTask` handles the generation process for chunks.
2use std::{
3    cmp::max,
4    future::{Future, poll_fn},
5    mem,
6    pin::Pin,
7    sync::{
8        Arc,
9        atomic::{AtomicBool, Ordering},
10    },
11    task::Poll,
12};
13
14use rayon::ThreadPool;
15use steel_utils::{ChunkPos, locks::SyncMutex};
16use tokio_util::sync::CancellationToken;
17
18use crate::chunk::{
19    chunk_holder::{ChunkHolder, ChunkSaveDependency},
20    chunk_map::ChunkMap,
21    chunk_pyramid::{GENERATION_PYRAMID, LOADING_PYRAMID},
22    status::ChunkStatus,
23};
24
25/// A pre-filled 2D cache of elements, efficient for async creation.
26pub struct StaticCache2D<T> {
27    min_x: i32,
28    min_z: i32,
29    size: i32,
30    /// Cache stored in row-major order (Z-then-X).
31    cache: Vec<T>,
32    /// Pins generation holders until the last user of this cache finishes.
33    save_dependencies: Option<Box<[ChunkSaveDependency]>>,
34}
35
36impl<T> StaticCache2D<T> {
37    /// Creates a `StaticCache2D` by populating it via a factory.
38    pub fn create<F>(center_x: i32, center_z: i32, radius: i32, factory: F) -> Self
39    where
40        F: Fn(i32, i32) -> T + Send + Sync + 'static,
41        T: Send + 'static,
42    {
43        let size = radius * 2 + 1;
44        let min_x = center_x - radius;
45        let min_z = center_z - radius;
46        let cap = (size * size) as usize;
47        let size_usize = size as usize;
48
49        let cache: Vec<T> = (0..cap)
50            .map(|index| {
51                let x_offset = (index % size_usize) as i32;
52                let z_offset = (index / size_usize) as i32;
53                factory(min_x + x_offset, min_z + z_offset)
54            })
55            .collect();
56
57        Self {
58            min_x,
59            min_z,
60            size,
61            cache,
62            save_dependencies: None,
63        }
64    }
65
66    /// Gets a reference to an element by world coordinates.
67    ///
68    /// # Panics
69    /// Panics if coordinates are out of bounds.
70    #[must_use]
71    pub fn get(&self, x: i32, z: i32) -> &T {
72        let Some(value) = self.try_get(x, z) else {
73            panic!(
74                "Out of bounds: ({x}, {z}) vs [({}, {}) to ({}, {})]",
75                self.min_x,
76                self.min_z,
77                self.min_x + self.size - 1,
78                self.min_z + self.size - 1
79            );
80        };
81        value
82    }
83
84    /// Gets a reference to an element by world coordinates.
85    #[must_use]
86    pub fn try_get(&self, x: i32, z: i32) -> Option<&T> {
87        let rel_x = x - self.min_x;
88        let rel_z = z - self.min_z;
89
90        if rel_x >= 0 && rel_x < self.size && rel_z >= 0 && rel_z < self.size {
91            let index = (rel_z * self.size + rel_x) as usize;
92            self.cache.get(index)
93        } else {
94            None
95        }
96    }
97}
98
99impl StaticCache2D<Arc<ChunkHolder>> {
100    fn pin_holders_for_generation(&mut self) {
101        self.save_dependencies = Some(
102            self.cache
103                .iter()
104                .map(ChunkHolder::add_save_dependency)
105                .collect::<Vec<_>>()
106                .into_boxed_slice(),
107        );
108    }
109}
110
111/// A pinned future representing a neighbor's readiness.
112pub type NeighborReady = Pin<Box<dyn Future<Output = Option<()>> + Send + Sync>>;
113
114/// A task responsible for driving a chunk to a target status.
115///
116/// It works in form of layers. Imagine a pyramid, to get to the top you first need to generate the base layer. And so on.
117/// This works in the same way but with chunk dependencies.
118///
119/// This is achieved using the Generation Pyramid and Loading Pyramid.
120///
121/// To make sure a chunk is only put through a stage once it uses an atomic with a CAS operation. Loading Pyramid must also be advanced with noop functions so this atomic can be driven forward.
122pub struct ChunkGenerationTask {
123    /// The chunk map associated with this task.
124    pub chunk_map: Arc<ChunkMap>,
125    /// The chunk position.
126    pub pos: ChunkPos,
127    /// The target generation status.
128    pub target_status: ChunkStatus,
129    /// The status scheduled for generation. Protected by a mutex for safe concurrent access.
130    pub scheduled_status: SyncMutex<Option<ChunkStatus>>,
131    /// Cancellation token — cancelled when this task should stop.
132    pub cancel_token: CancellationToken,
133    /// Cheap cancellation flag for scheduler-side filtering.
134    cancelled: AtomicBool,
135    /// Futures for neighbors. Protected by a mutex.
136    pub neighbor_ready: SyncMutex<Vec<NeighborReady>>,
137    /// Cache of required chunks.
138    pub cache: Arc<StaticCache2D<Arc<ChunkHolder>>>,
139    /// Holder for the chunk this task is targeting.
140    pub center_holder: Arc<ChunkHolder>,
141    /// Whether generation is required for this task.
142    pub needs_generation: AtomicBool,
143    /// The thread pool to use for generation.
144    pub thread_pool: Arc<ThreadPool>,
145}
146
147impl ChunkGenerationTask {
148    /// Creates a new generation task.
149    #[must_use]
150    #[inline]
151    #[expect(
152        clippy::missing_panics_doc,
153        reason = "panic is unreachable: ThreadPoolBuilder::build only fails on OS thread errors"
154    )]
155    pub fn new(
156        pos: ChunkPos,
157        target_status: ChunkStatus,
158        chunk_map: Arc<ChunkMap>,
159        thread_pool: Arc<ThreadPool>,
160        cancel_token: CancellationToken,
161    ) -> Self {
162        let worst_case_radius = GENERATION_PYRAMID
163            .get_step_to(target_status)
164            .accumulated_dependencies
165            .get_radius_of(ChunkStatus::Empty) as i32;
166
167        let chunk_map_clone = Arc::clone(&chunk_map);
168        let mut cache = StaticCache2D::create(pos.0.x, pos.0.y, worst_case_radius, move |x, y| {
169            chunk_map_clone
170                .chunks
171                .read_sync(&ChunkPos::new(x, y), |_, chunk_holder| {
172                    Arc::clone(chunk_holder)
173                })
174                .expect(
175                    "update_chunk_level installs a holder for every position granted a \
176                     loading level before its neighbors can be scheduled",
177                )
178        });
179        cache.pin_holders_for_generation();
180        let center_holder = Arc::clone(cache.get(pos.0.x, pos.0.y));
181
182        Self {
183            chunk_map,
184            pos,
185            target_status,
186            scheduled_status: SyncMutex::new(None),
187            cancel_token,
188            cancelled: AtomicBool::new(false),
189            neighbor_ready: SyncMutex::new(Vec::new()),
190            cache: Arc::new(cache),
191            center_holder,
192            needs_generation: AtomicBool::new(true),
193            thread_pool,
194        }
195    }
196
197    /// Cancels this task by triggering the cancellation token.
198    pub fn cancel(&self) {
199        if !self.cancelled.swap(true, Ordering::AcqRel) {
200            self.cancel_token.cancel();
201        }
202    }
203
204    /// Returns whether this task has been explicitly cancelled.
205    #[must_use]
206    pub fn is_cancelled(&self) -> bool {
207        self.cancelled.load(Ordering::Acquire)
208    }
209
210    /// Returns the holder for the chunk this task is targeting.
211    pub(crate) const fn center_holder(&self) -> &Arc<ChunkHolder> {
212        &self.center_holder
213    }
214
215    /// Schedules a chunk for a specific layer.
216    ///
217    /// # Panics
218    /// Panics if generation is required but not expected.
219    pub fn schedule_chunk_in_layer(
220        &self,
221        status: ChunkStatus,
222        needs_generation: bool,
223        chunk_holder: &Arc<ChunkHolder>,
224    ) -> bool {
225        let published_status = chunk_holder.published_status();
226
227        let generate;
228        if let Some(published_status) = published_status {
229            generate = status > published_status;
230        } else {
231            generate = true;
232        }
233
234        let pyramid = if generate {
235            &GENERATION_PYRAMID
236        } else {
237            &LOADING_PYRAMID
238        };
239
240        assert!(
241            !generate || needs_generation,
242            "Generation required but not expected for chunk load"
243        );
244
245        if let Some(future) = chunk_holder.apply_step(
246            pyramid.get_step_to(status),
247            &self.chunk_map,
248            &self.cache,
249            Arc::clone(&self.thread_pool),
250        ) {
251            self.neighbor_ready.lock().push(future);
252        } else {
253            self.cancel();
254        }
255
256        true
257    }
258
259    /// Schedules tasks for the current layer's neighbors.
260    pub fn schedule_layer(&self, status: ChunkStatus, needs_generation: bool) {
261        let radius = self.get_radius_for_layer(status, needs_generation);
262        // This for loop is inclusive, so if the radius is 0, we will only schedule the center chunk.
263        for x in (self.pos.0.x - radius)..=(self.pos.0.x + radius) {
264            for y in (self.pos.0.y - radius)..=(self.pos.0.y + radius) {
265                let chunk_holder = self.cache.get(x, y);
266                if self.is_cancelled()
267                    || !self.schedule_chunk_in_layer(status, needs_generation, chunk_holder)
268                {
269                    return;
270                }
271            }
272        }
273    }
274
275    const fn get_radius_for_layer(&self, status: ChunkStatus, needs_generation: bool) -> i32 {
276        let pyramid = if needs_generation {
277            &GENERATION_PYRAMID
278        } else {
279            &LOADING_PYRAMID
280        };
281
282        pyramid
283            .get_step_to(self.target_status)
284            .get_accumulated_radius_of(status) as i32
285    }
286
287    /// Schedules the next layer of generation dependencies.
288    ///
289    /// # Panics
290    /// Panics if the schedule is invalid.
291    pub fn schedule_next_layer(&self) {
292        let status_to_schedule = if self.scheduled_status.lock().is_none() {
293            ChunkStatus::Empty
294        } else if !self.needs_generation.load(Ordering::Relaxed)
295            && *self.scheduled_status.lock() == Some(ChunkStatus::Empty)
296            && !self.can_load_without_generation()
297        {
298            self.needs_generation.store(true, Ordering::Relaxed);
299            ChunkStatus::Empty
300        } else {
301            self.scheduled_status
302                .lock()
303                .expect("Scheduled status missing")
304                .next()
305                .expect("Next status missing")
306        };
307
308        self.schedule_layer(
309            status_to_schedule,
310            self.needs_generation.load(Ordering::Relaxed),
311        );
312        self.scheduled_status.lock().replace(status_to_schedule);
313    }
314
315    fn can_load_without_generation(&self) -> bool {
316        if self.target_status == ChunkStatus::Empty {
317            return true;
318        }
319        let center = self.cache.get(self.pos.0.x, self.pos.0.y);
320        let highest_generated_status = center.published_status();
321
322        if let Some(highest_status) = highest_generated_status {
323            if highest_status < self.target_status {
324                return false;
325            }
326
327            let dependencies = &LOADING_PYRAMID
328                .get_step_to(self.target_status)
329                .accumulated_dependencies;
330            let range = dependencies.get_radius() as i32;
331
332            for x in (self.pos.0.x - range)..=(self.pos.0.x + range) {
333                for z in (self.pos.0.y - range)..=(self.pos.0.y + range) {
334                    let distance = max((self.pos.0.x - x).abs(), (self.pos.0.y - z).abs()) as usize;
335                    if let Some(required_status) = dependencies.get(distance) {
336                        let neighbor = self.cache.get(x, z);
337                        let published = neighbor.published_status();
338                        if published < Some(required_status) {
339                            return false;
340                        }
341                    }
342                }
343            }
344            true
345        } else {
346            false
347        }
348    }
349
350    /// Runs the generation task loop.
351    pub async fn run(self: Arc<Self>) {
352        loop {
353            tokio::select! {
354                () = self.cancel_token.cancelled() => break,
355                () = self.wait_for_scheduled_layers() => {}
356            }
357
358            if *self.scheduled_status.lock() == Some(self.target_status) {
359                break;
360            }
361
362            self.schedule_next_layer();
363        }
364    }
365
366    /// Waits for all scheduled neighbor tasks to complete.
367    pub async fn wait_for_scheduled_layers(&self) {
368        // Collect all futures first to avoid locking the mutex during await
369        let mut futures: Vec<_> = {
370            let mut lock = self.neighbor_ready.lock();
371            mem::take(&mut *lock)
372        };
373
374        if futures.is_empty() {
375            return;
376        }
377
378        let mut failed = false;
379        poll_fn(|cx| {
380            let mut index = 0;
381            while index < futures.len() {
382                match futures[index].as_mut().poll(cx) {
383                    Poll::Ready(result) => {
384                        failed |= result.is_none();
385                        drop(futures.swap_remove(index));
386                    }
387                    Poll::Pending => index += 1,
388                }
389            }
390
391            if futures.is_empty() {
392                Poll::Ready(())
393            } else {
394                Poll::Pending
395            }
396        })
397        .await;
398
399        if failed {
400            self.cancel();
401        }
402    }
403}
steel_core/chunk/chunk_generation_task.rs

steel_core/chunk/
chunk_generation_task.rs
