fix(i3-per-workspace-layout): prevent creating a cascade of single-child containers

users/myuser/graphical/i3-per-workspace-layout/src/main.rs

@@ -1,5 +1,3 @@
-// Inspired by https://github.com/chmln/i3-auto-layout (MIT licensed)
-
 use anyhow::{anyhow, Context, Result};
 use clap::Parser;
 use log::{debug, info};
@@ -8,8 +6,9 @@ use std::collections::HashMap;
 use std::path::PathBuf;
 use tokio::{sync::mpsc, task::JoinHandle};
 use tokio_i3ipc::{
-    event::{Event, Subscribe, WorkspaceChange},
+    event::{Event, Subscribe, WindowChange, WorkspaceChange},
     msg::Msg,
+    reply::{Node, NodeType},
     I3,
 };
 use tokio_stream::StreamExt;
@@ -33,13 +32,32 @@ struct Cli {
 /// ```
 #[derive(Debug, Serialize, Deserialize)]
 struct Config {
-    /// Whether to force setting the layout even on filled workspaces
-    #[serde(default)] // false
-    force: bool,
     /// The workspace -> layout associations
     layouts: HashMap<String, String>,
 }
 
+fn find_workspace_for_window(tree: &Node, window_id: usize) -> Option<&Node> {
+    fn inner<'a>(node: &'a Node, window_id: usize, mut ret_workspace: Option<&'a Node>) -> Option<&'a Node> {
+        if node.node_type == NodeType::Workspace {
+            ret_workspace = Some(node);
+        }
+
+        if node.id == window_id {
+            return ret_workspace;
+        }
+
+        for child in &node.nodes {
+            if let Some(workspace) = inner(child, window_id, ret_workspace) {
+                return Some(workspace);
+            }
+        }
+
+        None
+    }
+
+    inner(tree, window_id, None)
+}
+
 #[tokio::main]
 async fn main() -> Result<()> {
     flexi_logger::Logger::try_with_env()?.start()?;
@@ -54,48 +72,45 @@ async fn main() -> Result<()> {
     let s_handle: JoinHandle<Result<()>> = tokio::spawn(async move {
         let mut event_listener = {
             let mut i3 = I3::connect().await?;
-            i3.subscribe([Subscribe::Workspace]).await?;
+            i3.subscribe([Subscribe::Workspace, Subscribe::Window]).await?;
             i3.listen()
         };
 
+        // Second connection to allow querying workspaces when needed.
+        let mut i3 = I3::connect().await?;
+
         info!("Waiting for workspace events...");
-        loop {
-            let Some(Ok(Event::Workspace(workspace_event))) = event_listener.next().await else { continue };
-
-            debug!(
-                "Got workspace event: name={:?}, change={:?}",
-                workspace_event.current.clone().and_then(|x| x.name),
-                workspace_event.change
-            );
-
-            if WorkspaceChange::Focus == workspace_event.change {
-                let workspace = workspace_event
+        while let Some(Ok(event)) = event_listener.next().await {
+            match event {
+                Event::Workspace(data) if data.change == WorkspaceChange::Focus => {
+                    let workspace = data
                         .current
                         .ok_or_else(|| anyhow!("Missing current field on workspace event. Is your i3 up-to-date?"))?;
 
-                // Only change the layout if the workspace is empty or force == true
-                if !workspace.nodes.is_empty() && !config.force {
-                    continue;
+                    let Some(cmd) = cmd_toggle_layout(&config, &workspace) else { continue };
+                    send.send(cmd).await.context("Failed to queue command for sending")?;
                 }
-
-                let Some(name) = workspace.name else { continue };
-                let Some(desired_layout) = config.layouts.get(&name) else { continue };
-
-                send.send(format!("[con_id={}] layout {}", workspace.id, desired_layout))
-                    .await
-                    .context("Failed to queue command for sending")?;
-
-                debug!("Changed layout of workspace {:?} to {}", &name, desired_layout);
+                Event::Window(data) if matches!(data.change, WindowChange::New | WindowChange::Move) => {
+                    let tree = i3.get_tree().await?;
+                    if let Some(workspace) = find_workspace_for_window(&tree, data.container.id) {
+                        let Some(cmd) = cmd_toggle_layout(&config, workspace) else { continue };
+                        send.send(cmd).await.context("Failed to queue command for sending")?;
+                    } else {
+                        debug!("Ignoring window without workspace {:?}", data.container.id);
                     }
                 }
+                _ => {}
+            }
+        }
+        Ok(())
     });
 
     let r_handle: JoinHandle<Result<()>> = tokio::spawn(async move {
         let mut i3 = I3::connect().await?;
-        loop {
-            let Some(cmd) = recv.recv().await else { continue };
+        while let Some(cmd) = recv.recv().await {
             i3.send_msg_body(Msg::RunCommand, cmd).await?;
         }
+        Ok(())
     });
 
     let (send, recv) = tokio::try_join!(s_handle, r_handle)?;
@@ -103,3 +118,36 @@ async fn main() -> Result<()> {
     debug!("Shutting down...");
     Ok(())
 }
+
+fn cmd_toggle_layout(config: &Config, workspace: &Node) -> Option<String> {
+    let mut con = workspace;
+    let name = workspace.name.as_ref()?;
+    let desired_layout = config.layouts.get(name)?;
+
+    // If the workspace already has a single child node that is a container,
+    // we want to change the layout of that one instead.
+    if workspace.nodes.len() == 1 && workspace.nodes[0].node_type == NodeType::Con {
+        con = &workspace.nodes[0];
+
+        // This command works very strangely, as it always targets the parent
+        // container of the specified container. So we now have to find the first child
+        // and operate on that instead... Wow.
+        // If the container is empty, we refuse to do anything.
+        if con.nodes.is_empty() {
+            return None;
+        }
+
+        con = &con.nodes[0];
+    } else {
+        // Strangely enough, setting a layout on the workspace container will create a new
+        // container inside of it. So if we haven't found a single container to modify,
+        // we can operate on the workspace.
+    }
+
+    debug!(
+        "Changing layout of workspace {:?} to {} (modifying con_id={})",
+        &name, desired_layout, con.id
+    );
+
+    Some(format!("[con_id={}] layout {}", con.id, desired_layout))
+}