From 293d87d31c7f27cb0c6096a58a8b3eb271faf1e5 Mon Sep 17 00:00:00 2001
From: jyn <github@jyn.dev>
Date: Sat, 30 Nov 2024 00:31:03 -0500
Subject: [PATCH] Add support for zsh shell completions

This also fixes `rr ls` not to list pre-installed files as traces.
---
 CMakeLists.txt            |  5 +++++
 scripts/rr_completion.zsh | 42 +++++++++++++++++++++++++++++++++++++++
 src/util.cc               |  8 ++++++--
 3 files changed, 53 insertions(+), 2 deletions(-)
 create mode 100644 scripts/rr_completion.zsh

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 124fbf69c29..66b47177614 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -845,6 +845,11 @@ install(PROGRAMS scripts/signal-rr-recording.sh
 install(PROGRAMS scripts/rr_completion
   DESTINATION ${CMAKE_INSTALL_DATADIR}/bash-completion/completions RENAME rr)
 
+# Note that this works fine when installing to /usr/local, but zsh by default doesn't autoload *any* completions in HOME,
+# so people will still have to manually set `FPATH=~/.local/share/zsh/site-functions`.
+install(PROGRAMS scripts/rr_completion.zsh
+  DESTINATION ${CMAKE_INSTALL_DATADIR}/zsh/site-functions RENAME _rr)
+
 set(RR_INSTALL_LIBS rrpreload rrpage rr_exec_stub)
 if(RTLD_AUDIT)
   set(RR_INSTALL_LIBS ${RR_INSTALL_LIBS} rraudit)
diff --git a/scripts/rr_completion.zsh b/scripts/rr_completion.zsh
new file mode 100644
index 00000000000..04daa57f05b
--- /dev/null
+++ b/scripts/rr_completion.zsh
@@ -0,0 +1,42 @@
+#compdef rr
+
+_rr() {
+  # allow overridding rr with another command (e.g. if multiple versions are installed)
+  zstyle -s :completion${curcontext}options command rr
+  : ${rr:=rr}
+
+  _rr_subcommands() {
+    $rr --list-commands | cut -s -d ' ' -f 3
+  }
+  _rr_traces() {
+    $rr ls | grep -v '^cpu_lock$'
+  }
+  _external_commands() {
+    _path_commands
+    _absolute_command_paths
+  }
+
+  _arguments -C \
+    '1:subcommand:($(_rr_subcommands))' \
+    '*::arg:->args'
+
+  case $state in
+    args) ;;
+    *) return;;
+  esac
+
+  # different subcommands can have different options. show the appropriate options for each.
+  # this is not ideal; `reply=` forces zsh to rerun `rr help` each time you hit tab.
+  # the alternative though is rewriting half the code in _arguments.
+  zstyle -e ':completion:*:*:rr:*:options' command \
+    'reply=( '${(q)service}' help ${words:#-*} )'
+
+  case $line[1] in
+    # complete a command, then delegate to that command's completion script
+    # -A means "don't use _normal until we've completed a non-option"
+    record) _arguments -A '-*' '1:command: _external_commands' '*:: :_normal -p $service' --;;
+    replay|rerun|ps|sources|traceinfo|pack|dump) _arguments '*:trace:($(_rr_traces))' --;;
+    help) _arguments ':subcommand:($(_rr_subcommands))' --;;
+    # TODO: complete `filename` and `explicit-sources`
+  esac
+}
diff --git a/src/util.cc b/src/util.cc
index 50ce31886af..7a7065da02a 100644
--- a/src/util.cc
+++ b/src/util.cc
@@ -2288,6 +2288,10 @@ bool remove_latest_trace_symlink() {
   return true;
 }
 
+static bool ends_with(std::string_view str, std::string_view suffix) {
+  return str.size() >= suffix.size() && str.compare(str.size()-suffix.size(), suffix.size(), suffix) == 0;
+}
+
 bool is_valid_trace_name(const string& entry, std::string* reason) {
   // filename corresponds to dirname
   const string name = filename(entry.c_str());
@@ -2310,9 +2314,9 @@ bool is_valid_trace_name(const string& entry, std::string* reason) {
     }
     return false;
   }
-  if (name == "cpu_lock") {
+  if (name == "cpu_lock" || name == "src" || ends_with(name, ".xml")) {
     if (reason) {
-      *reason = "Name cpu_lock is reserved";
+      *reason = "Name " + name + " is reserved";
     }
     return false;
   }