Rework def ref system to fix various bugs

Author: adriangb

src/build_context.rs

@@ -0,0 +1,99 @@
+use std::collections::hash_map::Entry;
+
+use pyo3::prelude::*;
+
+use ahash::AHashMap;
+
+use crate::build_tools::py_err;
+
+// An integer id for the reference
+type ReferenceId = usize;
+
+#[derive(Clone, Debug)]
+struct Definition<T> {
+    pub id: ReferenceId,
+    pub value: Option<T>,
+}
+
+#[derive(Clone, Debug)]
+pub struct DefinitionsBuilder<T> {
+    definitions: AHashMap<String, Definition<T>>,
+}
+
+impl<T: Clone + std::fmt::Debug> DefinitionsBuilder<T> {
+    pub fn new() -> Self {
+        Self {
+            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();
+        // 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
+            }
+        }
+    }
+
+    // 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_err!(format!("Duplicate ref: `{reference}`")),
+                None => Ok(entry.get().id),
+            },
+            Entry::Vacant(entry) => {
+                entry.insert(Definition {
+                    id: next_id,
+                    value: Some(value),
+                });
+                Ok(next_id)
+            }
+        }
+    }
+
+    // Retrieve an item definition using a ReferenceId
+    // Will raise an error if the definition for that reference does not yet exist
+    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_err!(format!(
+                    "Definitions error: no definition for ReferenceId `{reference_id}`"
+                ))
+            }
+        };
+        match def.value.as_ref() {
+            Some(v) => Ok(v),
+            None => py_err!(format!(
+                "Definitions error: attempted to use `{reference}` before it was filled"
+            )),
+        }
+    }
+
+    // Consume this Definitions into a vector of items, indexed by each items ReferenceId
+    pub fn build(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.into_iter() {
+            match def.value {
+                None => return py_err!(format!("Definitions error: definition {reference} was never filled")),
+                Some(v) => defs.push((def.id, v)),
+            }
+        }
+        defs.sort_by_key(|(id, _)| *id);
+        Ok(defs.into_iter().map(|(_, v)| v).collect())
+    }
+}

src/definitions.rs

@@ -10,8 +10,8 @@ use pyo3::prelude::*;
 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
 
 mod argument_markers;
-mod build_context;
 mod build_tools;
+mod definitions;
 mod errors;
 mod input;
 mod lazy_index_map;

src/lib.rs

@@ -1,10 +1,19 @@
 use ahash::AHashSet;
 
+type RecursionKey = (
+    // Identifier for the input object, e.g. the id() of a Python dict
+    usize,
+    // Identifier for the node we are traversing, e.g. the validator's id
+    // Generally only things that can be traversed multiple times, like a definition reference
+    // need to use the recursion guard, and those things should already have a natural node id
+    usize,
+);
+
 /// This is used to avoid cyclic references in input data causing recursive validation and a nasty segmentation fault.
 /// It's used in `validators/definition` to detect when a reference is reused within itself.
 #[derive(Debug, Clone, Default)]
 pub struct RecursionGuard {
-    ids: Option<AHashSet<usize>>,
+    ids: Option<AHashSet<RecursionKey>>,
     // see validators/definition::BACKUP_GUARD_LIMIT for details
     // depth could be a hashmap {validator_id => depth} but for simplicity and performance it's easier to just
     // use one number for all validators
@@ -13,14 +22,14 @@ pub struct RecursionGuard {
 
 impl RecursionGuard {
     // insert a new id into the set, return whether the set already had the id in it
-    pub fn contains_or_insert(&mut self, id: usize) -> bool {
+    pub fn contains_or_insert(&mut self, obj_id: usize, node_id: usize) -> bool {
         match self.ids {
             // https://doc.rust-lang.org/std/collections/struct.HashSet.html#method.insert
             // "If the set did not have this value present, `true` is returned."
-            Some(ref mut set) => !set.insert(id),
+            Some(ref mut set) => !set.insert((obj_id, node_id)),
             None => {
-                let mut set: AHashSet<usize> = AHashSet::with_capacity(10);
-                set.insert(id);
+                let mut set: AHashSet<RecursionKey> = AHashSet::with_capacity(10);
+                set.insert((obj_id, node_id));
                 self.ids = Some(set);
                 false
             }
@@ -37,10 +46,10 @@ impl RecursionGuard {
         self.depth -= 1;
     }
 
-    pub fn remove(&mut self, id: &usize) {
+    pub fn remove(&mut self, obj_id: usize, node_id: usize) {
         match self.ids {
             Some(ref mut set) => {
-                set.remove(id);
+                set.remove(&(obj_id, node_id));
             }
             None => unreachable!(),
         };