Clean up code and change to using output names

src/main.rs

@@ -1,13 +1,17 @@
 mod river_protocols;
 
+use river_protocols::zriver_seat_status_v1::ZriverSeatStatusV1;
 use river_protocols::{
-    zriver_output_status_v1, zriver_seat_status_v1, zriver_status_manager_v1::ZriverStatusManagerV1,
+    zriver_output_status_v1, zriver_seat_status_v1,
+    zriver_status_manager_v1::ZriverStatusManagerV1,
 };
-use serde::ser::{SerializeSeq, Serializer};
+use serde::ser::{SerializeSeq, SerializeStruct, Serializer};
 use serde::Serialize;
-use std::collections::BTreeMap;
+use wayland_client::protocol::wl_registry::WlRegistry;
-use wayland_client::protocol::{wl_output, wl_output::WlOutput, wl_seat, wl_seat::WlSeat};
+use std::collections::HashMap;
-use wayland_client::{Display, GlobalManager, Main};
+use wayland_client::protocol::{wl_output, wl_output::WlOutput, wl_seat,
+    wl_seat::WlSeat};
+use wayland_client::{Attached, DispatchData, Display, GlobalEvent, GlobalManager, Main};
 
 #[derive(Debug)]
 struct Flags {
@@ -17,7 +21,6 @@ struct Flags {
     viewstag: bool,
     layout: bool,
     mode: bool,
-    output: Option<String>,
     seat: Option<String>,
 }
 
@@ -30,34 +33,79 @@ impl Flags {
             viewstag: false,
             layout: false,
             mode: false,
-            output: None,
             seat: None,
         }
     }
 }
 
+#[derive(Clone)]
 struct Tags(u32);
 
-#[derive(Serialize)]
 struct Env {
-    #[serde(skip)]
     flags: Flags,
-    #[serde(skip_serializing_if = "Option::is_none")]
+    output_names: HashMap<u32, String>,
     layout: Option<String>,
-    #[serde(skip_serializing_if = "Option::is_none")]
     title: Option<String>,
-    #[serde(skip_serializing_if = "Option::is_none")]
     mode: Option<String>,
-    #[serde(skip_serializing_if = "Option::is_none")]
+    tags: Option<HashMap<u32, Tags>>,
-    tags: Option<BTreeMap<String, Tags>>,
+    urgency: Option<HashMap<u32, Tags>>,
-    #[serde(skip_serializing_if = "Option::is_none")]
+    viewstag: Option<HashMap<u32, Vec<u32>>>,
-    urgency: Option<BTreeMap<String, Tags>>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    viewstag: Option<BTreeMap<String, Vec<u32>>>,
-    #[serde(skip)]
     status_manager: Option<Main<ZriverStatusManagerV1>>,
 }
 
+impl Env {
+    fn translate_output_map<V: Clone>(
+        &self, input: &HashMap<u32, V>
+    ) -> HashMap<String, V> {
+        let mut output = HashMap::<String, V>::new();
+        output.reserve(input.len());
+        for (id, val) in input {
+            if let Some(name) = self.output_names.get(&id) {
+                output.insert(name.clone(), val.clone());
+            }
+        }
+        output
+    }
+}
+
+impl serde::Serialize for Env {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where S: Serializer {
+        let mut state = serializer.serialize_struct("Env", 6)?;
+        if let Some(layout) = self.layout.as_ref() {
+            state.serialize_field("layout", layout)?;
+        } else {
+            state.skip_field("layout")?;
+        }
+        if let Some(title) = self.title.as_ref() {
+            state.serialize_field("title", title)?;
+        } else {
+            state.skip_field("title")?;
+        }
+        if let Some(mode) = self.mode.as_ref() {
+            state.serialize_field("mode", mode)?;
+        } else {
+            state.skip_field("mode")?;
+        }
+        if let Some(tags) = self.tags.as_ref() {
+            state.serialize_field("tags", &self.translate_output_map(tags))?;
+        } else {
+            state.skip_field("tags")?;
+        }
+        if let Some(urgency) = self.urgency.as_ref() {
+            state.serialize_field("urgency", &self.translate_output_map(urgency))?;
+        } else {
+            state.skip_field("urgency")?;
+        }
+        if let Some(viewstag) = self.viewstag.as_ref() {
+            state.serialize_field("viewstag", &self.translate_output_map(viewstag))?;
+        } else {
+            state.skip_field("viewstag")?;
+        }
+        state.end()
+    }
+}
+
 impl Env {
     fn new() -> Env {
         let flags = configuration();
@@ -65,11 +113,12 @@ impl Env {
             title: None,
             layout: None,
             mode: None,
-            tags: flags.tags.then(BTreeMap::new),
+            tags: flags.tags.then(Default::default),
-            urgency: flags.urgency.then(BTreeMap::new),
+            urgency: flags.urgency.then(Default::default),
-            viewstag: flags.viewstag.then(BTreeMap::new),
+            viewstag: flags.viewstag.then(Default::default),
             status_manager: None,
             flags,
+            output_names: HashMap::default(),
         }
     }
 
@@ -86,6 +135,182 @@ impl Env {
     }
 }
 
