remove definitions vec

Author: davidhewitt

src/definitions.rs

@@ -3,16 +3,17 @@
 /// Unlike json schema we let you put definitions inline, not just in a single '#/$defs/' block or similar.
 /// We use DefinitionsBuilder to collect the references / definitions into a single vector
 /// and then get a definition from a reference using an integer id (just for performance of not using a HashMap)
-use std::collections::hash_map::Entry;
+use std::{
+    collections::hash_map::Entry,
+    fmt::Debug,
+    sync::{Arc, OnceLock},
+};
 
-use pyo3::prelude::*;
+use pyo3::{prelude::*, PyTraverseError, PyVisit};
 
 use ahash::AHashMap;
 
-use crate::build_tools::py_schema_err;
-
-// An integer id for the reference
-pub type ReferenceId = usize;
+use crate::{build_tools::py_schema_err, py_gc::PyGcTraverse};
 
 /// Definitions are validators and serializers that are
 /// shared by reference.
@@ -24,91 +25,154 @@ pub type ReferenceId = usize;
 /// They get indexed by a ReferenceId, which are integer identifiers
 /// that are handed out and managed by DefinitionsBuilder when the Schema{Validator,Serializer}
 /// gets build.
-pub type Definitions<T> = [T];
+#[derive(Clone)]
+pub struct Definitions<T>(AHashMap<Arc<String>, Definition<T>>);
 
-#[derive(Clone, Debug)]
-struct Definition<T> {
-    pub id: ReferenceId,
-    pub value: Option<T>,
+impl<T> Definitions<T> {
+    pub fn values(&self) -> impl Iterator<Item = &Definition<T>> {
+        self.0.values()
+    }
+}
+
+/// Internal type which contains a definition to be filled
+pub struct Definition<T>(Arc<OnceLock<T>>);
+
+impl<T> Definition<T> {
+    pub fn get(&self) -> Option<&T> {
+        self.0.get()
+    }
+}
+
+/// Reference to a definition.
+pub struct DefinitionRef<T> {
+    name: Arc<String>,
+    value: Definition<T>,
+}
+
+// DefinitionRef can always be cloned (#[derive(Clone)] would require T: Clone)
+impl<T> Clone for DefinitionRef<T> {
+    fn clone(&self) -> Self {
+        Self {
+            name: self.name.clone(),
+            value: self.value.clone(),
+        }
+    }
+}
+
+impl<T> DefinitionRef<T> {
+    pub fn id(&self) -> usize {
+        Arc::as_ptr(&self.value.0) as usize
+    }
+
+    pub fn name(&self) -> &str {
+        &self.name
+    }
+
+    pub fn get(&self) -> Option<&T> {
+        self.value.0.get()
+    }
+}
+
+impl<T: Debug> Debug for DefinitionRef<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        // To avoid possible infinite recursion from recursive definitions,
+        // a DefinitionRef just displays debug as its name
+        self.name.fmt(f)
+    }
+}
+
+impl<T: Debug> Debug for Definitions<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+impl<T> Clone for Definition<T> {
+    fn clone(&self) -> Self {
+        Self(self.0.clone())
+    }
+}
+
+impl<T: Debug> Debug for Definition<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self.0.get() {
+            Some(value) => value.fmt(f),
+            None => "...".fmt(f),
+        }
+    }
+}
+
+impl<T: PyGcTraverse> PyGcTraverse for DefinitionRef<T> {
+    fn py_gc_traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        if let Some(value) = self.value.0.get() {
+            value.py_gc_traverse(visit)?;
+        }
+        Ok(())
+    }
+}
+
+impl<T: PyGcTraverse> PyGcTraverse for Definitions<T> {
+    fn py_gc_traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
+        for value in self.0.values() {
+            if let Some(value) = value.0.get() {
+                value.py_gc_traverse(visit)?;
+            }
+        }
+        Ok(())
+    }
 }
 
 #[derive(Clone, Debug)]
 pub struct DefinitionsBuilder<T> {
-    definitions: AHashMap<String, Definition<T>>,
+    definitions: Definitions<T>,
 }
 
