diff --git a/Cargo.lock b/Cargo.lock index 25c9ffa5..233be45d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2376,6 +2376,7 @@ dependencies = [ "scroll 0.13.0", "serde", "serde_json", + "strsim", "thiserror", "toml 0.5.11", "winapi", diff --git a/Cargo.toml b/Cargo.toml index 4973003e..aee03e71 100644 --- a/Cargo.toml +++ b/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! diff --git a/cmd/stackmargin/src/lib.rs b/cmd/stackmargin/src/lib.rs index 021a60e0..77de1bd0 100644 --- a/cmd/stackmargin/src/lib.rs +++ b/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, +} #[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::>>()?, + ) + }; 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))?; diff --git a/humility-core/Cargo.toml b/humility-core/Cargo.toml index 3d79d595..19b4b640 100644 --- a/humility-core/Cargo.toml +++ b/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 diff --git a/humility-core/src/hubris.rs b/humility-core/src/hubris.rs index b8bfbc11..157f1617 100644 --- a/humility-core/src/hubris.rs +++ b/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 { + 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