Sentry positions

Tails8521 · Tails8521 · commit 7afd829b8f3d · 2020-10-11T16:18:45.000+02:00
diff --git a/src/demo_player.rs b/src/demo_player.rs
@@ -12,7 +12,7 @@ use std::{
     path::PathBuf,
 };
 
-use coldmaps::heatmap_analyser::{Class, HeatmapAnalyser, HeatmapAnalysis, PlayerState, Spawn, Team, UserId, UserInfo};
+use coldmaps::heatmap_analyser::{Class, HeatmapAnalyser, HeatmapAnalysis, PlayerState, Spawn, Team, UserId, UserInfo, handle_to_entity_index};
 use serde::Serialize;
 use tf_demo_parser::{
     demo::gamevent::GameEvent, demo::header::Header, demo::message::packetentities::EntityId, demo::message::packetentities::PacketEntity, demo::message::packetentities::PVS,
@@ -893,12 +893,4 @@ impl DemoAnalyzer {
         }
         Ok(())
     }
-}
-
-fn handle_to_entity_index(handle: i64) -> Option<NonZeroU32> {
-    let ret = handle as u32 & 0b111_1111_1111; // The rest of the bits is probably some kind of generational index
-    if ret == 2047 {
-        return None
-    }
-    NonZeroU32::new(ret)
 }
diff --git a/src/filters.rs b/src/filters.rs
@@ -72,11 +72,12 @@ impl Display for PropertyOperator {
 pub enum Property {
     Suicide,
     Posthumous,
-    DuringRound
+    DuringRound,
+    DiedToSentry,
 }
 
 impl Property {
-    pub const ALL: [Property; 3] = [Property::Suicide, Property::Posthumous, Property::DuringRound];
+    pub const ALL: [Property; 4] = [Property::Suicide, Property::Posthumous, Property::DuringRound, Property::DiedToSentry];
 }
 
 impl Default for Property {
@@ -91,6 +92,7 @@ impl Display for Property {
             Property::Suicide => write!(f, "Suicide"),
             Property::Posthumous => write!(f, "Posthumous"),
             Property::DuringRound => write!(f, "During round"),
+            Property::DiedToSentry => write!(f, "Died to sentry"),
         }
     }
 }
@@ -297,7 +299,8 @@ impl FilterTrait for PropertyFilter {
                 Some(PlayerEntity { state: PlayerState::Alive, .. }) => false,
                 _ => true,
             },
