Fix formatting and add score/team info to endofgame

analysis/src/endofgame.rs

@@ -1,7 +1,16 @@
+use std::collections::HashMap;
+
 #[derive(Debug, PartialEq)]
 pub struct EndOfGame {
     pub map: String,
     pub players: Vec<(PlayerInfo, PlayerStats)>,
+    pub teams: HashMap<u32, TeamInfo>,
+}
+
+#[derive(Debug, PartialEq)]
+pub struct TeamInfo {
+    pub name: String,
+    pub score: usize,
 }
 
 #[derive(Debug, PartialEq)]
@@ -28,7 +37,7 @@ pub fn parse(buf: &[u8]) -> Result<EndOfGame, ()> {
     let tmp = csdemo::Container::parse(buf).map_err(|e| ())?;
     let output = csdemo::parser::parse(
         csdemo::FrameIterator::parse(tmp.inner),
-        csdemo::parser::EntityFilter::disabled(),
+        csdemo::parser::EntityFilter::all(),
     )
     .map_err(|e| ())?;
 
@@ -69,9 +78,6 @@ pub fn parse(buf: &[u8]) -> Result<EndOfGame, ()> {
                             &mut player_life,
                         );
                     }
-                    csdemo::game_event::GameEvent::PlayerHurt(phurt) => {
-                        // println!("Untracked: {:?}", phurt);
-                    }
                     _ => {}
                 };
             }
@@ -79,6 +85,59 @@ pub fn parse(buf: &[u8]) -> Result<EndOfGame, ()> {
         };
     }
 
