1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

use ahash::HashMap;

use re_viewer_context::{test_context::TestContext, Contents, ViewClassExt, ViewerContext};

use crate::ViewportBlueprint;

/// Extension trait to [`TestContext`] for blueprint-related features.
pub trait TestContextExt {
    /// See docstring on the implementation below.
    fn setup_viewport_blueprint<R>(
        &mut self,
        setup_blueprint: impl FnOnce(&ViewerContext<'_>, &mut ViewportBlueprint) -> R,
    ) -> R;
}

impl TestContextExt for TestContext {
    /// Inspect or update the blueprint of a [`TestContext`].
    ///
    /// This helper works by deserializing the current blueprint, providing it to the provided
    /// closure, and saving it back to the blueprint store. The closure should call the appropriate
    /// methods of [`ViewportBlueprint`] to inspect and/or create views and containers as required.
    ///
    /// Each time [`TestContextExt::setup_viewport_blueprint`] is called, it entirely recomputes the
    /// "query results", i.e., the [`re_viewer_context::DataResult`]s that each view contains, based
    /// on the current content of the recording store.
    ///
    /// Important pre-requisite:
    /// - The view classes used by view must be already registered (see
    ///   [`TestContext::register_view_class`]).
    /// - The data store must be already populated for the views to have any content (see, e.g.,
    ///   [`TestContext::log_entity`]).
    ///
    fn setup_viewport_blueprint<R>(
        &mut self,
        setup_blueprint: impl FnOnce(&ViewerContext<'_>, &mut ViewportBlueprint) -> R,
    ) -> R {
        let mut setup_blueprint: Option<_> = Some(setup_blueprint);

        let mut result = None;

        egui::__run_test_ctx(|egui_ctx| {
            // We use `take` to ensure that the blueprint is setup only once, since egui forces
            // us to a `FnMut` closure.
            if let Some(setup_blueprint) = setup_blueprint.take() {
                self.run(egui_ctx, |ctx| {
                    let mut viewport_blueprint = ViewportBlueprint::try_from_db(
                        &self.blueprint_store,
                        &self.blueprint_query,
                    );
                    result = Some(setup_blueprint(ctx, &mut viewport_blueprint));
                    viewport_blueprint.save_to_blueprint_store(ctx);
                });

                self.handle_system_commands();

                // Reload the blueprint store and execute all view queries.
                let viewport_blueprint =
                    ViewportBlueprint::try_from_db(&self.blueprint_store, &self.blueprint_query);

                let maybe_visualizable_entities_per_visualizer = self
                    .view_class_registry
                    .maybe_visualizable_entities_for_visualizer_systems(
                        &self.recording_store.store_id(),
                    );
                let mut query_results = HashMap::default();

                self.run(egui_ctx, |ctx| {
                    viewport_blueprint.visit_contents(&mut |contents, _| {
                        if let Contents::View(view_id) = contents {
                            let view_blueprint = viewport_blueprint
                                .view(view_id)
                                .expect("view is known to exist");
                            let class_identifier = view_blueprint.class_identifier();

                            let visualizable_entities = ctx
                                .view_class_registry
                                .class(class_identifier)
                                .expect("The class must be registered beforehand")
                                .determine_visualizable_entities(
                                    &maybe_visualizable_entities_per_visualizer,
                                    ctx.recording(),
                                    &ctx.view_class_registry
                                        .new_visualizer_collection(class_identifier),
                                    &view_blueprint.space_origin,
                                );

                            let indicated_entities_per_visualizer = ctx
                                .view_class_registry
                                .indicated_entities_per_visualizer(&ctx.recording().store_id());

                            let mut data_query_result = view_blueprint.contents.execute_query(
                                ctx.store_context,
                                ctx.view_class_registry,
                                ctx.blueprint_query,
                                *view_id,
                                &visualizable_entities,
                            );

                            let resolver = view_blueprint.contents.build_resolver(
                                ctx.view_class_registry,
                                view_blueprint,
                                &maybe_visualizable_entities_per_visualizer,
                                &visualizable_entities,
                                &indicated_entities_per_visualizer,
                            );

                            resolver.update_overrides(
                                ctx.store_context.blueprint,
                                ctx.blueprint_query,
                                ctx.rec_cfg.time_ctrl.read().timeline(),
                                ctx.view_class_registry,
                                &mut data_query_result,
                                &mut self.view_states.lock(),
                            );

                            query_results.insert(*view_id, data_query_result);
                        }

                        true
                    });
                });

                self.query_results = query_results;
            }
        });

        result.expect("The `setup_closure` is expected to be called at least once")
    }
}