-            Property::DuringRound => death.during_round
+            Property::DuringRound => death.during_round,
+            Property::DiedToSentry => death.sentry_position.is_some(),
         };
         match self.op {
             PropertyOperator::IsPresent => ret,
diff --git a/src/heatmap.rs b/src/heatmap.rs
@@ -58,7 +58,8 @@ struct HeatMapParameters {
     top_y: f32,
     bottom_y: f32,
     radius: f32,
-    intensity: Option<f32>
+    intensity: Option<f32>,
+    use_sentry_position: bool,
 }
 
 #[derive(Debug)]
@@ -67,7 +68,7 @@ pub struct HeatMapGenerator {
 }
 
 impl HeatMapGenerator {
-    pub fn new(pos_x: f32, pos_y: f32, screen_width: u32, screen_height: u32, scale: f32, coords_type: CoordsType, radius: f32, intensity: Option<f32>) -> Self {
+    pub fn new(pos_x: f32, pos_y: f32, screen_width: u32, screen_height: u32, scale: f32, coords_type: CoordsType, radius: f32, intensity: Option<f32>, use_sentry_position: bool) -> Self {
         let screen_width = screen_width as f32;
         let screen_height = screen_height as f32;
         let aspect_ratio = screen_width / screen_height;
@@ -81,7 +82,8 @@ impl HeatMapGenerator {
                     top_y: pos_y + scale * LEVELOVERVIEW_SCALE_MULTIPLIER,
                     bottom_y: pos_y - scale * LEVELOVERVIEW_SCALE_MULTIPLIER,
                     radius,
-                    intensity
+                    intensity,
+                    use_sentry_position,
                 },
             },
             CoordsType::Console => Self {
@@ -93,7 +95,8 @@ impl HeatMapGenerator {
                     top_y: pos_y,
                     bottom_y: pos_y - scale * LEVELOVERVIEW_SCALE_MULTIPLIER * 2.0,
                     radius,
-                    intensity
+                    intensity,
+                    use_sentry_position,
                 },
             },
         }
@@ -118,9 +121,19 @@ impl HeatMapGenerator {
                 // LinSrgba::new(0.0, 0.0, 1.0, 1.0),
             ]);
             for death in deaths {
-                if let (Some(killer_entity), Some(victim_entity)) = (&death.killer_entity_state, &death.victim_entity_state) {
-                    let killer_coords = self.game_coords_to_screen_coords(killer_entity.position.x, killer_entity.position.y);
-                    let victim_coords = self.game_coords_to_screen_coords(victim_entity.position.x, victim_entity.position.y);
+                let killer_pos = if self.params.use_sentry_position {
+                    if let Some(sentry_position) = death.sentry_position {
+                        Some(sentry_position)
+                    } else {
+                        death.killer_entity_state.as_ref().map(|entity| entity.position)
+                    }
+                } else {
+                    death.killer_entity_state.as_ref().map(|entity| entity.position)
+                };
+                let victim_pos = death.victim_entity_state.as_ref().map(|entity| entity.position);
+                if let (Some(killer_pos), Some(victim_pos)) = (killer_pos, victim_pos) {
+                    let killer_coords = self.game_coords_to_screen_coords(killer_pos.x, killer_pos.y);
+                    let victim_coords = self.game_coords_to_screen_coords(victim_pos.x, victim_pos.y);
                     let points: Vec<((i32, i32), f32)> = line_drawing::XiaolinWu::<f32, i32>::new(killer_coords, victim_coords).collect();
 
                     // this is needed because the line drawing algorithm doesn't always go in the start-end order, we need to check what order was used and invert the gradient as needed
@@ -179,13 +192,19 @@ impl HeatMapGenerator {
         let radius = self.params.radius / 10.0;
         let pixels_iters = (radius * 2.0).ceil() as i32;
         for death in deaths {
-            let entity_state = match heatmap_type {
-                HeatmapType::VictimPosition => &death.victim_entity_state,
-                HeatmapType::KillerPosition => &death.killer_entity_state,
-                HeatmapType::Lines => unreachable!(),
+            let game_coords = match (heatmap_type, self.params.use_sentry_position) {
+                (HeatmapType::VictimPosition, _) => death.victim_entity_state.as_ref().map(|entity| entity.position),
+                (HeatmapType::KillerPosition, false) => death.killer_entity_state.as_ref().map(|entity| entity.position),
+                (HeatmapType::KillerPosition, true) => {
+                    if let Some(sentry_position) = death.sentry_position {
+                        Some(sentry_position)
+                    } else {
+                        death.killer_entity_state.as_ref().map(|entity| entity.position)
+                    }
+                }
+                (HeatmapType::Lines, _) => unreachable!(),
             };
-            if let Some(entity_state) = entity_state {
-                let game_coords = entity_state.position;
+            if let Some(game_coords) = game_coords {
                 let (x_f, y_f) = self.game_coords_to_screen_coords(game_coords.x, game_coords.y);
                 let x_i = x_f.round() as i32;
                 let y_i = y_f.round() as i32;
diff --git a/src/heatmap_analyser.rs b/src/heatmap_analyser.rs
@@ -1,9 +1,9 @@
 use num_enum::TryFromPrimitive;
 use serde::{Deserialize, Serialize};
-use std::collections::{BTreeMap, HashMap};
+use std::{collections::{BTreeMap, HashMap}, num::NonZeroU32};
 use std::convert::TryFrom;
 use std::str::FromStr;
-use tf_demo_parser::demo::gameevent_gen::{GameEvent, PlayerDeathEvent, PlayerSpawnEvent, TeamPlayRoundWinEvent};
+use tf_demo_parser::demo::{gameevent_gen::{GameEvent, PlayerDeathEvent, PlayerSpawnEvent, TeamPlayRoundWinEvent}, message::packetentities::PVS};
 use tf_demo_parser::demo::message::packetentities::{EntityId, PacketEntity};
 use tf_demo_parser::demo::message::usermessage::{ChatMessageKind, SayText2Message, UserMessage};
 use tf_demo_parser::demo::message::{Message, MessageType};
@@ -158,6 +158,7 @@ pub struct Death {
     pub tick: u32,
     pub round: u32,
     pub during_round: bool,
+    pub sentry_position: Option<Vector>
 }
 
 impl Death {
@@ -193,6 +194,7 @@ impl Death {
             victim_steamid: users.get(&victim).expect("Can't get victim").steam_id.clone(),
             victim_entity: event.victim_ent_index,
             victim_entity_state: None,
+            sentry_position: None
         }
     }
 }
@@ -259,6 +261,12 @@ pub struct PlayerEntity {
     pub state: PlayerState,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub enum OtherEntity {
+    Sentry { position: Option<Vector> },
+    SentryRocket { sentry: Option<EntityId> }
+}
+
 impl MessageHandler for HeatmapAnalyser {
     type Output = HeatmapAnalysis;
 
@@ -280,7 +288,15 @@ impl MessageHandler for HeatmapAnalyser {
             Message::UserMessage(message) => self.handle_user_message(&message, tick),
             Message::PacketEntities(message) => {
                 for entity in &message.entities {
-                    self.handle_entity(entity);
+                    if entity.pvs == PVS::Delete {
+                        let removed_entity = entity.entity_index;
+                        let _removed = self.state.other_entities.remove(&removed_entity);
+                    } else {
+                        self.handle_entity(entity);
+                    }
+                }
+                for removed_entity in &message.removed_entities {
+                    let _removed = self.state.other_entities.remove(removed_entity);
                 }
             }
             _ => unreachable!(),
@@ -318,6 +334,8 @@ impl HeatmapAnalyser {
             "CTFPlayer" => self.handle_player_entity(entity),
             "CTFPlayerResource" => self.handle_player_resource(entity),
             "CWorld" => self.handle_world_entity(entity),
+            "CObjectSentrygun" => self.handle_sentry_entity(entity),
+            "CTFProjectile_SentryRocket" => self.handle_sentry_rocket_entity(entity),
             _ => {}
         }
     }
@@ -384,6 +402,46 @@ impl HeatmapAnalyser {
         }
     }
 
+    fn handle_sentry_entity(&mut self, entity: &PacketEntity) {
+        let entry = self.state.other_entities.entry(entity.entity_index).or_insert_with(|| OtherEntity::Sentry {
+            position: None
+        });
+        let mut position = if let OtherEntity::Sentry { position } = *entry {
+            position
+        } else {
+            None
+        };
+        for prop in &entity.props {
+            match prop.definition.name.as_str() {
+                "m_vecOrigin" => position = Some(Vector::try_from(&prop.value).unwrap_or_default()),
+                _ => {}
+            }
+        }
+        *entry = OtherEntity::Sentry { position };
+    }
+
+    fn handle_sentry_rocket_entity(&mut self, entity: &PacketEntity) {
+        let entry = self.state.other_entities.entry(entity.entity_index).or_insert_with(|| OtherEntity::SentryRocket {
+            sentry: None
+        });
+        let mut sentry = if let OtherEntity::SentryRocket { sentry } = *entry {
+            sentry
+        } else {
+            None
+        };
+        for prop in &entity.props {
+            match prop.definition.name.as_str() {
+                "m_hOwnerEntity" => {
+                    let handle = i64::try_from(&prop.value).unwrap_or_default();
+                    let entity_id = handle_to_entity_index(handle);
+                    sentry = entity_id.map(|id| id.get().into());
+                }
+                _ => {}
+            }
+        }
+        *entry = OtherEntity::SentryRocket { sentry };
+    }
+
     fn handle_user_message(&mut self, message: &UserMessage, tick: u32) {
         if let UserMessage::SayText2(text_message) = message {
             if text_message.kind == ChatMessageKind::NameChange {
@@ -423,6 +481,30 @@ impl HeatmapAnalyser {
                 if let Some(victim_entity) = victim.entity_id {
                     death.victim_entity_state = Some(self.state.get_or_create_player_entity(victim_entity).clone());
                 }
+                match death.weapon.as_str() {
+                    "obj_sentrygun" | "obj_sentrygun2" | "obj_sentrygun3" | "obj_minisentry" => {
+                        if let Some(entity) = self.state.other_entities.get(&death.killer_entity.into()) {
+                            match entity {
+                                OtherEntity::Sentry { position } => {
+                                    death.sentry_position = *position;
+                                }
+                                OtherEntity::SentryRocket { sentry } => {
+                                    if let Some(sentry_entity) = sentry {
+                                        if let Some(entity) = self.state.other_entities.get(sentry_entity) {
+                                            match entity {
+                                                OtherEntity::Sentry { position } => {
+                                                    death.sentry_position = *position;
+                                                }
+                                                _ => {}
+                                            }
+                                        }
+                                    }
+                                }
+                            }
+                        }
+                    }
+                    _ => {}
+                }
                 self.state.deaths.push(death);
             }
             GameEvent::PlayerSpawn(event) => {
@@ -491,6 +573,7 @@ pub struct HeatmapAnalysis {
     pub in_round: bool,
 
     pub player_entities: Vec<PlayerEntity>,
+    pub other_entities: HashMap<EntityId, OtherEntity>,
     pub world: Option<World>,
 }
 
@@ -532,6 +615,7 @@ impl Default for HeatmapAnalysis {
                 player_entities.push(world);
                 player_entities
             },
+            other_entities: Default::default(),
             world: Default::default(),
         }
     }
@@ -568,3 +652,11 @@ impl HeatmapAnalysis {
         &mut self.player_entities[index]
     }
 }
+
+pub fn handle_to_entity_index(handle: i64) -> Option<NonZeroU32> {
+    let ret = handle as u32 & 0b111_1111_1111; // The rest of the bits is probably some kind of generational index
+    if ret == 2047 {
+        return None
+    }
+    NonZeroU32::new(ret)
+}
diff --git a/src/lib.rs b/src/lib.rs
@@ -73,9 +73,10 @@ pub fn generate_heatmap<'a>(
     scale: f32,
     coords_type: CoordsType,
     radius: f32,
-    intensity: Option<f32>
+    intensity: Option<f32>,
+    use_sentry_position: bool,
 ) -> ImageBuffer<Rgb<u8>, Vec<u8>> {
-    let heatmap_generator = heatmap::HeatMapGenerator::new(pos_x, pos_y, screen_width, screen_height, scale, coords_type, radius, intensity);
+    let heatmap_generator = heatmap::HeatMapGenerator::new(pos_x, pos_y, screen_width, screen_height, scale, coords_type, radius, intensity, use_sentry_position);
     heatmap_generator.generate_heatmap(heatmap_type, deaths, &mut image);
     image
 }
diff --git a/src/main.rs b/src/main.rs

Original file line number	Diff line number	Diff line change
`@@ -72,11 +72,12 @@ impl Display for PropertyOperator {`
`72`	`72`	`pub enum Property {`
`73`	`73`	`Suicide,`
`74`	`74`	`Posthumous,`
`75`		`- DuringRound`
	`75`	`+ DuringRound,`
	`76`	`+ DiedToSentry,`
`76`	`77`	`}`
`77`	`78`
`78`	`79`	`impl Property {`
`79`		`- pub const ALL: [Property; 3] = [Property::Suicide, Property::Posthumous, Property::DuringRound];`
	`80`	`+ pub const ALL: [Property; 4] = [Property::Suicide, Property::Posthumous, Property::DuringRound, Property::DiedToSentry];`
`80`	`81`	`}`
`81`	`82`
`82`	`83`	`impl Default for Property {`
`@@ -91,6 +92,7 @@ impl Display for Property {`
`91`	`92`	`Property::Suicide => write!(f, "Suicide"),`
`92`	`93`	`Property::Posthumous => write!(f, "Posthumous"),`
`93`	`94`	`Property::DuringRound => write!(f, "During round"),`
	`95`	`+ Property::DiedToSentry => write!(f, "Died to sentry"),`
`94`	`96`	`}`
`95`	`97`	`}`
`96`	`98`	`}`
`@@ -297,7 +299,8 @@ impl FilterTrait for PropertyFilter {`
`297`	`299`	`Some(PlayerEntity { state: PlayerState::Alive, .. }) => false,`
`298`	`300`	`_ => true,`
`299`	`301`	`},`
`300`		`- Property::DuringRound => death.during_round`
	`302`	`+ Property::DuringRound => death.during_round,`
	`303`	`+ Property::DiedToSentry => death.sentry_position.is_some(),`
`301`	`304`	`};`
`302`	`305`	`match self.op {`
`303`	`306`	`PropertyOperator::IsPresent => ret,`