Skip to main content

pylon_replication/
frame.rs

1//! The replication frame: what one subscriber receives per tick.
2//!
3//! Integers are LEB128 varints ([`crate::varint`]); signed values are
4//! zigzag varints. Positions travel quantized: `round(value / precision)`.
5//!
6//! ```text
7//! u8      version (1)
8//! u8      flags: bit 0 FULL: the client clears its table first
9//! f32 LE  precision
10//! varint  despawn count, then ids: the first absolute, each next as the
11//!         difference from the previous (ids ascend)
12//! varint  spawn count, then per entity (ids ascend, delta-coded):
13//!           id, x, y, z (zigzag, absolute quantized), components
14//! varint  update count, then per entity (ids ascend, delta-coded):
15//!           id, u8 mask (1 x, 2 y, 4 z, 8 components),
16//!           the changed axes (zigzag, quantized difference from the value
17//!           last sent), components when bit 8 is set
18//! components: varint count, then per component:
19//!           u8 id, varint (length + 1), bytes; length field 0 = removed
20//! ```
21//!
22//! A spawn carries every present component. An update carries the ones
23//! that changed or were removed since that subscriber was last updated.
24
25use std::collections::BTreeMap;
26
27use crate::varint;
28use crate::{ComponentId, EntityId};
29
30pub const VERSION: u8 = 1;
31pub const FLAG_FULL: u8 = 1;
32
33/// Bits of an update's mask byte.
34pub 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
41/// `value / precision`, rounded. Non-finite values quantize to 0.
42pub 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
54/// A component change inside an update: new bytes, or removed.
55pub type ComponentChange<'a> = (ComponentId, Option<&'a [u8]>);
56
57/// Encode an update body (everything after the id): the mask, the changed
58/// axes, and the component changes. Returns the body.
59pub 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/// Builds one frame. Add despawns in any order; add spawns and updates in
100/// ascending id order.
101#[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    /// Add a spawn. Ids must ascend across calls.
134    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    /// Add an update with a body from [`encode_update_body`]. Ids must
149    /// ascend across calls.
150    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    /// The frame bytes.
161    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// ---------------------------------------------------------------------------
193// Decoding
194// ---------------------------------------------------------------------------
195
196/// Why a frame did not decode or apply.
197#[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/// An entity as a client sees it.
209#[derive(Debug, Clone, Default, PartialEq)]
210pub struct ReplicaEntity {
211    /// The quantized position; deltas apply to it exactly.
212    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/// What one frame did.
223#[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/// A client's table of the entities it has been told about.
232#[derive(Debug, Clone, Default, PartialEq)]
233pub struct ReplicaTable {
234    pub entities: BTreeMap<EntityId, ReplicaEntity>,
235    pub precision: f32,
236}
237
238/// Longest id list a frame may declare, so a hostile count cannot make a
239/// decoder allocate without bound.
240const 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    /// Apply one frame. On an error the table may hold part of the frame;
252    /// the caller should treat the connection as broken.
253    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); // mask + two one-byte deltas
429    }
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()); // version 2
436        assert!(t.apply(&[1, 0, 0, 0, 0, 0]).is_err()); // precision 0
437        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        // A huge declared count.
444        assert!(t.apply(&head(&[0xff, 0xff, 0xff, 0xff, 0x0f])).is_err());
445        // An update for an entity the table never saw.
446        assert!(t.apply(&head(&[0, 0, 1, 7, 0])).is_err());
447        // Ids that do not ascend.
448        assert!(t.apply(&head(&[2, 5, 0, 0, 0])).is_err());
449        // Trailing bytes.
450        assert!(t.apply(&head(&[0, 0, 0, 9])).is_err());
451        // A component that runs past the end.
452        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}