| 1 | use std::cmp::Ordering; |
| 2 | |
| 3 | use ruma::OwnedRoomId; |
| 4 | |
| 5 | const FIRST: u8 = 0x20; |
| 6 | const LAST: u8 = 0x7e; |
| 7 | const BASE: u64 = (LAST - FIRST + 1) as u64; |
| 8 | const MAX_LEN: usize = 50; |
| 9 | |
| 10 | pub(crate) const EVENT_TYPE: &str = "org.matrix.msc3230.space_order"; |
| 11 | |
| 12 | pub(crate) fn valid(order: &str) -> bool { |
| 13 | (1..=MAX_LEN).contains(&order.len()) && order.bytes().all(|b| (FIRST..=LAST).contains(&b)) |
| 14 | } |
| 15 | |
| 16 | /// how msc3230 clients sort spaces: ordered ones by (order, room id), then the rest by room id |
| 17 | pub(crate) fn compare( |
| 18 | (a_id, a_order): (&OwnedRoomId, Option<&str>), |
| 19 | (b_id, b_order): (&OwnedRoomId, Option<&str>), |
| 20 | ) -> Ordering { |
| 21 | match (a_order, b_order) { |
| 22 | (Some(a), Some(b)) => a.cmp(b).then_with(|| a_id.cmp(b_id)), |
| 23 | (Some(_), None) => Ordering::Less, |
| 24 | (None, Some(_)) => Ordering::Greater, |
| 25 | (None, None) => a_id.cmp(b_id), |
| 26 | } |
| 27 | } |
| 28 | |
| 29 | /// the orders to write so `spaces`, given with their current valid orders, sort in |
| 30 | /// exactly this sequence. keeps the longest run of orders that already sorts right |
| 31 | pub(crate) fn plan(spaces: &[(OwnedRoomId, Option<String>)]) -> Vec<(OwnedRoomId, String)> { |
| 32 | let key = |index: usize| (&spaces[index].0, spaces[index].1.as_deref()); |
| 33 | let n = spaces.len(); |
| 34 | |
| 35 | let mut unordered_tail = n; |
| 36 | while unordered_tail > 0 |
| 37 | && spaces[unordered_tail - 1].1.is_none() |
| 38 | && (unordered_tail == n |
| 39 | || compare(key(unordered_tail - 1), key(unordered_tail)) == Ordering::Less) |
| 40 | { |
| 41 | unordered_tail -= 1; |
| 42 | } |
| 43 | |
| 44 | let mut chain = vec![0usize; n]; |
| 45 | let mut previous = vec![None; n]; |
| 46 | for i in (0..n).filter(|&i| spaces[i].1.is_some()) { |
| 47 | chain[i] = 1; |
| 48 | for j in 0..i { |
| 49 | if chain[j] > 0 && chain[j] + 1 > chain[i] && compare(key(j), key(i)) == Ordering::Less |
| 50 | { |
| 51 | chain[i] = chain[j] + 1; |
| 52 | previous[i] = Some(j); |
| 53 | } |
| 54 | } |
| 55 | } |
| 56 | |
| 57 | let mut longest: Option<usize> = None; |
| 58 | let mut best: Option<(usize, usize, Option<usize>)> = None; |
| 59 | for split in 0..=n { |
| 60 | if split > 0 && chain[split - 1] > longest.map_or(0, |end| chain[end]) { |
| 61 | longest = Some(split - 1); |
| 62 | } |
| 63 | if split < unordered_tail { |
| 64 | continue; |
| 65 | } |
| 66 | let cost = split - longest.map_or(0, |end| chain[end]); |
| 67 | if best.is_none_or(|(best_cost, _, _)| cost < best_cost) { |
| 68 | best = Some((cost, split, longest)); |
| 69 | } |
| 70 | } |
| 71 | let Some((_, split, end)) = best else { |
| 72 | return Vec::new(); |
| 73 | }; |
| 74 | |
| 75 | let mut kept = vec![false; n]; |
| 76 | let mut cursor = end; |
| 77 | while let Some(index) = cursor { |
| 78 | kept[index] = true; |
| 79 | cursor = previous[index]; |
| 80 | } |
| 81 | |
| 82 | let mut writes = Vec::new(); |
| 83 | let mut run_start = 0; |
| 84 | for index in 0..=split { |
| 85 | if index < split && !kept[index] { |
| 86 | continue; |
| 87 | } |
| 88 | if index > run_start { |
| 89 | let lo = run_start |
| 90 | .checked_sub(1) |
| 91 | .and_then(|i| spaces[i].1.as_deref()); |
| 92 | let hi = (index < split) |
| 93 | .then(|| spaces[index].1.as_deref()) |
| 94 | .flatten(); |
| 95 | let Some(orders) = between(lo, hi, index - run_start) else { |
| 96 | return rebalance(&spaces[..split]); |
| 97 | }; |
| 98 | writes.extend( |
| 99 | spaces[run_start..index] |
| 100 | .iter() |
| 101 | .map(|(id, _)| id.clone()) |
| 102 | .zip(orders), |
| 103 | ); |
| 104 | } |
| 105 | run_start = index + 1; |
| 106 | } |
| 107 | writes |
| 108 | } |
| 109 | |
| 110 | fn rebalance(spaces: &[(OwnedRoomId, Option<String>)]) -> Vec<(OwnedRoomId, String)> { |
| 111 | let orders = between(None, None, spaces.len()).unwrap_or_default(); |
| 112 | spaces |
| 113 | .iter() |
| 114 | .zip(orders) |
| 115 | .filter(|((_, current), order)| current.as_ref() != Some(order)) |
| 116 | .map(|((id, _), order)| (id.clone(), order)) |
| 117 | .collect() |
| 118 | } |
| 119 | |
| 120 | /// `count` evenly spread orders strictly between `lo` and `hi`, as short as they can be |
| 121 | fn between(lo: Option<&str>, hi: Option<&str>, count: usize) -> Option<Vec<String>> { |
| 122 | let lo = lo.unwrap_or("").as_bytes(); |
| 123 | (lo.len().max(1)..=MAX_LEN).find_map(|len| { |
| 124 | let lo = digits(lo, len); |
| 125 | let hi = match hi { |
| 126 | Some(hi) => digits(&hi.as_bytes()[..hi.len().min(len)], len), |
| 127 | None => { |
| 128 | let mut one = vec![0; len + 1]; |
| 129 | one[0] = 1; |
| 130 | one |
| 131 | } |
| 132 | }; |
| 133 | if hi <= lo { |
| 134 | return None; |
| 135 | } |
| 136 | let step = divide(&subtract(&hi, &lo), count as u64 + 1); |
| 137 | if step.iter().all(|&digit| digit == 0) { |
| 138 | return None; |
| 139 | } |
| 140 | let mut point = lo; |
| 141 | Some( |
| 142 | (0..count) |
| 143 | .map(|_| { |
| 144 | point = add(&point, &step); |
| 145 | let mut order: Vec<u8> = point[1..].iter().map(|&d| FIRST + d as u8).collect(); |
| 146 | while order.last() == Some(&FIRST) { |
| 147 | order.pop(); |
| 148 | } |
| 149 | String::from_utf8(order).expect("orders are ascii") |
| 150 | }) |
| 151 | .collect(), |
| 152 | ) |
| 153 | }) |
| 154 | } |
| 155 | |
| 156 | /// big-endian base-95 digits with one leading digit of headroom, zero padded to `len` |
| 157 | fn digits(order: &[u8], len: usize) -> Vec<u64> { |
| 158 | let mut out = vec![0; len + 1]; |
| 159 | for (slot, byte) in out[1..].iter_mut().zip(order) { |
| 160 | *slot = u64::from(byte - FIRST); |
| 161 | } |
| 162 | out |
| 163 | } |
| 164 | |
| 165 | fn subtract(a: &[u64], b: &[u64]) -> Vec<u64> { |
| 166 | let mut out = vec![0; a.len()]; |
| 167 | let mut borrow = 0; |
| 168 | for i in (0..a.len()).rev() { |
| 169 | let (value, under) = (a[i] + BASE - b[i] - borrow, a[i] < b[i] + borrow); |
| 170 | out[i] = value % BASE; |
| 171 | borrow = u64::from(under); |
| 172 | } |
| 173 | out |
| 174 | } |
| 175 | |
| 176 | fn add(a: &[u64], b: &[u64]) -> Vec<u64> { |
| 177 | let mut out = vec![0; a.len()]; |
| 178 | let mut carry = 0; |
| 179 | for i in (0..a.len()).rev() { |
| 180 | let sum = a[i] + b[i] + carry; |
| 181 | out[i] = sum % BASE; |
| 182 | carry = sum / BASE; |
| 183 | } |
| 184 | out |
| 185 | } |
| 186 | |
| 187 | fn divide(a: &[u64], divisor: u64) -> Vec<u64> { |
| 188 | let mut remainder = 0; |
| 189 | a.iter() |
| 190 | .map(|&digit| { |
| 191 | let current = remainder * BASE + digit; |
| 192 | remainder = current % divisor; |
| 193 | current / divisor |
| 194 | }) |
| 195 | .collect() |
| 196 | } |
| 197 | |
| 198 | #[cfg(test)] |
| 199 | mod tests { |
| 200 | use ruma::owned_room_id; |
| 201 | |
| 202 | use super::*; |
| 203 | |
| 204 | fn apply(spaces: &[(OwnedRoomId, Option<String>)]) -> (Vec<OwnedRoomId>, usize) { |
| 205 | let writes = plan(spaces); |
| 206 | let mut after: Vec<_> = spaces |
| 207 | .iter() |
| 208 | .map(|(id, order)| { |
| 209 | let written = writes.iter().find(|(w, _)| w == id).map(|(_, o)| o.clone()); |
| 210 | (id.clone(), written.or_else(|| order.clone())) |
| 211 | }) |
| 212 | .collect(); |
| 213 | assert!(after.iter().all(|(_, o)| o.as_deref().is_none_or(valid))); |
| 214 | after.sort_by(|a, b| compare((&a.0, a.1.as_deref()), (&b.0, b.1.as_deref()))); |
| 215 | (after.into_iter().map(|(id, _)| id).collect(), writes.len()) |
| 216 | } |
| 217 | |
| 218 | fn ids(spaces: &[(OwnedRoomId, Option<String>)]) -> Vec<OwnedRoomId> { |
| 219 | spaces.iter().map(|(id, _)| id.clone()).collect() |
| 220 | } |
| 221 | |
| 222 | #[test] |
| 223 | fn already_sorted_spaces_need_no_writes() { |
| 224 | let spaces = vec![ |
| 225 | (owned_room_id!("!z:x"), Some("a".to_owned())), |
| 226 | (owned_room_id!("!a:x"), None), |
| 227 | (owned_room_id!("!b:x"), None), |
| 228 | ]; |
| 229 | assert_eq!(apply(&spaces), (ids(&spaces), 0)); |
| 230 | } |
| 231 | |
| 232 | #[test] |
| 233 | fn moving_one_space_writes_one_order() { |
| 234 | let spaces = vec![ |
| 235 | (owned_room_id!("!c:x"), Some("c".to_owned())), |
| 236 | (owned_room_id!("!a:x"), Some("a".to_owned())), |
| 237 | (owned_room_id!("!b:x"), Some("b".to_owned())), |
| 238 | (owned_room_id!("!d:x"), Some("d".to_owned())), |
| 239 | ]; |
| 240 | assert_eq!(apply(&spaces), (ids(&spaces), 1)); |
| 241 | } |
| 242 | |
| 243 | #[test] |
| 244 | fn unordered_spaces_out_of_id_order_get_orders() { |
| 245 | let spaces = vec![ |
| 246 | (owned_room_id!("!c:x"), None), |
| 247 | (owned_room_id!("!a:x"), None), |
| 248 | (owned_room_id!("!b:x"), None), |
| 249 | ]; |
| 250 | assert_eq!(apply(&spaces), (ids(&spaces), 1)); |
| 251 | } |
| 252 | |
| 253 | #[test] |
| 254 | fn orders_with_no_room_between_them_rebalance() { |
| 255 | let spaces = vec![ |
| 256 | (owned_room_id!("!a:x"), Some("a".to_owned())), |
| 257 | (owned_room_id!("!c:x"), None), |
| 258 | (owned_room_id!("!b:x"), Some("a ".to_owned())), |
| 259 | ]; |
| 260 | assert_eq!(apply(&spaces).0, ids(&spaces)); |
| 261 | let prefix = "~".repeat(MAX_LEN - 1); |
| 262 | let spaces = vec![ |
| 263 | (owned_room_id!("!b:x"), Some(format!("{prefix}!"))), |
| 264 | (owned_room_id!("!a:x"), None), |
| 265 | (owned_room_id!("!c:x"), Some(format!("{prefix}\""))), |
| 266 | ]; |
| 267 | assert_eq!(apply(&spaces).0, ids(&spaces)); |
| 268 | } |
| 269 | } |