-impl<T: Clone + std::fmt::Debug> DefinitionsBuilder<T> {
+impl<T: std::fmt::Debug> DefinitionsBuilder<T> {
     pub fn new() -> Self {
         Self {
-            definitions: AHashMap::new(),
+            definitions: Definitions(AHashMap::new()),
         }
     }
 
     /// Get a ReferenceId for the given reference string.
-    // This ReferenceId can later be used to retrieve a definition
-    pub fn get_reference_id(&mut self, reference: &str) -> ReferenceId {
-        let next_id = self.definitions.len();
+    pub fn get_definition(&mut self, reference: &str) -> DefinitionRef<T> {
         // We either need a String copy or two hashmap lookups
         // Neither is better than the other
         // We opted for the easier outward facing API
-        match self.definitions.entry(reference.to_string()) {
-            Entry::Occupied(entry) => entry.get().id,
-            Entry::Vacant(entry) => {
-                entry.insert(Definition {
-                    id: next_id,
-                    value: None,
-                });
-                next_id
-            }
+        let name = Arc::new(reference.to_string());
+        let value = match self.definitions.0.entry(name.clone()) {
+            Entry::Occupied(entry) => entry.into_mut(),
+            Entry::Vacant(entry) => entry.insert(Definition(Arc::new(OnceLock::new()))),
+        };
+        DefinitionRef {
+            name,
+            value: value.clone(),
         }
     }
 
     /// Add a definition, returning the ReferenceId that maps to it
-    pub fn add_definition(&mut self, reference: String, value: T) -> PyResult<ReferenceId> {
-        let next_id = self.definitions.len();
-        match self.definitions.entry(reference.clone()) {
-            Entry::Occupied(mut entry) => match entry.get_mut().value.replace(value) {
-                Some(_) => py_schema_err!("Duplicate ref: `{}`", reference),
-                None => Ok(entry.get().id),
-            },
-            Entry::Vacant(entry) => {
-                entry.insert(Definition {
-                    id: next_id,
-                    value: Some(value),
-                });
-                Ok(next_id)
+    pub fn add_definition(&mut self, reference: String, value: T) -> PyResult<DefinitionRef<T>> {
+        let name = Arc::new(reference);
+        let value = match self.definitions.0.entry(name.clone()) {
+            Entry::Occupied(entry) => {
+                let definition = entry.into_mut();
+                match definition.0.set(value) {
+                    Ok(()) => definition.clone(),
+                    Err(_) => return py_schema_err!("Duplicate ref: `{}`", name),
+                }
             }
-        }
-    }
-
-    /// Retrieve an item definition using a ReferenceId
-    /// If the definition doesn't yet exist (as happens in recursive types) then we create it
-    /// At the end (in finish()) we check that there are no undefined definitions
-    pub fn get_definition(&self, reference_id: ReferenceId) -> PyResult<&T> {
-        let (reference, def) = match self.definitions.iter().find(|(_, def)| def.id == reference_id) {
-            Some(v) => v,
-            None => return py_schema_err!("Definitions error: no definition for ReferenceId `{}`", reference_id),
+            Entry::Vacant(entry) => entry.insert(Definition(Arc::new(OnceLock::from(value)))).clone(),
         };
-        match def.value.as_ref() {
-            Some(v) => Ok(v),
-            None => py_schema_err!(
-                "Definitions error: attempted to use `{}` before it was filled",
-                reference
-            ),
-        }
+        Ok(DefinitionRef { name, value })
     }
 
     /// Consume this Definitions into a vector of items, indexed by each items ReferenceId
-    pub fn finish(self) -> PyResult<Vec<T>> {
-        // We need to create a vec of defs according to the order in their ids
-        let mut defs: Vec<(usize, T)> = Vec::new();
-        for (reference, def) in self.definitions {
-            match def.value {
-                None => return py_schema_err!("Definitions error: definition {} was never filled", reference),
-                Some(v) => defs.push((def.id, v)),
+    pub fn finish(self) -> PyResult<Definitions<T>> {
+        for (reference, def) in &self.definitions.0 {
+            if def.0.get().is_none() {
+                return py_schema_err!("Definitions error: definition `{}` was never filled", reference);
             }
         }
-        defs.sort_by_key(|(id, _)| *id);
-        Ok(defs.into_iter().map(|(_, v)| v).collect())
+        Ok(self.definitions)
     }
 }

src/serializers/extra.rs

@@ -10,8 +10,6 @@ use serde::ser::Error;
 use super::config::SerializationConfig;
 use super::errors::{PydanticSerializationUnexpectedValue, UNEXPECTED_TYPE_SER_MARKER};
 use super::ob_type::ObTypeLookup;
-use super::shared::CombinedSerializer;
-use crate::definitions::Definitions;
 use crate::recursion_guard::RecursionGuard;
 
 /// this is ugly, would be much better if extra could be stored in `SerializationState`
@@ -48,7 +46,6 @@ impl SerializationState {
         Extra::new(
             py,
             mode,
-            &[],
             by_alias,
             &self.warnings,
             false,
@@ -72,7 +69,6 @@ impl SerializationState {
 #[cfg_attr(debug_assertions, derive(Debug))]
 pub(crate) struct Extra<'a> {
     pub mode: &'a SerMode,
-    pub definitions: &'a Definitions<CombinedSerializer>,
     pub ob_type_lookup: &'a ObTypeLookup,
     pub warnings: &'a CollectWarnings,
     pub by_alias: bool,
@@ -98,7 +94,6 @@ impl<'a> Extra<'a> {
     pub fn new(
         py: Python<'a>,
         mode: &'a SerMode,
-        definitions: &'a Definitions<CombinedSerializer>,
         by_alias: bool,
         warnings: &'a CollectWarnings,
         exclude_unset: bool,
@@ -112,7 +107,6 @@ impl<'a> Extra<'a> {
     ) -> Self {
         Self {
             mode,
-            definitions,
             ob_type_lookup: ObTypeLookup::cached(py),
             warnings,
             by_alias,
@@ -156,7 +150,6 @@ impl SerCheck {
 #[cfg_attr(debug_assertions, derive(Debug))]
 pub(crate) struct ExtraOwned {
     mode: SerMode,
-    definitions: Vec<CombinedSerializer>,
     warnings: CollectWarnings,
     by_alias: bool,
     exclude_unset: bool,
@@ -176,7 +169,6 @@ impl ExtraOwned {
     pub fn new(extra: &Extra) -> Self {
         Self {
             mode: extra.mode.clone(),
-            definitions: extra.definitions.to_vec(),
             warnings: extra.warnings.clone(),
             by_alias: extra.by_alias,
             exclude_unset: extra.exclude_unset,
@@ -196,7 +188,6 @@ impl ExtraOwned {
     pub fn to_extra<'py>(&'py self, py: Python<'py>) -> Extra<'py> {
         Extra {
             mode: &self.mode,
-            definitions: &self.definitions,
             ob_type_lookup: ObTypeLookup::cached(py),
             warnings: &self.warnings,
             by_alias: self.by_alias,

src/serializers/mod.rs

@@ -5,7 +5,7 @@ use pyo3::prelude::*;
 use pyo3::types::{PyBytes, PyDict};
 use pyo3::{PyTraverseError, PyVisit};
 
-use crate::definitions::DefinitionsBuilder;
+use crate::definitions::{Definitions, DefinitionsBuilder};
 use crate::py_gc::PyGcTraverse;
 
 use config::SerializationConfig;
@@ -30,7 +30,7 @@ mod type_serializers;
 #[derive(Debug)]
 pub struct SchemaSerializer {
     serializer: CombinedSerializer,
-    definitions: Vec<CombinedSerializer>,
+    definitions: Definitions<CombinedSerializer>,
     expected_json_size: AtomicUsize,
     config: SerializationConfig,
 }
@@ -54,7 +54,6 @@ impl SchemaSerializer {
         Extra::new(
             py,
             mode,
-            &self.definitions,
             by_alias,
             warnings,
             exclude_unset,
@@ -184,9 +183,7 @@ impl SchemaSerializer {
 
     fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
         self.serializer.py_gc_traverse(&visit)?;
-        for slot in &self.definitions {
-            slot.py_gc_traverse(&visit)?;
-        }
+        self.definitions.py_gc_traverse(&visit)?;
         Ok(())
     }
 }

src/serializers/shared.rs

@@ -13,7 +13,7 @@ use serde_json::ser::PrettyFormatter;
 
 use crate::build_tools::py_schema_err;
 use crate::build_tools::py_schema_error_type;
-use crate::definitions::{Definitions, DefinitionsBuilder};
+use crate::definitions::DefinitionsBuilder;
 use crate::py_gc::PyGcTraverse;
 use crate::tools::{py_err, SchemaDict};
 
@@ -293,7 +293,7 @@ pub(crate) trait TypeSerializer: Send + Sync + Clone + Debug {
     fn get_name(&self) -> &str;
 
     /// Used by union serializers to decide if it's worth trying again while allowing subclasses
-    fn retry_with_lax_check(&self, _definitions: &Definitions<CombinedSerializer>) -> bool {
+    fn retry_with_lax_check(&self) -> bool {
         false
     }
 

src/serializers/type_serializers/dataclass.rs

@@ -6,7 +6,7 @@ use std::borrow::Cow;
 use ahash::AHashMap;
 
 use crate::build_tools::{py_schema_error_type, ExtraBehavior};
-use crate::definitions::{Definitions, DefinitionsBuilder};
+use crate::definitions::DefinitionsBuilder;
 use crate::tools::SchemaDict;
 
 use super::{
@@ -179,7 +179,7 @@ impl TypeSerializer for DataclassSerializer {
         &self.name
     }
 
-    fn retry_with_lax_check(&self, _definitions: &Definitions<CombinedSerializer>) -> bool {
+    fn retry_with_lax_check(&self) -> bool {
         true
     }
 }