1use std::collections::BTreeMap;
26
27use crate::varint;
28use crate::{ComponentId, EntityId};
29
30pub const VERSION: u8 = 1;
31pub const FLAG_FULL: u8 = 1;
32
33pub mod mask {
35 pub const X: u8 = 1;
36 pub const Y: u8 = 2;
37 pub const Z: u8 = 4;
38 pub const COMPONENTS: u8 = 8;
39}
40
41pub fn quantize(value: f32, precision: f32) -> i64 {
43 if !value.is_finite() {
44 return 0;
45 }
46 let q = (value as f64 / precision as f64).round();
47 q.clamp(i64::MIN as f64, i64::MAX as f64) as i64
48}
49
50pub fn quantize3(pos: [f32; 3], precision: f32) -> [i64; 3] {
51 pos.map(|v| quantize(v, precision))
52}
53
54pub type ComponentChange<'a> = (ComponentId, Option<&'a [u8]>);
56
57pub fn encode_update_body(delta: [i64; 3], components: &[ComponentChange<'_>]) -> Vec<u8> {
60 let mut body = Vec::with_capacity(8);
61 let mut m = 0u8;
62 for (i, bit) in [mask::X, mask::Y, mask::Z].into_iter().enumerate() {
63 if delta[i] != 0 {
64 m |= bit;
65 }
66 }
67 if !components.is_empty() {
68 m |= mask::COMPONENTS;
69 }
70 body.push(m);
71 for d in delta {
72 if d != 0 {
73 varint::write_i64(&mut body, d);
74 }
75 }
76 if !components.is_empty() {
77 write_components(&mut body, components.iter().copied());
78 }
79 body
80}
81
82fn write_components<'a>(
83 out: &mut Vec<u8>,
84 components: impl ExactSizeIterator<Item = ComponentChange<'a>>,
85) {
86 varint::write_u64(out, components.len() as u64);
87 for (id, bytes) in components {
88 out.push(id);
89 match bytes {
90 Some(b) => {
91 varint::write_u64(out, b.len() as u64 + 1);
92 out.extend_from_slice(b);
93 }
94 None => varint::write_u64(out, 0),
95 }
96 }
97}
98
99#[derive(Debug)]
102pub struct FrameBuilder {
103 full: bool,
104 precision: f32,
105 despawns: Vec<EntityId>,
106 spawns: Vec<u8>,
107 spawn_count: u64,
108 last_spawn: Option<EntityId>,
109 updates: Vec<u8>,
110 update_count: u64,
111 last_update: Option<EntityId>,
112}
113
114impl FrameBuilder {
115 pub fn new(full: bool, precision: f32) -> Self {
116 Self {
117 full,
118 precision,
119 despawns: Vec::new(),
120 spawns: Vec::new(),
121 spawn_count: 0,
122 last_spawn: None,
123 updates: Vec::new(),
124 update_count: 0,
125 last_update: None,
126 }
127 }
128
129 pub fn despawn(&mut self, id: EntityId) {
130 self.despawns.push(id);
131 }
132
133 pub fn spawn<'a>(
135 &mut self,
136 id: EntityId,
137 q: [i64; 3],
138 components: impl ExactSizeIterator<Item = (ComponentId, &'a [u8])>,
139 ) {
140 write_delta_id(&mut self.spawns, &mut self.last_spawn, id);
141 for v in q {
142 varint::write_i64(&mut self.spawns, v);
143 }
144 write_components(&mut self.spawns, components.map(|(c, b)| (c, Some(b))));
145 self.spawn_count += 1;
146 }
147
148 pub fn update(&mut self, id: EntityId, body: &[u8]) {
151 write_delta_id(&mut self.updates, &mut self.last_update, id);
152 self.updates.extend_from_slice(body);
153 self.update_count += 1;
154 }
155
156 pub fn is_empty(&self) -> bool {
157 !self.full && self.despawns.is_empty() && self.spawn_count == 0 && self.update_count == 0
158 }
159
160 pub fn finish(mut self) -> Vec<u8> {
162 self.despawns.sort_unstable();
163 self.despawns.dedup();
164 let mut out = Vec::with_capacity(16 + self.spawns.len() + self.updates.len());
165 out.push(VERSION);
166 out.push(if self.full { FLAG_FULL } else { 0 });
167 out.extend_from_slice(&self.precision.to_le_bytes());
168 varint::write_u64(&mut out, self.despawns.len() as u64);
169 let mut last = None;
170 for id in &self.despawns {
171 write_delta_id(&mut out, &mut last, *id);
172 }
173 varint::write_u64(&mut out, self.spawn_count);
174 out.extend_from_slice(&self.spawns);
175 varint::write_u64(&mut out, self.update_count);
176 out.extend_from_slice(&self.updates);
177 out
178 }
179}
180
181fn write_delta_id(out: &mut Vec<u8>, last: &mut Option<EntityId>, id: EntityId) {
182 match *last {
183 Some(prev) => {
184 debug_assert!(id > prev, "ids must ascend");
185 varint::write_u64(out, id.wrapping_sub(prev));
186 }
187 None => varint::write_u64(out, id),
188 }
189 *last = Some(id);
190}
191
192#[derive(Debug, Clone, PartialEq, Eq)]
198pub struct DecodeError(pub String);
199
200impl std::fmt::Display for DecodeError {
201 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
202 f.write_str(&self.0)
203 }
204}
205
206impl std::error::Error for DecodeError {}
207
208#[derive(Debug, Clone, Default, PartialEq)]
210pub struct ReplicaEntity {
211 pub q: [i64; 3],
213 pub components: BTreeMap<ComponentId, Vec<u8>>,
214}
215
216impl ReplicaEntity {
217 pub fn pos(&self, precision: f32) -> [f32; 3] {
218 self.q.map(|v| (v as f64 * precision as f64) as f32)
219 }
220}
221
222#[derive(Debug, Clone, Default, PartialEq, Eq)]
224pub struct FrameSummary {
225 pub full: bool,
226 pub spawned: Vec<EntityId>,
227 pub updated: Vec<EntityId>,
228 pub despawned: Vec<EntityId>,
229}
230
231#[derive(Debug, Clone, Default, PartialEq)]
233pub struct ReplicaTable {
234 pub entities: BTreeMap<EntityId, ReplicaEntity>,
235 pub precision: f32,
236}
237
238const MAX_COUNT: u64 = 1 << 24;
241
242impl ReplicaTable {
243 pub fn new() -> Self {
244 Self::default()
245 }
246
247 pub fn pos(&self, id: EntityId) -> Option<[f32; 3]> {
248 self.entities.get(&id).map(|e| e.pos(self.precision))
249 }
250
251 pub fn apply(&mut self, frame: &[u8]) -> Result<FrameSummary, DecodeError> {
254 let err = |m: &str| DecodeError(m.to_string());
255 let mut b = frame;
256 let take = |b: &mut &[u8], n: usize| -> Result<Vec<u8>, DecodeError> {
257 if b.len() < n {
258 return Err(DecodeError("frame ends early".into()));
259 }
260 let (head, rest) = b.split_at(n);
261 *b = rest;
262 Ok(head.to_vec())
263 };
264 let head = take(&mut b, 6)?;
265 if head[0] != VERSION {
266 return Err(DecodeError(format!("replication version {}", head[0])));
267 }
268 let full = head[1] & FLAG_FULL != 0;
269 let precision = f32::from_le_bytes(head[2..6].try_into().unwrap());
270 if !(precision.is_finite() && precision > 0.0) {
271 return Err(err("bad precision"));
272 }
273 if full {
274 self.entities.clear();
275 }
276 self.precision = precision;
277 let mut summary = FrameSummary {
278 full,
279 ..FrameSummary::default()
280 };
281 let count = |b: &mut &[u8]| -> Result<u64, DecodeError> {
282 let n = varint::read_u64(b).ok_or_else(|| err("bad count"))?;
283 if n > MAX_COUNT {
284 return Err(err("count too large"));
285 }
286 Ok(n)
287 };
288 let next_id =
289 |b: &mut &[u8], last: &mut Option<EntityId>| -> Result<EntityId, DecodeError> {
290 let v = varint::read_u64(b).ok_or_else(|| err("bad id"))?;
291 let id = match *last {
292 Some(prev) => prev
293 .checked_add(v)
294 .filter(|_| v > 0)
295 .ok_or_else(|| err("ids do not ascend"))?,
296 None => v,
297 };
298 *last = Some(id);
299 Ok(id)
300 };
301
302 let n = count(&mut b)?;
303 let mut last = None;
304 for _ in 0..n {
305 let id = next_id(&mut b, &mut last)?;
306 self.entities.remove(&id);
307 summary.despawned.push(id);
308 }
309
310 let n = count(&mut b)?;
311 let mut last = None;
312 for _ in 0..n {
313 let id = next_id(&mut b, &mut last)?;
314 let mut q = [0i64; 3];
315 for v in &mut q {
316 *v = varint::read_i64(&mut b).ok_or_else(|| err("bad position"))?;
317 }
318 let mut entity = ReplicaEntity {
319 q,
320 components: BTreeMap::new(),
321 };
322 read_components(&mut b, &mut entity.components)?;
323 self.entities.insert(id, entity);
324 summary.spawned.push(id);
325 }
326
327 let n = count(&mut b)?;
328 let mut last = None;
329 for _ in 0..n {
330 let id = next_id(&mut b, &mut last)?;
331 let (&m, rest) = b.split_first().ok_or_else(|| err("frame ends early"))?;
332 b = rest;
333 let entity = self
334 .entities
335 .get_mut(&id)
336 .ok_or_else(|| DecodeError(format!("update for unknown entity {id}")))?;
337 for (i, bit) in [mask::X, mask::Y, mask::Z].into_iter().enumerate() {
338 if m & bit != 0 {
339 let d = varint::read_i64(&mut b).ok_or_else(|| err("bad delta"))?;
340 entity.q[i] = entity.q[i].wrapping_add(d);
341 }
342 }
343 if m & mask::COMPONENTS != 0 {
344 read_components(&mut b, &mut entity.components)?;
345 }
346 summary.updated.push(id);
347 }
348 if !b.is_empty() {
349 return Err(err("trailing bytes"));
350 }
351 Ok(summary)
352 }
353}
354
355fn read_components(
356 b: &mut &[u8],
357 into: &mut BTreeMap<ComponentId, Vec<u8>>,
358) -> Result<(), DecodeError> {
359 let err = |m: &str| DecodeError(m.to_string());
360 let n = varint::read_u64(b).ok_or_else(|| err("bad component count"))?;
361 if n > 256 {
362 return Err(err("more than 256 components"));
363 }
364 for _ in 0..n {
365 let (&id, rest) = b.split_first().ok_or_else(|| err("frame ends early"))?;
366 *b = rest;
367 let len = varint::read_u64(b).ok_or_else(|| err("bad component length"))?;
368 if len == 0 {
369 into.remove(&id);
370 continue;
371 }
372 let len = usize::try_from(len - 1).map_err(|_| err("component too large"))?;
373 if b.len() < len {
374 return Err(err("component runs past the end"));
375 }
376 let (value, rest) = b.split_at(len);
377 *b = rest;
378 into.insert(id, value.to_vec());
379 }
380 Ok(())
381}
382
383#[cfg(test)]
384mod tests {
385 use super::*;
386
387 #[test]
388 fn spawn_update_despawn_round_trip() {
389 let p = 0.1;
390 let mut f = FrameBuilder::new(true, p);
391 f.spawn(
392 3,
393 quantize3([1.0, 2.0, 0.0], p),
394 [(1u8, &b"hp"[..])].into_iter(),
395 );
396 f.spawn(10, quantize3([-5.0, 0.0, 0.0], p), std::iter::empty());
397 let mut t = ReplicaTable::new();
398 let s = t.apply(&f.finish()).unwrap();
399 assert!(s.full);
400 assert_eq!(s.spawned, vec![3, 10]);
401 assert_eq!(t.pos(3), Some([1.0, 2.0, 0.0]));
402 assert_eq!(t.entities[&3].components[&1], b"hp");
403
404 let mut f = FrameBuilder::new(false, p);
405 f.update(
406 3,
407 &encode_update_body([5, 0, 0], &[(1, None), (2, Some(b"new"))]),
408 );
409 f.update(10, &encode_update_body([0, -20, 0], &[]));
410 f.despawn(99);
411 let s = t.apply(&f.finish()).unwrap();
412 assert_eq!(s.updated, vec![3, 10]);
413 let pos3 = t.pos(3).unwrap();
414 assert!((pos3[0] - 1.5).abs() < 1e-6);
415 assert!(!t.entities[&3].components.contains_key(&1));
416 assert_eq!(t.entities[&3].components[&2], b"new");
417 assert_eq!(t.pos(10), Some([-5.0, -2.0, 0.0]));
418
419 let mut f = FrameBuilder::new(false, p);
420 f.despawn(10);
421 t.apply(&f.finish()).unwrap();
422 assert!(!t.entities.contains_key(&10));
423 }
424
425 #[test]
426 fn an_unchanged_entity_update_is_two_bytes_plus_the_id() {
427 let body = encode_update_body([3, -2, 0], &[]);
428 assert_eq!(body.len(), 3); }
430
431 #[test]
432 fn hostile_frames_are_refused_without_panicking() {
433 let mut t = ReplicaTable::new();
434 assert!(t.apply(&[]).is_err());
435 assert!(t.apply(&[2, 0, 0, 0, 0x80, 0x3f]).is_err()); assert!(t.apply(&[1, 0, 0, 0, 0, 0]).is_err()); let head = |extra: &[u8]| {
438 let mut v = vec![1, 0];
439 v.extend_from_slice(&1.0f32.to_le_bytes());
440 v.extend_from_slice(extra);
441 v
442 };
443 assert!(t.apply(&head(&[0xff, 0xff, 0xff, 0xff, 0x0f])).is_err());
445 assert!(t.apply(&head(&[0, 0, 1, 7, 0])).is_err());
447 assert!(t.apply(&head(&[2, 5, 0, 0, 0])).is_err());
449 assert!(t.apply(&head(&[0, 0, 0, 9])).is_err());
451 assert!(t.apply(&head(&[0, 1, 1, 0, 0, 0, 1, 4, 50, 1])).is_err());
453 }
454
455 #[test]
456 fn quantize_handles_bad_values() {
457 assert_eq!(quantize(f32::NAN, 0.1), 0);
458 assert_eq!(quantize(f32::INFINITY, 0.1), 0);
459 assert_eq!(quantize(f32::MAX, 1e-30), i64::MAX);
460 assert_eq!(quantize(-0.04, 0.1), 0);
461 assert_eq!(quantize(0.26, 0.1), 3);
462 }
463}