a_tree/
events.rs

1use crate::{
2    predicates::PredicateKind,
3    strings::{StringId, StringTable},
4};
5use itertools::Itertools;
6use rust_decimal::Decimal;
7use std::{
8    collections::HashMap,
9    fmt::{Display, Formatter},
10    ops::Index,
11};
12use thiserror::Error;
13
14#[derive(Error, PartialEq, Debug)]
15pub enum EventError {
16    #[error("attribute {0} has already been defined")]
17    AlreadyPresent(String),
18    #[error("event is missing some attributes")]
19    MissingAttributes,
20    #[error("ABE refers to non-existing attribute '{0:?}'")]
21    NonExistingAttribute(String),
22    #[error("{name:?}: wrong types => expected: {expected:?}, found: {actual:?}")]
23    WrongType {
24        name: String,
25        expected: AttributeKind,
26        actual: AttributeKind,
27    },
28    #[error("{name:?}: mismatching types => expected: {expected:?}, found: {actual:?}")]
29    MismatchingTypes {
30        name: String,
31        expected: AttributeKind,
32        actual: PredicateKind,
33    },
34}
35
36/// An [`Event`] builder
37///
38/// During the builder creation, it will set all the attributes to `undefined`. If some attributes
39/// are not assigned, they will be left `undefined`.
40#[derive(Debug)]
41pub struct EventBuilder<'atree> {
42    by_ids: Vec<AttributeValue>,
43    attributes: &'atree AttributeTable,
44    strings: &'atree StringTable,
45}
46
47impl<'atree> EventBuilder<'atree> {
48    pub(crate) fn new(attributes: &'atree AttributeTable, strings: &'atree StringTable) -> Self {
49        Self {
50            attributes,
51            strings,
52            by_ids: vec![AttributeValue::Undefined; attributes.len()],
53        }
54    }
55
56    /// Build the corresponding [`Event`].
57    ///
58    /// By default, the non-assigned attributes will be undefined.
59    ///
60    /// # Examples
61    ///
62    /// ```rust
63    /// use a_tree::{ATree, AttributeDefinition};
64    ///
65    /// let definitions = [
66    ///     AttributeDefinition::boolean("private"),
67    ///     AttributeDefinition::integer("exchange_id"),
68    ///     AttributeDefinition::string_list("deal_ids"),
69    /// ];
70    /// let atree = ATree::<u64>::new(&definitions).unwrap();
71    ///
72    /// let mut builder = atree.make_event();
73    /// builder.with_integer("exchange_id", 1).unwrap();
74    /// builder.with_boolean("private", false).unwrap();
75    ///
76    /// // The returned `Event` will have its `deal_ids` attribute `undefined` since it was not set
77    /// // by the builder
78    /// let event = builder.build().unwrap();
79    /// ```
80    pub fn build(self) -> Result<Event, EventError> {
81        Ok(Event(self.by_ids))
82    }
83
84    /// Set the specified boolean attribute.
85    ///
86    /// The specified attribute must exist within the [`crate::ATree`] and its type must be boolean.
87    pub fn with_boolean(&mut self, name: &str, value: bool) -> Result<(), EventError> {
88        self.add_value(name, AttributeKind::Boolean, || {
89            AttributeValue::Boolean(value)
90        })
91    }
92
93    /// Set the specified integer attribute.
94    ///
95    /// The specified attribute must exist within the [`crate::ATree`] and its type must be integer.
96    pub fn with_integer(&mut self, name: &str, value: i64) -> Result<(), EventError> {
97        self.add_value(name, AttributeKind::Integer, || {
98            AttributeValue::Integer(value)
99        })
100    }
101
102    /// Set the specified float attribute.
103    ///
104    /// The specified attribute must exist within the [`crate::ATree`] and its type must be float.
105    pub fn with_float(&mut self, name: &str, number: i64, scale: u32) -> Result<(), EventError> {
106        self.add_value(name, AttributeKind::Float, || {
107            AttributeValue::Float(Decimal::new(number, scale))
108        })
109    }
110
111    /// Set the specified string attribute.
112    ///
113    /// The specified attribute must exist within the [`crate::ATree`] and its type must be string.
114    pub fn with_string(&mut self, name: &str, value: &str) -> Result<(), EventError> {
115        self.add_value(name, AttributeKind::String, || {
116            let string_index = self.strings.get(value);
117            AttributeValue::String(string_index)
118        })
119    }
120
121    /// Set the specified list of integers attribute.
122    ///
123    /// The specified attribute must exist within the [`crate::ATree`] and its type must be a list
124    /// of integers.
125    pub fn with_integer_list(&mut self, name: &str, value: &[i64]) -> Result<(), EventError> {
126        self.add_value(name, AttributeKind::IntegerList, || {
127            let values = value.iter().sorted().unique().cloned().collect_vec();
128            AttributeValue::IntegerList(values)
129        })
130    }
131
132    /// Set the specified attribute to `undefined`.
133    ///
134    /// The specified attribute must exist within the [`crate::ATree`].
135    pub fn with_undefined(&mut self, name: &str) -> Result<(), EventError> {
136        let index = self
137            .attributes
138            .by_name(name)
139            .ok_or_else(|| EventError::NonExistingAttribute(name.to_string()))?;
140        self.by_ids[index.0] = AttributeValue::Undefined;
141        Ok(())
142    }
143
144    /// Set the specified string list attribute.
145    ///
146    /// The specified attribute must exist within the [`crate::ATree`] and its type must be a list
147    /// of strings.
148    pub fn with_string_list(&mut self, name: &str, values: &[&str]) -> Result<(), EventError> {
149        self.add_value(name, AttributeKind::StringList, || {
150            let values: Vec<_> = values
151                .iter()
152                .map(|v| self.strings.get(v))
153                .sorted()
154                .unique()
155                .collect();
156            AttributeValue::StringList(values)
157        })
158    }
159
160    fn add_value<F>(&mut self, name: &str, actual: AttributeKind, f: F) -> Result<(), EventError>
161    where
162        F: FnOnce() -> AttributeValue,
163    {
164        let index = self
165            .attributes
166            .by_name(name)
167            .ok_or_else(|| EventError::NonExistingAttribute(name.to_string()))?;
168        let expected = self.attributes.by_id(index);
169        if expected != actual {
170            return Err(EventError::WrongType {
171                name: name.to_owned(),
172                expected,
173                actual,
174            });
175        }
176        self.by_ids[index.0] = f();
177        Ok(())
178    }
179}
180
181/// An event that can be used by the [`crate::atree::ATree`] structure to match arbitrary boolean
182/// expressions
183#[derive(Clone, Debug)]
184pub struct Event(Vec<AttributeValue>);
185
186impl Index<AttributeId> for Event {
187    type Output = AttributeValue;
188
189    #[inline]
190    fn index(&self, index: AttributeId) -> &Self::Output {
191        &self.0[index.0]
192    }
193}
194
195#[derive(Clone, Debug)]
196pub enum AttributeValue {
197    Boolean(bool),
198    Integer(i64),
199    Float(Decimal),
200    String(StringId),
201    IntegerList(Vec<i64>),
202    StringList(Vec<StringId>),
203    Undefined,
204}
205
206#[derive(Clone, Debug)]
207pub struct AttributeTable {
208    by_names: HashMap<String, AttributeId>,
209    by_ids: Vec<AttributeKind>,
210}
211
212#[derive(Clone, Copy, Eq, Ord, PartialEq, PartialOrd, Debug, Hash)]
213pub struct AttributeId(usize);
214
215impl Display for AttributeId {
216    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
217        write!(formatter, "attribute({})", self.0)
218    }
219}
220
221impl AttributeTable {
222    pub fn new(definitions: &[AttributeDefinition]) -> Result<Self, EventError> {
223        let size = definitions.len();
224        let mut by_names = HashMap::with_capacity(size);
225        let mut by_ids = Vec::with_capacity(size);
226        for (i, definition) in definitions.iter().enumerate() {
227            let name = definition.name.to_owned();
228            if by_names.contains_key(&name) {
229                return Err(EventError::AlreadyPresent(name));
230            }
231
232            by_names.insert(name, AttributeId(i));
233            by_ids.push(definition.kind.clone());
234        }
235
236        Ok(Self { by_names, by_ids })
237    }
238
239    #[inline]
240    pub fn by_name(&self, name: &str) -> Option<AttributeId> {
241        self.by_names.get(name).cloned()
242    }
243
244    #[inline]
245    pub fn by_id(&self, id: AttributeId) -> AttributeKind {
246        self.by_ids[id.0].clone()
247    }
248
249    #[inline]
250    pub fn len(&self) -> usize {
251        self.by_ids.len()
252    }
253}
254
255/// The definition of an attribute that is usable by the [`crate::atree::ATree`]
256#[derive(Debug, Clone)]
257pub struct AttributeDefinition {
258    name: String,
259    kind: AttributeKind,
260}
261
262#[derive(Clone, PartialEq, Debug)]
263pub enum AttributeKind {
264    Boolean,
265    Integer,
266    Float,
267    String,
268    IntegerList,
269    StringList,
270}
271
272impl AttributeDefinition {
273    /// Create a boolean attribute definition.
274    pub fn boolean(name: &str) -> Self {
275        let kind = AttributeKind::Boolean;
276        Self {
277            name: name.to_owned(),
278            kind,
279        }
280    }
281
282    /// Create an integer attribute definition.
283    pub fn integer(name: &str) -> Self {
284        let kind = AttributeKind::Integer;
285        Self {
286            name: name.to_owned(),
287            kind,
288        }
289    }
290
291    /// Create a float attribute definition.
292    pub fn float(name: &str) -> Self {
293        let kind = AttributeKind::Float;
294        Self {
295            name: name.to_owned(),
296            kind,
297        }
298    }
299
300    /// Create a string attribute definition.
301    pub fn string(name: &str) -> Self {
302        let kind = AttributeKind::String;
303        Self {
304            name: name.to_owned(),
305            kind,
306        }
307    }
308
309    /// Create a list of integers attribute definition.
310    pub fn integer_list(name: &str) -> Self {
311        let kind = AttributeKind::IntegerList;
312        Self {
313            name: name.to_owned(),
314            kind,
315        }
316    }
317
318    /// Create a list of strings attribute definition.
319    pub fn string_list(name: &str) -> Self {
320        let kind = AttributeKind::StringList;
321        Self {
322            name: name.to_owned(),
323            kind,
324        }
325    }
326}
327
328#[cfg(test)]
329mod tests {
330    use super::*;
331
332    #[test]
333    fn can_create_an_attribute_table_with_no_attributes() {
334        assert!(AttributeTable::new(&[]).is_ok())
335    }
336
337    #[test]
338    fn can_create_an_attribute_table_with_some_attributes() {
339        let definitions = [
340            AttributeDefinition::boolean("private"),
341            AttributeDefinition::string_list("deals"),
342            AttributeDefinition::integer("exchange_id"),
343            AttributeDefinition::float("bidfloor"),
344            AttributeDefinition::string("country"),
345            AttributeDefinition::integer_list("segment_ids"),
346        ];
347
348        assert!(AttributeTable::new(&definitions).is_ok());
349    }
350
351    #[test]
352    fn return_an_error_on_duplicate_definitions() {
353        let definitions = [
354            AttributeDefinition::boolean("private"),
355            AttributeDefinition::string("country"),
356            AttributeDefinition::string_list("deals"),
357            AttributeDefinition::integer("exchange_id"),
358            AttributeDefinition::float("bidfloor"),
359            AttributeDefinition::integer("country"),
360            AttributeDefinition::integer_list("segment_ids"),
361        ];
362
363        assert!(AttributeTable::new(&definitions).is_err());
364    }
365
366    #[test]
367    fn can_add_a_boolean_attribute_value() {
368        let attributes = AttributeTable::new(&[AttributeDefinition::boolean("private")]).unwrap();
369        let strings = StringTable::new();
370        let mut event_builder = EventBuilder::new(&attributes, &strings);
371
372        let result = event_builder.with_boolean("private", true);
373
374        assert!(result.is_ok());
375    }
376
377    #[test]
378    fn can_add_an_integer_attribute_value() {
379        let attributes =
380            AttributeTable::new(&[AttributeDefinition::integer("exchange_id")]).unwrap();
381        let strings = StringTable::new();
382        let mut event_builder = EventBuilder::new(&attributes, &strings);
383
384        let result = event_builder.with_integer("exchange_id", 1);
385
386        assert!(result.is_ok());
387    }
388
389    #[test]
390    fn can_add_a_float_attribute_value() {
391        let attributes = AttributeTable::new(&[AttributeDefinition::float("bidfloor")]).unwrap();
392        let strings = StringTable::new();
393        let mut event_builder = EventBuilder::new(&attributes, &strings);
394
395        let result = event_builder.with_float("bidfloor", 1, 0);
396
397        assert!(result.is_ok());
398    }
399
400    #[test]
401    fn can_add_a_string_attribute_value() {
402        let attributes = AttributeTable::new(&[AttributeDefinition::string("country")]).unwrap();
403        let strings = StringTable::new();
404        let mut event_builder = EventBuilder::new(&attributes, &strings);
405
406        let result = event_builder.with_string("country", "US");
407
408        assert!(result.is_ok());
409    }
410
411    #[test]
412    fn can_add_an_integer_list_attribute_value() {
413        let attributes =
414            AttributeTable::new(&[AttributeDefinition::integer_list("segment_ids")]).unwrap();
415        let strings = StringTable::new();
416        let mut event_builder = EventBuilder::new(&attributes, &strings);
417
418        let result = event_builder.with_integer_list("segment_ids", &[1, 2, 3]);
419
420        assert!(result.is_ok());
421    }
422
423    #[test]
424    fn can_add_an_string_list_attribute_value() {
425        let attributes =
426            AttributeTable::new(&[AttributeDefinition::string_list("deal_ids")]).unwrap();
427        let strings = StringTable::new();
428        let mut event_builder = EventBuilder::new(&attributes, &strings);
429
430        let result = event_builder.with_string_list("deal_ids", &["deal-1", "deal-2"]);
431
432        assert!(result.is_ok());
433    }
434
435    #[test]
436    fn return_an_error_when_adding_a_non_existing_attribute() {
437        let attributes =
438            AttributeTable::new(&[AttributeDefinition::string_list("deal_ids")]).unwrap();
439        let strings = StringTable::new();
440        let mut event_builder = EventBuilder::new(&attributes, &strings);
441
442        let result = event_builder.with_boolean("non_existing", true);
443
444        assert!(matches!(result, Err(EventError::NonExistingAttribute(_))));
445    }
446
447    #[test]
448    fn can_create_an_event_with_no_attributes() {
449        let attributes = AttributeTable::new(&[]).unwrap();
450        let strings = StringTable::new();
451        let event_builder = EventBuilder::new(&attributes, &strings);
452
453        assert!(event_builder.build().is_ok());
454    }
455
456    #[test]
457    fn can_create_an_event_with_attributes() {
458        let attributes = AttributeTable::new(&[
459            AttributeDefinition::boolean("private"),
460            AttributeDefinition::string_list("deals"),
461            AttributeDefinition::integer("exchange_id"),
462            AttributeDefinition::float("bidfloor"),
463            AttributeDefinition::string("country"),
464            AttributeDefinition::integer_list("segment_ids"),
465        ])
466        .unwrap();
467        let strings = StringTable::new();
468        let mut builder = EventBuilder::new(&attributes, &strings);
469
470        assert!(builder.with_boolean("private", true).is_ok());
471        assert!(builder
472            .with_string_list("deals", &["deal-1", "deal-2"])
473            .is_ok());
474        assert!(builder.with_integer("exchange_id", 1).is_ok());
475        assert!(builder.with_float("bidfloor", 1, 0).is_ok());
476        assert!(builder.with_string("country", "US").is_ok());
477        assert!(builder.with_integer_list("segment_ids", &[1, 2, 3]).is_ok());
478
479        assert!(builder.build().is_ok());
480    }
481
482    #[test]
483    fn can_create_an_event_with_a_missing_attribute() {
484        let attributes = AttributeTable::new(&[AttributeDefinition::boolean("private")]).unwrap();
485        let strings = StringTable::new();
486        let event_builder = EventBuilder::new(&attributes, &strings);
487
488        assert!(event_builder.build().is_ok());
489    }
490
491    #[test]
492    fn return_an_error_when_trying_to_add_an_attribute_with_mismatched_type() {
493        let attributes = AttributeTable::new(&[AttributeDefinition::boolean("private")]).unwrap();
494        let strings = StringTable::new();
495        let mut event_builder = EventBuilder::new(&attributes, &strings);
496
497        let result = event_builder.with_integer("private", 1);
498
499        assert!(result.is_err());
500    }
501}
a_tree/events.rs

a_tree/
events.rs
