From 8c0c901fb1c508ba878aeaf371084daeebe52a24 Mon Sep 17 00:00:00 2001
From: Santo Cariotti <santo@dcariotti.me>
Date: Mon, 16 Oct 2023 14:16:25 +0200
Subject: Avoid `run_loop` flag

---
 src/main.rs  |  2 +-
 src/trace.rs | 44 +++++++++++++++++++++-----------------------
 src/ui.rs    | 14 +++++++++-----
 3 files changed, 31 insertions(+), 29 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 136c8c6..0656553 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -40,7 +40,7 @@ fn main() -> anyhow::Result<()> {
 
         ui.start(pid, &args)?;
     } else {
-        trace(pid, &args, true)?;
+        trace(pid, &args)?;
     }
 
     Ok(())
diff --git a/src/trace.rs b/src/trace.rs
index c097fba..e9ca481 100644
--- a/src/trace.rs
+++ b/src/trace.rs
@@ -41,7 +41,7 @@ pub fn attach(pid: Pid) -> anyhow::Result<()> {
 }
 
 /// Trace a process with `pid` ID and returns a list of `RegistersData`
-pub fn trace(pid: Pid, args: &Args, run_loop: bool) -> anyhow::Result<Vec<RegistersData>> {
+pub fn trace(pid: Pid, args: &Args) -> anyhow::Result<Vec<RegistersData>> {
     // First wait for the parent process
     _ = waitpid(pid, None)?;
 
@@ -55,30 +55,28 @@ pub fn trace(pid: Pid, args: &Args, run_loop: bool) -> anyhow::Result<Vec<Regist
 
     let mut lines: Vec<RegistersData> = Vec::new();
 
-    if run_loop {
-        // Since you have to do 2 syscalls (start and end) you have to alternate the print value,
-        // because it could be equals except for the `rax` register.
-        let mut have_to_print = true;
-
-        loop {
-            match trace_next(pid)? {
-                Some(reg) => {
-                    have_to_print ^= true;
-                    if have_to_print {
-                        if let Some(ref mut f) = f {
-                            writeln!(f, "{}", reg.output())?;
-                        }
-
-                        if args.no_tui {
-                            writeln!(io::stdout(), "{}", reg.output())?;
-                        }
-
-                        lines.push(reg);
+    // Since you have to do 2 syscalls (start and end) you have to alternate the print value,
+    // because it could be equals except for the `rax` register.
+    let mut have_to_print = true;
+
+    loop {
+        match trace_next(pid)? {
+            Some(reg) => {
+                have_to_print ^= true;
+                if have_to_print {
+                    if let Some(ref mut f) = f {
+                        writeln!(f, "{}", reg.output())?;
                     }
+
+                    if args.no_tui {
+                        writeln!(io::stdout(), "{}", reg.output())?;
+                    }
+
+                    lines.push(reg);
                 }
-                None => {
-                    break;
-                }
+            }
+            None => {
+                break;
             }
         }
     }
diff --git a/src/ui.rs b/src/ui.rs
index 8fd44b1..e6c8e1b 100644
--- a/src/ui.rs
+++ b/src/ui.rs
@@ -8,7 +8,7 @@ use crossterm::{
     terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen},
     ExecutableCommand,
 };
-use nix::unistd::Pid;
+use nix::{sys::wait::waitpid, unistd::Pid};
 use ratatui::{prelude::*, widgets::*};
 use std::io::{self, stdout};
 
@@ -66,12 +66,16 @@ impl UI {
         let mut have_to_trace = !args.command.is_some();
         let mut should_quit = false;
 
-        let lines = trace(pid, &args, args.command.is_some())?;
-        if lines.len() > 1 {
-            for line in lines {
-                self.add_line(line);
+        if args.command.is_some() {
+            let registers = trace(pid, &args)?;
+            for register in registers {
+                self.add_line(register);
             }
+        } else {
+            // First wait for the parent process
+            _ = waitpid(pid, None)?;
         }
+
         while !should_quit {
             if have_to_trace {
                 if let Some(reg) = trace_next(pid)? {
-- 
cgit v1.2.3-18-g5258