Merge pull request #49 from lun3x/48-embedded-indexes

Allow callers to use Borrowed key types, like stdlib collections
lun3x · Aug 15, 2024 · 241f214 · 241f214
2 parents 437b416 + 6c133c9
commit 241f214
Show file tree

Hide file tree

Showing 5 changed files with 114 additions and 18 deletions.
diff --git a/multi_index_map/tests/get_and_modify_mut.rs b/multi_index_map/tests/get_and_modify_mut.rs
@@ -1,8 +1,5 @@
 use multi_index_map::MultiIndexMap;
 
-#[derive(Hash, PartialEq, Eq, Clone)]
-struct TestNonPrimitiveType(u64);
-
 #[derive(MultiIndexMap, Clone, Debug)]
 #[multi_index_derive(Clone, Debug)]
 struct TestElement {

diff --git a/multi_index_map/tests/iter_after_modify.rs → ...ap/tests/hashed_non_unique_with_string.rs b/multi_index_map/tests/iter_after_modify.rs → ...ap/tests/hashed_non_unique_with_string.rs
@@ -69,3 +69,20 @@ fn iter_after_modify() {
         assert_eq!(it.next().unwrap().order_id, 4);
     }
 }
+
+#[test]
+fn get_by_borrowed_string() {
+    let o1 = Order {
+        order_id: 1,
+        timestamp: 111,
+        trader_name: "John".to_string(),
+    };
+
+    let mut map = MultiIndexOrderMap::default();
+
+    map.insert(o1);
+
+    let res = map.get_by_trader_name("John");
+    assert_eq!(res.len(), 1);
+    assert_eq!(res.first().unwrap().order_id, 1);
+}
diff --git a/multi_index_map/tests/hashed_unique.rs b/multi_index_map/tests/hashed_unique.rs
@@ -1,4 +1,4 @@
-use multi_index_map::{MultiIndexMap, UniquenessError};
+use multi_index_map::MultiIndexMap;
 
 #[derive(Hash, PartialEq, Eq, Clone)]
 struct TestNonPrimitiveType(u64);

diff --git a/multi_index_map/tests/update.rs b/multi_index_map/tests/update.rs
@@ -1,8 +1,5 @@
 use multi_index_map::MultiIndexMap;
 
-#[derive(Hash, PartialEq, Eq, Clone)]
-struct TestNonPrimitiveType(u64);
-
 #[derive(MultiIndexMap, PartialEq, Debug)]
 #[multi_index_derive(Debug)]
 struct TestElement {
@@ -12,6 +9,8 @@ struct TestElement {
     #[multi_index(hashed_unique)]
     field3: u32,
     field4: String,
+    #[multi_index(hashed_non_unique)]
+    field5: String,
 }
 
 #[test]
@@ -24,13 +23,15 @@ fn test_non_unique_update() {
                 field2: i as f64,
                 field3: i,
                 field4: i.to_string(),
+                field5: "42".to_string(),
             });
         } else {
             map.insert(TestElement {
                 field1: 37,
                 field2: i as f64,
                 field3: i,
                 field4: i.to_string(),
+                field5: "37".to_string(),
             });
         }
     }
@@ -51,6 +52,45 @@ fn test_non_unique_update() {
     }
 }
 
+#[test]
+fn test_non_unique_update_borrow() {
+    let mut map = MultiIndexTestElementMap::default();
+    for i in 0..10 {
+        if i % 2 == 0 {
+            map.insert(TestElement {
+                field1: 42,
+                field2: i as f64,
+                field3: i,
+                field4: i.to_string(),
+                field5: "42".to_string(),
+            });
+        } else {
+            map.insert(TestElement {
+                field1: 37,
+                field2: i as f64,
+                field3: i,
+                field4: i.to_string(),
+                field5: "37".to_string(),
+            });
+        }
+    }
+
+    let refs = map.update_by_field5("37", |field2, field4| {
+        *field2 = 99.0;
+        *field4 = "NinetyNine".to_string()
+    });
+    for r in refs.iter() {
+        assert_eq!(r.field2, 99.0);
+        assert_eq!(r.field4, "NinetyNine");
+    }
+
+    let refs = map.get_by_field1(&42);
+    for (i, r) in refs.iter().enumerate() {
+        assert_eq!(r.field2, i as f64 * 2.0);
+        assert_eq!(r.field4, (i * 2).to_string());
+    }
+}
+
 #[test]
 fn test_unique_update() {
     let mut map = MultiIndexTestElementMap::default();
@@ -61,13 +101,15 @@ fn test_unique_update() {
                 field2: i as f64,
                 field3: i,
                 field4: i.to_string(),
+                field5: "42".to_string(),
             });
         } else {
             map.insert(TestElement {
                 field1: 37,
                 field2: i as f64,
                 field3: i,
                 field4: i.to_string(),
+                field5: "37".to_string(),
             });
         }
     }
