use std::collections::HashMap;
use std::ffi::OsString;
use clap::ArgAction;
pub(crate) fn rewrite(root: &clap::Command, argv: &[OsString]) -> Vec<OsString> {
if argv.len() < 3 || argv[1].to_str() != Some("ktstr") {
return argv.to_vec();
}
let sub_name = match argv[2].to_str() {
Some(s) => s,
None => return argv.to_vec(),
};
let ktstr_cmd = match root.get_subcommands().find(|c| sub_matches(c, "ktstr")) {
Some(c) => c,
None => return argv.to_vec(),
};
let sub = match ktstr_cmd
.get_subcommands()
.find(|c| sub_matches(c, sub_name))
{
Some(c) => c,
None => return argv.to_vec(),
};
if !sub.get_arguments().any(clap::Arg::is_last_set) {
return argv.to_vec();
}
let (longs, shorts) = owned_flags(sub);
let (ktstr_bucket, pass_bucket) =
partition(&argv[3..], &longs, &shorts, rejected_flags(sub.get_name()));
let mut out = Vec::with_capacity(argv.len() + 1);
out.push(argv[0].clone());
out.push(argv[1].clone());
out.push(argv[2].clone());
out.extend(ktstr_bucket);
out.push(OsString::from("--"));
out.extend(pass_bucket);
out
}
fn sub_matches(c: &clap::Command, name: &str) -> bool {
c.get_name() == name || c.get_all_aliases().any(|a| a == name)
}
fn action_takes_value(action: &ArgAction) -> bool {
matches!(action, ArgAction::Set | ArgAction::Append)
}
fn owned_flags(sub: &clap::Command) -> (HashMap<String, bool>, HashMap<char, bool>) {
let mut longs: HashMap<String, bool> = HashMap::new();
let mut shorts: HashMap<char, bool> = HashMap::new();
for a in sub.get_arguments() {
let takes = action_takes_value(a.get_action());
if let Some(names) = a.get_long_and_visible_aliases() {
for l in names {
longs.insert(l.to_string(), takes);
}
}
if let Some(chs) = a.get_short_and_visible_aliases() {
for s in chs {
shorts.insert(s, takes);
}
}
}
longs.entry("help".to_string()).or_insert(false);
shorts.entry('h').or_insert(false);
(longs, shorts)
}
fn rejected_flags(sub_name: &str) -> &'static [&'static str] {
match sub_name {
"perf-delta" => &[
"a-scheduler",
"b-scheduler",
"a-topology",
"b-topology",
"dual-run",
],
_ => &[],
}
}
fn partition(
tail: &[OsString],
longs: &HashMap<String, bool>,
shorts: &HashMap<char, bool>,
rejected: &[&str],
) -> (Vec<OsString>, Vec<OsString>) {
let mut ktstr: Vec<OsString> = Vec::new();
let mut pass: Vec<OsString> = Vec::new();
let mut i = 0;
while i < tail.len() {
let tok = &tail[i];
match tok.to_str() {
Some("--") => {
pass.extend(tail[i + 1..].iter().cloned());
break;
}
Some(t) if t.starts_with("--") => {
let name = &t[2..];
let (flag, has_eq) = match name.split_once('=') {
Some((f, _)) => (f, true),
None => (name, false),
};
match longs.get(flag) {
Some(&takes_value) => {
ktstr.push(tok.clone());
if takes_value && !has_eq && i + 1 < tail.len() {
ktstr.push(tail[i + 1].clone());
i += 1;
}
}
None if rejected.contains(&flag) => ktstr.push(tok.clone()),
None => pass.push(tok.clone()),
}
}
Some(t) if t.starts_with('-') && t.len() > 1 => {
let first = t[1..].chars().next().expect("len > 1");
match shorts.get(&first) {
Some(&takes_value) => {
ktstr.push(tok.clone());
let glued_len = t.len() - 1 - first.len_utf8();
if takes_value && glued_len == 0 && i + 1 < tail.len() {
ktstr.push(tail[i + 1].clone());
i += 1;
}
}
None => pass.push(tok.clone()),
}
}
_ => pass.push(tok.clone()),
}
i += 1;
}
(ktstr, pass)
}