orml_utilities/
ordered_set.rs

1use frame_support::{traits::Get, BoundedVec, CloneNoBound, DefaultNoBound, PartialEqNoBound};
2use parity_scale_codec::{Decode, Encode, MaxEncodedLen};
3use scale_info::TypeInfo;
4#[cfg(feature = "std")]
5use sp_std::{fmt, prelude::*};
6
7/// An ordered set backed by `BoundedVec`
8#[derive(PartialEqNoBound, Eq, Encode, Decode, DefaultNoBound, CloneNoBound, TypeInfo, MaxEncodedLen)]
9#[codec(mel_bound())]
10#[scale_info(skip_type_params(S))]
11pub struct OrderedSet<T: Ord + Encode + Decode + MaxEncodedLen + Clone + Eq + PartialEq, S: Get<u32>>(
12	pub BoundedVec<T, S>,
13);
14
15impl<T: Ord + Encode + Decode + MaxEncodedLen + Clone + Eq + PartialEq, S: Get<u32>> OrderedSet<T, S> {
16	/// Create a new empty set
17	pub fn new() -> Self {
18		Self(BoundedVec::default())
19	}
20
21	/// Create a set from a `Vec`.
22	/// `v` will be sorted and dedup first.
23	pub fn from(bv: BoundedVec<T, S>) -> Self {
24		let mut v = bv.into_inner();
25		v.sort();
26		v.dedup();
27
28		Self::from_sorted_set(v.try_into().map_err(|_| ()).expect("Did not add any values"))
29	}
30
31	/// Create a set from a `Vec`.
32	/// Assume `v` is sorted and contain unique elements.
33	pub fn from_sorted_set(bv: BoundedVec<T, S>) -> Self {
34		Self(bv)
35	}
36
37	/// Insert an element.
38	/// Return true if insertion happened.
39	pub fn insert(&mut self, value: T) -> bool {
40		match self.0.binary_search(&value) {
41			Ok(_) => false,
42			Err(loc) => self.0.try_insert(loc, value).is_ok(),
43		}
44	}
45
46	/// Remove an element.
47	/// Return true if removal happened.
48	pub fn remove(&mut self, value: &T) -> bool {
49		match self.0.binary_search(value) {
50			Ok(loc) => {
51				self.0.remove(loc);
52				true
53			}
54			Err(_) => false,
55		}
56	}
57
58	/// Return if the set contains `value`
59	pub fn contains(&self, value: &T) -> bool {
60		self.0.binary_search(value).is_ok()
61	}
62
63	/// Clear the set
64	pub fn clear(&mut self) {
65		self.0 = BoundedVec::default();
66	}
67}
68
69impl<T: Ord + Encode + Decode + MaxEncodedLen + Clone + Eq + PartialEq, S: Get<u32>> From<BoundedVec<T, S>>
70	for OrderedSet<T, S>
71{
72	fn from(v: BoundedVec<T, S>) -> Self {
73		Self::from(v)
74	}
75}
76
77#[cfg(feature = "std")]
78impl<T, S> fmt::Debug for OrderedSet<T, S>
79where
80	T: Ord + Encode + Decode + MaxEncodedLen + Clone + Eq + PartialEq + fmt::Debug,
81	S: Get<u32>,
82{
83	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
84		f.debug_tuple("OrderedSet").field(&self.0).finish()
85	}
86}
87
88#[cfg(test)]
89mod tests {
90	use super::*;
91	use frame_support::parameter_types;
92	use sp_runtime::RuntimeDebug;
93
94	parameter_types! {
95		#[derive(PartialEq, Eq, RuntimeDebug)]
96		pub const Eight: u32 = 8;
97		#[derive(PartialEq, Eq, RuntimeDebug)]
98		pub const Five: u32 = 5;
99	}
100
101	#[test]
102	fn from() {
103		let v: BoundedVec<i32, Eight> = vec![4, 2, 3, 4, 3, 1].try_into().unwrap();
104		let set: OrderedSet<i32, Eight> = v.into();
105		assert_eq!(
106			set,
107			OrderedSet::<i32, Eight>::from(vec![1, 2, 3, 4].try_into().unwrap())
108		);
109	}
110
111	#[test]
112	fn insert() {
113		let mut set: OrderedSet<i32, Eight> = OrderedSet::new();
114		assert_eq!(set, OrderedSet::<i32, Eight>::from(vec![].try_into().unwrap()));
115
116		assert!(set.insert(1));
117		assert_eq!(set, OrderedSet::<i32, Eight>::from(vec![1].try_into().unwrap()));
118
119		assert!(set.insert(5));
120		assert_eq!(set, OrderedSet::<i32, Eight>::from(vec![1, 5].try_into().unwrap()));
121
122		assert!(set.insert(3));
123		assert_eq!(set, OrderedSet::<i32, Eight>::from(vec![1, 3, 5].try_into().unwrap()));
124
125		assert!(!set.insert(3));
126		assert_eq!(set, OrderedSet::<i32, Eight>::from(vec![1, 3, 5].try_into().unwrap()));
127	}
128
129	#[test]
130	fn remove() {
131		let mut set: OrderedSet<i32, Eight> = OrderedSet::from(vec![1, 2, 3, 4].try_into().unwrap());
132
133		assert!(!set.remove(&5));
134		assert_eq!(
135			set,
136			OrderedSet::<i32, Eight>::from(vec![1, 2, 3, 4].try_into().unwrap())
137		);
138
139		assert!(set.remove(&1));
140		assert_eq!(set, OrderedSet::<i32, Eight>::from(vec![2, 3, 4].try_into().unwrap()));
141
142		assert!(set.remove(&3));
143		assert_eq!(set, OrderedSet::<i32, Eight>::from(vec![2, 4].try_into().unwrap()));
144
145		assert!(!set.remove(&3));
146		assert_eq!(set, OrderedSet::<i32, Eight>::from(vec![2, 4].try_into().unwrap()));
147
148		assert!(set.remove(&4));
149		assert_eq!(set, OrderedSet::<i32, Eight>::from(vec![2].try_into().unwrap()));
150
151		assert!(set.remove(&2));
152		assert_eq!(set, OrderedSet::<i32, Eight>::from(vec![].try_into().unwrap()));
153
154		assert!(!set.remove(&2));
155		assert_eq!(set, OrderedSet::<i32, Eight>::from(vec![].try_into().unwrap()));
156	}
157
158	#[test]
159	fn contains() {
160		let set: OrderedSet<i32, Eight> = OrderedSet::from(vec![1, 2, 3, 4].try_into().unwrap());
161
162		assert!(!set.contains(&5));
163
164		assert!(set.contains(&1));
165
166		assert!(set.contains(&3));
167	}
168
169	#[test]
170	fn clear() {
171		let mut set: OrderedSet<i32, Eight> = OrderedSet::from(vec![1, 2, 3, 4].try_into().unwrap());
172		set.clear();
173		assert_eq!(set, OrderedSet::new());
174	}
175
176	#[test]
177	fn exceeding_max_size_should_fail() {
178		let mut set: OrderedSet<i32, Five> = OrderedSet::from(vec![1, 2, 3, 4, 5].try_into().unwrap());
179		let inserted = set.insert(6);
180
181		assert!(!inserted)
182	}
183}