tycho_collator/utils/
shard.rs

1use std::collections::VecDeque;
2
3use anyhow::{Context, Result, anyhow};
4use tycho_types::models::ShardIdent;
5
6#[derive(Debug, Clone, PartialEq)]
7pub enum SplitMergeAction {
8    Add(ShardIdent),
9    Split(ShardIdent),
10    Merge(ShardIdent, ShardIdent),
11}
12
13enum CalcSplitMergeStep<'a> {
14    CheckAction(
15        ShardIdent,
16        Option<Vec<&'a ShardIdent>>,
17        Option<SplitMergeAction>,
18    ),
19    DoAction(Vec<&'a ShardIdent>, SplitMergeAction),
20}
21
22/// Calculate the list of split/merge actions that are needed
23/// to move from the current shards set to a new
24pub fn calc_split_merge_actions(
25    from_current_shards: &[ShardIdent],
26    to_new_shards: Vec<&ShardIdent>,
27) -> Result<Vec<SplitMergeAction>> {
28    let workchain = to_new_shards
29        .first()
30        .map(|sid| sid.workchain())
31        .or_else(|| from_current_shards.first().map(|sid| sid.workchain()))
32        .context("`from_current_shards` or `to_new_shards` must not be empty")?;
33    let full_shard_id = ShardIdent::new_full(workchain);
34    let mut planned_actions = VecDeque::new();
35    if from_current_shards.is_empty() {
36        planned_actions.push_back(CalcSplitMergeStep::CheckAction(full_shard_id, None, None));
37    } else {
38        planned_actions.extend(
39            from_current_shards
40                .iter()
41                .map(|&sh| CalcSplitMergeStep::CheckAction(sh, None, None)),
42        );
43    }
44
45    let mut result_actions = vec![];
46
47    for new_shard_id in to_new_shards.iter() {
48        if from_current_shards.is_empty() {
49            result_actions.push(SplitMergeAction::Add(**new_shard_id));
50        }
51    }
52
53    let mut rest_to_shards = to_new_shards;
54    while let Some(next_planned_action) = planned_actions.pop_front() {
55        match next_planned_action {
56            CalcSplitMergeStep::CheckAction(from_shard_id, sub_to_shards_opt, action_opt) => {
57                if let Some(mut sub_to_shards) = sub_to_shards_opt {
58                    rest_to_shards = std::mem::take(&mut sub_to_shards);
59                }
60                let mut to_shards = std::mem::take(&mut rest_to_shards);
61                let mut child_to_shards = vec![];
62                for to_shard_id in to_shards.drain(..) {
63                    if &from_shard_id == to_shard_id {
64                        // do not need to split o merge
65                        if let Some(ref action) = action_opt {
66                            result_actions.push(action.clone());
67                        }
68                    } else if from_shard_id.is_ancestor_of(to_shard_id) {
69                        // need to split
70                        child_to_shards.push(to_shard_id);
71                    } else if to_shard_id.is_ancestor_of(&from_shard_id) {
72                        // need to merge
73                    } else {
74                        rest_to_shards.push(to_shard_id);
75                    }
76                }
77                if !child_to_shards.is_empty() {
78                    if let Some(ref action) = action_opt {
79                        result_actions.push(action.clone());
80                    }
81                    planned_actions.push_back(CalcSplitMergeStep::DoAction(
82                        child_to_shards,
83                        SplitMergeAction::Split(from_shard_id),
84                    ));
85                }
86            }
87            CalcSplitMergeStep::DoAction(child_to_shards, action) => match action {
88                SplitMergeAction::Split(from_shard_id) => {
89                    let (l_shard, r_shard) = from_shard_id.split().ok_or_else(|| {
90                        anyhow!(
91                            "Unable to split shard {}, MAX_SPLIT_DEPTH ({}) reached",
92                            from_shard_id,
93                            ShardIdent::MAX_SPLIT_DEPTH
94                        )
95                    })?;
96                    planned_actions.push_back(CalcSplitMergeStep::CheckAction(
97                        l_shard,
98                        Some(child_to_shards),
99                        Some(action.clone()),
100                    ));
101                    planned_actions.push_back(CalcSplitMergeStep::CheckAction(
102                        r_shard,
103                        None,
104                        Some(action),
105                    ));
106                }
107                SplitMergeAction::Merge(_from_shard_id_1, _from_shard_id_2) => {
108                    // do nothing
109                }
110                SplitMergeAction::Add(_) => {}
111            },
112        }
113    }
114
115    result_actions.dedup_by(|a, b| a == b);
116
117    Ok(result_actions)
118}
119
120#[cfg(test)]
121mod tests {
122    use tycho_types::models::ShardIdent;
123
124    use super::calc_split_merge_actions;
125    use crate::utils::shard::SplitMergeAction;
126
127    #[test]
128    fn test_calc_split_merge_actions() {
129        // BASECHAIN 0:80
130        let shard_80 = ShardIdent::new_full(0);
131
132        // split on 4 shards
133        let (shard_40, shard_c0) = shard_80.split().unwrap();
134        println!("full shard {shard_80}");
135        println!("shard {shard_80} split l {shard_40}");
136        println!("shard {shard_80} split r {shard_c0}");
137
138        let (shard_20, shard_60) = shard_40.split().unwrap();
139        println!("shard {shard_40} split l {shard_20}");
140        println!("shard {shard_40} split r {shard_60}");
141
142        let (shard_a0, shard_e0) = shard_c0.split().unwrap();
143        println!("shard {shard_c0} split l {shard_a0}");
144        println!("shard {shard_c0} split r {shard_e0}");
145
146        let shards_1_r = vec![&shard_80];
147        let shards_1_l = &[shard_80];
148        let actions = calc_split_merge_actions(&[], shards_1_r.clone()).unwrap();
149        println!("split/merge actions from [] to [1]: {actions:?}");
150        assert!(actions.contains(&SplitMergeAction::Add(shard_80)));
151
152        let shards_4_l = &[shard_20, shard_60, shard_a0, shard_e0];
153        let shards_4_r = vec![&shard_20, &shard_60, &shard_a0, &shard_e0];
154        let actions = calc_split_merge_actions(&[], shards_4_r.clone()).unwrap();
155        println!("split/merge actions from [] to [4]: {actions:?}");
156        assert!(actions.contains(&SplitMergeAction::Add(shard_20)));
157        assert!(actions.contains(&SplitMergeAction::Add(shard_60)));
158        assert!(actions.contains(&SplitMergeAction::Add(shard_a0)));
159        assert!(actions.contains(&SplitMergeAction::Add(shard_e0)));
160
161        let actions = calc_split_merge_actions(shards_1_l, shards_4_r.clone()).unwrap();
162        println!("split/merge actions from [1] to [4]: {actions:?}");
163        assert!(actions.contains(&SplitMergeAction::Split(shard_80)));
164        assert!(actions.contains(&SplitMergeAction::Split(shard_40)));
165        assert!(actions.contains(&SplitMergeAction::Split(shard_c0)));
166
167        let shards_2_l = &[shard_40, shard_c0];
168        let actions = calc_split_merge_actions(shards_2_l, shards_4_r.clone()).unwrap();
169        println!("split/merge actions from [2] to [4]: {actions:?}");
170        assert!(actions.contains(&SplitMergeAction::Split(shard_40)));
171        assert!(actions.contains(&SplitMergeAction::Split(shard_c0)));
172
173        let shards_3_r = vec![&shard_40, &shard_a0, &shard_e0];
174        let shards_3_l = &[shard_40, shard_a0, shard_e0];
175        let actions = calc_split_merge_actions(shards_2_l, shards_3_r.clone()).unwrap();
176        println!("split/merge actions from [2] to [3]: {actions:?}");
177        assert!(actions.contains(&SplitMergeAction::Split(shard_c0)));
178
179        let actions = calc_split_merge_actions(shards_3_l, shards_4_r.clone()).unwrap();
180        println!("split/merge actions from [3] to [4]: {actions:?}");
181        assert!(actions.contains(&SplitMergeAction::Split(shard_40)));
182
183        let actions = calc_split_merge_actions(shards_4_l, shards_4_r.clone()).unwrap();
184        println!("split/merge actions from [4] to [4]: {actions:?}");
185        assert!(actions.is_empty());
186
187        // BASECHAIN 1:80
188        let shard_80 = ShardIdent::new_full(1);
189
190        // split on 4 shards
191        let (shard_40, shard_c0) = shard_80.split().unwrap();
192        println!("full shard {shard_80}");
193        println!("shard {shard_80} split l {shard_40}");
194        println!("shard {shard_80} split r {shard_c0}");
195
196        let (shard_20, shard_60) = shard_40.split().unwrap();
197        println!("shard {shard_40} split l {shard_20}");
198        println!("shard {shard_40} split r {shard_60}");
199
200        let (shard_a0, shard_e0) = shard_c0.split().unwrap();
201        println!("shard {shard_c0} split l {shard_a0}");
202        println!("shard {shard_c0} split r {shard_e0}");
203
204        let shards_1_r = vec![&shard_80];
205        let shards_1_l = &[shard_80];
206        let actions = calc_split_merge_actions(&[], shards_1_r.clone()).unwrap();
207        println!("split/merge actions from [] to [1]: {actions:?}");
208        assert!(actions.contains(&SplitMergeAction::Add(shard_80)));
209
210        let shards_4_l = &[shard_20, shard_60, shard_a0, shard_e0];
211        let shards_4_r = vec![&shard_20, &shard_60, &shard_a0, &shard_e0];
212        let actions = calc_split_merge_actions(&[], shards_4_r.clone()).unwrap();
213        println!("split/merge actions from [] to [4]: {actions:?}");
214        assert!(actions.contains(&SplitMergeAction::Add(shard_20)));
215        assert!(actions.contains(&SplitMergeAction::Add(shard_60)));
216        assert!(actions.contains(&SplitMergeAction::Add(shard_a0)));
217        assert!(actions.contains(&SplitMergeAction::Add(shard_e0)));
218
219        let actions = calc_split_merge_actions(shards_1_l, shards_4_r.clone()).unwrap();
220        println!("split/merge actions from [1] to [4]: {actions:?}");
221        assert!(actions.contains(&SplitMergeAction::Split(shard_80)));
222        assert!(actions.contains(&SplitMergeAction::Split(shard_40)));
223        assert!(actions.contains(&SplitMergeAction::Split(shard_c0)));
224
225        let actions = calc_split_merge_actions(shards_4_l, shards_4_r.clone()).unwrap();
226        println!("split/merge actions from [4] to [4]: {actions:?}");
227        assert!(actions.is_empty());
228
229        // should error on empty input
230        let res = calc_split_merge_actions(&[], vec![]);
231        assert!(res.is_err());
232    }
233}
tycho_collator/utils/shard.rs

tycho_collator/utils/
shard.rs
