steel_core/command/builtins/
locate.rs1use std::{sync::Arc, time::Instant};
4
5use steel_utils::{BlockPos, Identifier, translations};
6use text_components::{
7 Modifier, TextComponent,
8 format::Color,
9 interactivity::{ClickEvent, HoverEvent},
10};
11
12use super::super::{
13 brigadier::{CommandNodeBuilder, CommandSyntaxError},
14 execution::{
15 CommandResultSuspension, CommandResultSuspensionPoll, CommandSource, SteelArgumentType,
16 SteelCommandContext, SteelCommandRuntime, StructureOrTagKey, argument, literal,
17 },
18 registration::CommandRegistration,
19};
20use crate::{
21 chunk::{
22 chunk_request::{ChunkRequest, ChunkRequestHandle, ChunkRequestState, ChunkTicketKind},
23 status::ChunkStatus,
24 },
25 world::World,
26 worldgen::{
27 generator::ChunkGenerator,
28 structure::{StructureLocateCandidate, StructureLocatePlan, squared_distance},
29 },
30};
31
32const MAX_STRUCTURE_SEARCH_RADIUS: i32 = 100;
33
34pub(super) fn registration() -> CommandRegistration<CommandSource> {
35 CommandRegistration::new(Identifier::vanilla_static("locate"), |_| command())
36}
37
38fn command() -> CommandNodeBuilder<CommandSource, SteelCommandRuntime> {
39 literal("locate").then(
40 literal("structure").then(
41 argument("structure", SteelArgumentType::structure_or_tag_key())
42 .executes_suspended(start_structure_search),
43 ),
44 )
45 }
48
49fn start_structure_search(
50 context: &SteelCommandContext<CommandSource>,
51) -> Result<LocateStructureSearch, CommandSyntaxError> {
52 let query = context.structure_or_tag_key("structure")?;
53 let Some(structures) = query.resolve() else {
54 return Err(invalid_structure(query));
55 };
56 if structures.is_empty() {
57 return Err(structure_not_found(query));
58 }
59
60 let world = context.source().world();
61 let Some(structure_generator) = world
62 .chunk_map
63 .world_gen_context
64 .generator
65 .structure_generator()
66 else {
67 return Err(structure_not_found(query));
68 };
69 let structure_keys = structures
70 .iter()
71 .map(|structure| structure.key.clone())
72 .collect::<Vec<_>>();
73 let Some(plan) = structure_generator.locate_plan_for_structures(&structure_keys) else {
74 return Err(structure_not_found(query));
75 };
76 if plan.is_empty() {
77 return Err(structure_not_found(query));
78 }
79
80 Ok(LocateStructureSearch {
81 source: context.source().clone(),
82 world: Arc::clone(world),
83 query: query.clone(),
84 plan,
85 origin: BlockPos::from(context.source().position()),
86 phase: LocatePhase::Start,
87 pending: None,
88 candidates: Vec::new(),
89 best: None,
90 random_radius: 0,
91 started_at: Instant::now(),
92 })
93}
94
95enum LocatePhase {
96 Start,
97 WaitingRings,
98 RandomSpread,
99 WaitingRandomSpread,
100}
101
102struct LocatedStructure {
103 candidate: StructureLocateCandidate,
104 found_structure: Identifier,
105 distance_sqr: i64,
106}
107
108struct LocateStructureSearch {
109 source: CommandSource,
110 world: Arc<World>,
111 query: StructureOrTagKey,
112 plan: StructureLocatePlan,
113 origin: BlockPos,
114 phase: LocatePhase,
115 pending: Option<ChunkRequestHandle>,
116 candidates: Vec<StructureLocateCandidate>,
117 best: Option<LocatedStructure>,
118 random_radius: i32,
119 started_at: Instant,
120}
121
122impl CommandResultSuspension for LocateStructureSearch {
123 fn poll(&mut self) -> CommandResultSuspensionPoll {
124 loop {
125 match self.phase {
126 LocatePhase::Start => {
127 self.candidates = self.plan.ring_candidates(self.origin);
128 if self.candidates.is_empty() {
129 self.phase = LocatePhase::RandomSpread;
130 continue;
131 }
132 self.pending = Some(self.request_current_candidates());
133 self.phase = LocatePhase::WaitingRings;
134 return CommandResultSuspensionPoll::Pending;
135 }
136 LocatePhase::WaitingRings => match self.poll_pending_request() {
137 PendingRequest::Pending => return CommandResultSuspensionPoll::Pending,
138 PendingRequest::Cancelled => return Self::cancelled_result(),
139 PendingRequest::Ready => {
140 self.best = self.first_valid_candidate();
141 self.clear_request();
142
143 if self.best.is_some() && !self.plan.has_random_spread() {
144 return self.success_result();
145 }
146
147 self.phase = LocatePhase::RandomSpread;
148 }
149 },
150 LocatePhase::RandomSpread => {
151 if self.random_radius > MAX_STRUCTURE_SEARCH_RADIUS {
152 return self.finished_result();
153 }
154
155 self.candidates = self
156 .plan
157 .random_spread_candidates_at_radius(self.origin, self.random_radius);
158 self.random_radius += 1;
159
160 if self.candidates.is_empty() {
161 continue;
162 }
163
164 self.pending = Some(self.request_current_candidates());
165 self.phase = LocatePhase::WaitingRandomSpread;
166 return CommandResultSuspensionPoll::Pending;
167 }
168 LocatePhase::WaitingRandomSpread => match self.poll_pending_request() {
169 PendingRequest::Pending => return CommandResultSuspensionPoll::Pending,
170 PendingRequest::Cancelled => return Self::cancelled_result(),
171 PendingRequest::Ready => {
172 if self.update_best_after_random_radius() {
173 return self.success_result();
174 }
175
176 self.clear_request();
177 self.phase = LocatePhase::RandomSpread;
178 }
179 },
180 }
181 }
182 }
183
184 fn cancel(&mut self) {
185 if let Some(pending) = &mut self.pending {
186 pending.cancel();
187 }
188 }
189}
190
191impl LocateStructureSearch {
192 fn request_current_candidates(&self) -> ChunkRequestHandle {
193 let positions = self
194 .candidates
195 .iter()
196 .map(|candidate| candidate.chunk_pos)
197 .collect();
198 self.world.chunk_map.request_chunks(ChunkRequest {
199 status: ChunkStatus::StructureStarts,
200 positions,
201 ticket_kind: ChunkTicketKind::StructureLocate,
202 })
203 }
204
205 fn poll_pending_request(&self) -> PendingRequest {
206 let Some(pending) = &self.pending else {
207 return PendingRequest::Cancelled;
208 };
209
210 match pending.poll() {
211 ChunkRequestState::Pending { .. } => PendingRequest::Pending,
212 ChunkRequestState::Ready => PendingRequest::Ready,
213 ChunkRequestState::Cancelled => PendingRequest::Cancelled,
214 }
215 }
216
217 fn clear_request(&mut self) {
218 self.pending = None;
219 self.candidates.clear();
220 }
221
222 fn first_valid_candidate(&self) -> Option<LocatedStructure> {
223 self.candidates.iter().copied().find_map(|candidate| {
224 self.generated_structure_at_candidate(candidate)
225 .map(|found_structure| LocatedStructure {
226 candidate,
227 found_structure,
228 distance_sqr: squared_distance(candidate.locate_pos, self.origin),
229 })
230 })
231 }
232
233 fn update_best_after_random_radius(&mut self) -> bool {
234 let mut best = self.best.take();
235 let mut current_scan = None;
236 let mut found_current_scan = false;
237 let mut found_in_this_radius = false;
238
239 for candidate in &self.candidates {
240 if current_scan != Some(candidate.scan_id()) {
241 current_scan = Some(candidate.scan_id());
242 found_current_scan = false;
243 }
244
245 if found_current_scan {
246 continue;
247 }
248
249 let Some(found_structure) = self.generated_structure_at_candidate(*candidate) else {
250 continue;
251 };
252 found_current_scan = true;
253 found_in_this_radius = true;
254 let located = LocatedStructure {
255 candidate: *candidate,
256 found_structure,
257 distance_sqr: squared_distance(candidate.locate_pos, self.origin),
258 };
259 if best
260 .as_ref()
261 .is_none_or(|current| located.distance_sqr < current.distance_sqr)
262 {
263 best = Some(located);
264 }
265 }
266
267 self.best = best;
268 found_in_this_radius
269 }
270
271 fn generated_structure_at_candidate(
272 &self,
273 candidate: StructureLocateCandidate,
274 ) -> Option<Identifier> {
275 let holder = self
276 .world
277 .chunk_map
278 .chunks
279 .read_sync(&candidate.chunk_pos, |_, holder| Arc::clone(holder))?;
280 let chunk = holder.try_chunk(ChunkStatus::StructureStarts)?;
281 let starts = chunk.structure_starts();
282 let structures = self.plan.structures_for_candidate(candidate)?;
283 structures.iter().find_map(|structure| {
284 starts
285 .get(structure)
286 .is_some_and(|start| !start.pieces.is_empty())
287 .then(|| structure.clone())
288 })
289 }
290
291 fn finished_result(&self) -> CommandResultSuspensionPoll {
292 if self.best.is_some() {
293 self.success_result()
294 } else {
295 CommandResultSuspensionPoll::Ready(Err(structure_not_found(&self.query)))
296 }
297 }
298
299 fn success_result(&self) -> CommandResultSuspensionPoll {
300 let Some(best) = &self.best else {
301 return CommandResultSuspensionPoll::Ready(Err(structure_not_found(&self.query)));
302 };
303 let pos = best.candidate.locate_pos;
304 let distance = horizontal_distance(self.origin, pos);
305 let structure_name = self.query.found_name(&best.found_structure);
306 self.source.send_success(
307 &locate_success_component(structure_name.clone(), pos, distance),
308 false,
309 );
310 tracing::info!(
311 "Locating element {} took {} ms",
312 structure_name,
313 self.started_at.elapsed().as_millis()
314 );
315 CommandResultSuspensionPoll::Ready(Ok(distance))
316 }
317
318 fn cancelled_result() -> CommandResultSuspensionPoll {
319 CommandResultSuspensionPoll::Ready(Err(CommandSyntaxError::dynamic(
320 "Structure search was cancelled",
321 )))
322 }
323}
324
325enum PendingRequest {
326 Pending,
327 Ready,
328 Cancelled,
329}
330
331fn invalid_structure(query: &StructureOrTagKey) -> CommandSyntaxError {
332 CommandSyntaxError::dynamic(
333 translations::COMMANDS_LOCATE_STRUCTURE_INVALID
334 .message([TextComponent::from(query.as_printable())])
335 .component(),
336 )
337}
338
339fn structure_not_found(query: &StructureOrTagKey) -> CommandSyntaxError {
340 CommandSyntaxError::dynamic(
341 translations::COMMANDS_LOCATE_STRUCTURE_NOT_FOUND
342 .message([TextComponent::from(query.as_printable())])
343 .component(),
344 )
345}
346
347fn horizontal_distance(a: BlockPos, b: BlockPos) -> i32 {
348 let dx = b.0.x.wrapping_sub(a.0.x);
349 let dz = b.0.z.wrapping_sub(a.0.z);
350 let squared = dx.wrapping_mul(dx).wrapping_add(dz.wrapping_mul(dz));
351 (f64::from(squared as f32).sqrt() as f32).floor() as i32
352}
353
354fn locate_success_component(structure_name: String, pos: BlockPos, distance: i32) -> TextComponent {
355 translations::COMMANDS_LOCATE_STRUCTURE_SUCCESS
356 .message([
357 TextComponent::from(structure_name),
358 locate_coordinates_component(pos),
359 TextComponent::from(distance.to_string()),
360 ])
361 .component()
362}
363
364fn locate_coordinates_component(pos: BlockPos) -> TextComponent {
365 let displayed_y = "~";
366 TextComponent::plain("[")
367 .add_child(
368 translations::CHAT_COORDINATES
369 .message([
370 TextComponent::from(pos.0.x.to_string()),
371 TextComponent::from(displayed_y),
372 TextComponent::from(pos.0.z.to_string()),
373 ])
374 .component(),
375 )
376 .add_child(TextComponent::plain("]"))
377 .color(Color::Green)
378 .hover_event(HoverEvent::show_text(
379 &translations::CHAT_COORDINATES_TOOLTIP,
380 ))
381 .click_event(ClickEvent::suggest_command(format!(
382 "/tp @s {} {} {}",
383 pos.0.x, displayed_y, pos.0.z
384 )))
385}
386
387#[cfg(test)]
388mod tests {
389 use super::super::create_dispatcher;
390 use super::*;
391 use crate::command::{
392 brigadier::{CommandDispatcher, NodeId},
393 execution::SteelCommandRuntime,
394 };
395 use steel_registry::init_vanilla_registry;
396
397 type Dispatcher = CommandDispatcher<CommandSource, SteelCommandRuntime>;
398
399 fn child(dispatcher: &Dispatcher, parent: NodeId, name: &str) -> NodeId {
400 let Some(children) = dispatcher.children(parent) else {
401 panic!("parent node should exist");
402 };
403 let Some(child) = children.iter().copied().find(|child| {
404 dispatcher
405 .node(*child)
406 .is_some_and(|node| node.name() == name)
407 }) else {
408 panic!("child {name} should exist");
409 };
410 child
411 }
412
413 #[test]
414 fn locate_graph_exposes_only_the_supported_typed_structure_branch() {
415 init_vanilla_registry();
416 let Ok(dispatcher) = create_dispatcher() else {
417 panic!("built-in commands should register");
418 };
419 let locate = child(&dispatcher, dispatcher.root(), "locate");
420 let structure = child(&dispatcher, locate, "structure");
421 let target = child(&dispatcher, structure, "structure");
422
423 assert_eq!(
424 dispatcher
425 .node(target)
426 .and_then(|node| node.argument_type()),
427 Some(&SteelArgumentType::structure_or_tag_key())
428 );
429 let Some(target_node) = dispatcher.node(target) else {
430 panic!("locate structure argument should exist");
431 };
432 assert!(target_node.is_executable());
433 assert!(dispatcher.children(target).is_some_and(<[_]>::is_empty));
434 assert_eq!(dispatcher.children(locate).map(<[_]>::len), Some(1));
435 }
436
437 #[test]
438 fn locate_coordinates_component_matches_vanilla_interactivity() {
439 let component = locate_coordinates_component(BlockPos::new(12, 0, -34));
440
441 assert_eq!(component.format.color, Some(Color::Green));
442 assert!(matches!(
443 component.interactions.click,
444 Some(ClickEvent::SuggestCommand { ref command })
445 if command.as_ref() == "/tp @s 12 ~ -34"
446 ));
447 assert!(matches!(
448 component.interactions.hover,
449 Some(HoverEvent::ShowText { .. })
450 ));
451 }
452
453 #[test]
454 fn horizontal_distance_matches_vanillas_wrapping_int_and_float_math() {
455 assert_eq!(
456 horizontal_distance(BlockPos::new(0, 0, 0), BlockPos::new(3, 100, 4)),
457 5
458 );
459 assert_eq!(
460 horizontal_distance(
461 BlockPos::new(-30_000_000, 0, 0),
462 BlockPos::new(30_000_000, 0, 0)
463 ),
464 36_907
465 );
466 }
467}