Skip to main content

dynamis_world/world/
query_pool.rs

1use dynamis_layout::{MAX_HITS_PER_QUERY, QueryResultHeader, QueryResultRecord};
2use dynamis_model::BodyHandle;
3use std::collections::VecDeque;
4
5#[derive(Clone, Copy, Debug, PartialEq, Eq)]
6pub struct QueryHandle {
7    pub batch: u64,
8    pub index: u32,
9}
10
11#[derive(Clone, Copy, Debug, PartialEq)]
12pub struct QueryHit {
13    pub body: BodyHandle,
14    pub collider: u32,
15    pub distance: f32,
16    pub point: [f32; 3],
17    pub normal: [f32; 3],
18    pub step: u64,
19}
20
21pub(crate) struct QueryOutcome {
22    pub hits: Vec<Vec<QueryHit>>,
23    pub overflow: Vec<bool>,
24}
25
26struct QueryBatch {
27    batch: u64,
28    step: u64,
29    width: usize,
30    outcome: Option<QueryOutcome>,
31}
32
33pub(crate) struct QueryPool {
34    batches: VecDeque<QueryBatch>,
35}
36
37impl QueryPool {
38    pub(crate) fn new() -> Self {
39        Self {
40            batches: VecDeque::new(),
41        }
42    }
43
44    pub(crate) fn submit(&mut self, batch: u64, step: u64, width: usize) {
45        self.batches.push_back(QueryBatch {
46            batch,
47            step,
48            width,
49            outcome: None,
50        });
51    }
52
53    pub(crate) fn hit(&self, handle: QueryHandle) -> Option<QueryHit> {
54        self.hits(handle).first().copied()
55    }
56
57    pub(crate) fn hits(&self, handle: QueryHandle) -> &[QueryHit] {
58        &self.outcome(handle).hits[handle.index as usize]
59    }
60
61    pub(crate) fn overflow(&self, handle: QueryHandle) -> bool {
62        self.outcome(handle).overflow[handle.index as usize]
63    }
64
65    fn outcome(&self, handle: QueryHandle) -> &QueryOutcome {
66        self.batch(handle)
67            .and_then(|batch| batch.outcome.as_ref())
68            .expect("query outcome read before it arrived")
69    }
70
71    pub(crate) fn is_current(&self, handle: QueryHandle) -> bool {
72        self.batch(handle).is_some()
73    }
74
75    pub(crate) fn is_ready(&self, handle: QueryHandle) -> bool {
76        self.batch(handle)
77            .is_some_and(|batch| batch.outcome.is_some())
78    }
79
80    fn batch(&self, handle: QueryHandle) -> Option<&QueryBatch> {
81        self.batches
82            .iter()
83            .find(|batch| batch.batch == handle.batch)
84    }
85
86    pub(crate) fn collect(&mut self, batch_id: u64, bytes: &[u8]) {
87        let batch = self
88            .batches
89            .iter_mut()
90            .find(|batch| batch.batch == batch_id)
91            .unwrap_or_else(|| panic!("no query batch is registered for batch {batch_id}"));
92        let step = batch.step;
93        let records = dynamis_layout::decode::<QueryResultRecord>(bytes);
94        let mut hits = vec![Vec::new(); batch.width];
95        let mut overflow = vec![false; batch.width];
96        for (index, result) in records.iter().take(batch.width).enumerate() {
97            let QueryResultHeader {
98                count,
99                overflow: spilled,
100                ..
101            } = result.header;
102            assert!(
103                count <= MAX_HITS_PER_QUERY,
104                "GPU query result exceeds the hit lane count"
105            );
106            hits[index] = result.hits[..count as usize]
107                .iter()
108                .map(|record| QueryHit {
109                    body: BodyHandle {
110                        id: record.body_id,
111                        generation: record.body_generation,
112                    },
113                    collider: record.collider_index,
114                    distance: record.distance,
115                    point: record.point,
116                    normal: record.normal,
117                    step,
118                })
119                .collect();
120            hits[index].sort_unstable_by(|left, right| left.distance.total_cmp(&right.distance));
121            overflow[index] = spilled != 0;
122        }
123        batch.outcome = Some(QueryOutcome { hits, overflow });
124        while self
125            .batches
126            .front()
127            .is_some_and(|oldest| oldest.outcome.is_some() && oldest.step < step)
128        {
129            self.batches.pop_front();
130        }
131    }
132}