From b0814b8a834487d0ddfe070ab1fdb8e199d4d5c2 Mon Sep 17 00:00:00 2001
From: aNNiMON <annimon119@gmail.com>
Date: Wed, 24 Jul 2024 19:54:04 +0300
Subject: [PATCH] Show found processes as tree in Filtered mode

---
 src/main.rs | 45 +++++++++++++++++++++++++++------------------
 1 file changed, 27 insertions(+), 18 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index a0c3f28..92d90cd 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,6 @@
 use itertools::Itertools;
 use processes::{terminate_process, ProcessInfo};
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::process::ExitCode;
 use std::{env, fmt::Display};
 
@@ -18,11 +18,6 @@ enum Mode {
     Help,
 }
 
-// pkill                 => all unique processed, excl. svchost
-// pkill --all           => all processes, incl. svchost
-// pkill <regex>         => filtered processes
-// pkill --pid name      => PID of the first occurence
-// pkill --kill PID/name => terminate process by its PID or name
 fn main() -> ExitCode {
     match parse_mode() {
         Mode::All => {
@@ -30,7 +25,11 @@ fn main() -> ExitCode {
             ExitCode::SUCCESS
         }
         Mode::Tree => {
-            show_processes_tree();
+            let processes: Vec<ProcessInfo> = processes::list_processes()
+                .into_iter()
+                .sorted_by_key(|p| p.parent_pid)
+                .collect();
+            show_processes_tree(processes);
             ExitCode::SUCCESS
         }
         Mode::Minimal => {
@@ -38,6 +37,8 @@ fn main() -> ExitCode {
                 "svchost",
                 "registry",
                 "[system process]",
+                "system",
+                "memory compression",
                 "smss.exe",
                 "csrss.exe",
                 "lsass.exe",
@@ -55,11 +56,24 @@ fn main() -> ExitCode {
             ExitCode::SUCCESS
         }
         Mode::Filtered(filter) => {
-            show_processes_table(
-                processes::list_processes()
-                    .iter()
-                    .filter(|p| p.name.to_ascii_lowercase().contains(&filter))
-                    .sorted_by_key(|p| &p.name),
+            let all_processes = processes::list_processes();
+            let filtered_pids: HashSet<u32> = all_processes
+                .clone()
+                .into_iter()
+                .flat_map(|p| {
+                    if p.name.to_ascii_lowercase().contains(&filter) {
+                        vec![p.pid, p.parent_pid]
+                    } else {
+                        vec![]
+                    }
+                })
+                .collect();
+            show_processes_tree(
+                all_processes
+                    .into_iter()
+                    .filter(|p| filtered_pids.contains(&p.pid))
+                    .sorted_by_key(|p| p.name.clone())
+                    .collect(),
             );
             ExitCode::SUCCESS
         }
@@ -145,12 +159,7 @@ fn parse_mode() -> Mode {
     mode
 }
 
-fn show_processes_tree() {
-    let processes: Vec<ProcessInfo> = processes::list_processes()
-        .into_iter()
-        .sorted_by_key(|p| p.parent_pid)
-        .collect();
-
+fn show_processes_tree(processes: Vec<ProcessInfo>) {
     let by_pid: HashMap<u32, ProcessInfo> =
         processes.clone().into_iter().map(|p| (p.pid, p)).collect();
     let mut tree: HashMap<u32, Vec<ProcessInfo>> = HashMap::new();