@@ -83,7 +125,8 @@ fn test_unique_update() {
             field1: 42,
             field2: 99.0,
             field3: 0,
-            field4: "NinetyNine".to_string()
+            field4: "NinetyNine".to_string(),
+            field5: "42".to_string()
         })
     );
 
@@ -95,7 +138,8 @@ fn test_unique_update() {
             field1: 37,
             field2: 1.0,
             field3: 1,
-            field4: 1.to_string()
+            field4: 1.to_string(),
+            field5: "37".to_string()
         })
     );
 }
diff --git a/multi_index_map_derive/src/generators.rs b/multi_index_map_derive/src/generators.rs
@@ -334,21 +334,39 @@ fn generate_field_getter(
     field_idents: &FieldIdents,
     field_info: &FieldInfo,
     element_name: &Ident,
+    ordering: &Ordering,
     uniqueness: &Uniqueness,
 ) -> proc_macro2::TokenStream {
     let getter_name = format_ident!("get_by_{}", &field_idents.name);
     let index_name = &field_idents.index_name;
     let field_vis = &field_info.vis;
     let field_type = &field_info.ty;
 
+    let key_bounds = match ordering {
+        Ordering::Hashed => quote! {
+            Q: ::std::hash::Hash + Eq + ?Sized
+        },
+        Ordering::Ordered => quote! {
+            Q: Ord + ?Sized
+        },
+    };
+
     match uniqueness {
         Uniqueness::Unique => quote! {
-            #field_vis fn #getter_name(&self, key: &#field_type) -> Option<&#element_name> {
+            #field_vis fn #getter_name<Q>(&self, key: &Q) -> Option<&#element_name>
+            where
+                #field_type: ::std::borrow::Borrow<Q>,
+                #key_bounds,
+            {
                 Some(&self._store[*self.#index_name.get(key)?])
             }
         },
         Uniqueness::NonUnique => quote! {
-            #field_vis fn #getter_name(&self, key: &#field_type) -> Vec<&#element_name> {
+            #field_vis fn #getter_name<Q>(&self, key: &Q) -> Vec<&#element_name>
+            where
+                #field_type: ::std::borrow::Borrow<Q>,
+                #key_bounds,
+            {
                 if let Some(idxs) = self.#index_name.get(key) {
                     let mut elem_refs = Vec::with_capacity(idxs.len());
                     for idx in idxs {
@@ -472,6 +490,7 @@ fn generate_field_updater(
     field_idents: &FieldIdents,
     field_info: &FieldInfo,
     element_name: &Ident,
+    ordering: &Ordering,
     uniqueness: &Uniqueness,
     unindexed_types: &[&Type],
     unindexed_idents: &[&Ident],
@@ -482,25 +501,42 @@ fn generate_field_updater(
     let field_type = &field_info.ty;
     let field_name_str = &field_info.str;
 
+    let key_bounds = match ordering {
+        Ordering::Hashed => quote! {
+            Q: ::std::hash::Hash + Eq + ?Sized
+        },
+        Ordering::Ordered => quote! {
+            Q: Ord + ?Sized
+        },
+    };
+
     match uniqueness {
         Uniqueness::Unique => quote! {
-            #field_vis fn #updater_name(
+            #field_vis fn #updater_name<Q>(
                 &mut self,
-                key: &#field_type,
+                key: &Q,
                 f: impl FnOnce(#(&mut #unindexed_types,)*)
-            ) -> Option<&#element_name> {
+            ) -> Option<&#element_name>
+            where
+                #field_type: ::std::borrow::Borrow<Q>,
+                #key_bounds,
+            {
                 let idx = *self.#index_name.get(key)?;
                 let elem = &mut self._store[idx];
                 f(#(&mut elem.#unindexed_idents,)*);
                 Some(elem)
             }
         },
         Uniqueness::NonUnique => quote! {
-            #field_vis fn #updater_name(
+            #field_vis fn #updater_name<Q>(
                 &mut self,
-                key: &#field_type,
+                key: &Q,
                 mut f: impl FnMut(#(&mut #unindexed_types,)*)
-            ) -> Vec<&#element_name> {
+            ) -> Vec<&#element_name>
+            where
+                #field_type: ::std::borrow::Borrow<Q>,
+                #key_bounds,
+            {
                 let empty = ::std::collections::BTreeSet::<usize>::new();
                 let idxs = match self.#index_name.get(key) {
                     Some(container) => container,
@@ -663,7 +699,8 @@ pub(crate) fn generate_accessors<'a>(
                 str: &idents.name.to_string(),
             };
 
-            let getter = generate_field_getter(idents, &field_info, element_name, uniqueness);
+            let getter =
+                generate_field_getter(idents, &field_info, element_name, ordering, uniqueness);
 
             let mut_getter =
                 generate_field_mut_getter(idents, &field_info, element_name, uniqueness);
@@ -675,6 +712,7 @@ pub(crate) fn generate_accessors<'a>(
                 idents,
                 &field_info,
                 element_name,
+                ordering,
                 uniqueness,
                 &unindexed_types,
                 &unindexed_idents,