tf_demo_parser/demo/message/
packetentities.rs

1use bitbuffer::{
2    BitRead, BitReadSized, BitReadStream, BitWrite, BitWriteSized, BitWriteStream, Endianness,
3    LittleEndian,
4};
5use serde::{Deserialize, Serialize};
6use serde_repr::{Deserialize_repr, Serialize_repr};
7use std::borrow::Cow;
8
9use crate::demo::message::stringtable::log_base2;
10use crate::demo::packet::datatable::{ClassId, SendTable};
11use crate::demo::parser::{Encode, ParseBitSkip};
12use crate::demo::sendprop::{SendProp, SendPropIdentifier, SendPropValue};
13use crate::{Parse, ParseError, ParserState, ReadResult, Result, Stream};
14use parse_display::{Display, FromStr};
15use std::cmp::{min, Ordering};
16use std::collections::HashSet;
17
18use crate::demo::data::ServerTick;
19use itertools::Either;
20use std::fmt;
21#[cfg(feature = "trace")]
22use tracing::trace;
23
24#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
25#[derive(
26    Debug,
27    Copy,
28    Clone,
29    PartialEq,
30    Eq,
31    Hash,
32    Serialize,
33    Deserialize,
34    Display,
35    Ord,
36    PartialOrd,
37    FromStr,
38    Default,
39)]
40pub struct EntityId(u32);
41
42impl From<u32> for EntityId {
43    fn from(num: u32) -> Self {
44        EntityId(num)
45    }
46}
47
48impl From<EntityId> for u32 {
49    fn from(id: EntityId) -> Self {
50        id.0
51    }
52}
53
54impl From<usize> for EntityId {
55    fn from(num: usize) -> Self {
56        EntityId(num as u32)
57    }
58}
59
60impl From<EntityId> for usize {
61    fn from(id: EntityId) -> Self {
62        id.0 as usize
63    }
64}
65
66impl PartialEq<u32> for EntityId {
67    fn eq(&self, other: &u32) -> bool {
68        self.0 == *other
69    }
70}
71
72impl PartialOrd<u32> for EntityId {
73    fn partial_cmp(&self, other: &u32) -> Option<Ordering> {
74        self.0.partial_cmp(other)
75    }
76}
77
78#[cfg_attr(feature = "schema", derive(schemars::JsonSchema_repr))]
79#[derive(
80    BitRead, BitWrite, Clone, Copy, Debug, PartialEq, Eq, Serialize_repr, Deserialize_repr,
81)]
82#[discriminant_bits = 2]
83#[repr(u8)]
84pub enum UpdateType {
85    Preserve = 0b00,
86    Leave = 0b01,
87    Enter = 0b10,
88    Delete = 0b11,
89}
90
91#[cfg_attr(feature = "schema", derive(schemars::JsonSchema_repr))]
92#[derive(Debug, Eq, PartialEq, Copy, Clone, Serialize_repr, Deserialize_repr, Default)]
93#[repr(u8)]
94pub enum BaselineIndex {
95    #[default]
96    First = 0,
97    Second = 1,
98}
99
100impl BaselineIndex {
101    pub fn other(self) -> Self {
102        match self {
103            BaselineIndex::First => BaselineIndex::Second,
104            BaselineIndex::Second => BaselineIndex::First,
105        }
106    }
107}
108
109impl From<bool> for BaselineIndex {
110    fn from(value: bool) -> Self {
111        match value {
112            false => BaselineIndex::First,
113            true => BaselineIndex::Second,
114        }
115    }
116}
117
118impl<E: Endianness> BitRead<'_, E> for BaselineIndex {
119    fn read(stream: &mut BitReadStream<'_, E>) -> ReadResult<Self> {
120        bool::read(stream).map(BaselineIndex::from)
121    }
122
123    unsafe fn read_unchecked(stream: &mut BitReadStream<'_, E>, end: bool) -> ReadResult<Self> {
124        bool::read_unchecked(stream, end).map(BaselineIndex::from)
125    }
126
127    fn bit_size() -> Option<usize> {
128        Some(1)
129    }
130}
131
132impl<E: Endianness> BitWrite<E> for BaselineIndex {
133    fn write(&self, stream: &mut BitWriteStream<E>) -> ReadResult<()> {
134        let val = match self {
135            BaselineIndex::First => false,
136            BaselineIndex::Second => true,
137        };
138        bool::write(&val, stream)
139    }
140}
141
142#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
143#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
144pub struct PacketEntity {
145    pub server_class: ClassId,
146    pub entity_index: EntityId,
147    pub props: Vec<SendProp>,
148    pub in_pvs: bool,
149    pub update_type: UpdateType,
150    pub serial_number: u32,
151    pub delay: Option<f32>,
152    pub delta: Option<ServerTick>,
153    pub baseline_index: BaselineIndex,
154}
155
156impl fmt::Display for PacketEntity {
157    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
158        writeln!(f, "{}({}) {{", self.entity_index, self.server_class)?;
159        for child in self.props.iter() {
160            writeln!(f, "\t{}", child)?;
161        }
162        write!(f, "}}")
163    }
164}
165
166impl PacketEntity {
167    fn mut_prop_by_identifier(&mut self, index: &SendPropIdentifier) -> Option<&mut SendProp> {
168        self.props.iter_mut().find(|prop| prop.identifier == *index)
169    }
170
171    pub fn get_prop_by_identifier(
172        &self,
173        index: &SendPropIdentifier,
174        parser_state: &ParserState,
175    ) -> Option<SendProp> {
176        self.props(parser_state)
177            .find(|prop| prop.identifier == *index)
178    }
179
180    pub fn apply_update(&mut self, props: &[SendProp]) {
181        for prop in props {
182            match self.mut_prop_by_identifier(&prop.identifier) {
183                Some(existing_prop) => existing_prop.value = prop.value.clone(),
184                None => self.props.push(prop.clone()),
185            }
186        }
187    }
188
189    pub fn get_prop_by_name(
190        &self,
191        table_name: &str,
192        name: &str,
193        parser_state: &ParserState,
194    ) -> Option<SendProp> {
195        let identifier = SendPropIdentifier::new(table_name, name);
196        self.get_prop_by_identifier(&identifier, parser_state)
197    }
198
199    pub fn get_baseline_props<'a>(&self, parser_state: &'a ParserState) -> Cow<'a, [SendProp]> {
200        let Some(send_table) = parser_state.send_tables.get(usize::from(self.server_class)) else {
201            return Cow::default();
202        };
203        parser_state
204            .get_baseline(
205                self.baseline_index,
206                self.entity_index,
207                self.server_class,
208                send_table,
209                self.delta.is_some(),
210            )
211            .unwrap_or_default()
212    }
213
214    pub fn props<'a>(
215        &'a self,
216        parser_state: &'a ParserState,
217    ) -> impl Iterator<Item = SendProp> + 'a {
218        if self.update_type == UpdateType::Enter {
219            let mut found_props = HashSet::<SendPropIdentifier>::new();
220            let props = self.props.iter().cloned();
221            #[allow(clippy::unnecessary_to_owned)]
222            let baseline_props = self
223                .get_baseline_props(parser_state)
224                .into_owned()
225                .into_iter();
226            Either::Left(props.chain(baseline_props).filter(move |prop| {
227                let found = found_props.contains(&prop.identifier);
228                found_props.insert(prop.identifier);
229                !found
230            }))
231        } else {
232            Either::Right(self.props.iter().cloned())
233        }
234    }
235}
236
237fn read_bit_var<'a, T: BitReadSized<'a, LittleEndian>>(stream: &mut Stream<'a>) -> ReadResult<T> {
238    let ty: u8 = stream.read_sized(2)?;
239
240    let bits = match ty {
241        0 => 4,
242        1 => 8,
243        2 => 12,
244        _ => 32,
245    };
246    stream.read_sized(bits)
247}
248
249fn write_bit_var(var: u32, stream: &mut BitWriteStream<LittleEndian>) -> ReadResult<()> {
250    let (ty, bits): (u8, usize) = if var >= 2u32.pow(12) {
251        (3, 32)
252    } else if var >= 2u32.pow(8) {
253        (2, 12)
254    } else if var >= 2u32.pow(4) {
255        (1, 8)
256    } else {
257        (0, 4)
258    };
259
260    ty.write_sized(stream, 2)?;
261    var.write_sized(stream, bits)
262}
263
264#[test]
265fn test_bit_var_roundtrip() {
266    use bitbuffer::{BitReadBuffer, BitReadStream};
267
268    fn bit_var_normal(val: u32) {
269        let mut data = Vec::with_capacity(16);
270        let pos = {
271            let mut write = BitWriteStream::new(&mut data, LittleEndian);
272            write_bit_var(val, &mut write).unwrap();
273            write.bit_len()
274        };
275        let mut read = BitReadStream::new(BitReadBuffer::new(&data, LittleEndian));
276        assert_eq!(val, read_bit_var::<u32>(&mut read).unwrap());
277        assert_eq!(pos, read.pos());
278    }
279    bit_var_normal(0);
280    bit_var_normal(1);
281    bit_var_normal(24);
282    bit_var_normal(1234);
283    bit_var_normal(12345);
284    bit_var_normal(123456);
285    bit_var_normal(1234567);
286    bit_var_normal(12345678);
287    bit_var_normal(123456789);
288    bit_var_normal(u32::MAX);
289
290    for i in 0..31 {
291        bit_var_normal(2u32.pow(i));
292        bit_var_normal(2u32.pow(i) - 1);
293        bit_var_normal(2u32.pow(i) + 1);
294    }
295}
296
297#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
298#[derive(Debug, PartialEq, Serialize, Deserialize, Clone, Default)]
299pub struct PacketEntitiesMessage {
300    pub entities: Vec<PacketEntity>,
301    pub removed_entities: Vec<EntityId>,
302    pub max_entries: u16,
303    pub delta: Option<ServerTick>,
304    pub base_line: BaselineIndex,
305    pub updated_base_line: bool,
306}
307
308fn get_send_table(state: &ParserState, class: ClassId) -> Result<&SendTable> {
309    state
310        .send_tables
311        .get(usize::from(class))
312        .ok_or(ParseError::UnknownServerClass(class))
313}
314
315fn get_entity_for_update(
316    state: &ParserState,
317    entity_index: EntityId,
318    update_type: UpdateType,
319    delta: Option<ServerTick>,
320) -> Result<PacketEntity> {
321    let class_id = state
322        .entity_classes
323        .get(&entity_index)
324        .copied()
325        .ok_or(ParseError::UnknownEntity(entity_index))?;
326
327    Ok(PacketEntity {
328        server_class: class_id,
329        entity_index,
330        props: Vec::with_capacity(8),
331        in_pvs: false,
332        update_type,
333        serial_number: 0,
334        delay: None,
335        delta,
336        baseline_index: BaselineIndex::First,
337    })
338}
339
340impl Parse<'_> for PacketEntitiesMessage {
341    fn parse(stream: &mut Stream, state: &ParserState) -> Result<Self> {
342        let max_entries = stream.read_sized(11)?;
343        let delta: Option<ServerTick> = stream.read()?;
344        let base_line = stream.read()?;
345        let updated_entries: u16 = stream.read_sized(11)?;
346        let length: u32 = stream.read_sized(20)?;
347        let updated_base_line = stream.read()?;
348
349        let mut data = stream.read_bits(length as usize)?;
350
351        let mut entities = Vec::with_capacity(min(updated_entries, 128) as usize);
352        let mut removed_entities = Vec::new();
353
354        let mut last_index: i32 = -1;
355
356        for _ in 0..updated_entries {
357            let diff: u32 = read_bit_var(&mut data)?;
358            let index = last_index.saturating_add(diff as i32).saturating_add(1);
359            if index >= 2048 {
360                return Err(ParseError::InvalidDemo("invalid entity index"));
361            }
362            let entity_index = EntityId::from(index as u32);
363
364            let update_type = data.read()?;
365
366            if update_type == UpdateType::Enter {
367                let mut entity =
368                    Self::read_enter(&mut data, entity_index, state, base_line, delta)?;
369                let send_table = get_send_table(state, entity.server_class)?;
370                Self::read_update(&mut data, send_table, &mut entity.props, entity_index)?;
371
372                entities.push(entity);
373            } else if update_type == UpdateType::Preserve {
374                let mut entity = get_entity_for_update(state, entity_index, update_type, delta)?;
375                let send_table = get_send_table(state, entity.server_class)?;
376
377                Self::read_update(&mut data, send_table, &mut entity.props, entity_index)?;
378                entity.in_pvs = true;
379
380                entities.push(entity);
381            } else if state.entity_classes.contains_key(&entity_index) {
382                let entity = get_entity_for_update(state, entity_index, update_type, delta)?;
383                entities.push(entity);
384            } else {
385                entities.push(PacketEntity {
386                    server_class: 0.into(),
387                    entity_index,
388                    props: vec![],
389                    in_pvs: false,
390                    update_type,
391                    serial_number: 0,
392                    delay: None,
393                    delta,
394                    baseline_index: BaselineIndex::First,
395                });
396            }
397
398            last_index = index;
399        }
400
401        if delta.is_some() {
402            while data.read()? {
403                removed_entities.push(data.read_sized::<u32>(11)?.into())
404            }
405        }
406
407        Ok(PacketEntitiesMessage {
408            entities,
409            removed_entities,
410            max_entries,
411            delta,
412            base_line,
413            updated_base_line,
414        })
415    }
416}
417
418impl Encode for PacketEntitiesMessage {
419    fn encode(&self, stream: &mut BitWriteStream<LittleEndian>, state: &ParserState) -> Result<()> {
420        self.max_entries.write_sized(stream, 11)?;
421        self.delta.is_some().write(stream)?;
422        if let Some(delta) = self.delta {
423            delta.write(stream)?;
424        }
425        self.base_line.write(stream)?;
426        self.entities.len().write_sized(stream, 11)?;
427
428        stream.reserve_int(20, |stream| {
429            self.updated_base_line.write(stream)?;
430
431            let length_start = stream.bit_len();
432
433            let mut last_index: i32 = -1;
434
435            for entity in self.entities.iter() {
436                let diff = entity.entity_index.0 as i32 - last_index - 1;
437                write_bit_var(diff as u32, stream)?;
438                last_index = entity.entity_index.0 as i32;
439
440                entity.update_type.write(stream)?;
441
442                let send_table = get_send_table(state, entity.server_class)?;
443                match entity.update_type {
444                    UpdateType::Enter => {
445                        Self::write_enter(entity, stream, state)?;
446                        Self::write_update(&entity.props, stream, send_table, entity.entity_index)?;
447                    }
448                    UpdateType::Preserve => {
449                        Self::write_update(&entity.props, stream, send_table, entity.entity_index)?;
450                    }
451                    _ => {}
452                }
453            }
454
455            if self.delta.is_some() {
456                for removed in self.removed_entities.iter() {
457                    true.write(stream)?;
458                    removed.0.write_sized(stream, 11)?;
459                }
460                false.write(stream)?;
461            }
462
463            let length_end = stream.bit_len();
464
465            Ok((length_end - length_start) as u64)
466        })
467    }
468}
469
470impl PacketEntitiesMessage {
471    fn read_enter(
472        stream: &mut Stream,
473        entity_index: EntityId,
474        state: &ParserState,
475        baseline_index: BaselineIndex,
476        delta: Option<ServerTick>,
477    ) -> Result<PacketEntity> {
478        let bits = log_base2(state.server_classes.len()) + 1;
479        let class_index: ClassId = stream.read_sized::<u16>(bits as usize)?.into();
480
481        let serial = stream.read_sized(10)?;
482
483        Ok(PacketEntity {
484            server_class: class_index,
485            entity_index,
486            props: vec![],
487            in_pvs: true,
488            update_type: UpdateType::Enter,
489            serial_number: serial,
490            delay: None,
491            delta,
492            baseline_index,
493        })
494    }
495
496    fn write_enter(
497        entity: &PacketEntity,
498        stream: &mut BitWriteStream<LittleEndian>,
499        state: &ParserState,
500    ) -> Result<()> {
501        let bits = log_base2(state.server_classes.len()) + 1;
502        u16::from(entity.server_class).write_sized(stream, bits as usize)?;
503        entity.serial_number.write_sized(stream, 10)?;
504
505        Ok(())
506    }
507
508    pub fn read_update(
509        stream: &mut Stream,
510        send_table: &SendTable,
511        props: &mut Vec<SendProp>,
512        entity_index: EntityId,
513    ) -> Result<()> {
514        let mut index: i32 = -1;
515
516        #[cfg(feature = "trace")]
517        trace!(entity_index = display(entity_index), "reading update");
518        #[cfg(not(feature = "trace"))]
519        let _ = entity_index;
520
521        while stream.read()? {
522            let diff: u32 = read_bit_var(stream)?;
523            index = index.saturating_add(diff as i32).saturating_add(1);
524
525            match send_table.flattened_props.get(index as usize) {
526                Some(definition) => {
527                    let value = SendPropValue::parse(stream, &definition.parse_definition)?;
528
529                    #[cfg(feature = "trace")]
530                    trace!(
531                        entity_index = display(entity_index),
532                        index = display(index),
533                        value = debug(&value),
534                        definition = display(definition.identifier),
535                        "reading prop"
536                    );
537                    props.push(SendProp {
538                        index: index as u32,
539                        identifier: definition.identifier,
540                        value,
541                    });
542                }
543                None => {
544                    return Err(ParseError::PropIndexOutOfBounds {
545                        index,
546                        prop_count: send_table.flattened_props.len(),
547                        table: send_table.name.to_string(),
548                    });
549                }
550            }
551        }
552
553        Ok(())
554    }
555
556    pub fn write_update<'a, Props: IntoIterator<Item = &'a SendProp>>(
557        props: Props,
558        stream: &mut BitWriteStream<LittleEndian>,
559        send_table: &SendTable,
560        _entity_index: EntityId,
561    ) -> Result<()> {
562        let mut last_index: i32 = -1;
563
564        let mut props: Vec<&SendProp> = props.into_iter().collect();
565        props.sort_by(|a, b| a.index.cmp(&b.index));
566
567        for prop in props {
568            true.write(stream)?;
569
570            let index = prop.index as usize;
571
572            let definition =
573                send_table
574                    .flattened_props
575                    .get(index)
576                    .ok_or(ParseError::PropIndexOutOfBounds {
577                        index: index as i32,
578                        prop_count: send_table.flattened_props.len(),
579                        table: send_table.name.to_string(),
580                    })?;
581            write_bit_var((index as i32 - last_index - 1) as u32, stream)?;
582            last_index = index as i32;
583            prop.value.encode(stream, &definition.parse_definition)?;
584        }
585        false.write(stream)?;
586        Ok(())
587    }
588}
589
590impl ParseBitSkip<'_> for PacketEntitiesMessage {
591    fn parse_skip(stream: &mut Stream, _state: &ParserState) -> Result<()> {
592        stream.skip_bits(11)?;
593        if stream.read()? {
594            stream.skip_bits(32)?;
595        }
596        stream.skip_bits(12)?;
597        let length: u32 = stream.read_sized(20)?;
598        stream
599            .skip_bits(length as usize + 1)
600            .map_err(ParseError::from)
601    }
602}
603
604#[test]
605fn test_packet_entitier_message_roundtrip() {
606    use crate::demo::packet::datatable::{SendTable, SendTableName, ServerClass, ServerClassName};
607    use crate::demo::sendprop::{FloatDefinition, SendPropDefinition, SendPropParseDefinition};
608
609    let mut state = ParserState::new(24, |_| false, false);
610    state.server_classes = vec![
611        ServerClass {
612            id: ClassId::from(0),
613            name: ServerClassName::from("class1"),
614            data_table: SendTableName::from("table1"),
615        },
616        ServerClass {
617            id: ClassId::from(1),
618            name: ServerClassName::from("class2"),
619            data_table: SendTableName::from("table2"),
620        },
621    ];
622    state.send_tables = vec![
623        SendTable {
624            name: SendTableName::from("table1"),
625            needs_decoder: false,
626            flattened_props: vec![],
627        },
628        SendTable {
629            name: SendTableName::from("table2"),
630            needs_decoder: false,
631            flattened_props: vec![
632                SendPropDefinition {
633                    identifier: SendPropIdentifier::new("table2", "prop1"),
634                    parse_definition: SendPropParseDefinition::Int {
635                        changes_often: false,
636                        bit_count: 8,
637                    },
638                },
639                SendPropDefinition {
640                    identifier: SendPropIdentifier::new("table2", "prop2"),
641                    parse_definition: SendPropParseDefinition::String {
642                        changes_often: false,
643                    },
644                },
645                SendPropDefinition {
646                    identifier: SendPropIdentifier::new("table2", "prop3"),
647                    parse_definition: SendPropParseDefinition::Float {
648                        changes_often: false,
649                        definition: FloatDefinition::Coord,
650                    },
651                },
652            ],
653        },
654    ];
655    state
656        .entity_classes
657        .insert(EntityId::from(4u32), ClassId::from(1));
658    crate::test_roundtrip_encode(
659        PacketEntitiesMessage {
660            entities: vec![],
661            removed_entities: vec![],
662            max_entries: 0,
663            delta: None,
664            base_line: BaselineIndex::First,
665            updated_base_line: false,
666        },
667        &state,
668    );
669    crate::test_roundtrip_encode(
670        PacketEntitiesMessage {
671            entities: vec![PacketEntity {
672                server_class: ClassId::from(0),
673                entity_index: Default::default(),
674                props: vec![],
675                in_pvs: true,
676                update_type: UpdateType::Enter,
677                serial_number: 0,
678                delay: None,
679                delta: None,
680                baseline_index: BaselineIndex::First,
681            }],
682            removed_entities: vec![],
683            max_entries: 4,
684            delta: None,
685            base_line: BaselineIndex::First,
686            updated_base_line: false,
687        },
688        &state,
689    );
690    crate::test_roundtrip_encode(
691        PacketEntitiesMessage {
692            entities: vec![
693                PacketEntity {
694                    server_class: ClassId::from(0),
695                    entity_index: EntityId::from(0u32),
696                    props: vec![],
697                    in_pvs: true,
698                    update_type: UpdateType::Enter,
699                    serial_number: 0,
700                    delay: None,
701                    delta: None,
702                    baseline_index: BaselineIndex::First,
703                },
704                PacketEntity {
705                    server_class: ClassId::from(1),
706                    entity_index: EntityId::from(4u32),
707                    props: vec![
708                        SendProp {
709                            index: 0,
710                            identifier: SendPropIdentifier::new("table2", "prop1"),
711                            value: SendPropValue::Integer(4),
712                        },
713                        SendProp {
714                            index: 2,
715                            identifier: SendPropIdentifier::new("table2", "prop3"),
716                            value: SendPropValue::Float(1.0),
717                        },
718                    ],
719                    in_pvs: true,
720                    update_type: UpdateType::Preserve,
721                    serial_number: 0,
722                    delay: None,
723                    delta: None,
724                    baseline_index: BaselineIndex::First,
725                },
726                PacketEntity {
727                    server_class: ClassId::from(1),
728                    entity_index: EntityId::from(5u32),
729                    delta: None,
730                    baseline_index: BaselineIndex::First,
731                    props: vec![
732                        SendProp {
733                            index: 0,
734                            identifier: SendPropIdentifier::new("table2", "prop1"),
735                            value: SendPropValue::Integer(4),
736                        },
737                        SendProp {
738                            index: 2,
739                            identifier: SendPropIdentifier::new("table2", "prop3"),
740                            value: SendPropValue::Float(1.0),
741                        },
742                    ],
743                    in_pvs: true,
744                    update_type: UpdateType::Enter,
745                    serial_number: 0,
746                    delay: None,
747                },
748            ],
749            removed_entities: vec![],
750            max_entries: 4,
751            delta: None,
752            base_line: BaselineIndex::First,
753            updated_base_line: false,
754        },
755        &state,
756    );
757}