xref: /third_party/rust/crates/clap/src/mkeymap.rs (revision 19625d8c) 
1use std::iter::Iterator;
2use std::ops::Index;
3
4use crate::builder::OsStr;
5use crate::Arg;
6use crate::INTERNAL_ERROR_MSG;
7
8#[derive(PartialEq, Eq, Debug, Clone)]
9pub(crate) struct Key {
10    key: KeyType,
11    index: usize,
12}
13
14#[derive(Default, PartialEq, Eq, Debug, Clone)]
15pub(crate) struct MKeyMap {
16    /// All of the arguments.
17    args: Vec<Arg>,
18
19    // Cache part:
20    /// Will be set after `_build()`.
21    keys: Vec<Key>,
22}
23
24#[derive(Debug, PartialEq, Eq, Hash, Clone)]
25pub(crate) enum KeyType {
26    Short(char),
27    Long(OsStr),
28    Position(usize),
29}
30
31impl KeyType {
32    pub(crate) fn is_position(&self) -> bool {
33        matches!(self, KeyType::Position(_))
34    }
35}
36
37impl PartialEq<usize> for KeyType {
38    fn eq(&self, rhs: &usize) -> bool {
39        match self {
40            KeyType::Position(x) => x == rhs,
41            _ => false,
42        }
43    }
44}
45
46impl PartialEq<&str> for KeyType {
47    fn eq(&self, rhs: &&str) -> bool {
48        match self {
49            KeyType::Long(l) => l == rhs,
50            _ => false,
51        }
52    }
53}
54
55impl PartialEq<str> for KeyType {
56    fn eq(&self, rhs: &str) -> bool {
57        match self {
58            KeyType::Long(l) => l == rhs,
59            _ => false,
60        }
61    }
62}
63
64impl PartialEq<OsStr> for KeyType {
65    fn eq(&self, rhs: &OsStr) -> bool {
66        match self {
67            KeyType::Long(l) => l == rhs,
68            _ => false,
69        }
70    }
71}
72
73impl PartialEq<char> for KeyType {
74    fn eq(&self, rhs: &char) -> bool {
75        match self {
76            KeyType::Short(c) => c == rhs,
77            _ => false,
78        }
79    }
80}
81
82impl MKeyMap {
83    /// If any arg has corresponding key in this map, we can search the key with
84    /// u64(for positional argument), char(for short flag), &str and OsString
85    /// (for long flag)
86    pub(crate) fn contains<K>(&self, key: K) -> bool
87    where
88        KeyType: PartialEq<K>,
89    {
90        self.keys.iter().any(|x| x.key == key)
91    }
92
93    /// Push an argument in the map.
94    pub(crate) fn push(&mut self, new_arg: Arg) {
95        self.args.push(new_arg);
96    }
97
98    /// Find the arg have corresponding key in this map, we can search the key
99    /// with u64(for positional argument), char(for short flag), &str and
100    /// OsString (for long flag)
101    pub(crate) fn get<K: ?Sized>(&self, key: &K) -> Option<&Arg>
102    where
103        KeyType: PartialEq<K>,
104    {
105        self.keys
106            .iter()
107            .find(|k| &k.key == key)
108            .map(|k| &self.args[k.index])
109    }
110
111    /// Return iterators of all keys.
112    pub(crate) fn keys(&self) -> impl Iterator<Item = &KeyType> {
113        self.keys.iter().map(|x| &x.key)
114    }
115
116    /// Return iterators of all args.
117    pub(crate) fn args(&self) -> impl Iterator<Item = &Arg> {
118        self.args.iter()
119    }
120
121    /// Return mutable iterators of all args.
122    pub(crate) fn args_mut(&mut self) -> impl Iterator<Item = &mut Arg> {
123        self.args.iter_mut()
124    }
125
126    /// We need a lazy build here since some we may change args after creating
127    /// the map, you can checkout who uses `args_mut`.
128    pub(crate) fn _build(&mut self) {
129        // There will be at least as many keys as args, so that is a good starting point
130        self.keys.reserve(self.args.len());
131        for (i, arg) in self.args.iter().enumerate() {
132            append_keys(&mut self.keys, arg, i);
133        }
134    }
135
136    /// Remove an arg in the graph by Id, usually used by `mut_arg`. Return
137    /// `Some(arg)` if removed.
138    pub(crate) fn remove_by_name(&mut self, name: &str) -> Option<Arg> {
139        self.args
140            .iter()
141            .position(|arg| arg.id == name)
142            // since it's a cold function, using this wouldn't hurt much
143            .map(|i| self.args.remove(i))
144    }
145}
146
147impl Index<&'_ KeyType> for MKeyMap {
148    type Output = Arg;
149
150    fn index(&self, key: &KeyType) -> &Self::Output {
151        self.get(key).expect(INTERNAL_ERROR_MSG)
152    }
153}
154
155/// Generate key types for an specific Arg.
156fn append_keys(keys: &mut Vec<Key>, arg: &Arg, index: usize) {
157    if let Some(pos_index) = arg.index {
158        let key = KeyType::Position(pos_index);
159        keys.push(Key { key, index });
160    } else {
161        if let Some(short) = arg.short {
162            let key = KeyType::Short(short);
163            keys.push(Key { key, index });
164        }
165        if let Some(long) = arg.long.clone() {
166            let key = KeyType::Long(long.into());
167            keys.push(Key { key, index });
168        }
169
170        for (short, _) in arg.short_aliases.iter() {
171            let key = KeyType::Short(*short);
172            keys.push(Key { key, index });
173        }
174        for (long, _) in arg.aliases.iter() {
175            let key = KeyType::Long(long.into());
176            keys.push(Key { key, index });
177        }
178    }
179}
180