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
130
131

use ahash::HashMap;

use re_types::components::Color;
use re_viewer_context::{Annotations, QueryContext, ResolvedAnnotationInfos};

use crate::clamped_or_nothing;

/// Process [`Color`] components using annotations and default colors.
pub fn process_color_slice<'a>(
    ctx: &QueryContext<'_>,
    fallback_provider: &'a dyn re_viewer_context::TypedComponentFallbackProvider<Color>,
    num_instances: usize,
    annotation_infos: &'a ResolvedAnnotationInfos,
    colors: &'a [Color],
) -> Vec<egui::Color32> {
    re_tracing::profile_function_if!(10_000 < num_instances);

    if let Some(last_color) = colors.last() {
        // If we have colors we can ignore the annotation infos/contexts.

        if colors.len() == num_instances {
            // Common happy path
            colors.iter().map(|c| egui::Color32::from(*c)).collect()
        } else if colors.len() == 1 {
            // Common happy path
            vec![egui::Color32::from(*last_color); num_instances]
        } else {
            let colors = clamped_or_nothing(colors, num_instances);
            colors.map(|c| egui::Color32::from(*c)).collect()
        }
    } else {
        match annotation_infos {
            ResolvedAnnotationInfos::Same(count, annotation_info) => {
                re_tracing::profile_scope!("no colors, same annotation");
                let color = annotation_info
                    .color()
                    .unwrap_or_else(|| fallback_provider.fallback_for(ctx).into());
                vec![color; *count]
            }
            ResolvedAnnotationInfos::Many(annotation_info) => {
                re_tracing::profile_scope!("no-colors, many annotations");
                let fallback = fallback_provider.fallback_for(ctx).into();
                annotation_info
                    .iter()
                    .map(|annotation_info| annotation_info.color().unwrap_or(fallback))
                    .collect()
            }
        }
    }
}

pub type Keypoints = HashMap<
    (re_types::components::ClassId, i64),
    HashMap<re_types::datatypes::KeypointId, glam::Vec3>,
>;

/// Resolves all annotations and keypoints for the given entity view.
pub fn process_annotation_and_keypoint_slices(
    latest_at: re_log_types::TimeInt,
    num_instances: usize,
    positions: impl Iterator<Item = glam::Vec3>,
    keypoint_ids: &[re_types::components::KeypointId],
    class_ids: &[re_types::components::ClassId],
    annotations: &Annotations,
) -> (ResolvedAnnotationInfos, Keypoints) {
    re_tracing::profile_function!();

    let mut keypoints: Keypoints = HashMap::default();

    // No need to process annotations if we don't have class-ids
    if class_ids.is_empty() {
        let resolved_annotation = annotations
            .resolved_class_description(None)
            .annotation_info();

        return (
            ResolvedAnnotationInfos::Same(num_instances, resolved_annotation),
            keypoints,
        );
    };

    let class_ids = clamped_or_nothing(class_ids, num_instances);

    if keypoint_ids.is_empty() {
        let annotation_info = class_ids
            .map(|&class_id| {
                let class_description = annotations.resolved_class_description(Some(class_id));
                class_description.annotation_info()
            })
            .collect();

        (
            ResolvedAnnotationInfos::Many(annotation_info),
            Default::default(),
        )
    } else {
        let keypoint_ids = clamped_or_nothing(keypoint_ids, num_instances);
        let annotation_info = itertools::izip!(positions, keypoint_ids, class_ids)
            .map(|(position, keypoint_id, &class_id)| {
                let class_description = annotations.resolved_class_description(Some(class_id));

                keypoints
                    .entry((class_id, latest_at.as_i64()))
                    .or_default()
                    .insert(keypoint_id.0, position);
                class_description.annotation_info_with_keypoint(keypoint_id.0)
            })
            .collect();

        (ResolvedAnnotationInfos::Many(annotation_info), keypoints)
    }
}

/// Resolves all annotations for the given entity view.
pub fn process_annotation_slices(
    latest_at: re_log_types::TimeInt,
    num_instances: usize,
    class_ids: &[re_types::components::ClassId],
    annotations: &Annotations,
) -> ResolvedAnnotationInfos {
    let (annotations, _keypoints) = process_annotation_and_keypoint_slices(
        latest_at,
        num_instances,
        std::iter::empty(), // positions are only needed for keypoint lookup
        &[],
        class_ids,
        annotations,
    );

    annotations
}