tenant.openagents/omega

1use crate::Channel;
2use client::ChannelId;
3use collections::BTreeMap;
4use rpc::proto;
5use std::sync::Arc;
6
7#[derive(Default, Debug)]
8pub struct ChannelIndex {
9    channels_ordered: Vec<ChannelId>,
10    channels_by_id: BTreeMap<ChannelId, Arc<Channel>>,
11}
12
13impl ChannelIndex {
14    pub fn by_id(&self) -> &BTreeMap<ChannelId, Arc<Channel>> {
15        &self.channels_by_id
16    }
17
18    pub fn ordered_channels(&self) -> &[ChannelId] {
19        &self.channels_ordered
20    }
21
22    pub fn clear(&mut self) {
23        self.channels_ordered.clear();
24        self.channels_by_id.clear();
25    }
26
27    /// Delete the given channels from this index.
28    pub fn delete_channels(&mut self, channels: &[ChannelId]) {
29        self.channels_by_id
30            .retain(|channel_id, _| !channels.contains(channel_id));
31        self.channels_ordered
32            .retain(|channel_id| !channels.contains(channel_id));
33    }
34
35    pub fn bulk_insert(&mut self) -> ChannelPathsInsertGuard<'_> {
36        ChannelPathsInsertGuard {
37            channels_ordered: &mut self.channels_ordered,
38            channels_by_id: &mut self.channels_by_id,
39        }
40    }
41}
42
43/// A guard for ensuring that the paths index maintains its sort and uniqueness
44/// invariants after a series of insertions
45#[derive(Debug)]
46pub struct ChannelPathsInsertGuard<'a> {
47    channels_ordered: &'a mut Vec<ChannelId>,
48    channels_by_id: &'a mut BTreeMap<ChannelId, Arc<Channel>>,
49}
50
51impl ChannelPathsInsertGuard<'_> {
52    pub fn insert(&mut self, channel_proto: proto::Channel) -> bool {
53        let mut ret = false;
54        let parent_path = channel_proto
55            .parent_path
56            .iter()
57            .map(|cid| ChannelId(*cid))
58            .collect();
59        if let Some(existing_channel) = self.channels_by_id.get_mut(&ChannelId(channel_proto.id)) {
60            let existing_channel = Arc::make_mut(existing_channel);
61
62            ret = existing_channel.visibility != channel_proto.visibility()
63                || existing_channel.name != channel_proto.name
64                || existing_channel.parent_path != parent_path
65                || existing_channel.channel_order != channel_proto.channel_order;
66
67            existing_channel.visibility = channel_proto.visibility();
68            existing_channel.name = channel_proto.name.into();
69            existing_channel.parent_path = parent_path;
70            existing_channel.channel_order = channel_proto.channel_order;
71        } else {
72            self.channels_by_id.insert(
73                ChannelId(channel_proto.id),
74                Arc::new(Channel {
75                    id: ChannelId(channel_proto.id),
76                    visibility: channel_proto.visibility(),
77                    name: channel_proto.name.into(),
78                    parent_path,
79                    channel_order: channel_proto.channel_order,
80                }),
81            );
82            self.insert_root(ChannelId(channel_proto.id));
83        }
84        ret
85    }
86
87    fn insert_root(&mut self, channel_id: ChannelId) {
88        self.channels_ordered.push(channel_id);
89    }
90}
91
92impl Drop for ChannelPathsInsertGuard<'_> {
93    fn drop(&mut self) {
94        self.channels_ordered.sort_by(|a, b| {
95            let a = channel_path_sorting_key(*a, self.channels_by_id);
96            let b = channel_path_sorting_key(*b, self.channels_by_id);
97            a.cmp(b)
98        });
99        self.channels_ordered.dedup();
100    }
101}
102
103fn channel_path_sorting_key(
104    id: ChannelId,
105    channels_by_id: &BTreeMap<ChannelId, Arc<Channel>>,
106) -> impl Iterator<Item = (i32, ChannelId)> {
107    let (parent_path, order_and_id) =
108        channels_by_id
109            .get(&id)
110            .map_or((&[] as &[_], None), |channel| {
111                (
112                    channel.parent_path.as_slice(),
113                    Some((channel.channel_order, channel.id)),
114                )
115            });
116    parent_path
117        .iter()
118        .filter_map(|id| Some((channels_by_id.get(id)?.channel_order, *id)))
119        .chain(order_and_id)
120}
121