+    let mut teams = HashMap::new();
+    let mut entity_to_team = HashMap::new();
+    let mut entity_to_name = HashMap::<_, String>::new();
+    for tick_state in output.entity_states.ticks {
+        for state in tick_state.states {
+            if state.class.as_ref() != "CCSTeam" {
+                continue;
+            }
+
+            let team = match state
+                .get_prop("CCSTeam.m_iTeamNum")
+                .map(|p| p.value.as_u32())
+                .flatten()
+            {
+                Some(team) => {
+                    entity_to_team.insert(state.id, team.clone());
+                    team
+                }
+                None => match entity_to_team.get(&state.id) {
+                    Some(t) => t.clone(),
+                    None => continue,
+                },
+            };
+
+            let name = match entity_to_name.get(&state.id) {
+                Some(n) => n.to_owned(),
+                None => match state
+                    .get_prop("CCSTeam.m_szTeamname")
+                    .map(|p| match &p.value {
+                        csdemo::parser::Variant::String(v) => Some(v.to_owned()),
+                        _ => None,
+                    })
+                    .flatten()
+                {
+                    Some(n) => {
+                        entity_to_name.insert(state.id, n.clone());
+                        n
+                    }
+                    None => continue,
+                },
+            };
+
+            if let Some(score) = state
+                .get_prop("CCSTeam.m_iScore")
+                .map(|p| p.value.as_i32())
+                .flatten()
+                .map(|v| v as usize)
+            {
+                teams.insert(team, TeamInfo { name, score });
+            }
+        }
+    }
+
     let mut players: Vec<_> = player_stats
         .into_iter()
         .filter_map(|(id, stats)| {
@@ -100,7 +159,11 @@ pub fn parse(buf: &[u8]) -> Result<EndOfGame, ()> {
 
     let map = header.map_name().to_owned();
 
-    Ok(EndOfGame { map, players })
+    Ok(EndOfGame {
+        map,
+        players,
+        teams,
+    })
 }
 
 fn player_death(

analysis/src/heatmap.rs

@@ -34,7 +34,7 @@ impl HeatMap {
         }
 
         self.max_y = self.max_y.max(y);
-            let row = self.rows.get_mut(y ).unwrap();
+        let row = self.rows.get_mut(y).unwrap();
 
         if row.len() <= x {
             row.resize(x + 1, 0);
@@ -43,7 +43,6 @@ impl HeatMap {
         self.max_x = self.max_x.max(x);
         let cell = row.get_mut(x).unwrap();
 
-
         *cell += 1;
 
         self.max_value = self.max_value.max(*cell);
@@ -87,26 +86,26 @@ pub fn parse(config: &Config, buf: &[u8]) -> Result<HeatMapOutput, ()> {
     .map_err(|e| ())?;
 
     let pawn_ids = {
-        let mut tmp = std::collections::HashMap::<PawnID,_>::new();
+        let mut tmp = std::collections::HashMap::<PawnID, _>::new();
 
         for event in output.events.iter() {
             let entry = match event {
                 csdemo::DemoEvent::GameEvent(ge) => match ge.as_ref() {
-                    csdemo::game_event::GameEvent::PlayerSpawn(pspawn) => match pspawn.userid_pawn.as_ref() {
+                    csdemo::game_event::GameEvent::PlayerSpawn(pspawn) => {
+                        match pspawn.userid_pawn.as_ref() {
                             Some(csdemo::RawValue::I32(v)) => {
                                 Some((PawnID::from(*v), pspawn.userid.unwrap()))
                             }
-                        _ => {
+                            _ => None,
-                            None
+                        }
-                        },
+                    }
-                    },
                     _ => None,
                 },
                 _ => None,
             };
 
             if let Some((pawn, userid)) = entry {
-                if let Some(previous) = tmp.insert(pawn, userid){
+                if let Some(previous) = tmp.insert(pawn, userid) {
                     assert_eq!(previous, userid);
                 }
             }
@@ -162,15 +161,24 @@ fn process_tick(
         .filter(|s| matches!(s.class.as_ref(), "CCSPlayerPawn" | "CCSTeam"))
     {
         if entity_state.class.as_ref() == "CCSTeam" {
-            let raw_team_name = match entity_state.get_prop("CCSTeam.m_szTeamname").map(|p| match &p.value {
+            let raw_team_name = match entity_state
+                .get_prop("CCSTeam.m_szTeamname")
+                .map(|p| match &p.value {
                     csdemo::parser::Variant::String(v) => Some(v),
                     _ => None,
-            }).flatten() {
+                })
+                .flatten()
+            {
                 Some(n) => n,
                 None => continue,
             };
 
-            for prop in entity_state.props.iter().filter(|p| p.prop_info.prop_name.as_ref() == "CCSTeam.m_aPawns").filter_map(|p| p.value.as_u32().map(|v| PawnID::from(v))) {
+            for prop in entity_state
+                .props
+                .iter()
+                .filter(|p| p.prop_info.prop_name.as_ref() == "CCSTeam.m_aPawns")
+                .filter_map(|p| p.value.as_u32().map(|v| PawnID::from(v)))
+            {
                 teams.insert(prop, raw_team_name.clone());
             }
 
@@ -187,35 +195,67 @@ fn process_tick(
             None => continue,
         };
 
-        let _inner_guard =
+        let _inner_guard = tracing::trace_span!("Entity", entity_id=?entity_state.id).entered();
-            tracing::trace_span!("Entity", entity_id=?entity_state.id).entered();
 
-        let x_cell = match entity_state.get_prop("CCSPlayerPawn.CBodyComponentBaseAnimGraph.m_cellX").map(|prop| prop.value.as_u32()).flatten() {
+        let x_cell = match entity_state
+            .get_prop("CCSPlayerPawn.CBodyComponentBaseAnimGraph.m_cellX")
+            .map(|prop| prop.value.as_u32())
+            .flatten()
+        {
             Some(c) => c,
             None => player_cells.get(&user_id).map(|(x, _, _)| *x).unwrap_or(0),
         };
-        let y_cell = match entity_state.get_prop("CCSPlayerPawn.CBodyComponentBaseAnimGraph.m_cellY").map(|prop| prop.value.as_u32()).flatten() {
+        let y_cell = match entity_state
+            .get_prop("CCSPlayerPawn.CBodyComponentBaseAnimGraph.m_cellY")
+            .map(|prop| prop.value.as_u32())
+            .flatten()
+        {
             Some(c) => c,
             None => player_cells.get(&user_id).map(|(_, y, _)| *y).unwrap_or(0),
         };
-        let z_cell = match entity_state.get_prop("CCSPlayerPawn.CBodyComponentBaseAnimGraph.m_cellZ").map(|prop| prop.value.as_u32()).flatten() {
+        let z_cell = match entity_state
+            .get_prop("CCSPlayerPawn.CBodyComponentBaseAnimGraph.m_cellZ")
+            .map(|prop| prop.value.as_u32())
+            .flatten()
+        {
             Some(c) => c,
             None => player_cells.get(&user_id).map(|(_, _, z)| *z).unwrap_or(0),
         };
 
         player_cells.insert(user_id, (x_cell, y_cell, z_cell));
 
-        let x_coord = match entity_state.get_prop("CCSPlayerPawn.CBodyComponentBaseAnimGraph.m_vecX").map(|prop| prop.value.as_f32()).flatten() {
+        let x_coord = match entity_state
+            .get_prop("CCSPlayerPawn.CBodyComponentBaseAnimGraph.m_vecX")
+            .map(|prop| prop.value.as_f32())
+            .flatten()
+        {
             Some(c) => c,
-            None => player_position.get(&user_id).map(|(x, _, _)| *x).unwrap_or(0.0),
+            None => player_position
+                .get(&user_id)
+                .map(|(x, _, _)| *x)
+                .unwrap_or(0.0),
         };
-        let y_coord = match entity_state.get_prop("CCSPlayerPawn.CBodyComponentBaseAnimGraph.m_vecY").map(|prop| prop.value.as_f32()).flatten() {
+        let y_coord = match entity_state
+            .get_prop("CCSPlayerPawn.CBodyComponentBaseAnimGraph.m_vecY")
+            .map(|prop| prop.value.as_f32())
+            .flatten()
+        {
             Some(c) => c,
-            None => player_position.get(&user_id).map(|(_, y, _)| *y).unwrap_or(0.0),
+            None => player_position
+                .get(&user_id)
+                .map(|(_, y, _)| *y)
+                .unwrap_or(0.0),
         };
-        let z_coord = match entity_state.get_prop("CCSPlayerPawn.CBodyComponentBaseAnimGraph.m_vecZ").map(|prop| prop.value.as_f32()).flatten() {
+        let z_coord = match entity_state
+            .get_prop("CCSPlayerPawn.CBodyComponentBaseAnimGraph.m_vecZ")
+            .map(|prop| prop.value.as_f32())
+            .flatten()
+        {
             Some(c) => c,
-            None => player_position.get(&user_id).map(|(_, _, z)| *z).unwrap_or(0.0),
+            None => player_position
+                .get(&user_id)
+                .map(|(_, _, z)| *z)
+                .unwrap_or(0.0),
         };
 
         player_position.insert(user_id, (x_coord, y_coord, z_coord));
@@ -224,9 +264,9 @@ fn process_tick(
         assert!(y_coord >= 0.0);
         assert!(z_coord >= 0.0);
 
-        let x_cell_coord = ((x_cell as f32 * (1 << 9) as f32)) as f32;
+        let x_cell_coord = (x_cell as f32 * (1 << 9) as f32) as f32;
-        let y_cell_coord = ((y_cell as f32 * (1 << 9) as f32)) as f32;
+        let y_cell_coord = (y_cell as f32 * (1 << 9) as f32) as f32;
-        let z_cell_coord = ((z_cell as f32 * (1 << 9) as f32)) as f32;
+        let z_cell_coord = (z_cell as f32 * (1 << 9) as f32) as f32;
 
         let x_coord = x_cell_coord + x_coord;
         let y_coord = y_cell_coord + y_coord;
@@ -265,7 +305,9 @@ fn process_tick(
 
         // tracing::trace!("Coord (X, Y, Z): {:?} -> {:?}", (x_coord, y_coord, z_coord), (x_cell, y_cell));
 
-        let heatmap = heatmaps.entry((user_id.clone(), team)).or_insert(HeatMap::new(config.cell_size));
+        let heatmap = heatmaps
+            .entry((user_id.clone(), team))
+            .or_insert(HeatMap::new(config.cell_size));
         heatmap.increment(x_cell, y_cell);
     }
 }
@@ -276,7 +318,7 @@ impl core::fmt::Display for HeatMap {
 
         for row in self.rows.iter() {
             for cell in row.iter().copied() {
-                write!(f, "{: ^width$} ", cell, width=size)?;
+                write!(f, "{: ^width$} ", cell, width = size)?;
             }
             writeln!(f)?;
         }
@@ -288,22 +330,45 @@ impl core::fmt::Display for HeatMap {
 impl HeatMap {
     pub fn coords(&self) -> ((f32, f32), (f32, f32)) {
         (
-            (self.min_x as f32 * self.block_size - MAX_COORD, self.max_x as f32 * self.block_size - MAX_COORD),
+            (
-            (self.min_y as f32 * self.block_size - MAX_COORD, self.max_y as f32 * self.block_size - MAX_COORD)
+                self.min_x as f32 * self.block_size - MAX_COORD,
+                self.max_x as f32 * self.block_size - MAX_COORD,
+            ),
+            (
+                self.min_y as f32 * self.block_size - MAX_COORD,
+                self.max_y as f32 * self.block_size - MAX_COORD,
+            ),
         )
     }
 
     pub fn as_image(&self) -> image::RgbImage {
         use colors_transform::Color;
 
-        let mut buffer = image::RgbImage::new((self.max_x - self.min_x) as u32 + 1, (self.max_y - self.min_y) as u32 + 1);
+        let mut buffer = image::RgbImage::new(
+            (self.max_x - self.min_x) as u32 + 1,
+            (self.max_y - self.min_y) as u32 + 1,
+        );
 
         for (y, row) in self.rows.iter().rev().enumerate() {
-            for (x, cell) in row.iter().copied().chain(core::iter::repeat(0)).enumerate().take(self.max_x - self.min_x) {
+            for (x, cell) in row
-                let scaled = (1.0/(1.0 + (cell as f32))) * 240.0;
+                .iter()
+                .copied()
+                .chain(core::iter::repeat(0))
+                .enumerate()
+                .take(self.max_x - self.min_x)
+            {
+                let scaled = (1.0 / (1.0 + (cell as f32))) * 240.0;
                 let raw_rgb = colors_transform::Hsl::from(scaled, 100.0, 50.0).to_rgb();
 
-                buffer.put_pixel(x as u32, y as u32, image::Rgb([raw_rgb.get_red() as u8, raw_rgb.get_green() as u8, raw_rgb.get_blue() as u8]))
+                buffer.put_pixel(
+                    x as u32,
+                    y as u32,
+                    image::Rgb([
+                        raw_rgb.get_red() as u8,
+                        raw_rgb.get_green() as u8,
+                        raw_rgb.get_blue() as u8,
+                    ]),
+                )
             }
         }
 
@@ -352,18 +417,12 @@ mod tests {
         assert_eq!(input.max_y, 3);
 
         assert_eq!(
-            &vec![
+            &vec![vec![], vec![], vec![0, 0, 1], vec![0, 0, 0, 1]],
-                vec![],
-                vec![],
-                vec![0, 0, 1],
-                vec![0, 0, 0, 1]
-            ],
             &input.rows
         );
 
         input.fit(2.0..10.0, 2.0..10.0);
 
-        
         assert_eq!(
             &vec![
                 vec![0, 0, 0, 0],
@@ -388,23 +447,12 @@ mod tests {
         assert_eq!(input.max_y, 3);
 
         assert_eq!(
-            &vec![
+            &vec![vec![], vec![], vec![0, 0, 1], vec![0, 0, 0, 1]],
-                vec![],
-                vec![],
-                vec![0, 0, 1],
-                vec![0, 0, 0, 1]
-            ],
             &input.rows
         );
 
         input.fit(6.0..10.0, 6.0..10.0);
 
-        assert_eq!(
+        assert_eq!(&vec![vec![1, 0], vec![0, 0]], &input.rows);
-            &vec![
-                vec![1, 0],
-                vec![0, 0]
-            ],
-            &input.rows
-        );
     }
 }

analysis/src/perround.rs

@@ -60,28 +60,29 @@ pub struct Round {
 pub enum RoundEvent {
     BombPlanted,
     BombDefused,
-    Kill {
+    Kill { attacker: u64, died: u64 },
-        attacker: u64,
-        died: u64,
-    },
 }
 
 #[derive(Debug)]
 pub struct PerRound {
-    pub rounds: Vec<Round>
+    pub rounds: Vec<Round>,
 }
 
 pub fn parse(buf: &[u8]) -> Result<PerRound, ()> {
     let tmp = csdemo::Container::parse(buf).map_err(|e| ())?;
     let output = csdemo::parser::parse(
         csdemo::FrameIterator::parse(tmp.inner),
-        csdemo::parser::EntityFilter::all()
+        csdemo::parser::EntityFilter::all(),
-    ).map_err(|e| ())?;
+    )
+    .map_err(|e| ())?;
 
     let mut rounds: Vec<Round> = Vec::new();
     for tick in output.entity_states.ticks.iter() {
         for state in tick.states.iter() {
-            let round_start_count = state.get_prop("CCSGameRulesProxy.CCSGameRules.m_nRoundStartCount").map(|v| v.value.as_u32()).flatten();
+            let round_start_count = state
+                .get_prop("CCSGameRulesProxy.CCSGameRules.m_nRoundStartCount")
+                .map(|v| v.value.as_u32())
+                .flatten();
             if let Some(round_start_count) = round_start_count {
                 if rounds.len() < (round_start_count - 1) as usize {
                     rounds.push(Round {
@@ -93,21 +94,32 @@ pub fn parse(buf: &[u8]) -> Result<PerRound, ()> {
                 }
             }
 
-            let round_end_count = state.get_prop("CCSGameRulesProxy.CCSGameRules.m_nRoundEndCount").map(|v| v.value.as_u32()).flatten();
+            let round_end_count = state
+                .get_prop("CCSGameRulesProxy.CCSGameRules.m_nRoundEndCount")
+                .map(|v| v.value.as_u32())
+                .flatten();
             if let Some(round_end_count) = round_end_count {
                 if rounds.len() == (round_end_count - 1) as usize {
                     rounds.last_mut().unwrap().end = tick.tick;
                 }
             }
 
-            let total_rounds_played = state.get_prop("CCSGameRulesProxy.CCSGameRules.m_totalRoundsPlayed").map(|v| v.value.as_i32()).flatten();
+            let total_rounds_played = state
+                .get_prop("CCSGameRulesProxy.CCSGameRules.m_totalRoundsPlayed")
+                .map(|v| v.value.as_i32())
+                .flatten();
             if let Some(total_rounds_played) = total_rounds_played {
                 debug_assert_eq!(total_rounds_played, rounds.len() as i32);
             }
 
-
             if state.class.as_ref() == "CCSGameRulesProxy" {
-                let round_win_reason = state.get_prop("CCSGameRulesProxy.CCSGameRules.m_eRoundWinReason").map(|p| p.value.as_i32()).flatten().map(|v| ROUND_WIN_REASON.get(&v)).flatten().filter(|r| !matches!(r, WinReason::StillInProgress));
+                let round_win_reason = state
+                    .get_prop("CCSGameRulesProxy.CCSGameRules.m_eRoundWinReason")
+                    .map(|p| p.value.as_i32())
+                    .flatten()
+                    .map(|v| ROUND_WIN_REASON.get(&v))
+                    .flatten()
+                    .filter(|r| !matches!(r, WinReason::StillInProgress));
                 if let Some(round_win_reason) = round_win_reason {
                     rounds.last_mut().unwrap().winreason = round_win_reason.clone();
                 }
@@ -138,12 +150,8 @@ pub fn parse(buf: &[u8]) -> Result<PerRound, ()> {
                 }
 
                 let event = match ge.as_ref() {
-                    csdemo::game_event::GameEvent::BombPlanted(planted) => {
+                    csdemo::game_event::GameEvent::BombPlanted(planted) => RoundEvent::BombPlanted,
-                        RoundEvent::BombPlanted
+                    csdemo::game_event::GameEvent::BombDefused(defused) => RoundEvent::BombDefused,
-                    }
-                    csdemo::game_event::GameEvent::BombDefused(defused) => {
-                        RoundEvent::BombDefused
-                    }
                     csdemo::game_event::GameEvent::PlayerDeath(death) => {
                         let died = match death.userid {
                             Some(d) => d,
@@ -171,7 +179,5 @@ pub fn parse(buf: &[u8]) -> Result<PerRound, ()> {
         };
     }
 
-    Ok(PerRound {
+    Ok(PerRound { rounds })
-        rounds,
-    })
 }

analysis/tests/endofgame.rs

@@ -11,6 +11,24 @@ fn endofgame_nuke() {
 
     let expected = endofgame::EndOfGame {
         map: "de_nuke".to_owned(),
+        teams: [
+            (
+                3,
+                endofgame::TeamInfo {
+                    name: "CT".to_owned(),
+                    score: 8,
+                },
+            ),
+            (
+                2,
+                endofgame::TeamInfo {
+                    name: "TERRORIST".to_owned(),
+                    score: 13,
+                },
+            ),
+        ]
+        .into_iter()
+        .collect(),
         players: vec![
             (
                 endofgame::PlayerInfo {
@@ -199,8 +217,5 @@ fn endofgame_nuke() {
         ],
     };
 
-    // TODO
-    // Add stats for rest of players
-
     assert_eq!(result, expected);
 }