1use std::collections::{BTreeMap, BTreeSet};
4use std::ops::Range;
5
6#[derive(Debug, Clone, Eq, Hash, Ord, PartialEq, PartialOrd)]
8pub struct ExecutionGroupId(String);
9
10impl ExecutionGroupId {
11 pub fn new(id: impl Into<String>) -> Result<Self, ExecutionGraphError> {
13 let id = id.into();
14 if id.trim().is_empty() {
15 return Err(ExecutionGraphError::EmptyGroupId);
16 }
17 Ok(Self(id))
18 }
19
20 pub fn as_str(&self) -> &str {
22 &self.0
23 }
24}
25
26impl std::fmt::Display for ExecutionGroupId {
27 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28 formatter.write_str(&self.0)
29 }
30}
31
32#[derive(Debug, Clone, Eq, PartialEq)]
34pub struct ExecutionGroupSpec {
35 id: String,
36 dependencies: Vec<String>,
37}
38
39impl ExecutionGroupSpec {
40 pub fn root(id: impl Into<String>) -> Self {
42 Self {
43 id: id.into(),
44 dependencies: Vec::new(),
45 }
46 }
47
48 pub fn with_dependencies(
50 id: impl Into<String>,
51 dependencies: impl IntoIterator<Item = impl Into<String>>,
52 ) -> Self {
53 Self {
54 id: id.into(),
55 dependencies: dependencies.into_iter().map(Into::into).collect(),
56 }
57 }
58
59 pub fn id(&self) -> &str {
61 &self.id
62 }
63
64 pub fn dependencies(&self) -> &[String] {
66 &self.dependencies
67 }
68}
69
70#[derive(Debug, Clone, Eq, PartialEq)]
72pub struct ExecutionGraph {
73 groups: Vec<ExecutionGroupSpec>,
74 dependencies: Vec<Vec<usize>>,
75 dependents: Vec<Vec<usize>>,
76 execution_order: Vec<usize>,
77 output: usize,
78}
79
80#[derive(Debug, Clone, Copy, Eq, PartialEq)]
82pub struct ExecutionUnitAddress {
83 group: usize,
84 index: usize,
85}
86
87impl ExecutionUnitAddress {
88 pub const fn group(self) -> usize {
90 self.group
91 }
92
93 pub const fn index(self) -> usize {
95 self.index
96 }
97
98 pub const fn with_index(self, index: usize) -> Self {
100 Self {
101 group: self.group,
102 index,
103 }
104 }
105}
106
107#[derive(Debug, Clone, Eq, PartialEq)]
109pub struct ExecutionUnitLayout {
110 group_ids: Vec<ExecutionGroupId>,
111 group_ranges: Vec<Range<usize>>,
112 addresses: Vec<ExecutionUnitAddress>,
113}
114
115impl ExecutionUnitLayout {
116 pub fn new(
118 graph: &ExecutionGraph,
119 group_unit_counts: impl IntoIterator<Item = usize>,
120 ) -> Result<Self, ExecutionUnitLayoutError> {
121 let counts = group_unit_counts.into_iter().collect::<Vec<_>>();
122 if counts.len() != graph.groups().len() {
123 return Err(ExecutionUnitLayoutError::GroupCountMismatch {
124 graph_groups: graph.groups().len(),
125 declared_groups: counts.len(),
126 });
127 }
128 let group_ids = graph
129 .groups()
130 .iter()
131 .map(|group| {
132 ExecutionGroupId::new(group.id().to_owned())
133 .expect("validated execution graph has non-empty group identifiers")
134 })
135 .collect();
136 let mut group_ranges = Vec::with_capacity(counts.len());
137 let mut addresses = Vec::new();
138 for (group, count) in counts.into_iter().enumerate() {
139 let start = addresses.len();
140 let end = start
141 .checked_add(count)
142 .ok_or(ExecutionUnitLayoutError::UnitCountOverflow)?;
143 addresses.reserve(count);
144 addresses.extend((0..count).map(|index| ExecutionUnitAddress { group, index }));
145 group_ranges.push(start..end);
146 }
147 Ok(Self {
148 group_ids,
149 group_ranges,
150 addresses,
151 })
152 }
153
154 pub fn len(&self) -> usize {
156 self.addresses.len()
157 }
158
159 pub fn is_empty(&self) -> bool {
161 self.addresses.is_empty()
162 }
163
164 pub fn group_count(&self) -> usize {
166 self.group_ranges.len()
167 }
168
169 pub fn group_id(&self, group: usize) -> Option<&ExecutionGroupId> {
171 self.group_ids.get(group)
172 }
173
174 pub fn group_range(&self, group: usize) -> Option<Range<usize>> {
176 self.group_ranges.get(group).cloned()
177 }
178
179 pub fn address(&self, ordinal: usize) -> Option<ExecutionUnitAddress> {
181 self.addresses.get(ordinal).copied()
182 }
183
184 pub fn ordinal(&self, group: usize, index: usize) -> Option<usize> {
186 let range = self.group_ranges.get(group)?;
187 (index < range.len()).then_some(range.start + index)
188 }
189}
190
191#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
193pub enum ExecutionUnitLayoutError {
194 #[error(
196 "execution graph contains {graph_groups} groups but the architecture declared {declared_groups} group counts"
197 )]
198 GroupCountMismatch {
199 graph_groups: usize,
201 declared_groups: usize,
203 },
204 #[error("execution-unit count overflowed usize")]
206 UnitCountOverflow,
207}
208
209impl ExecutionGraph {
210 pub fn new(
212 groups: Vec<ExecutionGroupSpec>,
213 output: impl AsRef<str>,
214 ) -> Result<Self, ExecutionGraphError> {
215 if groups.is_empty() {
216 return Err(ExecutionGraphError::EmptyGraph);
217 }
218 let mut by_id = BTreeMap::new();
219 for (index, group) in groups.iter().enumerate() {
220 if group.id.trim().is_empty() {
221 return Err(ExecutionGraphError::EmptyGroupId);
222 }
223 if by_id.insert(group.id.clone(), index).is_some() {
224 return Err(ExecutionGraphError::DuplicateGroup(group.id.clone()));
225 }
226 }
227 let output_name = output.as_ref();
228 let output = by_id
229 .get(output_name)
230 .copied()
231 .ok_or_else(|| ExecutionGraphError::UnknownOutput(output_name.to_owned()))?;
232 let mut dependencies = Vec::with_capacity(groups.len());
233 let mut dependents = vec![Vec::new(); groups.len()];
234 let mut indegree = vec![0usize; groups.len()];
235 for (index, group) in groups.iter().enumerate() {
236 let mut seen = BTreeSet::new();
237 let mut resolved = Vec::with_capacity(group.dependencies.len());
238 for dependency in &group.dependencies {
239 let dependency_index = by_id.get(dependency).copied().ok_or_else(|| {
240 ExecutionGraphError::UnknownDependency {
241 group: group.id.clone(),
242 dependency: dependency.clone(),
243 }
244 })?;
245 if dependency_index == index {
246 return Err(ExecutionGraphError::SelfDependency(group.id.clone()));
247 }
248 if !seen.insert(dependency_index) {
249 return Err(ExecutionGraphError::DuplicateDependency {
250 group: group.id.clone(),
251 dependency: dependency.clone(),
252 });
253 }
254 resolved.push(dependency_index);
255 dependents[dependency_index].push(index);
256 }
257 indegree[index] = resolved.len();
258 dependencies.push(resolved);
259 }
260 let mut ready = indegree
261 .iter()
262 .enumerate()
263 .filter_map(|(index, °ree)| (degree == 0).then_some(index))
264 .collect::<BTreeSet<_>>();
265 let mut execution_order = Vec::with_capacity(groups.len());
266 while let Some(index) = ready.pop_first() {
267 execution_order.push(index);
268 for &dependent in &dependents[index] {
269 indegree[dependent] -= 1;
270 if indegree[dependent] == 0 {
271 ready.insert(dependent);
272 }
273 }
274 }
275 if execution_order.len() != groups.len() {
276 return Err(ExecutionGraphError::Cycle);
277 }
278 let mut contributes = BTreeSet::new();
279 let mut pending = vec![output];
280 while let Some(index) = pending.pop() {
281 if contributes.insert(index) {
282 pending.extend(dependencies[index].iter().copied());
283 }
284 }
285 if contributes.len() != groups.len() {
286 let disconnected = groups
287 .iter()
288 .enumerate()
289 .filter_map(|(index, group)| {
290 (!contributes.contains(&index)).then_some(group.id.clone())
291 })
292 .collect();
293 return Err(ExecutionGraphError::Disconnected { disconnected });
294 }
295 Ok(Self {
296 groups,
297 dependencies,
298 dependents,
299 execution_order,
300 output,
301 })
302 }
303
304 pub fn chain(
306 ids: impl IntoIterator<Item = impl Into<String>>,
307 ) -> Result<Self, ExecutionGraphError> {
308 let ids = ids.into_iter().map(Into::into).collect::<Vec<String>>();
309 let output = ids.last().cloned().ok_or(ExecutionGraphError::EmptyGraph)?;
310 let groups = ids
311 .iter()
312 .enumerate()
313 .map(|(index, id)| match index.checked_sub(1) {
314 Some(previous) => Self::group_with_dependency(id.clone(), ids[previous].clone()),
315 None => ExecutionGroupSpec::root(id.clone()),
316 })
317 .collect();
318 Self::new(groups, output)
319 }
320
321 fn group_with_dependency(id: String, dependency: String) -> ExecutionGroupSpec {
322 ExecutionGroupSpec::with_dependencies(id, [dependency])
323 }
324
325 pub fn groups(&self) -> &[ExecutionGroupSpec] {
327 &self.groups
328 }
329
330 pub fn group_index(&self, id: &str) -> Option<usize> {
332 self.groups.iter().position(|group| group.id() == id)
333 }
334
335 pub fn execution_order(&self) -> &[usize] {
337 &self.execution_order
338 }
339
340 pub fn dependencies(&self, group: usize) -> Option<&[usize]> {
342 self.dependencies.get(group).map(Vec::as_slice)
343 }
344
345 pub fn dependents(&self, group: usize) -> Option<&[usize]> {
347 self.dependents.get(group).map(Vec::as_slice)
348 }
349
350 pub const fn output(&self) -> usize {
352 self.output
353 }
354
355 pub fn consumer_counts(&self) -> Vec<usize> {
357 let mut counts = vec![0; self.groups.len()];
358 for dependencies in &self.dependencies {
359 for &dependency in dependencies {
360 counts[dependency] += 1;
361 }
362 }
363 counts
364 }
365}
366
367#[derive(Debug, Clone, Copy, Eq, PartialEq)]
369pub enum ReadyGroupState {
370 Pending,
372 Ordered,
374 Failed,
376 Blocked,
378}
379
380#[derive(Debug)]
381struct ExecutionGroupReadySet<'a> {
382 graph: &'a ExecutionGraph,
383 remaining_dependencies: Vec<usize>,
384 states: Vec<ReadyGroupState>,
385 ready: BTreeSet<usize>,
386}
387
388#[derive(Debug)]
390pub struct ExecutionGroupSchedule<'a> {
391 graph: &'a ExecutionGraph,
392 ready: ExecutionGroupReadySet<'a>,
393 started: Vec<bool>,
394 remaining_consumers: Vec<usize>,
395}
396
397impl<'a> ExecutionGroupSchedule<'a> {
398 pub fn new(graph: &'a ExecutionGraph) -> Self {
400 Self {
401 graph,
402 ready: ExecutionGroupReadySet::new(graph),
403 started: vec![false; graph.groups.len()],
404 remaining_consumers: graph.consumer_counts(),
405 }
406 }
407
408 pub fn startable_groups(&self) -> impl Iterator<Item = usize> + '_ {
410 self.ready
411 .ready_groups()
412 .filter(|&group| !self.started[group])
413 }
414
415 pub fn dependencies(&self, group: usize) -> Result<&[usize], ExecutionScheduleError> {
417 self.graph
418 .dependencies(group)
419 .ok_or(ExecutionScheduleError::UnknownGroup {
420 group,
421 count: self.started.len(),
422 })
423 }
424
425 pub fn started(&mut self, group: usize) -> Result<Vec<usize>, ExecutionScheduleError> {
428 let count = self.started.len();
429 let started = self
430 .started
431 .get_mut(group)
432 .ok_or(ExecutionScheduleError::UnknownGroup { group, count })?;
433 if *started {
434 return Err(ExecutionScheduleError::AlreadyStarted { group });
435 }
436 if !self.ready.ready.contains(&group) {
437 return Err(ExecutionScheduleError::DependenciesPending { group });
438 }
439 *started = true;
440 let mut releasable = Vec::new();
441 for &dependency in &self.graph.dependencies[group] {
442 self.remaining_consumers[dependency] -= 1;
443 if self.remaining_consumers[dependency] == 0 {
444 releasable.push(dependency);
445 }
446 }
447 Ok(releasable)
448 }
449
450 pub fn ordered(&mut self, group: usize) -> Result<(), ExecutionScheduleError> {
452 match self.started.get(group).copied() {
453 None => Err(ExecutionScheduleError::UnknownGroup {
454 group,
455 count: self.started.len(),
456 }),
457 Some(false) => Err(ExecutionScheduleError::NotStarted { group }),
458 Some(true) if self.ready.state(group) == Some(ReadyGroupState::Pending) => {
459 self.ready.ordered(group);
460 Ok(())
461 }
462 Some(true) => Err(ExecutionScheduleError::AlreadyOrdered { group }),
463 }
464 }
465
466 pub fn fail(&mut self, group: usize) -> Result<(), ExecutionScheduleError> {
468 if group >= self.started.len() {
469 return Err(ExecutionScheduleError::UnknownGroup {
470 group,
471 count: self.started.len(),
472 });
473 }
474 self.ready.fail(group);
475 Ok(())
476 }
477
478 pub fn state(&self, group: usize) -> Option<ReadyGroupState> {
480 self.ready.state(group)
481 }
482}
483
484#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
486pub enum ExecutionScheduleError {
487 #[error("execution group {group} is outside the {count}-group schedule")]
489 UnknownGroup {
490 group: usize,
492 count: usize,
494 },
495 #[error("execution group {group} was already started")]
497 AlreadyStarted {
498 group: usize,
500 },
501 #[error("execution group {group} still has unordered dependencies")]
503 DependenciesPending {
504 group: usize,
506 },
507 #[error("execution group {group} was ordered before it started")]
509 NotStarted {
510 group: usize,
512 },
513 #[error("execution group {group} was already ordered or closed")]
515 AlreadyOrdered {
516 group: usize,
518 },
519}
520
521impl<'a> ExecutionGroupReadySet<'a> {
522 fn new(graph: &'a ExecutionGraph) -> Self {
523 let remaining_dependencies = graph.dependencies.iter().map(Vec::len).collect::<Vec<_>>();
524 let ready = remaining_dependencies
525 .iter()
526 .enumerate()
527 .filter_map(|(group, &remaining)| (remaining == 0).then_some(group))
528 .collect();
529 Self {
530 graph,
531 remaining_dependencies,
532 states: vec![ReadyGroupState::Pending; graph.groups.len()],
533 ready,
534 }
535 }
536
537 fn ready_groups(&self) -> impl Iterator<Item = usize> + '_ {
538 self.ready.iter().copied()
539 }
540
541 fn ordered(&mut self, group: usize) {
542 debug_assert_eq!(self.states[group], ReadyGroupState::Pending);
543 self.ready.remove(&group);
544 self.states[group] = ReadyGroupState::Ordered;
545 for &dependent in &self.graph.dependents[group] {
546 if self.states[dependent] != ReadyGroupState::Pending {
547 continue;
548 }
549 self.remaining_dependencies[dependent] -= 1;
550 if self.remaining_dependencies[dependent] == 0 {
551 self.ready.insert(dependent);
552 }
553 }
554 }
555
556 fn fail(&mut self, group: usize) {
557 self.close_subgraph(group, ReadyGroupState::Failed);
558 }
559
560 fn close_subgraph(&mut self, group: usize, state: ReadyGroupState) {
561 let mut pending = vec![(group, state)];
562 while let Some((group, state)) = pending.pop() {
563 if self.states[group] != ReadyGroupState::Pending {
564 continue;
565 }
566 self.ready.remove(&group);
567 self.states[group] = state;
568 pending.extend(
569 self.graph.dependents[group]
570 .iter()
571 .copied()
572 .map(|dependent| (dependent, ReadyGroupState::Blocked)),
573 );
574 }
575 }
576
577 fn state(&self, group: usize) -> Option<ReadyGroupState> {
578 self.states.get(group).copied()
579 }
580}
581
582#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
584pub enum ExecutionGraphError {
585 #[error("execution-group graph must contain at least one group")]
587 EmptyGraph,
588 #[error("execution-group identifiers must not be empty")]
590 EmptyGroupId,
591 #[error("duplicate execution-group identifier {0:?}")]
593 DuplicateGroup(String),
594 #[error("execution-group graph output {0:?} does not exist")]
596 UnknownOutput(String),
597 #[error("execution group {group:?} depends on unknown group {dependency:?}")]
599 UnknownDependency {
600 group: String,
602 dependency: String,
604 },
605 #[error("execution group {0:?} cannot depend on itself")]
607 SelfDependency(String),
608 #[error("execution group {group:?} repeats dependency {dependency:?}")]
610 DuplicateDependency {
611 group: String,
613 dependency: String,
615 },
616 #[error("execution-group graph contains a dependency cycle")]
618 Cycle,
619 #[error("execution groups do not contribute to the graph output: {disconnected:?}")]
621 Disconnected {
622 disconnected: Vec<String>,
624 },
625}
626
627#[cfg(test)]
628mod tests {
629 use super::*;
630
631 #[test]
632 fn graph_order_is_stable_and_dependency_driven() {
633 let graph = ExecutionGraph::new(
634 vec![
635 ExecutionGroupSpec::root("image"),
636 ExecutionGroupSpec::root("audio"),
637 ExecutionGroupSpec::with_dependencies("text", ["image", "audio"]),
638 ],
639 "text",
640 )
641 .unwrap();
642 assert_eq!(graph.execution_order(), &[0, 1, 2]);
643 assert_eq!(graph.dependencies(2), Some([0, 1].as_slice()));
644
645 let mut ready = ExecutionGroupReadySet::new(&graph);
646 assert_eq!(ready.ready_groups().collect::<Vec<_>>(), vec![0, 1]);
647 ready.ordered(1);
648 assert_eq!(ready.ready_groups().collect::<Vec<_>>(), vec![0]);
649 ready.ordered(0);
650 assert_eq!(ready.ready_groups().collect::<Vec<_>>(), vec![2]);
651 }
652
653 #[test]
654 fn schedule_releases_dependency_outputs_after_their_final_consumer_starts() {
655 let graph = ExecutionGraph::new(
656 vec![
657 ExecutionGroupSpec::root("root"),
658 ExecutionGroupSpec::with_dependencies("left", ["root"]),
659 ExecutionGroupSpec::with_dependencies("right", ["root"]),
660 ExecutionGroupSpec::with_dependencies("output", ["left", "right"]),
661 ],
662 "output",
663 )
664 .unwrap();
665 let mut schedule = ExecutionGroupSchedule::new(&graph);
666 assert_eq!(schedule.startable_groups().collect::<Vec<_>>(), vec![0]);
667 assert!(schedule.started(1).is_err());
668 assert!(schedule.started(0).unwrap().is_empty());
669 schedule.ordered(0).unwrap();
670 assert_eq!(schedule.startable_groups().collect::<Vec<_>>(), vec![1, 2]);
671 assert!(schedule.started(1).unwrap().is_empty());
672 assert_eq!(schedule.started(2).unwrap(), vec![0]);
673 schedule.ordered(1).unwrap();
674 schedule.ordered(2).unwrap();
675 assert_eq!(schedule.started(3).unwrap(), vec![1, 2]);
676 assert!(schedule.ordered(3).is_ok());
677 assert_eq!(schedule.state(3), Some(ReadyGroupState::Ordered));
678 }
679
680 #[test]
681 fn execution_unit_layout_preserves_group_major_residency_order() {
682 let graph = ExecutionGraph::new(
683 vec![
684 ExecutionGroupSpec::root("vision"),
685 ExecutionGroupSpec::with_dependencies("text", ["vision"]),
686 ],
687 "text",
688 )
689 .unwrap();
690 let layout = ExecutionUnitLayout::new(&graph, [2, 3]).unwrap();
691
692 assert_eq!(graph.group_index("vision"), Some(0));
693 assert_eq!(graph.group_index("text"), Some(1));
694 assert_eq!(graph.group_index("missing"), None);
695 assert_eq!(layout.len(), 5);
696 assert_eq!(layout.group_count(), 2);
697 assert_eq!(layout.group_id(0).unwrap().as_str(), "vision");
698 assert_eq!(layout.group_id(1).unwrap().as_str(), "text");
699 assert_eq!(layout.group_range(0), Some(0..2));
700 assert_eq!(layout.group_range(1), Some(2..5));
701 assert_eq!(layout.address(3).unwrap().group(), 1);
702 assert_eq!(layout.address(3).unwrap().index(), 1);
703 assert_eq!(layout.ordinal(1, 2), Some(4));
704 assert_eq!(layout.ordinal(0, 2), None);
705 }
706
707 #[test]
708 fn execution_unit_layout_rejects_graph_count_drift() {
709 let graph = ExecutionGraph::chain(["vision", "text"]).unwrap();
710 assert_eq!(
711 ExecutionUnitLayout::new(&graph, [2]).unwrap_err(),
712 ExecutionUnitLayoutError::GroupCountMismatch {
713 graph_groups: 2,
714 declared_groups: 1,
715 }
716 );
717 }
718
719 #[test]
720 fn invalid_graphs_fail_closed() {
721 let groups = vec![
722 ExecutionGroupSpec::with_dependencies("left", ["right"]),
723 ExecutionGroupSpec::with_dependencies("right", ["left"]),
724 ];
725 assert_eq!(
726 ExecutionGraph::new(groups, "right"),
727 Err(ExecutionGraphError::Cycle)
728 );
729 }
730}