feature: option to have process tree entries be collapsed by default (#1770)

* Add option to have process tree collapsed by default

* Fix collapse logic

* format

* tweak how it's done

* oops

* slight tweaks to the no-children collapse logic

* update schema

---------

Co-authored-by: ceres <ceres.bezuidenhout@trintel.co.za>
Co-authored-by: Bucket-Bucket-Bucket <107044719+Bucket-Bucket-Bucket@users.noreply.github.com>

CHANGELOG.md

@@ -30,6 +30,7 @@ That said, these are more guidelines rather than hardset rules, though the proje
 - [#1717](https://github.com/ClementTsang/bottom/pull/1717): Support delete key (fn + delete on macOS) to kill processes.
 - [#1306](https://github.com/ClementTsang/bottom/pull/1306): Support using left/right key to collapse/expand process trees respectively.
 - [#1767](https://github.com/ClementTsang/bottom/pull/1767): Add a virtual memory column for processes.
+- [#1770](https://github.com/ClementTsang/bottom/pull/1770) (originally [#1627](https://github.com/ClementTsang/bottom/pull/1627)): Add option to have process tree entries be collapsed by default.
 
 ### Bug Fixes
 

docs/content/configuration/command-line-options.md

@@ -37,6 +37,7 @@ see information on these options by running `btm -h`, or run `btm --help` to dis
 | `-T, --tree`                | Makes the process widget use tree mode by default.                                            |
 | `-n, --unnormalized_cpu`    | Show process CPU% usage without averaging over the number of CPU cores.                       |
 | `-W, --whole_word`          | Enables whole-word matching by default while searching.                                       |
+| `--tree_collapse`           | Collapse process tree by default.                                                             |
 
 ## Temperature Options
 

docs/content/configuration/config-file/flags.md

@@ -51,3 +51,4 @@ each time:
 | `memory_legend`              | String (one of ["none", "top-left", "top", "top-right", "left", "right", "bottom-left", "bottom", "bottom-right"]) | Where to place the legend for the memory widget.                                                                             |
 | `network_legend`             | String (one of ["none", "top-left", "top", "top-right", "left", "right", "bottom-left", "bottom", "bottom-right"]) | Where to place the legend for the network widget.                                                                            |
 | `average_cpu_row`            | Boolean                                                                                                            | Moves the average CPU usage entry to its own row when using basic mode.                                                      |
+| `tree_collapse`              | Boolean                                                                                                            | Collapse process tree by default.                                                                                            |

schema/nightly/bottom.json

@@ -2,7 +2,7 @@
   "$id": "https://github.com/ClementTsang/bottom/blob/main/schema/nightly/bottom.json",
   "$schema": "https://json-schema.org/draft/2020-12/schema",
   "title": "Schema for bottom's config file (nightly)",
-  "description": "https://clementtsang.github.io/bottom/nightly/configuration/config-file",
+  "description": "https://bottom.pages.dev/nightly/configuration/config-file/",
   "type": "object",
   "properties": {
     "cpu": {
@@ -493,6 +493,12 @@
             "null"
           ]
         },
+        "tree_collapse": {
+          "type": [
+            "boolean",
+            "null"
+          ]
+        },
         "unnormalized_cpu": {
           "type": [
             "boolean",
@@ -710,6 +716,8 @@
         "GPU%",
         "Mem",
         "Mem%",
+        "Memory",
+        "Memory%",
         "Name",
         "PID",
         "R/s",
@@ -721,7 +729,13 @@
         "TRead",
         "TWrite",
         "Time",
+        "Total Read",
+        "Total Write",
         "User",
+        "Virt",
+        "VirtMem",
+        "Virtual",
+        "Virtual Memory",
         "W/s",
         "Wps",
         "Write"

src/app.rs

@@ -14,6 +14,7 @@ pub use states::*;
 use unicode_segmentation::{GraphemeCursor, UnicodeSegmentation};
 
 use crate::canvas::dialogs::process_kill_dialog::ProcessKillDialog;
+use crate::widgets::TreeCollapsed;
 use crate::{
     canvas::components::time_graph::LegendPosition,
     constants,
@@ -62,6 +63,7 @@ pub struct AppConfigFields {
     pub network_use_binary_prefix: bool,
     pub retention_ms: u64,
     pub dedicated_average_row: bool,
+    pub default_tree_collapse: bool,
 }
 
 /// For filtering out information
@@ -423,9 +425,9 @@ impl App {
                     proc_widget_state.force_rerender_and_update();
                 }
                 ProcWidgetMode::Normal => {
-                    proc_widget_state.mode = ProcWidgetMode::Tree {
+                    proc_widget_state.mode = ProcWidgetMode::Tree(TreeCollapsed::new(
-                        collapsed_pids: Default::default(),
+                        self.app_config_fields.default_tree_collapse,
-                    };
+                    ));
                     proc_widget_state.force_rerender_and_update();
                 }
                 ProcWidgetMode::Grouped => {}

src/options.rs

@@ -227,6 +227,7 @@ pub(crate) fn init_app(args: BottomArgs, config: Config) -> Result<(App, BottomL
     #[cfg(any(target_os = "linux", target_os = "macos", target_os = "freebsd"))]
     let is_advanced_kill = !(is_flag_enabled!(disable_advanced_kill, args.process, config));
     let process_memory_as_value = is_flag_enabled!(process_memory_as_value, args.process, config);
+    let is_default_tree_collapsed = is_flag_enabled!(tree_collapse, args.process, config);
 
     // For CPU
     let default_cpu_selection = get_default_cpu_selection(args, config);
@@ -306,6 +307,7 @@ pub(crate) fn init_app(args: BottomArgs, config: Config) -> Result<(App, BottomL
         network_use_binary_prefix,
         retention_ms,
         dedicated_average_row: get_dedicated_avg_row(config),
+        default_tree_collapse: is_default_tree_collapsed,
     };
 
     let table_config = ProcTableConfig {
@@ -383,9 +385,7 @@ pub(crate) fn init_app(args: BottomArgs, config: Config) -> Result<(App, BottomL
                             let mode = if is_grouped {
                                 ProcWidgetMode::Grouped
                             } else if is_default_tree {
-                                ProcWidgetMode::Tree {
+                                ProcWidgetMode::Tree(TreeCollapsed::new(is_default_tree_collapsed))
-                                    collapsed_pids: Default::default(),
-                                }
                             } else {
                                 ProcWidgetMode::Normal
                             };

src/options/args.rs

@@ -371,6 +371,14 @@ pub struct ProcessArgs {
         alias = "whole-word"
     )]
     pub whole_word: bool,
+
+    #[arg(
+        long,
+        action = ArgAction::SetTrue,
+        help = "Collapse process tree by default.",
+        alias = "tree-collapse"
+    )]
+    pub tree_collapse: bool,
 }
 
 /// Temperature arguments/config options.

src/options/config/flags.rs

@@ -45,4 +45,5 @@ pub(crate) struct FlagConfig {
     pub(crate) enable_cache_memory: Option<bool>,
     pub(crate) retention: Option<StringOrNum>,
     pub(crate) average_cpu_row: Option<bool>, // FIXME: This makes no sense outside of basic mode, add a basic mode config section.
+    pub(crate) tree_collapse: Option<bool>,
 }

src/widgets/process_table.rs

@@ -60,9 +60,83 @@ impl ProcessSearchState {
     }
 }
 
+/// Whether to expand or collapse by default.
 #[derive(Clone, Debug, Eq, PartialEq)]
-pub enum ProcWidgetMode {
+pub(crate) enum TreeCollapsed {
-    Tree { collapsed_pids: HashSet<Pid> },
+    DefaultCollapse { expanded_pids: HashSet<Pid> },
+    DefaultExpand { collapsed_pids: HashSet<Pid> },
+}
+
+impl TreeCollapsed {
+    /// Creates a new [`TreeCollapsed`].
+    pub(crate) fn new(default_collapsed: bool) -> Self {
+        if default_collapsed {
+            TreeCollapsed::DefaultCollapse {
+                expanded_pids: HashSet::new(),
+            }
+        } else {
+            TreeCollapsed::DefaultExpand {
+                collapsed_pids: HashSet::new(),
+            }
+        }
+    }
+
+    /// Check whether the given PID is collapsed.
+    pub(crate) fn is_collapsed(&self, pid: Pid) -> bool {
+        match self {
+            TreeCollapsed::DefaultCollapse { expanded_pids } => !expanded_pids.contains(&pid),
+            TreeCollapsed::DefaultExpand { collapsed_pids } => collapsed_pids.contains(&pid),
+        }
+    }
+
+    /// Collapse the given PID.
+    pub(crate) fn collapse(&mut self, pid: Pid) {
+        match self {
+            TreeCollapsed::DefaultCollapse { expanded_pids } => {
+                expanded_pids.remove(&pid);
+            }
+            TreeCollapsed::DefaultExpand { collapsed_pids } => {
+                collapsed_pids.insert(pid);
+            }
+        }
+    }
+
+    /// Expand the given PID.
+    pub(crate) fn expand(&mut self, pid: Pid) {
+        match self {
+            TreeCollapsed::DefaultCollapse { expanded_pids } => {
+                expanded_pids.insert(pid);
+            }
+            TreeCollapsed::DefaultExpand { collapsed_pids } => {
+                collapsed_pids.remove(&pid);
+            }
+        }
+    }
+
+    /// Toggle the given PID.
+    pub(crate) fn toggle(&mut self, pid: Pid) {
+        match self {
+            TreeCollapsed::DefaultCollapse { expanded_pids } => {
+                if expanded_pids.contains(&pid) {
+                    expanded_pids.remove(&pid);
+                } else {
+                    expanded_pids.insert(pid);
+                }
+            }
+            TreeCollapsed::DefaultExpand { collapsed_pids } => {
+                if collapsed_pids.contains(&pid) {
+                    collapsed_pids.remove(&pid);
+                } else {
+                    collapsed_pids.insert(pid);
+                }
+            }
+        }
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub(crate) enum ProcWidgetMode {
+    Tree(TreeCollapsed),
     Grouped,
     Normal,
 }
@@ -132,7 +206,7 @@ pub enum ProcWidgetColumn {
 // This is temporary. Switch back to `ProcColumn` later!
 
 pub struct ProcWidgetState {
-    pub mode: ProcWidgetMode,
+    pub(crate) mode: ProcWidgetMode,
 
     /// The state of the search box.
     pub proc_search: ProcessSearchState,
@@ -200,7 +274,7 @@ impl ProcWidgetState {
         DataTable::new_sortable(columns, props, styling)
     }
 
-    pub fn new(
+    pub(crate) fn new(
         config: &AppConfigFields, mode: ProcWidgetMode, table_config: ProcTableConfig,
         colours: &Styles, config_columns: &Option<IndexSet<ProcWidgetColumn>>,
     ) -> Self {
@@ -404,16 +478,14 @@ impl ProcWidgetState {
             ProcWidgetMode::Grouped | ProcWidgetMode::Normal => {
                 self.get_normal_data(&stored_data.process_data.process_harvest)
             }
-            ProcWidgetMode::Tree { collapsed_pids } => {
+            ProcWidgetMode::Tree(collapse) => self.get_tree_data(collapse, stored_data),
-                self.get_tree_data(collapsed_pids, stored_data)
-            }
         };
         self.table.set_data(data);
         self.force_update_data = false;
     }
 
     fn get_tree_data(
-        &self, collapsed_pids: &HashSet<Pid>, stored_data: &StoredData,
+        &self, collapsed: &TreeCollapsed, stored_data: &StoredData,
     ) -> Vec<ProcWidgetData> {
         const BRANCH_END: char = '└';
         const BRANCH_SPLIT: char = '├';
@@ -572,8 +644,9 @@ impl ProcWidgetState {
             let disabled = !kept_pids.contains(&process.pid);
             let is_last = *siblings_left == 0;
 
-            if collapsed_pids.contains(&process.pid) {
+            if collapsed.is_collapsed(process.pid) {
                 let mut summed_process = process.clone();
+                let mut has_children = false;
 
                 if let Some(children_pids) = filtered_tree.get(&process.pid) {
                     let mut sum_queue = children_pids
@@ -596,13 +669,27 @@ impl ProcWidgetState {
                             }));
                         }
                     }
+
+                    has_children = !children_pids.is_empty();
                 }
 
-                let prefix = if prefixes.is_empty() {
+                // This is so that if an entry is "collapsed" but there are no children, avoid drawing the "+".
-                    "+ ".to_string()
+                let prefix = if has_children {
+                    if prefixes.is_empty() {
+                        "+ ".to_string()
+                    } else {
+                        format!(
+                            "{}{}{} + ",
+                            prefixes.join(""),
+                            if is_last { BRANCH_END } else { BRANCH_SPLIT },
+                            BRANCH_HORIZONTAL
+                        )
+                    }
+                } else if prefixes.is_empty() {
+                    String::default()
                 } else {
                     format!(
-                        "{}{}{} + ",
+                        "{}{}{} ",
                         prefixes.join(""),
                         if is_last { BRANCH_END } else { BRANCH_SPLIT },
                         BRANCH_HORIZONTAL
@@ -839,35 +926,27 @@ impl ProcWidgetState {
     }
 
     pub fn collapse_current_tree_branch_entry(&mut self) {
-        if let ProcWidgetMode::Tree { collapsed_pids } = &mut self.mode {
+        if let ProcWidgetMode::Tree(collapsed) = &mut self.mode {
             if let Some(process) = self.table.current_item() {
-                let pid = process.pid;
+                collapsed.collapse(process.pid);
-
-                collapsed_pids.insert(pid);
                 self.force_data_update();
             }
         }
     }
 
     pub fn expand_current_tree_branch_entry(&mut self) {
-        if let ProcWidgetMode::Tree { collapsed_pids } = &mut self.mode {
+        if let ProcWidgetMode::Tree(collapsed) = &mut self.mode {
             if let Some(process) = self.table.current_item() {
-                let pid = process.pid;
+                collapsed.expand(process.pid);
-
-                collapsed_pids.remove(&pid);
                 self.force_data_update();
             }
         }
     }
 
     pub fn toggle_current_tree_branch_entry(&mut self) {
-        if let ProcWidgetMode::Tree { collapsed_pids } = &mut self.mode {
+        if let ProcWidgetMode::Tree(collapsed) = &mut self.mode {
             if let Some(process) = self.table.current_item() {
-                let pid = process.pid;
+                collapsed.toggle(process.pid);
-
-                if !collapsed_pids.remove(&pid) {
-                    collapsed_pids.insert(pid);
-                }
                 self.force_data_update();
             }
         }
@@ -1539,4 +1618,44 @@ mod test {
         state.toggle_command();
         assert_eq!(get_columns(&state.table), original_columns);
     }
+
+    /// Sanity test to ensure tree collapse logic works, both when enabled-by-default or disabled-by-default.
+    #[test]
+    fn test_tree_collapse() {
+        {
+            let mut collapsed_by_default = TreeCollapsed::new(true);
+
+            assert!(collapsed_by_default.is_collapsed(1));
+
+            collapsed_by_default.collapse(1);
+            assert!(collapsed_by_default.is_collapsed(1));
+
+            collapsed_by_default.expand(1);
+            assert!(!collapsed_by_default.is_collapsed(1));
+
+            collapsed_by_default.toggle(1);
+            assert!(collapsed_by_default.is_collapsed(1));
+
+            collapsed_by_default.toggle(1);
+            assert!(!collapsed_by_default.is_collapsed(1));
+        }
+
+        {
+            let mut expanded_by_default = TreeCollapsed::new(false);
+
+            assert!(!expanded_by_default.is_collapsed(1));
+
+            expanded_by_default.collapse(1);
+            assert!(expanded_by_default.is_collapsed(1));
+
+            expanded_by_default.expand(1);
+            assert!(!expanded_by_default.is_collapsed(1));
+
+            expanded_by_default.toggle(1);
+            assert!(expanded_by_default.is_collapsed(1));
+
+            expanded_by_default.toggle(1);
+            assert!(!expanded_by_default.is_collapsed(1));
+        }
+    }
 }