Skip to main content

steel_core/chunk/
chunk_generation_task.rs

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