wip

specta-jsdoc/src/lib.rs

@@ -79,7 +79,7 @@ impl Language for JSDoc {
     type Error = specta_typescript::ExportError; // TODO: Custom error type
 
     // TODO: Make this properly export JSDoc
-    fn export(&self, type_map: TypeMap) -> Result<String, Self::Error> {
+    fn export(&self, _type_map: &TypeMap) -> Result<String, Self::Error> {
         todo!("Coming soon...");
         // let mut out = self.0.header.to_string();
         // if !self.0.remove_default_header {

specta-typescript/src/typescript.rs

@@ -132,7 +132,7 @@ impl Typescript {
 impl Language for Typescript {
     type Error = ExportError;
 
-    fn export(&self, type_map: TypeMap) -> Result<String, Self::Error> {
+    fn export(&self, type_map: &TypeMap) -> Result<String, Self::Error> {
         let mut out = self.header.to_string();
         if !self.remove_default_header {
             out += "// This file has been generated by Specta. DO NOT EDIT.\n\n";
@@ -142,7 +142,7 @@ impl Language for Typescript {
             return Err(ExportError::DuplicateTypeName(ty_name, l0, l1));
         }
 
-        for (_, ty) in type_map.iter() {
+        for (_, ty) in type_map.into_iter() {
             is_valid_ty(&ty.inner, &type_map)?;
 
             out += &export_named_datatype(self, ty, &type_map)?;

specta-util/src/export.rs

@@ -5,20 +5,24 @@ use std::{
 
 use specta::{datatype::NamedDataType, NamedType, SpectaID, TypeMap};
 
-use crate::TypeCollection;
-
 // Global type store for collecting custom types to export.
 static TYPES: OnceLock<Mutex<HashMap<SpectaID, fn(&mut TypeMap) -> NamedDataType>>> =
     OnceLock::new();
 
 /// Get the global type store containing all registered types.
-pub fn export() -> TypeCollection {
-    let type_map = TYPES
+pub fn export() -> TypeMap {
+    // TODO: Make `TYPES` should just hold a `TypeMap` directly???
+    let types = TYPES
         .get_or_init(Default::default)
         .lock()
         .unwrap_or_else(PoisonError::into_inner);
 
-    TypeCollection::from_raw(type_map.clone())
+    let mut map = TypeMap::default();
+    for (id, export) in types.iter() {
+        let dt = export(&mut map);
+        map.insert(*id, dt);
+    }
+    map
 }
 
 #[doc(hidden)]

specta-util/src/lib.rs

@@ -19,10 +19,8 @@ pub use export::export;
 #[cfg_attr(docsrs, doc(cfg(feature = "serde")))]
 mod selection;
 mod static_types;
-mod type_collection;
 
 #[cfg(feature = "serde")]
 #[cfg_attr(docsrs, doc(cfg(feature = "serde")))]
 // pub use selection::selection;
 pub use static_types::{Any, Unknown};
-pub use type_collection::TypeCollection;

specta/src/internal.rs

@@ -333,8 +333,8 @@ pub fn detect_duplicate_type_names(
 ) -> Vec<(Cow<'static, str>, ImplLocation, ImplLocation)> {
     let mut errors = Vec::new();
 
-    let mut map = HashMap::with_capacity(type_map.len());
-    for (sid, dt) in type_map.iter() {
+    let mut map = HashMap::with_capacity(type_map.into_iter().len());
+    for (sid, dt) in type_map.into_iter() {
         if let Some(ext) = &dt.ext {
             if let Some((existing_sid, existing_impl_location)) =
                 map.insert(dt.name.clone(), (sid, ext.impl_location))

specta/src/language.rs

@@ -11,7 +11,7 @@ pub trait Language {
     type Error: std::error::Error + From<std::io::Error>;
 
     /// TODO
-    fn export(&self, type_map: TypeMap) -> Result<String, Self::Error>;
+    fn export(&self, type_map: &TypeMap) -> Result<String, Self::Error>;
 
     /// TODO
     // TODO: Not sure I love this here but it's for Tauri Specta.
@@ -22,7 +22,7 @@ pub trait Language {
 impl<T: Language> Language for &T {
     type Error = T::Error;
 
-    fn export(&self, type_map: TypeMap) -> Result<String, Self::Error> {
+    fn export(&self, type_map: &TypeMap) -> Result<String, Self::Error> {
         (*self).export(type_map)
     }
 

specta/src/type_map.rs

@@ -1,8 +1,16 @@
-use std::{collections::BTreeMap, fmt};
+use std::{
+    borrow::Borrow,
+    collections::{btree_map, BTreeMap},
+    fmt,
+    path::Path,
+};
 
-use crate::{datatype::NamedDataType, SpectaID};
+use crate::{datatype::NamedDataType, Language, NamedType, SpectaID};
 
-/// A map used to store the types "discovered" while exporting a type.
+/// Define a set of types which can be exported together.
+///
+/// While exporting a type will add all of the types it depends on to the collection.
+/// You can construct your own collection to easily export a set of types together.
 #[derive(Default, Clone, PartialEq)]
 pub struct TypeMap {
     // `None` indicates that the entry is a placeholder. It was reference and we are currently working out it's definition.
@@ -18,6 +26,46 @@ impl fmt::Debug for TypeMap {
 }
 
 impl TypeMap {
+    /// Register a type with the collection.
+    pub fn register<T: NamedType>(&mut self) -> &mut Self {
+        let def = T::definition_named_data_type(self);
+        self.map.insert(T::sid(), Some(def));
+        self
+    }
+
+    /// Insert a type into the collection.
+    /// You should prefer to use `TypeMap::register` as it ensures all invariants are met.
+    ///
+    /// When using this method it's the responsibility of the caller to:
+    ///  - Ensure the `SpectaID` and `NamedDataType` are correctly matched.
+    ///  - Ensure the same `TypeMap` was used when calling `NamedType::definition_named_data_type`.
+    /// Not honoring these rules will result in a broken collection.
+    pub fn insert(&mut self, sid: SpectaID, def: NamedDataType) -> &mut Self {
+        self.map.insert(sid, Some(def));
+        self
+    }
+
+    /// Join another type collection into this one.
+    pub fn extend(&mut self, collection: impl Borrow<Self>) -> &mut Self {
+        self.map
+            .extend(collection.borrow().map.iter().map(|(k, v)| (*k, v.clone())));
+        self
+    }
+
+    /// TODO
+    pub fn export<L: Language>(&self, language: L) -> Result<String, L::Error> {
+        language.export(self)
+    }
+
+    /// TODO
+    pub fn export_to<L: Language>(
+        &self,
+        language: L,
+        path: impl AsRef<Path>,
+    ) -> Result<(), L::Error> {
+        std::fs::write(path, self.export(language)?).map_err(Into::into)
+    }
+
     #[track_caller]
     pub fn get(&self, sid: SpectaID) -> Option<&NamedDataType> {
         #[allow(clippy::bind_instead_of_map)]
@@ -33,42 +81,35 @@ impl TypeMap {
             }
         })
     }
+}
 
-    pub fn insert(&mut self, sid: SpectaID, dt: NamedDataType) {
-        self.map.insert(sid, Some(dt));
-    }
+impl<'a> IntoIterator for &'a TypeMap {
+    type Item = (SpectaID, &'a NamedDataType);
+    type IntoIter = TypeMapInterator<'a>;
 
-    pub fn is_empty(&self) -> bool {
-        self.map.is_empty()
+    fn into_iter(self) -> Self::IntoIter {
+        TypeMapInterator(self.map.iter())
     }
+}
 
-    pub fn len(&self) -> usize {
-        self.map.len()
-    }
+// Sealed
+pub struct TypeMapInterator<'a>(btree_map::Iter<'a, SpectaID, Option<NamedDataType>>);
 
-    pub fn contains_key(&self, sid: SpectaID) -> bool {
-        self.map.contains_key(&sid)
-    }
+impl<'a> ExactSizeIterator for TypeMapInterator<'a> {}
 
-    pub fn remove(&mut self, sid: SpectaID) -> Option<NamedDataType> {
-        self.map.remove(&sid).flatten()
-    }
+impl<'a> Iterator for TypeMapInterator<'a> {
+    type Item = (SpectaID, &'a NamedDataType);
 
-    pub fn append(&mut self, type_map: &mut TypeMap) {
-        self.map.append(&mut type_map.map);
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            let (sid, ndt) = self.0.next()?;
+            if let Some(ndt) = ndt {
+                return Some((*sid, ndt));
+            }
+        }
     }
 
-    // TODO: It would be nice if this would a proper `Iterator` or `IntoIterator` implementation!
-    pub fn iter(&self) -> impl Iterator<Item = (SpectaID, &NamedDataType)> {
-        #[allow(clippy::unnecessary_filter_map)]
-        self.map.iter().filter_map(|(sid, ndt)| match ndt {
-            Some(ndt) => Some((*sid, ndt)),
-            None => {
-                #[cfg(debug_assertions)]
-                unreachable!("specta: `TypeMap::into_iter` found a type placeholder!");
-                #[cfg(not(debug_assertions))]
-                None
-            }
-        })
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (0, Some(self.0.len()))
     }
 }