re_viewer_context/
query_context.rs

1use std::sync::LazyLock;
2
3use ahash::HashMap;
4use slotmap::SlotMap;
5use smallvec::SmallVec;
6
7use re_log_types::{EntityPath, EntityPathHash};
8
9use crate::{
10    DataResult, StoreContext, ViewContext, ViewId, ViewState, ViewerContext, blueprint_timeline,
11};
12
13slotmap::new_key_type! {
14    /// Identifier for a [`DataResultNode`]
15    pub struct DataResultHandle;
16}
17
18/// Context for a latest-at query in a specific view.
19// TODO(andreas) this is centered around latest-at queries. Does it have to be? Makes sense for UI, but that means it won't scale much into Visualizer queriers.
20// This is currently used only for fallback providers, but the expectation is that we're using this more widely as the primary context object
21// in all places where we query a specific entity in a specific view.
22pub struct QueryContext<'a> {
23    pub view_ctx: &'a ViewContext<'a>,
24
25    /// Target entity path which is lacking the component and needs a fallback.
26    ///
27    /// For editing overrides/defaults, this is the path to the store entity where they override/default is used.
28    /// For view properties this is the path that stores the respective view property archetype.
29    pub target_entity_path: &'a re_log_types::EntityPath,
30
31    /// Archetype name in which context the component is needed.
32    ///
33    /// View properties always have an archetype context, but overrides/defaults may not.
34    pub archetype_name: Option<re_types::ArchetypeName>,
35
36    /// Query which didn't yield a result for the component at the target entity path.
37    pub query: &'a re_chunk_store::LatestAtQuery,
38}
39
40impl QueryContext<'_> {
41    #[inline]
42    pub fn viewer_ctx(&self) -> &ViewerContext<'_> {
43        self.view_ctx.viewer_ctx
44    }
45
46    #[inline]
47    pub fn store_ctx(&self) -> &StoreContext<'_> {
48        self.view_ctx.viewer_ctx.store_context
49    }
50
51    #[inline]
52    pub fn render_ctx(&self) -> &re_renderer::RenderContext {
53        self.view_ctx.viewer_ctx.global_context.render_ctx
54    }
55
56    #[inline]
57    pub fn egui_ctx(&self) -> &egui::Context {
58        self.view_ctx.viewer_ctx.global_context.egui_ctx
59    }
60
61    #[inline]
62    pub fn recording(&self) -> &re_entity_db::EntityDb {
63        self.view_ctx.recording()
64    }
65
66    #[inline]
67    pub fn view_state(&self) -> &dyn ViewState {
68        self.view_ctx.view_state
69    }
70}
71
72/// The result of executing a single data query for a specific view.
73#[derive(Debug)]
74pub struct DataQueryResult {
75    /// The [`DataResultTree`] for the query
76    pub tree: DataResultTree,
77
78    /// The number of entities that matched the query, including those that are not visualizable.
79    pub num_matching_entities: usize,
80
81    /// Of the matched queries, the number of entities that are visualizable by any given visualizer.
82    ///
83    /// This does *not* take into account the actual selection of visualizers
84    /// which may be an explicit none for any given entity.
85    pub num_visualized_entities: usize,
86
87    /// Latest-at results for all component defaults in this view.
88    pub component_defaults: re_query::LatestAtResults,
89}
90
91impl Default for DataQueryResult {
92    fn default() -> Self {
93        Self {
94            tree: Default::default(),
95            num_matching_entities: 0,
96            num_visualized_entities: 0,
97            component_defaults: re_query::LatestAtResults {
98                entity_path: "<defaults>".into(),
99                query: re_chunk_store::LatestAtQuery::latest(blueprint_timeline()),
100                compound_index: (re_chunk::TimeInt::STATIC, re_chunk::RowId::ZERO),
101                components: Default::default(),
102            },
103        }
104    }
105}
106
107impl DataQueryResult {
108    #[inline]
109    pub fn is_empty(&self) -> bool {
110        self.tree.is_empty()
111    }
112
113    #[inline]
114    pub fn result_for_entity(&self, path: &EntityPath) -> Option<&DataResult> {
115        self.tree
116            .lookup_result_by_path(path)
117            .filter(|result| !result.tree_prefix_only)
118    }
119}
120
121impl Clone for DataQueryResult {
122    fn clone(&self) -> Self {
123        re_tracing::profile_function!();
124        Self {
125            tree: self.tree.clone(),
126            num_matching_entities: self.num_matching_entities,
127            num_visualized_entities: self.num_visualized_entities,
128            component_defaults: self.component_defaults.clone(),
129        }
130    }
131}
132
133/// A hierarchical tree of [`DataResult`]s
134#[derive(Clone, Default, Debug)]
135pub struct DataResultTree {
136    data_results: SlotMap<DataResultHandle, DataResultNode>,
137    // TODO(jleibs): Decide if we really want to compute this per-query.
138    // at the moment we only look up a single path per frame for the selection panel. It's probably
139    // less over-head to just walk the tree once instead of pre-computing an entire map we use for
140    // a single lookup.
141    data_results_by_path: HashMap<EntityPathHash, DataResultHandle>,
142    root_handle: Option<DataResultHandle>,
143}
144
145/// A single node in the [`DataResultTree`]
146#[derive(Clone, Debug)]
147pub struct DataResultNode {
148    pub data_result: DataResult,
149    pub children: SmallVec<[DataResultHandle; 4]>,
150}
151
152impl DataResultTree {
153    pub fn new(
154        data_results: SlotMap<DataResultHandle, DataResultNode>,
155        root_handle: Option<DataResultHandle>,
156    ) -> Self {
157        re_tracing::profile_function!();
158        let data_results_by_path = data_results
159            .iter()
160            .map(|(handle, node)| (node.data_result.entity_path.hash(), handle))
161            .collect();
162
163        Self {
164            data_results,
165            data_results_by_path,
166            root_handle,
167        }
168    }
169
170    pub fn root_handle(&self) -> Option<DataResultHandle> {
171        self.root_handle
172    }
173
174    pub fn root_node(&self) -> Option<&DataResultNode> {
175        self.data_results.get(self.root_handle?)
176    }
177
178    /// Depth-first traversal of the tree, calling `visitor` on each result.
179    ///
180    /// Stops traversing a branch if `visitor` returns `false`.
181    pub fn visit<'a>(&'a self, visitor: &mut impl FnMut(&'a DataResultNode) -> bool) {
182        if let Some(root_handle) = self.root_handle {
183            self.visit_recursive(root_handle, visitor);
184        }
185    }
186
187    /// Depth-first traversal of the tree, calling `visitor` on each result, starting from a
188    /// specific node.
189    ///
190    /// Stops traversing a branch if `visitor` returns `false`.
191    pub fn visit_from_node<'a>(
192        &'a self,
193        node: &DataResultNode,
194        visitor: &mut impl FnMut(&'a DataResultNode) -> bool,
195    ) {
196        if let Some(handle) = self
197            .data_results_by_path
198            .get(&node.data_result.entity_path.hash())
199        {
200            self.visit_recursive(*handle, visitor);
201        }
202    }
203
204    /// Depth-first search of a node based on the provided predicate.
205    ///
206    /// If a `staring_node` is provided, the search starts at that node. Otherwise, it starts at the
207    /// root node.
208    pub fn find_node_by(
209        &self,
210        starting_node: Option<&DataResultNode>,
211        predicate: impl Fn(&DataResultNode) -> bool,
212    ) -> Option<&DataResultNode> {
213        let mut result = None;
214
215        let node = starting_node.or_else(|| self.root_node())?;
216        self.visit_from_node(node, &mut |node| {
217            if predicate(node) {
218                result = Some(node);
219            }
220
221            // keep recursing until we find something
222            result.is_none()
223        });
224        result
225    }
226
227    /// Look up a [`DataResult`] in the tree based on its handle.
228    #[inline]
229    pub fn lookup_result(&self, handle: DataResultHandle) -> Option<&DataResult> {
230        self.data_results.get(handle).map(|node| &node.data_result)
231    }
232
233    /// Look up a [`DataResultNode`] in the tree based on its handle.
234    #[inline]
235    pub fn lookup_node(&self, handle: DataResultHandle) -> Option<&DataResultNode> {
236        self.data_results.get(handle)
237    }
238
239    /// Look up a [`DataResultNode`] in the tree based on its handle.
240    #[inline]
241    pub fn lookup_node_mut(&mut self, handle: DataResultHandle) -> Option<&mut DataResultNode> {
242        self.data_results.get_mut(handle)
243    }
244
245    /// Look up a [`DataResultNode`] in the tree based on an [`EntityPath`].
246    #[inline]
247    pub fn lookup_node_by_path(&self, path: &EntityPath) -> Option<&DataResultNode> {
248        self.data_results_by_path
249            .get(&path.hash())
250            .and_then(|handle| self.lookup_node(*handle))
251    }
252
253    /// Look up a [`DataResult`] in the tree based on an [`EntityPath`].
254    #[inline]
255    pub fn lookup_result_by_path(&self, path: &EntityPath) -> Option<&DataResult> {
256        self.data_results_by_path
257            .get(&path.hash())
258            .and_then(|handle| self.lookup_result(*handle))
259    }
260
261    #[inline]
262    pub fn is_empty(&self) -> bool {
263        self.data_results_by_path.is_empty()
264    }
265
266    fn visit_recursive<'a>(
267        &'a self,
268        handle: DataResultHandle,
269        visitor: &mut impl FnMut(&'a DataResultNode) -> bool,
270    ) {
271        if let Some(result) = self.data_results.get(handle)
272            && visitor(result)
273        {
274            for child in &result.children {
275                self.visit_recursive(*child, visitor);
276            }
277        }
278    }
279}
280
281static EMPTY_QUERY: LazyLock<DataQueryResult> = LazyLock::new(Default::default);
282
283impl ViewerContext<'_> {
284    pub fn lookup_query_result(&self, id: ViewId) -> &DataQueryResult {
285        self.query_results.get(&id).unwrap_or_else(|| {
286            if cfg!(debug_assertions) {
287                re_log::warn!("Tried looking up a query that doesn't exist: {:?}", id);
288            } else {
289                re_log::debug!("Tried looking up a query that doesn't exist: {:?}", id);
290            }
291            &EMPTY_QUERY
292        })
293    }
294}
re_viewer_context/query_context.rs

re_viewer_context/
query_context.rs
