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