1 files changed, 50 insertions, 52 deletions

diff --git a/src/session.rs b/src/session.rs
index a7ed261..aad299a 100644
--- a/src/session.rs
+++ b/src/session.rs
@@ -1,40 +1,59 @@
 use std::{
-    collections::{hash_map::Entry, HashMap, HashSet},
+    collections::{hash_map::Entry, HashMap},
     fmt::Display,
     io::{BufWriter, Write},
     iter::IntoIterator,
     time::Duration,
 };
 
+use clap::Args;
 use serde::{Deserialize, Serialize};
 
+pub trait SessionWriter {
+    type Error: Display;
+    type Writer: Write;
+
+    fn writer(&self) -> Result<Self::Writer, std::io::Error>;
+    fn format(&self, session: &Session) -> Result<String, Self::Error>;
+    fn filter(&self, session: &Session) -> bool;
+}
+
+#[derive(Debug, Clone, Args)]
+#[group(skip)]
+pub struct Config {
+    /// Enable sorting
+    #[arg(long, default_value_t = true)]
+    pub sort: bool,
+}
+
 #[derive(Debug, Default)]
 pub struct Sessions {
     inner: HashMap<String, State>,
-    exclude: Vec<String>,
+    sort: bool,
 }
 
 impl Sessions {
-    pub fn new(exclude: Vec<String>) -> Self {
+    pub fn new(Config { sort }: Config) -> Self {
         Self {
-            exclude,
+            sort,
             ..Default::default()
         }
     }
 
-    pub fn sorted(self) -> Vec<Session> {
-        let mut sessions: Vec<Session> = self.into_iter().map(Session::from).collect();
-        sessions.sort();
-        sessions
-    }
+    pub fn write_sessions<W: SessionWriter>(&self, writer: W) -> std::io::Result<()> {
+        let mut buf_writer = BufWriter::new(writer.writer()?);
+
+        let mut sessions: Vec<Session> = self.inner.iter().map(Session::from).collect();
 
-    pub fn write_into<W: Write>(&self, writer: W) -> std::io::Result<()> {
-        let mut buf_writer = BufWriter::new(writer);
+        if self.sort {
+            sessions.sort();
+        }
 
-        for session in self.inner.iter().map(Session::from) {
-            match ron::to_string(&session) {
-                Ok(ser) => writeln!(buf_writer, "{ser}")?,
-                Err(err) => tracing::warn!(%err, "Failed to serialize session"),
+        for session in sessions {
+            match writer.format(&session) {
+                Ok(fmt) if writer.filter(&session) => writeln!(buf_writer, "{fmt}")?,
+                Err(err) => tracing::warn!(%err, "Failed to format session"),
+                _ => tracing::debug!(%session, "Skipping filtered session"),
             }
         }
 
@@ -45,11 +64,6 @@ impl Sessions {
         let span = tracing::trace_span!("Entry", ?item);
         let _guard = span.enter();
 
-        if self.exclude.contains(&item.name) {
-            tracing::debug!(item.name, "Skipping excluded item");
-            return;
-        }
-
         match self.inner.entry(item.name) {
             Entry::Occupied(mut occupied) if &item.state > occupied.get() => {
                 tracing::trace!(?occupied, new_value=?item.state, "New entry is more recent, replacing");
@@ -66,20 +80,6 @@ impl Sessions {
     }
 }
 
-impl IntoIterator for Sessions {
-    type Item = Session;
-
-    type IntoIter = std::collections::hash_set::IntoIter<Self::Item>;
-
-    fn into_iter(self) -> Self::IntoIter {
-        self.inner
-            .into_iter()
-            .map(Into::into)
-            .collect::<HashSet<Session>>()
-            .into_iter()
-    }
-}
-
 impl Extend<Session> for Sessions {
     fn extend<T: IntoIterator<Item = Session>>(&mut self, iter: T) {
         for item in iter {
@@ -98,12 +98,6 @@ pub enum State {
     LocalHost,
 }
 
-#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
-pub struct Session {
-    pub state: State,
-    pub name: String,
-}
-
 mod epoch_timestamp {
     use std::time::Duration;
 
@@ -124,30 +118,34 @@ mod epoch_timestamp {
     }
 }
 
-impl Display for Session {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}", self.name)
-    }
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
+pub struct Session {
+    pub state: State,
+    pub name: String,
 }
 
-impl From<&str> for Session {
-    fn from(value: &str) -> Self {
+impl Session {
+    pub fn discover(name: impl Into<String>) -> Self {
         Self {
+            name: name.into(),
             state: State::Discovered,
-            name: value.to_owned(),
         }
     }
-}
 
-impl From<String> for Session {
-    fn from(name: String) -> Self {
+    pub fn localhost(name: impl Into<String>) -> Self {
         Self {
-            state: State::Discovered,
-            name,
+            name: name.into(),
+            state: State::LocalHost,
         }
     }
 }
 
+impl Display for Session {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.name)
+    }
+}
+
 impl From<(String, State)> for Session {
     fn from((name, state): (String, State)) -> Self {
         Self { state, name }