Cleanup type collection

oscartbeaumont · Apr 30, 2024 · 029e296 · 029e296
1 parent a542d55
commit 029e296
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 6 deletions.
diff --git a/src/type_collection.rs b/src/type_collection.rs
@@ -1,4 +1,7 @@
-use std::collections::HashMap;
+use std::{
+ borrow::{Borrow, Cow},
+ collections::HashMap,
+};
 
 use crate::{NamedDataType, NamedType, SpectaID, TypeMap};
 
@@ -18,8 +21,8 @@ impl Default for TypeCollection {
 
 impl TypeCollection {
  /// Join another type collection into this one.
- pub fn join(&mut self, collection: TypeCollection) -> &mut Self {
- self.types.extend(collection.types);
+ pub fn extend(&mut self, collection: impl Borrow<TypeCollection>) -> &mut Self {
+ self.types.extend(collection.borrow().types.iter());
  self
  }
 
@@ -31,7 +34,7 @@ impl TypeCollection {
  }
 
  /// Export all the types in the collection.
- pub fn export(&mut self, mut type_map: &mut crate::TypeMap) {
+ pub fn export(&self, mut type_map: &mut crate::TypeMap) {
  for (sid, export) in self.types.iter() {
  let dt = export(&mut type_map);
  type_map.insert(*sid, dt);

diff --git a/tests/type_collection.rs b/tests/type_collection.rs
@@ -35,10 +35,16 @@ fn type_collection_merge() {
  let mut type_map = TypeMap::default();
  TypeCollection::default()
  .register::<D>()
- .join(a)
- .join(b)
+ .extend(a)
+ .extend(b)
  .export(&mut type_map);
  assert_eq!(type_map.len(), 4);
+
+ // Check it compile with any valid arg
+ TypeCollection::default()
+ .extend(&TypeCollection::default())
+ .extend(&mut TypeCollection::default())
+ .extend(TypeCollection::default());
 }
 
 #[test]