Skip to main content

dbsp/dynamic/
weight.rs

1use crate::{
2    DBWeight,
3    algebra::{AddAssignByRef, AddByRef, HasZero},
4    declare_trait_object, declare_typed_trait_object,
5    dynamic::data::DataTyped,
6};
7
8use super::{Data, DataTrait, DataTraitTyped};
9
10/// A trait for types that can be used as weights in DBSP.
11///
12/// Such types must implement the plus operation and have a neutral element wrt addition (zero).
13pub trait Weight: Data {
14    /// Set the weight to zero.
15    fn set_zero(&mut self);
16
17    /// Check if the weight is zero.
18    fn is_zero(&self) -> bool;
19
20    /// Add the value of `rhs` to `self`, store the result in `result`.
21    ///
22    /// # Safety
23    ///
24    /// `rhs` and `result` must be valid pointers to values of type `Self`.
25    unsafe fn raw_add(&self, rhs: *const u8, result: *mut u8);
26
27    /// Add the value of `rhs` to `self`, store the result in `self`.
28    ///
29    /// # Safety
30    ///
31    /// `rhs` must be a valid pointer to a value of type `Self`
32    unsafe fn raw_add_assign(&mut self, rhs: *const u8);
33}
34
35impl<T: DBWeight> Weight for T {
36    fn set_zero(&mut self) {
37        *self = HasZero::zero();
38    }
39
40    fn is_zero(&self) -> bool {
41        HasZero::is_zero(self)
42    }
43
44    unsafe fn raw_add(&self, rhs: *const u8, result: *mut u8) {
45        unsafe { *(result as *mut Self) = AddByRef::add_by_ref(self, &*(rhs as *const Self)) }
46    }
47
48    unsafe fn raw_add_assign(&mut self, rhs: *const u8) {
49        unsafe { AddAssignByRef::add_assign_by_ref(self, &*(rhs as *const Self)) }
50    }
51}
52
53/// Strongly typed weight.
54///
55/// A trait that includes the concrete weight type in its type signature.
56/// See [`WeightTraitTyped`].
57pub trait WeightTyped: Weight + DataTyped {}
58
59impl<T> WeightTyped for T where T: Weight + DataTyped {}
60
61/// A trait for trait objects that represent weights.
62pub trait WeightTrait: DataTrait + Weight {
63    /// Add `rhs` to `self` and store the result in `result`.
64    fn add(&self, rhs: &Self, result: &mut Self) {
65        debug_assert_eq!(self.as_any().type_id(), rhs.as_any().type_id());
66        debug_assert_eq!(self.as_any().type_id(), result.as_any().type_id());
67
68        unsafe {
69            Weight::raw_add(
70                self,
71                rhs as *const Self as *const u8,
72                result as *mut Self as *mut u8,
73            )
74        }
75    }
76
77    /// Add `rhs` to `self`, store the result in `self`.
78    fn add_assign(&mut self, rhs: &Self) {
79        debug_assert_eq!(self.as_any().type_id(), rhs.as_any().type_id());
80
81        unsafe { Weight::raw_add_assign(self, rhs as *const Self as *const u8) }
82    }
83}
84
85impl<Trait: ?Sized> WeightTrait for Trait where Trait: DataTrait + Weight {}
86
87/// A trait for trait objects that include the concrete weight type in
88/// their type signatures.
89///
90/// Such trait objects can be dereferenced into their concrete types.
91///
92/// See [`DynZWeight`](`crate::DynZWeight`) for an example of such a trait object.
93pub trait WeightTraitTyped: WeightTrait + DataTraitTyped {}
94
95impl<Trait: ?Sized> WeightTraitTyped for Trait where Trait: WeightTrait + DataTraitTyped {}
96
97declare_trait_object!(DynWeight<> = dyn Weight<>
98where
99);
100
101declare_typed_trait_object!(DynWeightTyped<T> = dyn WeightTyped<Type = T>
102    [T]
103where
104    T: DBWeight,
105);