tf_demo_parser/demo/parser/
analyser.rs

1use crate::demo::data::{DemoTick, ServerTick};
2use crate::demo::gameevent_gen::{
3    GameEvent, PlayerDeathEvent, PlayerSpawnEvent, TeamPlayRoundWinEvent,
4};
5use crate::demo::message::packetentities::EntityId;
6use crate::demo::message::usermessage::{
7    ChatMessageKind, HudTextLocation, SayText2Message, TextMessage, UserMessage,
8};
9use crate::demo::message::{Message, MessageType};
10use crate::demo::packet::stringtable::StringTableEntry;
11use crate::demo::parser::handler::{BorrowMessageHandler, MessageHandler};
12use crate::demo::vector::Vector;
13use crate::{ParserState, ReadResult, Stream};
14use bitbuffer::{BitWrite, BitWriteStream, Endianness};
15use num_enum::TryFromPrimitive;
16use parse_display::{Display, FromStr};
17use serde::de::Error;
18use serde::{ser::SerializeMap, Deserialize, Deserializer, Serialize, Serializer};
19use std::collections::{BTreeMap, HashMap};
20use std::convert::TryFrom;
21
22#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
23pub struct ChatMessage {
24    pub kind: ChatMessageKind,
25    pub from: String,
26    pub text: String,
27    pub tick: DemoTick,
28}
29
30impl ChatMessage {
31    pub fn from_message(message: &SayText2Message, tick: DemoTick) -> Self {
32        ChatMessage {
33            kind: message.kind,
34            from: message
35                .from
36                .as_ref()
37                .map(|s| s.to_string())
38                .unwrap_or_default(),
39            text: message.plain_text(),
40            tick,
41        }
42    }
43
44    pub fn from_text(message: &TextMessage, tick: DemoTick) -> Self {
45        ChatMessage {
46            kind: ChatMessageKind::Empty,
47            from: String::new(),
48            text: message.plain_text(),
49            tick,
50        }
51    }
52}
53
54#[derive(
55    Debug,
56    Clone,
57    Serialize,
58    Deserialize,
59    Copy,
60    PartialEq,
61    Eq,
62    Hash,
63    TryFromPrimitive,
64    Default,
65    Display,
66)]
67#[serde(rename_all = "lowercase")]
68#[repr(u8)]
69pub enum Team {
70    #[default]
71    Other = 0,
72    Spectator = 1,
73    Red = 2,
74    Blue = 3,
75}
76
77impl Team {
78    pub fn new<U>(number: U) -> Self
79    where
80        u8: TryFrom<U>,
81    {
82        Team::try_from(u8::try_from(number).unwrap_or_default()).unwrap_or_default()
83    }
84
85    pub fn is_player(&self) -> bool {
86        *self == Team::Red || *self == Team::Blue
87    }
88}
89
90#[derive(
91    Debug, Clone, Serialize, Copy, PartialEq, Eq, Hash, TryFromPrimitive, Display, FromStr, Default,
92)]
93#[display(style = "lowercase")]
94#[serde(rename_all = "lowercase")]
95#[repr(u8)]
96pub enum Class {
97    #[default]
98    Other = 0,
99    Scout = 1,
100    Sniper = 2,
101    Soldier = 3,
102    Demoman = 4,
103    Medic = 5,
104    Heavy = 6,
105    Pyro = 7,
106    Spy = 8,
107    Engineer = 9,
108}
109
110impl<'de> Deserialize<'de> for Class {
111    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
112    where
113        D: Deserializer<'de>,
114    {
115        #[derive(Deserialize, Debug)]
116        #[serde(untagged)]
117        enum IntOrStr<'a> {
118            Int(u8),
119            Str(&'a str),
120        }
121
122        let raw = IntOrStr::deserialize(deserializer)?;
123        match raw {
124            IntOrStr::Int(class) => Class::try_from_primitive(class).map_err(D::Error::custom),
125            IntOrStr::Str(class) if class.len() == 1 => {
126                Class::try_from_primitive(class.parse().map_err(D::Error::custom)?)
127                    .map_err(D::Error::custom)
128            }
129            IntOrStr::Str(class) => class.parse().map_err(D::Error::custom),
130        }
131    }
132}
133
134#[test]
135fn test_class_deserialize() {
136    assert_eq!(Class::Scout, serde_json::from_str(r#""scout""#).unwrap());
137    assert_eq!(Class::Scout, serde_json::from_str(r#""1""#).unwrap());
138    assert_eq!(Class::Scout, serde_json::from_str("1").unwrap());
139}
140
141impl Class {
142    pub fn new<U>(number: U) -> Self
143    where
144        u8: TryFrom<U>,
145    {
146        Class::try_from(u8::try_from(number).unwrap_or_default()).unwrap_or_default()
147    }
148}
149
150#[derive(Default, Debug, Eq, PartialEq, Deserialize, Clone)]
151#[serde(from = "HashMap<Class, u16>")]
152pub struct ClassList([u16; 10]);
153
154impl ClassList {
155    /// Get an iterator for all classes played and the number of spawn on the class
156    pub fn iter(&self) -> impl Iterator<Item = (Class, u16)> + '_ {
157        self.0
158            .iter()
159            .copied()
160            .enumerate()
161            .map(|(class, count)| (Class::new(class), count))
162            .filter(|(_, count)| *count > 0)
163    }
164
165    /// Get an iterator for all classes played and the number of spawn on the class, sorted by the number of spawns
166    pub fn sorted(&self) -> impl Iterator<Item = (Class, u16)> {
167        let mut classes = self.iter().collect::<Vec<(Class, u16)>>();
168        classes.sort_by(|a, b| a.1.cmp(&b.1).reverse());
169        classes.into_iter()
170    }
171
172    pub fn get(&self, class: Class) -> u16 {
173        // class number is always in bounds
174        #[allow(clippy::indexing_slicing)]
175        self.0[class as u8 as usize]
176    }
177
178    pub fn get_mut(&mut self, class: Class) -> &mut u16 {
179        // class number is always in bounds
180        #[allow(clippy::indexing_slicing)]
181        &mut self.0[class as u8 as usize]
182    }
183}
184
185#[test]
186fn test_classlist_sorted() {
187    let list = ClassList([0, 1, 5, 0, 0, 3, 0, 0, 0, 0]);
188    assert_eq!(
189        list.sorted().collect::<Vec<_>>(),
190        &[(Class::Sniper, 5), (Class::Medic, 3), (Class::Scout, 1)]
191    )
192}
193
194impl Serialize for ClassList {
195    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
196    where
197        S: Serializer,
198    {
199        let count = self.0.iter().filter(|c| **c > 0).count();
200        let mut classes = serializer.serialize_map(Some(count))?;
201        for (class, count) in self.0.iter().copied().enumerate() {
202            if count > 0 {
203                classes.serialize_entry(&class, &count)?;
204            }
205        }
206
207        classes.end()
208    }
209}
210
211impl From<HashMap<Class, u16>> for ClassList {
212    fn from(map: HashMap<Class, u16>) -> Self {
213        let mut classes = ClassList::default();
214
215        for (class, count) in map.into_iter() {
216            *classes.get_mut(class) = count;
217        }
218
219        classes
220    }
221}
222
223#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
224#[derive(
225    Debug,
226    Clone,
227    Serialize,
228    Deserialize,
229    Copy,
230    PartialEq,
231    Eq,
232    Hash,
233    Ord,
234    PartialOrd,
235    Default,
236    Display,
237)]
238pub struct UserId(u16);
239
240impl<E: Endianness> BitWrite<E> for UserId {
241    fn write(&self, stream: &mut BitWriteStream<E>) -> ReadResult<()> {
242        (self.0 as u32).write(stream)
243    }
244}
245
246impl From<u32> for UserId {
247    fn from(int: u32) -> Self {
248        UserId(int as u16)
249    }
250}
251
252impl From<u16> for UserId {
253    fn from(int: u16) -> Self {
254        UserId(int)
255    }
256}
257
258impl From<UserId> for u16 {
259    fn from(id: UserId) -> Self {
260        id.0
261    }
262}
263
264impl From<UserId> for u32 {
265    fn from(id: UserId) -> Self {
266        id.0 as u32
267    }
268}
269
270impl PartialEq<u16> for UserId {
271    fn eq(&self, other: &u16) -> bool {
272        self.0 == *other
273    }
274}
275
276#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
277pub struct Spawn {
278    pub user: UserId,
279    pub class: Class,
280    pub team: Team,
281    pub tick: DemoTick,
282}
283
284impl Spawn {
285    pub fn from_event(event: &PlayerSpawnEvent, tick: DemoTick) -> Self {
286        Spawn {
287            user: UserId::from(event.user_id),
288            class: Class::new(event.class),
289            team: Team::new(event.team),
290            tick,
291        }
292    }
293}
294
295#[derive(Debug, Clone, Serialize, Deserialize)]
296#[serde(rename_all = "camelCase")]
297pub struct UserInfo {
298    pub classes: ClassList,
299    pub name: String,
300    pub user_id: UserId,
301    pub steam_id: String,
302    #[serde(skip)]
303    pub entity_id: EntityId,
304    pub team: Team,
305}
306
307impl From<crate::demo::data::UserInfo> for UserInfo {
308    fn from(info: crate::demo::data::UserInfo) -> Self {
309        UserInfo {
310            classes: ClassList::default(),
311            name: info.player_info.name,
312            user_id: info.player_info.user_id,
313            steam_id: info.player_info.steam_id,
314            entity_id: info.entity_id,
315            team: Team::default(),
316        }
317    }
318}
319
320impl PartialEq for UserInfo {
321    fn eq(&self, other: &UserInfo) -> bool {
322        self.classes == other.classes
323            && self.name == other.name
324            && self.user_id == other.user_id
325            && self.steam_id == other.steam_id
326            && self.team == other.team
327    }
328}
329
330#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
331pub struct Death {
332    pub weapon: String,
333    pub victim: UserId,
334    pub assister: Option<UserId>,
335    pub killer: UserId,
336    pub tick: DemoTick,
337}
338
339impl Death {
340    pub fn from_event(event: &PlayerDeathEvent, tick: DemoTick) -> Self {
341        let assister = if event.assister < (16 * 1024) {
342            Some(UserId::from(event.assister))
343        } else {
344            None
345        };
346        Death {
347            assister,
348            tick,
349            killer: UserId::from(event.attacker),
350            weapon: event.weapon.to_string(),
351            victim: UserId::from(event.user_id),
352        }
353    }
354}
355
356#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
357pub struct Round {
358    pub winner: Team,
359    pub length: f32,
360    pub end_tick: DemoTick,
361}
362
363impl Round {
364    pub fn from_event(event: &TeamPlayRoundWinEvent, tick: DemoTick) -> Self {
365        Round {
366            winner: Team::new(event.team),
367            length: event.round_time,
368            end_tick: tick,
369        }
370    }
371}
372
373#[derive(Default, Debug, Serialize, Deserialize, PartialEq)]
374pub struct World {
375    pub boundary_min: Vector,
376    pub boundary_max: Vector,
377}
378
379#[derive(Default, Debug, Serialize, Deserialize, PartialEq)]
380pub struct Analyser {
381    state: MatchState,
382    pause_start: Option<DemoTick>,
383    user_id_map: HashMap<EntityId, UserId>,
384}
385
386#[derive(Default, Debug, Serialize, Deserialize, PartialEq)]
387pub struct Pause {
388    from: DemoTick,
389    to: DemoTick,
390}
391
392impl MessageHandler for Analyser {
393    type Output = MatchState;
394
395    fn does_handle(message_type: MessageType) -> bool {
396        matches!(
397            message_type,
398            MessageType::GameEvent
399                | MessageType::UserMessage
400                | MessageType::ServerInfo
401                | MessageType::NetTick
402                | MessageType::SetPause
403        )
404    }
405
406    fn handle_message(&mut self, message: &Message, tick: DemoTick, _parser_state: &ParserState) {
407        match message {
408            Message::NetTick(msg) => {
409                if self.state.start_tick == 0 {
410                    self.state.start_tick = msg.tick;
411                }
412            }
413            Message::ServerInfo(message) => {
414                self.state.interval_per_tick = message.interval_per_tick
415            }
416            Message::GameEvent(message) => self.handle_event(&message.event, tick),
417            Message::UserMessage(message) => self.handle_user_message(message, tick),
418            Message::SetPause(message) => {
419                if message.pause {
420                    self.pause_start = Some(tick);
421                } else {
422                    let start = self.pause_start.unwrap_or_default();
423                    self.state.pauses.push(Pause {
424                        from: start,
425                        to: tick,
426                    })
427                }
428            }
429            _ => {}
430        }
431    }
432
433    fn handle_string_entry(
434        &mut self,
435        table: &str,
436        index: usize,
437        entry: &StringTableEntry,
438        _parser_state: &ParserState,
439    ) {
440        if table == "userinfo" {
441            let _ = self.parse_user_info(
442                index,
443                entry.text.as_ref().map(|s| s.as_ref()),
444                entry.extra_data.as_ref().map(|data| data.data.clone()),
445            );
446        }
447    }
448
449    fn into_output(self, _state: &ParserState) -> Self::Output {
450        self.state
451    }
452}
453
454impl BorrowMessageHandler for Analyser {
455    fn borrow_output(&self, _state: &ParserState) -> &Self::Output {
456        &self.state
457    }
458}
459
460impl Analyser {
461    pub fn new() -> Self {
462        Self::default()
463    }
464
465    fn handle_user_message(&mut self, message: &UserMessage, tick: DemoTick) {
466        match message {
467            UserMessage::SayText2(text_message) => {
468                if text_message.kind == ChatMessageKind::NameChange {
469                    if let Some(from) = text_message.from.clone() {
470                        self.change_name(from.into(), text_message.plain_text());
471                    }
472                } else {
473                    self.state
474                        .chat
475                        .push(ChatMessage::from_message(text_message, tick));
476                }
477            }
478            UserMessage::Text(text_message) => {
479                if text_message.location == HudTextLocation::PrintTalk {
480                    self.state
481                        .chat
482                        .push(ChatMessage::from_text(text_message, tick));
483                }
484            }
485            _ => {}
486        }
487    }
488
489    fn change_name(&mut self, from: String, to: String) {
490        if let Some(user) = self.state.users.values_mut().find(|user| user.name == from) {
491            user.name = to;
492        }
493    }
494
495    fn handle_event(&mut self, event: &GameEvent, tick: DemoTick) {
496        const WIN_REASON_TIME_LIMIT: u8 = 6;
497
498        match event {
499            GameEvent::PlayerDeath(event) => self.state.deaths.push(Death::from_event(event, tick)),
500            GameEvent::PlayerSpawn(event) => {
501                let spawn = Spawn::from_event(event, tick);
502                if let Some(user_state) = self.state.users.get_mut(&spawn.user) {
503                    *user_state.classes.get_mut(spawn.class) += 1;
504                    user_state.team = spawn.team;
505                }
506            }
507            GameEvent::TeamPlayRoundWin(event) => {
508                if event.win_reason != WIN_REASON_TIME_LIMIT {
509                    self.state.rounds.push(Round::from_event(event, tick))
510                }
511            }
512            _ => {}
513        }
514    }
515
516    fn parse_user_info(
517        &mut self,
518        index: usize,
519        text: Option<&str>,
520        data: Option<Stream>,
521    ) -> ReadResult<()> {
522        if let Some(user_info) =
523            crate::demo::data::UserInfo::parse_from_string_table(index as u16, text, data)?
524        {
525            self.state
526                .users
527                .entry(user_info.player_info.user_id)
528                .and_modify(|info| {
529                    info.entity_id = user_info.entity_id;
530                })
531                .or_insert_with(|| user_info.into());
532        }
533
534        Ok(())
535    }
536}
537
538#[derive(Debug, Serialize, Deserialize, PartialEq, Default)]
539#[serde(rename_all = "camelCase")]
540pub struct MatchState {
541    pub chat: Vec<ChatMessage>,
542    pub users: BTreeMap<UserId, UserInfo>,
543    pub deaths: Vec<Death>,
544    pub rounds: Vec<Round>,
545    pub start_tick: ServerTick,
546    pub interval_per_tick: f32,
547    pub pauses: Vec<Pause>,
548}
tf_demo_parser/demo/parser/analyser.rs

tf_demo_parser/demo/parser/
analyser.rs