+fn handle_zriver_status_manager_v1(
+    status_manager: Main<ZriverStatusManagerV1>, mut env: DispatchData
+) {
+    if let Some(env) = env.get::<Env>() {
+        env.status_manager = Some(status_manager);
+    }
+}
+
+fn handle_river_seat_status(
+    _status: Main<ZriverSeatStatusV1>, event: zriver_seat_status_v1::Event,
+    mut env: DispatchData
+) {
+    match event {
+        zriver_seat_status_v1::Event::FocusedView { title } => {
+            if let Some(env) = env.get::<Env>() {
+                if env.flags.title {
+                    env.title = Some(title);
+                }
+            }
+        }
+        zriver_seat_status_v1::Event::Mode { name } => {
+            if let Some(env) = env.get::<Env>() {
+                if env.flags.mode {
+                    env.mode = Some(name);
+                }
+            }
+        }
+        _ => {}
+    }
+}
+
+fn handle_seat_event(
+    seat: Main<WlSeat>, event: wl_seat::Event, mut env: DispatchData
+) {
+    match event {
+        wl_seat::Event::Name { name } => {
+            if let Some(env) = env.get::<Env>() {
+                if (env.flags.title || env.flags.mode) && (env.flags.seat.is_none()
+                    || name.eq(env.flags.seat.as_ref().unwrap())) {
+                        if let Some(status_manager) = &env.status_manager {
+                            let seat_status =
+                                status_manager.get_river_seat_status(&seat);
+                            seat_status.quick_assign(handle_river_seat_status);
+                        }
+                    }
+            }
+        }
+        _ => {}
+    }
+}
+
+fn handle_river_output_status(
+    output: &Main<WlOutput>, event: zriver_output_status_v1::Event,
+    mut env: DispatchData
+) {
+    if let Some(env) = env.get::<Env>() {
+        let output_id = output.as_ref().id();
+        match event {
+            zriver_output_status_v1::Event::FocusedTags {
+                tags: focused_tags,
+            } => {
+                if let Some(tags) = &mut env.tags {
+                    if let Some(inner_value) = tags.get_mut(&output_id) {
+                        (*inner_value) = Tags(focused_tags);
+                    } else {
+                        tags.insert(output_id, Tags(focused_tags));
+                    }
+                }
+            }
+            zriver_output_status_v1::Event::ViewTags {
+                tags,
+            } => {
+                if let Some(viewstag) = &mut env.viewstag {
+                    let tags: Vec<u32> = tags[0..]
+                        .chunks(4)
+                        .map(|s| {
+                            let buf = [s[0], s[1], s[2], s[3]];
+                            let tagmask =
+                                u32::from_le_bytes(buf);
+                            for i in 0..32 {
+                                if 1 << i == tagmask {
+                                    return 1 + i;
+                                }
+                            }
+                            0
+                        })
+                        .collect();
+                    if let Some(inner_value) = viewstag.get_mut(&output_id) {
+                        (*inner_value) = tags;
+                    } else {
+                        viewstag.insert(output_id, tags);
+                    }
+                }
+            }
+            zriver_output_status_v1::Event::UrgentTags {
+                tags,
+            } => {
+                if let Some(urgency) = &mut env.urgency {
+                    if let Some(inner_value) = urgency.get_mut(&output_id) {
+                        (*inner_value) = Tags(tags);
+                    } else {
+                        urgency .insert(output_id, Tags(tags));
+                    }
+                }
+            }
+            zriver_output_status_v1::Event::LayoutName {
+                name,
+            } => {
+                env.layout = Some(name);
+            }
+            _ => {}
+        }
+    }
+    
+}
+
+fn handle_output_event(
+    output: Main<WlOutput>, event: wl_output::Event, mut env: DispatchData
+) {
+    match event {
+        wl_output::Event::Geometry {
+            x: _,
+            y: _,
+            physical_width: _,
+            physical_height: _,
+            subpixel: _,
+            make: _,
+            model: _,
+            transform: _,
+        } => {
+            if let Some(env) = env.get::<Env>() {
+                if let Some(status_manager) = &env.status_manager {
+                    let output_status =
+                        status_manager.get_river_output_status(&output);
+                    output_status.quick_assign(move |_, event, env| {
+                        handle_river_output_status(&output, event, env);
+                    });
+                }
+            }
+        }
+        wl_output::Event::Name {
+            name,
+        } => {
+            if let Some(env) = env.get::<Env>() {
+                env.output_names.insert(output.as_ref().id(), name);
+            }
+        }
+        _ => {}
+    }
+}
+
+fn global_manager_callback(
+    event: GlobalEvent, registry: Attached<WlRegistry>, mut env: DispatchData
+) {
+    let mut global_filter = wayland_client::global_filter!(
+        [ZriverStatusManagerV1, 1, handle_zriver_status_manager_v1],
+        [WlSeat, 7, |seat: Main<WlSeat>, _env: DispatchData| {
+            seat.quick_assign(handle_seat_event);
+        }],
+        [WlOutput, 3, |output: Main<WlOutput>, _env: DispatchData| {
+            output.quick_assign(handle_output_event);
+        }]
+    );
+
+    match event {
+        GlobalEvent::New { id: _, interface: _, version: _ } => {
+            global_filter(event, registry, env);
+        }
+        GlobalEvent::Removed { id, interface: _ } => {
+            if let Some(env) = env.get::<Env>() {
+                env.output_names.remove(&id);
+            }
+        }
+    }
+}
+
 fn main() {
     let mut env = Env::new();
 
@@ -95,149 +320,7 @@ fn main() {
 
     GlobalManager::new_with_cb(
         &attached_display,
-        wayland_client::global_filter!(
+        global_manager_callback,
-            [
-                ZriverStatusManagerV1,
-                1,
-                |status_manager: Main<ZriverStatusManagerV1>, mut env: DispatchData| {
-                    if let Some(env) = env.get::<Env>() {
-                        env.status_manager = Some(status_manager);
-                    }
-                }
-            ],
-            [WlSeat, 7, |seat: Main<WlSeat>, _env: DispatchData| {
-                seat.quick_assign(move |seat, event, mut env| match event {
-                    wl_seat::Event::Name { name } => {
-                        if let Some(env) = env.get::<Env>() {
-                            if (env.flags.title || env.flags.mode)
-                                && (env.flags.seat.is_none()
-                                    || name.eq(env.flags.seat.as_ref().unwrap()))
-                            {
-                                if let Some(status_manager) = &env.status_manager {
-                                    let seat_status = status_manager.get_river_seat_status(&seat);
-                                    seat_status.quick_assign(
-                                        move |_, event, mut env| match event {
-                                            zriver_seat_status_v1::Event::FocusedView { title } => {
-                                                if let Some(env) = env.get::<Env>() {
-                                                    if env.flags.title {
-                                                        env.title = Some(title);
-                                                    }
-                                                }
-                                            }
-                                            zriver_seat_status_v1::Event::Mode { name } => {
-                                                if let Some(env) = env.get::<Env>() {
-                                                    if env.flags.mode {
-                                                        env.mode = Some(name);
-                                                    }
-                                                }
-                                            }
-                                            _ => {}
-                                        },
-                                    );
-                                }
-                            }
-                        }
-                    }
-                    _ => {}
-                });
-            }],
-            [WlOutput, 3, |output: Main<WlOutput>, _env: DispatchData| {
-                output.quick_assign(move |output, event, mut env| match event {
-                    wl_output::Event::Geometry {
-                        x: _,
-                        y: _,
-                        physical_width: _,
-                        physical_height: _,
-                        subpixel: _,
-                        mut make,
-                        model: _,
-                        transform: _,
-                    } => {
-                        if let Some(env) = env.get::<Env>() {
-                            if env.flags.output.is_none()
-                                || env.flags.output.as_ref().unwrap().eq(&make)
-                            {
-                                if let Some(status_manager) = &env.status_manager {
-                                    make = make.replace(' ', "").to_string();
-                                    let output_status =
-                                        status_manager.get_river_output_status(&output);
-                                    output_status.quick_assign(move |_, event, mut env| {
-                                        if let Some(env) = env.get::<Env>() {
-                                            match event {
-                                                zriver_output_status_v1::Event::FocusedTags {
-                                                    tags: focused_tags,
-                                                } => {
-                                                    if let Some(tags) = &mut env.tags {
-                                                        if let Some(inner_value) =
-                                                            tags.get_mut(&make)
-                                                        {
-                                                            (*inner_value) = Tags(focused_tags);
-                                                        } else {
-                                                            tags.insert(
-                                                                make.clone(),
-                                                                Tags(focused_tags),
-                                                            );
-                                                        }
-                                                    }
-                                                }
-                                                zriver_output_status_v1::Event::ViewTags {
-                                                    tags,
-                                                } => {
-                                                    if let Some(viewstag) = &mut env.viewstag {
-                                                        let tags: Vec<u32> = tags[0..]
-                                                            .chunks(4)
-                                                            .map(|s| {
-                                                                let buf = [s[0], s[1], s[2], s[3]];
-                                                                let tagmask =
-                                                                    u32::from_le_bytes(buf);
-                                                                for i in 0..32 {
-                                                                    if 1 << i == tagmask {
-                                                                        return 1 + i;
-                                                                    }
-                                                                }
-                                                                0
-                                                            })
-                                                            .collect();
-                                                        if let Some(inner_value) =
-                                                            viewstag.get_mut(&make)
-                                                        {
-                                                            (*inner_value) = tags;
-                                                        } else {
-                                                            viewstag.insert(make.clone(), tags);
-                                                        }
-                                                    }
-                                                }
-                                                zriver_output_status_v1::Event::UrgentTags {
-                                                    tags,
-                                                } => {
-                                                    if let Some(urgency) = &mut env.urgency {
-                                                        if let Some(inner_value) =
-                                                            urgency.get_mut(&make)
-                                                        {
-                                                            (*inner_value) = Tags(tags);
-                                                        } else {
-                                                            urgency
-                                                                .insert(make.clone(), Tags(tags));
-                                                        }
-                                                    }
-                                                }
-                                                zriver_output_status_v1::Event::LayoutName {
-                                                    name,
-                                                } => {
-                                                    env.layout = Some(name);
-                                                }
-                                                _ => {}
-                                            }
-                                        }
-                                    });
-                                }
-                            }
-                        }
-                    }
-                    _ => {}
-                });
-            }]
-        ),
     );
 
     loop {
@@ -263,7 +346,6 @@ fn configuration() -> Flags {
         match args.next() {
             Some(flag) => match flag.as_str() {
                 "--seat" | "-s" => default.seat = args.next(),
-                "--output" | "-o" => default.output = args.next(),
                 "--urgency" | "-u" => default.urgency = true,
                 "--title" | "-w" => default.title = true,
                 "--mode" | "-m" => default.mode = true,
@@ -272,14 +354,13 @@ fn configuration() -> Flags {
                 "--views-tag" | "-vt" => default.viewstag = true,
                 "--help" | "-h" => {
                     print!("Usage: ristate [option]\n\n");
-                    print!("  --tag | -t 			the focused tag\n");
+                    print!("  --tag | -t                the focused tag\n");
-                    print!("  --title | -w	   	 	the title of the focused view\n");
+                    print!("  --title | -w              the title of the focused view\n");
-                    print!("  --mode | -m	   	 	the current input mode\n");
+                    print!("  --mode | -m               the current input mode\n");
-                    print!("  --urgency | -u 		urgent tag\n");
+                    print!("  --urgency | -u            urgent tag\n");
-                    print!("  --views-tag | -vt    		the tag of all views\n");
+                    print!("  --views-tag | -vt         the tag of all views\n");
-                    print!("  --seat | -s <string>  	select the seat\n");
+                    print!("  --seat | -s <string>      select the seat\n");
-                    print!("  --output | -o <string> 	select the output\n");
+                    print!("  --layout | -l <string>    display the name of the layout\n");
-                    print!("  --layout | -l <string> 	display the name of the layout\n");
                     std::process::exit(0);
                 }
                 _ => {}