diff --git a/src/deps/tests.rs b/src/deps/tests.rs index 8d02881..5ff8006 100644 --- a/src/deps/tests.rs +++ b/src/deps/tests.rs @@ -176,7 +176,7 @@ fn test_deps_tuple_large_tuples() { let node_ids = (&handles[0], &handles[1], &handles[2], &handles[3]).to_node_ids(); assert_eq!(node_ids.len(), 4); for (i, node_id) in node_ids.iter().enumerate().take(4) { - assert_eq!(*node_id, NodeId(i)); + assert_eq!(*node_id, NodeId(i as u32)); } // Test 8-tuple @@ -193,7 +193,7 @@ fn test_deps_tuple_large_tuples() { .to_node_ids(); assert_eq!(node_ids.len(), 8); for (i, node_id) in node_ids.iter().enumerate().take(8) { - assert_eq!(*node_id, NodeId(i)); + assert_eq!(*node_id, NodeId(i as u32)); } } diff --git a/src/error.rs b/src/error.rs index fa9674c..53fb983 100644 --- a/src/error.rs +++ b/src/error.rs @@ -6,19 +6,16 @@ #[derive(Debug, Clone, PartialEq)] pub enum DagError { /// Invalid dependency: task does not exist - InvalidDependency { task_id: usize }, + InvalidDependency { task_id: u32 }, /// Type mismatch in task dependencies TypeMismatch { expected: &'static str, found: &'static str, }, /// Task panicked during execution - TaskPanicked { - task_id: usize, - panic_message: String, - }, + TaskPanicked { task_id: u32, panic_message: String }, /// Result not found for task - ResultNotFound { task_id: usize }, + ResultNotFound { task_id: u32 }, /// Attempted to run the DAG while it's already running ConcurrentExecution, } diff --git a/src/runner.rs b/src/runner.rs index 59364da..b70f37b 100644 --- a/src/runner.rs +++ b/src/runner.rs @@ -8,6 +8,7 @@ use std::any::Any; use std::collections::{HashMap, HashSet, VecDeque}; use std::future::Future; +use std::hash::{BuildHasher, Hasher}; use std::panic::AssertUnwindSafe; use std::pin::Pin; use std::sync::{ @@ -40,6 +41,35 @@ impl<'a> Drop for RunGuard<'a> { } } +#[derive(Default, Clone)] +pub(crate) struct PassThroughHasher { + hash: u64, +} + +impl Hasher for PassThroughHasher { + fn finish(&self) -> u64 { + self.hash + } + + fn write_u32(&mut self, i: u32) { + self.hash = i as u64; + } + + fn write(&mut self, _bytes: &[u8]) { + panic!("PassThroughHasher used on invalid type"); + } +} + +impl BuildHasher for PassThroughHasher { + type Hasher = PassThroughHasher; + + fn build_hasher(&self) -> Self::Hasher { + PassThroughHasher::default() + } +} + +pub(crate) type PassThroughHashMap = HashMap; + /// Build and execute a typed DAG of tasks. /// /// A `DagRunner` is the main orchestrator for building and executing a directed acyclic graph @@ -100,10 +130,11 @@ impl<'a> Drop for RunGuard<'a> { /// Arc enables efficient sharing during fanout without cloning data. pub struct DagRunner { pub(crate) nodes: Mutex>>>, - pub(crate) outputs: Mutex>>, - pub(crate) edges: Mutex>>, // node -> dependencies - pub(crate) dependents: Mutex>>, // node -> tasks that depend on it - pub(crate) next_id: Mutex, + pub(crate) outputs: + Mutex>>, + pub(crate) edges: Mutex>>, // node -> dependencies + pub(crate) dependents: Mutex>>, // node -> tasks that depend on it + pub(crate) next_id: Mutex, pub(crate) run_lock: AtomicBool, // Ensures only one run() at a time } @@ -126,9 +157,9 @@ impl DagRunner { pub fn new() -> Self { Self { nodes: Mutex::new(Vec::new()), - outputs: Mutex::new(HashMap::new()), - edges: Mutex::new(HashMap::new()), - dependents: Mutex::new(HashMap::new()), + outputs: Mutex::new(HashMap::default()), + edges: Mutex::new(HashMap::default()), + dependents: Mutex::new(HashMap::default()), next_id: Mutex::new(0), run_lock: AtomicBool::new(false), } @@ -314,16 +345,18 @@ impl DagRunner { // Build topological layers let layers = self.compute_layers()?; + let total_tasks = layers.iter().map(|l| l.len()).sum::(); + #[cfg(feature = "tracing")] debug!( layer_count = layers.len(), - total_tasks = layers.iter().map(|l| l.len()).sum::(), - "computed topological layers" + total_tasks, "computed topological layers" ); let edges = self.edges.lock().clone(); - let mut outputs: HashMap> = HashMap::new(); + let mut outputs: PassThroughHashMap> = + HashMap::with_capacity_and_hasher(total_tasks, PassThroughHasher::default()); let mut first_error = None; for layer in layers { @@ -368,7 +401,7 @@ impl DagRunner { // Take ownership of the node let node = { let mut nodes_lock = self.nodes.lock(); - nodes_lock[node_id.0].take() + nodes_lock[node_id.0 as usize].take() }; if let Some(node) = node { @@ -432,7 +465,7 @@ impl DagRunner { // Take ownership of the node let node = { let mut nodes_lock = self.nodes.lock(); - nodes_lock[node_id.0].take() + nodes_lock[node_id.0 as usize].take() }; if let Some(node) = node { let dependencies: Vec<_> = edges[&node_id] @@ -586,7 +619,7 @@ impl DagRunner { #[cfg(feature = "tracing")] debug!("computing topological layers"); - let mut in_degree: HashMap = HashMap::new(); + let mut in_degree: PassThroughHashMap = HashMap::default(); let mut layers = Vec::new(); // Calculate in-degrees: for each node, count how many dependencies it has diff --git a/src/types.rs b/src/types.rs index 1b2feea..2ebf2ae 100644 --- a/src/types.rs +++ b/src/types.rs @@ -8,7 +8,7 @@ use crate::task::Task; /// Opaque node identifier #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] -pub struct NodeId(pub usize); +pub struct NodeId(pub u32); /// Opaque, typed token for a node's output. ///