Merge branch 'gamemode-filters'

1 files changed, 28 insertions, 1 deletions

diff --git a/src/filters/log.rs b/src/filters/log.rs
index ac166e1..c878b78 100644
--- a/src/filters/log.rs
+++ b/src/filters/log.rs
@@ -9,7 +9,7 @@ use super::{
 
 use std::collections::HashSet;
 
-use evtclib::Encounter;
+use evtclib::{Encounter, GameMode};
 
 use chrono::{DateTime, Datelike, Utc, Weekday};
 
@@ -37,6 +37,33 @@ pub fn encounter(bosses: HashSet<Encounter>) -> Box<dyn LogFilter> {
 }
 
 #[derive(Debug, Clone)]
+struct GameModeFilter(HashSet<GameMode>);
+
+impl Filter<EarlyLogResult, LogResult> for GameModeFilter {
+    fn filter_early(&self, early_log: &EarlyLogResult) -> Inclusion {
+        let encounter_id = early_log.evtc.header.combat_id;
+        // Special WvW encounter
+        const GENERIC_ENCOUNTER_ID: u16 = 1;
+        let mode = if encounter_id == GENERIC_ENCOUNTER_ID {
+            Some(GameMode::WvW)
+        } else {
+            Encounter::from_header_id(encounter_id).map(Encounter::game_mode)
+        };
+        mode.map(|m| self.0.contains(&m).into())
+            .unwrap_or(Inclusion::Exclude)
+    }
+
+    fn filter(&self, log: &LogResult) -> bool {
+        log.game_mode.map(|m| self.0.contains(&m)).unwrap_or(false)
+    }
+}
+
+/// A [`LogFilter`] that only accepts logs with the given game mode.
+pub fn game_mode(game_modes: HashSet<GameMode>) -> Box<dyn LogFilter> {
+    Box::new(GameModeFilter(game_modes))
+}
+
+#[derive(Debug, Clone)]
 struct OutcomeFilter(HashSet<FightOutcome>);
 
 impl Filter<EarlyLogResult, LogResult> for OutcomeFilter {