1use std::collections::{BTreeMap, BTreeSet, VecDeque};
4
5#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
7pub enum ComponentDomain {
8 TokenIds,
10 PatchMatrix,
12 AudioFeatures,
14 HiddenStates,
16 Logits,
18 TargetStates,
20}
21
22#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
24pub enum ComponentKind {
25 StaticText,
27 Vision,
29 Audio,
31 Assembly,
33 Decoder,
35 Prediction,
37 Assistant,
39 OutputProjection,
41}
42
43#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
45pub enum ComponentResidencyClass {
46 Static,
48 Media,
50 Decoder,
52 Experts,
54 Draft,
56}
57
58#[derive(Debug, Clone, Eq, PartialEq)]
60pub struct ComponentSpec {
61 pub id: String,
63 pub kind: ComponentKind,
65 pub external_inputs: Vec<ComponentDomain>,
67 pub dependencies: Vec<String>,
69 pub dependency_inputs: Vec<ComponentDomain>,
71 pub output: ComponentDomain,
73 pub residency: ComponentResidencyClass,
75}
76
77#[derive(Debug, Clone, Eq, PartialEq)]
79pub struct ComponentGraph {
80 units: Vec<ComponentSpec>,
81 execution_order: Vec<usize>,
82 outputs: Vec<usize>,
83}
84
85impl ComponentGraph {
86 pub fn new(
88 units: Vec<ComponentSpec>,
89 outputs: impl IntoIterator<Item = impl AsRef<str>>,
90 ) -> Result<Self, ComponentGraphError> {
91 if units.is_empty() {
92 return Err(ComponentGraphError::Empty);
93 }
94 let mut by_id = BTreeMap::new();
95 for (index, unit) in units.iter().enumerate() {
96 if unit.id.trim().is_empty() {
97 return Err(ComponentGraphError::EmptyId);
98 }
99 if unit.dependencies.len() != unit.dependency_inputs.len() {
100 return Err(ComponentGraphError::DependencyArity {
101 unit: unit.id.clone(),
102 dependencies: unit.dependencies.len(),
103 domains: unit.dependency_inputs.len(),
104 });
105 }
106 if by_id.insert(unit.id.clone(), index).is_some() {
107 return Err(ComponentGraphError::DuplicateId(unit.id.clone()));
108 }
109 }
110 let mut edges = vec![Vec::new(); units.len()];
111 let mut indegree = vec![0_usize; units.len()];
112 for (unit_index, unit) in units.iter().enumerate() {
113 let mut seen = BTreeSet::new();
114 for (slot, dependency) in unit.dependencies.iter().enumerate() {
115 let dependency_index = by_id.get(dependency).copied().ok_or_else(|| {
116 ComponentGraphError::UnknownDependency {
117 unit: unit.id.clone(),
118 dependency: dependency.clone(),
119 }
120 })?;
121 if !seen.insert(dependency_index) {
122 return Err(ComponentGraphError::DuplicateDependency {
123 unit: unit.id.clone(),
124 dependency: dependency.clone(),
125 });
126 }
127 let actual = units[dependency_index].output;
128 let expected = unit.dependency_inputs[slot];
129 if actual != expected {
130 return Err(ComponentGraphError::DomainMismatch {
131 unit: unit.id.clone(),
132 dependency: dependency.clone(),
133 expected,
134 actual,
135 });
136 }
137 edges[dependency_index].push(unit_index);
138 indegree[unit_index] += 1;
139 }
140 }
141 let mut ready = indegree
142 .iter()
143 .enumerate()
144 .filter_map(|(index, degree)| (*degree == 0).then_some(index))
145 .collect::<VecDeque<_>>();
146 let mut execution_order = Vec::with_capacity(units.len());
147 while let Some(index) = ready.pop_front() {
148 execution_order.push(index);
149 for dependent in &edges[index] {
150 indegree[*dependent] -= 1;
151 if indegree[*dependent] == 0 {
152 ready.push_back(*dependent);
153 }
154 }
155 }
156 if execution_order.len() != units.len() {
157 return Err(ComponentGraphError::Cycle);
158 }
159 let outputs = outputs
160 .into_iter()
161 .map(|output| {
162 let output = output.as_ref();
163 by_id
164 .get(output)
165 .copied()
166 .ok_or_else(|| ComponentGraphError::UnknownOutput(output.to_owned()))
167 })
168 .collect::<Result<Vec<_>, _>>()?;
169 if outputs.is_empty() {
170 return Err(ComponentGraphError::NoOutputs);
171 }
172 if outputs.iter().copied().collect::<BTreeSet<_>>().len() != outputs.len() {
173 return Err(ComponentGraphError::DuplicateOutput);
174 }
175 Ok(Self {
176 units,
177 execution_order,
178 outputs,
179 })
180 }
181
182 pub fn units(&self) -> &[ComponentSpec] {
184 &self.units
185 }
186
187 pub fn execution_order(&self) -> impl Iterator<Item = &ComponentSpec> {
189 self.execution_order.iter().map(|index| &self.units[*index])
190 }
191
192 pub fn outputs(&self) -> impl Iterator<Item = &ComponentSpec> {
194 self.outputs.iter().map(|index| &self.units[*index])
195 }
196
197 pub fn residency_count(&self, class: ComponentResidencyClass) -> usize {
199 self.units
200 .iter()
201 .filter(|unit| unit.residency == class)
202 .count()
203 }
204}
205
206#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
208pub enum ComponentGraphError {
209 #[error("component graph cannot be empty")]
211 Empty,
212 #[error("component identity cannot be empty")]
214 EmptyId,
215 #[error("component {0:?} was declared more than once")]
217 DuplicateId(String),
218 #[error("component {unit:?} has {dependencies} dependencies but {domains} dependency domains")]
220 DependencyArity {
221 unit: String,
223 dependencies: usize,
225 domains: usize,
227 },
228 #[error("component {unit:?} depends on unknown component {dependency:?}")]
230 UnknownDependency {
231 unit: String,
233 dependency: String,
235 },
236 #[error("component {unit:?} repeats dependency {dependency:?}")]
238 DuplicateDependency {
239 unit: String,
241 dependency: String,
243 },
244 #[error("component {unit:?} expects {expected:?} from {dependency:?}, got {actual:?}")]
246 DomainMismatch {
247 unit: String,
249 dependency: String,
251 expected: ComponentDomain,
253 actual: ComponentDomain,
255 },
256 #[error("component graph contains a dependency cycle")]
258 Cycle,
259 #[error("component graph requires at least one output")]
261 NoOutputs,
262 #[error("component graph repeats an output")]
264 DuplicateOutput,
265 #[error("unknown component output {0:?}")]
267 UnknownOutput(String),
268}
269
270#[cfg(test)]
271mod tests {
272 use super::*;
273
274 fn unit(id: &str, output: ComponentDomain) -> ComponentSpec {
275 ComponentSpec {
276 id: id.into(),
277 kind: ComponentKind::StaticText,
278 external_inputs: vec![],
279 dependencies: vec![],
280 dependency_inputs: vec![],
281 output,
282 residency: ComponentResidencyClass::Static,
283 }
284 }
285
286 #[test]
287 fn component_graph_validates_domains_and_preserves_outputs() {
288 let embedding = ComponentSpec {
289 external_inputs: vec![ComponentDomain::TokenIds],
290 ..unit("embedding", ComponentDomain::HiddenStates)
291 };
292 let decoder = ComponentSpec {
293 id: "decoder.0".into(),
294 kind: ComponentKind::Decoder,
295 dependencies: vec!["embedding".into()],
296 dependency_inputs: vec![ComponentDomain::HiddenStates],
297 output: ComponentDomain::HiddenStates,
298 residency: ComponentResidencyClass::Decoder,
299 external_inputs: vec![],
300 };
301 let graph = ComponentGraph::new(vec![decoder, embedding], ["decoder.0"]).unwrap();
302 assert_eq!(
303 graph
304 .execution_order()
305 .map(|unit| unit.id.as_str())
306 .collect::<Vec<_>>(),
307 ["embedding", "decoder.0"]
308 );
309 assert_eq!(graph.residency_count(ComponentResidencyClass::Decoder), 1);
310 }
311
312 #[test]
313 fn component_graph_rejects_domain_mismatch_and_cycles() {
314 let mut decoder = unit("decoder", ComponentDomain::HiddenStates);
315 decoder.dependencies = vec!["tokens".into()];
316 decoder.dependency_inputs = vec![ComponentDomain::HiddenStates];
317 let tokens = unit("tokens", ComponentDomain::TokenIds);
318 assert!(matches!(
319 ComponentGraph::new(vec![tokens, decoder], ["decoder"]),
320 Err(ComponentGraphError::DomainMismatch { .. })
321 ));
322 let mut left = unit("left", ComponentDomain::HiddenStates);
323 left.dependencies = vec!["right".into()];
324 left.dependency_inputs = vec![ComponentDomain::HiddenStates];
325 let mut right = unit("right", ComponentDomain::HiddenStates);
326 right.dependencies = vec!["left".into()];
327 right.dependency_inputs = vec![ComponentDomain::HiddenStates];
328 assert_eq!(
329 ComponentGraph::new(vec![left, right], ["left"]).unwrap_err(),
330 ComponentGraphError::Cycle
331 );
332 }
333}