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    #[inline]
161    fn add_value<F>(&mut self, name: &str, actual: AttributeKind, f: F) -> Result<(), EventError>
162    where
163        F: FnOnce() -> AttributeValue,
164    {
165        let index = self
166            .attributes
167            .by_name(name)
168            .ok_or_else(|| EventError::NonExistingAttribute(name.to_string()))?;
169        let expected = self.attributes.by_id(index);
170        if expected != actual {
171            return Err(EventError::WrongType {
172                name: name.to_owned(),
173                expected,
174                actual,
175            });
176        }
177        self.by_ids[index.0] = f();
178        Ok(())
179    }
180}
181
182/// An event that can be used by the [`crate::atree::ATree`] structure to match arbitrary boolean
183/// expressions
184#[derive(Clone, Debug)]
185pub struct Event(Vec<AttributeValue>);
186
187impl Index<AttributeId> for Event {
188    type Output = AttributeValue;
189
190    #[inline]
191    fn index(&self, index: AttributeId) -> &Self::Output {
192        &self.0[index.0]
193    }
194}
195
196#[derive(Clone, Debug)]
197pub enum AttributeValue {
198    Boolean(bool),
199    Integer(i64),
200    Float(Decimal),
201    String(StringId),
202    IntegerList(Vec<i64>),
203    StringList(Vec<StringId>),
204    Undefined,
205}
206
207#[derive(Clone, Debug)]
208pub struct AttributeTable {
209    by_names: HashMap<String, AttributeId>,
210    by_ids: Vec<AttributeKind>,
211}
212
213#[derive(Clone, Copy, Eq, Ord, PartialEq, PartialOrd, Debug, Hash)]
214pub struct AttributeId(usize);
215
216impl Display for AttributeId {
217    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
218        write!(formatter, "attribute({})", self.0)
219    }
220}
221
222impl AttributeTable {
223    pub fn new(definitions: &[AttributeDefinition]) -> Result<Self, EventError> {
224        let size = definitions.len();
225        let mut by_names = HashMap::with_capacity(size);
226        let mut by_ids = Vec::with_capacity(size);
227        for (i, definition) in definitions.iter().enumerate() {
228            let name = definition.name.to_owned();
229            if by_names.contains_key(&name) {
230                return Err(EventError::AlreadyPresent(name));
231            }
232
233            by_names.insert(name, AttributeId(i));
234            by_ids.push(definition.kind.clone());
235        }
236
237        Ok(Self { by_names, by_ids })
238    }
239
240    #[inline]
241    pub fn by_name(&self, name: &str) -> Option<AttributeId> {
242        self.by_names.get(name).cloned()
243    }
244
245    #[inline]
246    pub fn by_id(&self, id: AttributeId) -> AttributeKind {
247        self.by_ids[id.0].clone()
248    }
249
250    #[inline]
251    pub fn len(&self) -> usize {
252        self.by_ids.len()
253    }
254}
255
256/// The definition of an attribute that is usable by the [`crate::atree::ATree`]
257#[derive(Debug, Clone)]
258pub struct AttributeDefinition {
259    name: String,
260    kind: AttributeKind,
261}
262
263#[derive(Clone, PartialEq, Debug)]
264pub enum AttributeKind {
265    Boolean,
266    Integer,
267    Float,
268    String,
269    IntegerList,
270    StringList,
271}
272
273impl AttributeDefinition {
274    /// Create a boolean attribute definition.
275    pub fn boolean(name: &str) -> Self {
276        let kind = AttributeKind::Boolean;
277        Self {
278            name: name.to_owned(),
279            kind,
280        }
281    }
282
283    /// Create an integer attribute definition.
284    pub fn integer(name: &str) -> Self {
285        let kind = AttributeKind::Integer;
286        Self {
287            name: name.to_owned(),
288            kind,
289        }
290    }
291
292    /// Create a float attribute definition.
293    pub fn float(name: &str) -> Self {
294        let kind = AttributeKind::Float;
295        Self {
296            name: name.to_owned(),
297            kind,
298        }
299    }
300
301    /// Create a string attribute definition.
302    pub fn string(name: &str) -> Self {
303        let kind = AttributeKind::String;
304        Self {
305            name: name.to_owned(),
306            kind,
307        }
308    }
309
310    /// Create a list of integers attribute definition.
311    pub fn integer_list(name: &str) -> Self {
312        let kind = AttributeKind::IntegerList;
313        Self {
314            name: name.to_owned(),
315            kind,
316        }
317    }
318
319    /// Create a list of strings attribute definition.
320    pub fn string_list(name: &str) -> Self {
321        let kind = AttributeKind::StringList;
322        Self {
323            name: name.to_owned(),
324            kind,
325        }
326    }
327}
328
329#[cfg(test)]
330mod tests {
331    use super::*;
332
333    #[test]
334    fn can_create_an_attribute_table_with_no_attributes() {
335        assert!(AttributeTable::new(&[]).is_ok())
336    }
337
338    #[test]
339    fn can_create_an_attribute_table_with_some_attributes() {
340        let definitions = [
341            AttributeDefinition::boolean("private"),
342            AttributeDefinition::string_list("deals"),
343            AttributeDefinition::integer("exchange_id"),
344            AttributeDefinition::float("bidfloor"),
345            AttributeDefinition::string("country"),
346            AttributeDefinition::integer_list("segment_ids"),
347        ];
348
349        assert!(AttributeTable::new(&definitions).is_ok());
350    }
351
352    #[test]
353    fn return_an_error_on_duplicate_definitions() {
354        let definitions = [
355            AttributeDefinition::boolean("private"),
356            AttributeDefinition::string("country"),
357            AttributeDefinition::string_list("deals"),
358            AttributeDefinition::integer("exchange_id"),
359            AttributeDefinition::float("bidfloor"),
360            AttributeDefinition::integer("country"),
361            AttributeDefinition::integer_list("segment_ids"),
362        ];
363
364        assert!(AttributeTable::new(&definitions).is_err());
365    }
366
367    #[test]
368    fn can_add_a_boolean_attribute_value() {
369        let attributes = AttributeTable::new(&[AttributeDefinition::boolean("private")]).unwrap();
370        let strings = StringTable::new();
371        let mut event_builder = EventBuilder::new(&attributes, &strings);
372
373        let result = event_builder.with_boolean("private", true);
374
375        assert!(result.is_ok());
376    }
377
378    #[test]
379    fn can_add_an_integer_attribute_value() {
380        let attributes =
381            AttributeTable::new(&[AttributeDefinition::integer("exchange_id")]).unwrap();
382        let strings = StringTable::new();
383        let mut event_builder = EventBuilder::new(&attributes, &strings);
384
385        let result = event_builder.with_integer("exchange_id", 1);
386
387        assert!(result.is_ok());
388    }
389
390    #[test]
391    fn can_add_a_float_attribute_value() {
392        let attributes = AttributeTable::new(&[AttributeDefinition::float("bidfloor")]).unwrap();
393        let strings = StringTable::new();
394        let mut event_builder = EventBuilder::new(&attributes, &strings);
395
396        let result = event_builder.with_float("bidfloor", 1, 0);
397
398        assert!(result.is_ok());
399    }
400
401    #[test]
402    fn can_add_a_string_attribute_value() {
403        let attributes = AttributeTable::new(&[AttributeDefinition::string("country")]).unwrap();
404        let strings = StringTable::new();
405        let mut event_builder = EventBuilder::new(&attributes, &strings);
406
407        let result = event_builder.with_string("country", "US");
408
409        assert!(result.is_ok());
410    }
411
412    #[test]
413    fn can_add_an_integer_list_attribute_value() {
414        let attributes =
415            AttributeTable::new(&[AttributeDefinition::integer_list("segment_ids")]).unwrap();
416        let strings = StringTable::new();
417        let mut event_builder = EventBuilder::new(&attributes, &strings);
418
419        let result = event_builder.with_integer_list("segment_ids", &[1, 2, 3]);
420
421        assert!(result.is_ok());
422    }
423
424    #[test]
425    fn can_add_an_string_list_attribute_value() {
426        let attributes =
427            AttributeTable::new(&[AttributeDefinition::string_list("deal_ids")]).unwrap();
428        let strings = StringTable::new();
429        let mut event_builder = EventBuilder::new(&attributes, &strings);
430
431        let result = event_builder.with_string_list("deal_ids", &["deal-1", "deal-2"]);
432
433        assert!(result.is_ok());
434    }
435
436    #[test]
437    fn return_an_error_when_adding_a_non_existing_attribute() {
438        let attributes =
439            AttributeTable::new(&[AttributeDefinition::string_list("deal_ids")]).unwrap();
440        let strings = StringTable::new();
441        let mut event_builder = EventBuilder::new(&attributes, &strings);
442
443        let result = event_builder.with_boolean("non_existing", true);
444
445        assert!(matches!(result, Err(EventError::NonExistingAttribute(_))));
446    }
447
448    #[test]
449    fn can_create_an_event_with_no_attributes() {
450        let attributes = AttributeTable::new(&[]).unwrap();
451        let strings = StringTable::new();
452        let event_builder = EventBuilder::new(&attributes, &strings);
453
454        assert!(event_builder.build().is_ok());
455    }
456
457    #[test]
458    fn can_create_an_event_with_attributes() {
459        let attributes = AttributeTable::new(&[
460            AttributeDefinition::boolean("private"),
461            AttributeDefinition::string_list("deals"),
462            AttributeDefinition::integer("exchange_id"),
463            AttributeDefinition::float("bidfloor"),
464            AttributeDefinition::string("country"),
465            AttributeDefinition::integer_list("segment_ids"),
466        ])
467        .unwrap();
468        let strings = StringTable::new();
469        let mut builder = EventBuilder::new(&attributes, &strings);
470
471        assert!(builder.with_boolean("private", true).is_ok());
472        assert!(builder
473            .with_string_list("deals", &["deal-1", "deal-2"])
474            .is_ok());
475        assert!(builder.with_integer("exchange_id", 1).is_ok());
476        assert!(builder.with_float("bidfloor", 1, 0).is_ok());
477        assert!(builder.with_string("country", "US").is_ok());
478        assert!(builder.with_integer_list("segment_ids", &[1, 2, 3]).is_ok());
479
480        assert!(builder.build().is_ok());
481    }
482
483    #[test]
484    fn can_create_an_event_with_a_missing_attribute() {
485        let attributes = AttributeTable::new(&[AttributeDefinition::boolean("private")]).unwrap();
486        let strings = StringTable::new();
487        let event_builder = EventBuilder::new(&attributes, &strings);
488
489        assert!(event_builder.build().is_ok());
490    }
491
492    #[test]
493    fn return_an_error_when_trying_to_add_an_attribute_with_mismatched_type() {
494        let attributes = AttributeTable::new(&[AttributeDefinition::boolean("private")]).unwrap();
495        let strings = StringTable::new();
496        let mut event_builder = EventBuilder::new(&attributes, &strings);
497
498        let result = event_builder.with_integer("private", 1);
499
500        assert!(result.is_err());
501    }
502}
a_tree/events.rs

a_tree/
events.rs
