Allow humility stackmargin to take task names as arguments

Cargo.toml

@@ -251,6 +251,7 @@ srec = "0.2"
 strum = "0.25"
 strum_macros = "0.25"
 syn = "1.0"
+strsim = "0.10"
 tempfile = "3.3"
 termimad = "0.21"
 termios = "0.3" # not usable on windows!

cmd/stackmargin/src/lib.rs

@@ -30,19 +30,37 @@
 use anyhow::{Result, bail};
 use clap::{CommandFactory, Parser};
 use humility::hubris::*;
-use humility_cli::ExecutionContext;
+use humility_cli::{ExecutionContext, Subcommand};
 use humility_cmd::{Archive, Attach, Command, CommandKind, Validate};
-use std::convert::TryInto;
+use std::{collections::HashSet, convert::TryInto};
 
 #[derive(Parser, Debug)]
 #[clap(name = "stackmargin", about = env!("CARGO_PKG_DESCRIPTION"))]
-struct StackmarginArgs {}
+struct StackmarginArgs {
+    /// Tasks to check (leave empty to check all tasks)
+    #[clap(multiple_occurrences = true)]
+    tasks: Vec<String>,
+}
 
 #[rustfmt::skip::macros(println, bail)]
 fn stackmargin(context: &mut ExecutionContext) -> Result<()> {
+    let Subcommand::Other(subargs) = context.cli.cmd.as_ref().unwrap();
+    let subargs = StackmarginArgs::try_parse_from(subargs)?;
+
     let core = &mut **context.core.as_mut().unwrap();
     let hubris = context.archive.as_ref().unwrap();
 
+    let valid_tasks = if subargs.tasks.is_empty() {
+        None
+    } else {
+        Some(
+            subargs
+                .tasks
+                .iter()
+                .map(|t| hubris.try_lookup_task(t))
+                .collect::<Result<HashSet<HubrisTask>>>()?,
+        )
+    };
     let regions = hubris.regions(core)?;
 
     let (base, size) = hubris.task_table(core)?;
@@ -91,6 +109,12 @@ fn stackmargin(context: &mut ExecutionContext) -> Result<()> {
         {
             continue;
         }
+        if valid_tasks
+            .as_ref()
+            .is_some_and(|tasks| !tasks.contains(&HubrisTask::Task(i)))
+        {
+            continue;
+        }
 
         let module = hubris.lookup_module(HubrisTask::Task(i))?;
 

humility-core/Cargo.toml

@@ -13,9 +13,9 @@ fallible-iterator.workspace = true
 gimli.workspace = true
 goblin.workspace = true
 hubpack.workspace = true
-humility_load_derive.workspace = true
 humility-arch-arm.workspace = true
 humility-log.workspace = true
+humility_load_derive.workspace = true
 humpty.workspace = true
 idol.workspace = true
 indexmap.workspace = true
@@ -33,6 +33,7 @@ rustc-demangle.workspace = true
 scroll.workspace = true
 serde.workspace = true
 serde_json.workspace = true
+strsim = { workspace = true }
 thiserror.workspace = true
 toml.workspace = true
 zerocopy.workspace = true

humility-core/src/hubris.rs

@@ -1980,6 +1980,39 @@ impl HubrisArchive {
         self.tasks.get(name)
     }
 
+    /// Tries to look up a task by name, returning an error with similar names
+    pub fn try_lookup_task(&self, name: &str) -> Result<HubrisTask> {
+        self.tasks
+            .get(name)
+            .cloned()
+            .ok_or_else(|| self.task_name_suggestion(name))
+    }
+
+    fn task_name_suggestion(&self, name: &str) -> anyhow::Error {
+        // Suggest only for very small differences
+        // High number can result in inaccurate suggestions for short queries e.g. `rls`
+        const MAX_DISTANCE: usize = 3;
+
+        let mut scored: Vec<_> = self
+            .tasks
+            .keys()
+            .filter_map(|s| {
+                let distance = strsim::damerau_levenshtein(name, s);
+                if distance <= MAX_DISTANCE {
+                    Some((distance, s))
+                } else {
+                    None
+                }
+            })
+            .collect();
+        scored.sort();
+        let mut out = format!("'{name}' is not a valid task name.");
+        if let Some((_, s)) = scored.first() {
+            out.push_str(&format!(" Did you mean '{s}'?"));
+        }
+        anyhow!("{out}")
+    }
+
     pub fn task_name(&self, index: usize) -> Option<&str> {
         let index = HubrisTask::Task(index as u32);
         // TODO this is super gross but we don't have the inverse of the tasks