message type testing

src/demo/message/mod.rs

@@ -2,6 +2,8 @@ use enum_primitive_derive::Primitive;
 use num_traits::FromPrimitive;
 
 pub use generated::*;
+use serde::{Deserialize, Serialize};
+use serde_repr::{Deserialize_repr, Serialize_repr};
 
 use crate::{Parse, ParseError, ParserState, Result, Stream};
 use crate::demo::message::bspdecal::*;
@@ -24,7 +26,8 @@ pub mod tempentities;
 pub mod usermessage;
 pub mod voice;
 
-#[derive(Primitive, Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Primitive, Debug, Clone, Copy, PartialEq, Eq, Serialize_repr, Deserialize_repr)]
+#[repr(u8)]
 pub enum MessageType {
     Empty = 0,
     File = 2,

src/demo/parser/handler.rs

@@ -1,8 +1,8 @@
 use crate::demo::message::{Message, MessageType};
 use crate::demo::packet::datatable::{SendTable, ServerClass};
-use crate::demo::packet::stringtable::{StringTable, StringTableEntry};
 use crate::demo::packet::Packet;
-use crate::demo::parser::analyser::{Analyser};
+use crate::demo::packet::stringtable::{StringTable, StringTableEntry};
+use crate::demo::parser::analyser::Analyser;
 use crate::ParserState;
 
 pub trait MessageHandler {
@@ -17,7 +17,6 @@ pub trait MessageHandler {
     fn get_output(self, state: ParserState) -> Self::Output;
 }
 
-#[derive(Default)]
 pub struct DemoHandler<T: MessageHandler> {
     tick: u32,
     string_table_names: Vec<String>,
@@ -33,7 +32,7 @@ impl DemoHandler<Analyser> {
 
 impl<T: MessageHandler> DemoHandler<T> {
     pub fn with_analyser(analyser: T) -> Self {
-        let state_handler = ParserState::new();
+        let state_handler = ParserState::new(T::does_handle);
 
         DemoHandler {
             tick: 0,

src/demo/parser/messagetypeanalyser.rs

@@ -0,0 +1,46 @@
+use std::collections::HashMap;
+
+use serde::{Deserialize, Serialize};
+use serde_repr::{Deserialize_repr, Serialize_repr};
+
+use crate::{ParserState, ReadResult, Stream};
+use crate::demo::gameevent_gen::{
+    GameEvent, PlayerDeathEvent, PlayerSpawnEvent, TeamPlayRoundWinEvent,
+};
+use crate::demo::message::{Message, MessageType};
+use crate::demo::message::packetentities::EntityId;
+use crate::demo::message::usermessage::{ChatMessageKind, SayText2Message, UserMessage};
+use crate::demo::packet::PacketType;
+use crate::demo::packet::stringtable::StringTableEntry;
+use crate::demo::parser::handler::MessageHandler;
+use crate::demo::vector::Vector;
+
+pub struct MessageTypeAnalyser {
+    packet_types: Vec<MessageType>
+}
+
+impl MessageHandler for MessageTypeAnalyser {
+    type Output = Vec<MessageType>;
+
+    fn does_handle(message_type: MessageType) -> bool {
+        true
+    }
+
+    fn handle_message(&mut self, message: Message, tick: u32) {
+        self.packet_types.push(message.get_message_type())
+    }
+
+    fn handle_string_entry(&mut self, table: &String, _index: usize, entry: &StringTableEntry) {}
+
+    fn get_output(self, state: ParserState) -> Self::Output {
+        self.packet_types
+    }
+}
+
+impl MessageTypeAnalyser {
+    pub fn new() -> Self {
+        MessageTypeAnalyser {
+            packet_types: Vec::with_capacity(1024)
+        }
+    }
+}

src/demo/parser/mod.rs

@@ -8,10 +8,12 @@ pub use crate::demo::parser::analyser::MatchState;
 use crate::demo::parser::handler::{DemoHandler, MessageHandler};
 pub use crate::demo::parser::state::ParserState;
 use crate::Stream;
+pub use self::messagetypeanalyser::MessageTypeAnalyser;
 
 mod analyser;
 mod handler;
 mod state;
+mod messagetypeanalyser;
 
 /// Errors that can occur during parsing
 #[derive(Debug)]

src/demo/parser/state.rs

@@ -7,10 +7,10 @@ use crate::demo::message::packetentities::EntityId;
 use crate::demo::message::stringtable::StringTableMeta;
 use crate::demo::packet::datatable::{SendTable, ServerClass};
 use crate::demo::packet::stringtable::StringTableEntry;
-use crate::demo::parser::handler::{MessageHandler};
+use crate::demo::parser::analyser::Analyser;
+use crate::demo::parser::handler::MessageHandler;
 use crate::demo::sendprop::SendProp;
 use crate::Stream;
-use crate::demo::parser::analyser::Analyser;
 
 #[derive(Default)]
 pub struct DemoMeta {
@@ -19,7 +19,6 @@ pub struct DemoMeta {
     pub interval_per_tick: f32,
 }
 
-#[derive(Default)]
 pub struct ParserState {
     pub static_baselines: HashMap<u32, StaticBaseline>,
     pub event_definitions: HashMap<GameEventTypeId, GameEventDefinition>,
@@ -29,6 +28,7 @@ pub struct ParserState {
     pub server_classes: Vec<ServerClass>,
     pub instance_baselines: [HashMap<EntityId, Vec<SendProp>>; 2],
     pub demo_meta: DemoMeta,
+    analyser_handles: fn(message_type: MessageType) -> bool,
 }
 
 pub struct StaticBaseline {
@@ -48,8 +48,18 @@ impl StaticBaseline {
 }
 
 impl ParserState {
-    pub fn new() -> Self {
+    pub fn new(analyser_handles: fn(message_type: MessageType) -> bool) -> Self {
-        ParserState::default()
+        ParserState {
+            static_baselines: HashMap::new(),
+            event_definitions: HashMap::new(),
+            string_tables: Vec::new(),
+            entity_classes: HashMap::new(),
+            send_tables: HashMap::new(),
+            server_classes: Vec::new(),
+            instance_baselines: [HashMap::new(), HashMap::new()],
+            demo_meta: DemoMeta::default(),
+            analyser_handles,
+        }
     }
 
     pub fn handle_data_table(
@@ -68,7 +78,7 @@ impl ParserState {
     }
 
     pub fn should_parse_message(&self, message_type: MessageType) -> bool {
-        Self::does_handle(message_type) || Analyser::does_handle(message_type)
+        Self::does_handle(message_type) || (self.analyser_handles)(message_type)
     }
 }
 

src/lib.rs

@@ -1,10 +1,11 @@
-pub use crate::demo::{
-    parser::{
-        DemoParser, Parse, ParseError, ParserState, Result, MatchState
-    },
-    Demo, Stream,
-};
-
 pub use bitstream_reader::Result as ReadResult;
 
+pub use crate::demo::{
+    Demo,
+    message::MessageType, parser::{
+        DemoParser, MatchState, MessageTypeAnalyser, Parse, ParseError, ParserState, Result,
+    },
+    Stream,
+};
+
 mod demo;

tests/tests.rs

@@ -1,7 +1,7 @@
 use std::fs;
 use pretty_assertions::assert_eq;
 
-use tf_demo_parser::{Demo, DemoParser, MatchState};
+use tf_demo_parser::{Demo, DemoParser, MatchState, MessageTypeAnalyser, MessageType};
 
 fn snapshot_test(input_file: &str, snapshot_file: &str) {
     let file = fs::read(input_file).expect("Unable to read file");
@@ -12,6 +12,23 @@ fn snapshot_test(input_file: &str, snapshot_file: &str) {
     assert_eq!(expected, state);
 }
 
+fn test_message_types(input_file: &str, snapshot_file: &str) {
+    let file = fs::read(input_file).expect("Unable to read file");
+    let demo = Demo::new(file);
+    let (_, message_types) = DemoParser::parse_with_analyser(demo.get_stream(), MessageTypeAnalyser::new()).unwrap();
+
+    let expected: Vec<MessageType> = serde_json::from_slice(fs::read(snapshot_file).expect("Unable to read file").as_slice()).unwrap();
+    assert_eq!(expected, message_types);
+}
+
+fn dump_message_types(input_file: &str, snapshot_file: &str) {
+    let file = fs::read(input_file).expect("Unable to read file");
+    let demo = Demo::new(file);
+    let (_, message_types) = DemoParser::parse_with_analyser(demo.get_stream(), MessageTypeAnalyser::new()).unwrap();
+
+    fs::write(snapshot_file, serde_json::to_vec(&message_types).unwrap()).unwrap();
+}
+
 #[test]
 fn snapshot_test_small() {
     snapshot_test("data/small.dem", "data/small.json");
@@ -26,3 +43,8 @@ fn snapshot_test_gully() {
 fn snapshot_test_comp() {
     snapshot_test("data/comp.dem", "data/comp.json");
 }
+
+#[test]
+fn message_type_test_comp() {
+    dump_message_types("data/comp.dem", "data/comp_message_types.json");
+}