Another project
1use std::collections::BTreeMap;
2
3use bone_types::{ParameterIndex, ParentIndex, ResidualIndex};
4
5use crate::system::ConstraintSystem;
6
7#[derive(Clone, Debug, PartialEq, Eq)]
8pub struct Component {
9 parameters: Vec<ParameterIndex>,
10 residual_parents: Vec<ParentIndex>,
11}
12
13impl Component {
14 #[must_use]
15 pub fn parameters(&self) -> &[ParameterIndex] {
16 &self.parameters
17 }
18
19 #[must_use]
20 pub fn residual_parents(&self) -> &[ParentIndex] {
21 &self.residual_parents
22 }
23
24 #[must_use]
25 pub fn residual_rows(&self, system: &ConstraintSystem) -> Vec<ResidualIndex> {
26 let offsets = system.row_offsets();
27 self.residual_parents
28 .iter()
29 .flat_map(|&parent| {
30 let start = offsets[parent.as_usize()].value();
31 (0..system.residual_at_parent(parent).rows()).map(move |o| {
32 let Ok(ov) = u32::try_from(o) else {
33 unreachable!("residual row count is a small constant")
34 };
35 ResidualIndex::new(start + ov)
36 })
37 })
38 .collect()
39 }
40}
41
42#[derive(Clone, Debug, PartialEq, Eq)]
43pub struct Decomposition {
44 components: Vec<Component>,
45}
46
47impl Decomposition {
48 #[must_use]
49 pub fn components(&self) -> &[Component] {
50 &self.components
51 }
52}
53
54#[must_use]
55pub fn decompose(system: &ConstraintSystem) -> Decomposition {
56 let n = system.parameter_count();
57 let dsu_final = system.residuals().iter().fold(Dsu::new(n), |dsu, r| {
58 let ps = r.parameters();
59 match ps.split_first() {
60 None => dsu,
61 Some((&first, rest)) => rest
62 .iter()
63 .fold(dsu, |acc, &p| acc.union(first.as_usize(), p.as_usize())),
64 }
65 });
66 let roots: Vec<usize> = (0..n).map(|p| dsu_final.find(p)).collect();
67 let param_groups: BTreeMap<usize, Vec<ParameterIndex>> =
68 (0..n).fold(BTreeMap::new(), |mut acc, p| {
69 let Ok(pv) = u32::try_from(p) else {
70 unreachable!("parameter count stays within ParameterIndex range")
71 };
72 acc.entry(roots[p])
73 .or_default()
74 .push(ParameterIndex::new(pv));
75 acc
76 });
77 let parent_groups: BTreeMap<usize, Vec<ParentIndex>> = system
78 .residuals()
79 .iter()
80 .enumerate()
81 .fold(BTreeMap::new(), |mut acc, (pi, r)| {
82 if let Some(&first) = r.parameters().first() {
83 let Ok(piv) = u32::try_from(pi) else {
84 unreachable!("parent count fits in u32 via ResidualIndex")
85 };
86 acc.entry(roots[first.as_usize()])
87 .or_default()
88 .push(ParentIndex::new(piv));
89 }
90 acc
91 });
92 let mut keyed: Vec<(usize, Component)> = param_groups
93 .into_iter()
94 .map(|(root, parameters)| {
95 let residual_parents = parent_groups.get(&root).cloned().unwrap_or_default();
96 let min_param = parameters
97 .first()
98 .copied()
99 .map_or(usize::MAX, ParameterIndex::as_usize);
100 (
101 min_param,
102 Component {
103 parameters,
104 residual_parents,
105 },
106 )
107 })
108 .collect();
109 keyed.sort_by_key(|(min_param, _)| *min_param);
110 Decomposition {
111 components: keyed.into_iter().map(|(_, c)| c).collect(),
112 }
113}
114
115#[derive(Clone, Debug)]
116struct Dsu {
117 parent: Vec<usize>,
118}
119
120impl Dsu {
121 fn new(n: usize) -> Self {
122 Self {
123 parent: (0..n).collect(),
124 }
125 }
126
127 fn find(&self, x: usize) -> usize {
128 std::iter::successors(Some(x), |&i| {
129 let p = self.parent[i];
130 (p != i).then_some(p)
131 })
132 .fold(x, |_, i| i)
133 }
134
135 fn union(self, a: usize, b: usize) -> Self {
136 let ra = self.find(a);
137 let rb = self.find(b);
138 if ra == rb {
139 return self;
140 }
141 let (lo, hi) = if ra < rb { (ra, rb) } else { (rb, ra) };
142 let Self { mut parent } = self;
143 parent[hi] = lo;
144 Self { parent }
145 }
146}