matrix_sdk/event_handler/
maps.rs

1// Copyright 2022 The Matrix.org Foundation C.I.C.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::{
16    borrow::Borrow,
17    collections::{BTreeMap, btree_map},
18};
19
20use ruma::{OwnedRoomId, RoomId};
21
22use super::{
23    EventHandlerFn, EventHandlerHandle, EventHandlerWrapper, HandlerKind, StaticEventTypePart,
24};
25
26#[derive(Default)]
27pub(super) struct EventHandlerMaps {
28    by_kind: BTreeMap<HandlerKind, Vec<EventHandlerWrapper>>,
29    by_kind_type: BTreeMap<KindTypeWrap, Vec<EventHandlerWrapper>>,
30    by_kind_type_prefix: BTreeMap<HandlerKind, BTreeMap<&'static str, Vec<EventHandlerWrapper>>>,
31    by_kind_roomid: BTreeMap<KindRoomId, Vec<EventHandlerWrapper>>,
32    by_kind_type_roomid: BTreeMap<KindTypeRoomIdWrap, Vec<EventHandlerWrapper>>,
33    by_kind_type_prefix_roomid:
34        BTreeMap<KindRoomId, BTreeMap<&'static str, Vec<EventHandlerWrapper>>>,
35}
36
37impl EventHandlerMaps {
38    pub fn add(&mut self, handle: EventHandlerHandle, handler_fn: Box<EventHandlerFn>) {
39        let wrapper = EventHandlerWrapper { handler_id: handle.handler_id, handler_fn };
40
41        match Key::new(handle) {
42            Key::Kind(key) => {
43                self.by_kind.entry(key).or_default().push(wrapper);
44            }
45            Key::KindType(key) => {
46                self.by_kind_type.entry(key).or_default().push(wrapper);
47            }
48            Key::KindTypePrefix(outer_key, inner_key) => {
49                self.by_kind_type_prefix
50                    .entry(outer_key)
51                    .or_default()
52                    .entry(inner_key)
53                    .or_default()
54                    .push(wrapper);
55            }
56            Key::KindRoomId(key) => {
57                self.by_kind_roomid.entry(key).or_default().push(wrapper);
58            }
59            Key::KindTypeRoomId(key) => {
60                self.by_kind_type_roomid.entry(key).or_default().push(wrapper)
61            }
62            Key::KindTypePrefixRoomId(outer_key, inner_key) => self
63                .by_kind_type_prefix_roomid
64                .entry(outer_key)
65                .or_default()
66                .entry(inner_key)
67                .or_default()
68                .push(wrapper),
69        }
70    }
71
72    pub fn get_handlers<'a>(
73        &'a self,
74        ev_kind: HandlerKind,
75        ev_type: &str,
76        room_id: Option<&'a RoomId>,
77    ) -> impl Iterator<Item = (EventHandlerHandle, &'a EventHandlerFn)> + 'a {
78        // Use get_key_value instead of just get to be able to access the event_type
79        // from the BTreeMap key as &'static str, required for EventHandlerHandle.
80        let kind_kv = self.by_kind.get_key_value(&ev_kind).map(|(_, handlers)| (None, handlers));
81        let kind_type_kv = self
82            .by_kind_type
83            .get_key_value(&KindType { ev_kind, ev_type })
84            .map(|(key, handlers)| (Some(StaticEventTypePart::Full(key.0.ev_type)), handlers));
85        let kind_type_prefix_kv = self
86            .by_kind_type_prefix
87            .get_key_value(&ev_kind)
88            .and_then(|(_, inner_map)| {
89                inner_map.iter().find(|(ev_type_prefix, _)| ev_type.starts_with(**ev_type_prefix))
90            })
91            .map(|(ev_type_prefix, handlers)| {
92                (Some(StaticEventTypePart::Prefix(ev_type_prefix)), handlers)
93            });
94        let maybe_roomid_kvs = room_id
95            .map(|r_id| {
96                let room_id = r_id.to_owned();
97                let kind_roomid_kv = self
98                    .by_kind_roomid
99                    .get_key_value(&KindRoomId { ev_kind, room_id })
100                    .map(|(_, handlers)| (None, handlers));
101
102                let room_id = r_id.to_owned();
103                let kind_type_roomid_kv = self
104                    .by_kind_type_roomid
105                    .get_key_value(&KindTypeRoomId { ev_kind, ev_type, room_id })
106                    .map(|(key, handlers)| {
107                        (Some(StaticEventTypePart::Full(key.0.ev_type)), handlers)
108                    });
109
110                let room_id = r_id.to_owned();
111                let kind_type_prefix_roomid_kv = self
112                    .by_kind_type_prefix_roomid
113                    .get_key_value(&KindRoomId { ev_kind, room_id })
114                    .and_then(|(_, inner_map)| {
115                        inner_map
116                            .iter()
117                            .find(|(ev_type_prefix, _)| ev_type.starts_with(**ev_type_prefix))
118                    })
119                    .map(|(ev_type_prefix, handlers)| {
120                        (Some(StaticEventTypePart::Prefix(ev_type_prefix)), handlers)
121                    });
122
123                [kind_roomid_kv, kind_type_roomid_kv, kind_type_prefix_roomid_kv]
124            })
125            .into_iter()
126            .flatten();
127
128        [kind_kv, kind_type_kv, kind_type_prefix_kv]
129            .into_iter()
130            .chain(maybe_roomid_kvs)
131            .flatten()
132            .flat_map(move |(ev_type, handlers)| {
133                handlers.iter().map(move |wrap| {
134                    let handle = EventHandlerHandle {
135                        ev_kind,
136                        ev_type,
137                        room_id: room_id.map(ToOwned::to_owned),
138                        handler_id: wrap.handler_id,
139                    };
140
141                    (handle, &*wrap.handler_fn)
142                })
143            })
144    }
145
146    pub fn remove(&mut self, handle: EventHandlerHandle) {
147        // Generic closure would be a lot nicer
148        fn remove_entry<K: Ord>(
149            map: &mut BTreeMap<K, Vec<EventHandlerWrapper>>,
150            key: K,
151            handler_id: u64,
152        ) {
153            if let btree_map::Entry::Occupied(mut o) = map.entry(key) {
154                let v = o.get_mut();
155                v.retain(|e| e.handler_id != handler_id);
156
157                if v.is_empty() {
158                    o.remove();
159                }
160            }
161        }
162        fn remove_nested_entry<K1: Ord, K2: Ord>(
163            map: &mut BTreeMap<K1, BTreeMap<K2, Vec<EventHandlerWrapper>>>,
164            outer_key: K1,
165            inner_key: K2,
166            handler_id: u64,
167        ) {
168            if let btree_map::Entry::Occupied(mut o) = map.entry(outer_key) {
169                let v = o.get_mut();
170                remove_entry(v, inner_key, handler_id);
171
172                if v.is_empty() {
173                    o.remove();
174                }
175            }
176        }
177
178        let handler_id = handle.handler_id;
179        match Key::new(handle) {
180            Key::Kind(key) => {
181                remove_entry(&mut self.by_kind, key, handler_id);
182            }
183            Key::KindType(key) => {
184                remove_entry(&mut self.by_kind_type, key, handler_id);
185            }
186            Key::KindTypePrefix(outer_key, inner_key) => {
187                remove_nested_entry(
188                    &mut self.by_kind_type_prefix,
189                    outer_key,
190                    inner_key,
191                    handler_id,
192                );
193            }
194            Key::KindRoomId(key) => {
195                remove_entry(&mut self.by_kind_roomid, key, handler_id);
196            }
197            Key::KindTypeRoomId(key) => {
198                remove_entry(&mut self.by_kind_type_roomid, key, handler_id);
199            }
200            Key::KindTypePrefixRoomId(outer_key, inner_key) => {
201                remove_nested_entry(
202                    &mut self.by_kind_type_prefix_roomid,
203                    outer_key,
204                    inner_key,
205                    handler_id,
206                );
207            }
208        }
209    }
210
211    #[cfg(test)]
212    pub fn len(&self) -> usize {
213        self.by_kind_type.len() + self.by_kind_type_roomid.len()
214    }
215}
216
217enum Key {
218    Kind(HandlerKind),
219    KindType(KindTypeWrap),
220    KindTypePrefix(HandlerKind, &'static str),
221    KindRoomId(KindRoomId),
222    KindTypeRoomId(KindTypeRoomIdWrap),
223    KindTypePrefixRoomId(KindRoomId, &'static str),
224}
225
226impl Key {
227    fn new(handle: EventHandlerHandle) -> Self {
228        let EventHandlerHandle { ev_kind, ev_type, room_id, .. } = handle;
229        match (ev_type, room_id) {
230            (None, None) => Key::Kind(ev_kind),
231            (Some(ev_type), None) => match ev_type {
232                StaticEventTypePart::Full(ev_type) => {
233                    Key::KindType(KindTypeWrap(KindType { ev_kind, ev_type }))
234                }
235                StaticEventTypePart::Prefix(ev_type_prefix) => {
236                    Key::KindTypePrefix(ev_kind, ev_type_prefix)
237                }
238            },
239            (None, Some(room_id)) => Key::KindRoomId(KindRoomId { ev_kind, room_id }),
240            (Some(ev_type), Some(room_id)) => match ev_type {
241                StaticEventTypePart::Full(ev_type) => {
242                    Key::KindTypeRoomId(KindTypeRoomIdWrap(KindTypeRoomId {
243                        ev_kind,
244                        ev_type,
245                        room_id,
246                    }))
247                }
248                StaticEventTypePart::Prefix(ev_type_prefix) => {
249                    Key::KindTypePrefixRoomId(KindRoomId { ev_kind, room_id }, ev_type_prefix)
250                }
251            },
252        }
253    }
254}
255
256#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
257struct KindTypeWrap(KindType<'static>);
258
259#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
260struct KindType<'a> {
261    ev_kind: HandlerKind,
262    ev_type: &'a str,
263}
264
265#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
266struct KindRoomId {
267    ev_kind: HandlerKind,
268    room_id: OwnedRoomId,
269}
270
271#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
272struct KindTypeRoomIdWrap(KindTypeRoomId<'static>);
273
274#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
275struct KindTypeRoomId<'a> {
276    ev_kind: HandlerKind,
277    ev_type: &'a str,
278    room_id: OwnedRoomId,
279}
280
281// These lifetime-generic impls are what makes it possible to obtain a
282// &'static str event type from get_key_value in call_event_handlers.
283impl<'a> Borrow<KindType<'a>> for KindTypeWrap {
284    fn borrow(&self) -> &KindType<'a> {
285        &self.0
286    }
287}
288
289impl<'a> Borrow<KindTypeRoomId<'a>> for KindTypeRoomIdWrap {
290    fn borrow(&self) -> &KindTypeRoomId<'a> {
291        &self.0
292    }
293}
matrix_sdk/event_handler/maps.rs

matrix_sdk/event_handler/
maps.rs
