redis_driver/resp/
command_args.rs

1use crate::resp::CommandArg;
2use smallvec::{smallvec, SmallVec};
3use std::{
4    collections::{BTreeMap, BTreeSet, HashMap, HashSet},
5    hash::BuildHasher,
6    iter::{once, Once},
7    ops::Deref,
8};
9
10#[derive(Debug, Clone)]
11pub enum CommandArgs {
12    Empty,
13    Single(CommandArg),
14    Array2([CommandArg; 2]),
15    Array3([CommandArg; 3]),
16    Array4([CommandArg; 4]),
17    Array5([CommandArg; 5]),
18    Vec(SmallVec<[CommandArg; 10]>),
19}
20
21impl CommandArgs {
22    #[must_use]
23    #[inline]
24    pub fn arg<A>(self, args: A) -> Self
25    where
26        A: IntoArgs,
27    {
28        args.into_args(self)
29    }
30
31    #[must_use]
32    #[inline]
33    pub fn arg_if<A>(self, condition: bool, arg: A) -> Self
34    where
35        A: IntoArgs,
36    {
37        if condition {
38            arg.into_args(self)
39        } else {
40            self
41        }
42    }
43
44    #[must_use]
45    #[inline]
46    pub fn len(&self) -> usize {
47        match self {
48            CommandArgs::Empty => 0,
49            CommandArgs::Single(_) => 1,
50            CommandArgs::Array2(_) => 2,
51            CommandArgs::Array3(_) => 3,
52            CommandArgs::Array4(_) => 4,
53            CommandArgs::Array5(_) => 5,
54            CommandArgs::Vec(v) => v.len(),
55        }
56    }
57
58    #[must_use]
59    #[inline]
60    pub fn is_empty(&self) -> bool {
61        self.len() == 0
62    }
63}
64
65impl Default for CommandArgs {
66    fn default() -> Self {
67        CommandArgs::Empty
68    }
69}
70
71impl<'a> IntoIterator for &'a CommandArgs {
72    type Item = &'a CommandArg;
73    type IntoIter = CommandArgsIterator<'a>;
74
75    fn into_iter(self) -> Self::IntoIter {
76        match self {
77            CommandArgs::Empty => CommandArgsIterator::Empty,
78            CommandArgs::Single(s) => CommandArgsIterator::Single(Some(s)),
79            CommandArgs::Array2(a) => CommandArgsIterator::Iter(a.iter()),
80            CommandArgs::Array3(a) => CommandArgsIterator::Iter(a.iter()),
81            CommandArgs::Array4(a) => CommandArgsIterator::Iter(a.iter()),
82            CommandArgs::Array5(a) => CommandArgsIterator::Iter(a.iter()),
83            CommandArgs::Vec(a) => CommandArgsIterator::Iter(a.iter()),
84        }
85    }
86}
87
88#[derive(Clone)]
89pub enum CommandArgsIterator<'a> {
90    Empty,
91    Single(Option<&'a CommandArg>),
92    Iter(std::slice::Iter<'a, CommandArg>),
93}
94
95impl<'a> Iterator for CommandArgsIterator<'a> {
96    type Item = &'a CommandArg;
97
98    fn next(&mut self) -> Option<Self::Item> {
99        match self {
100            CommandArgsIterator::Empty => None,
101            CommandArgsIterator::Single(s) => s.take(),
102            CommandArgsIterator::Iter(i) => i.next(),
103        }
104    }
105}
106
107impl Deref for CommandArgs {
108    type Target = [CommandArg];
109
110    fn deref(&self) -> &Self::Target {
111        match self {
112            CommandArgs::Empty => &[],
113            CommandArgs::Single(s) => std::slice::from_ref(s),
114            CommandArgs::Array2(a) => a,
115            CommandArgs::Array3(a) => a,
116            CommandArgs::Array4(a) => a,
117            CommandArgs::Array5(a) => a,
118            CommandArgs::Vec(v) => v,
119        }
120    }
121}
122
123/// Types compatible with command args
124pub trait IntoArgs {
125    fn into_args(self, args: CommandArgs) -> CommandArgs;
126    fn num_args(&self) -> usize {
127        unimplemented!()
128    }
129}
130
131impl<T> IntoArgs for T
132where
133    T: Into<CommandArg>,
134{
135    #[inline]
136    fn into_args(self, args: CommandArgs) -> CommandArgs {
137        match args {
138            CommandArgs::Empty => CommandArgs::Single(self.into()),
139            CommandArgs::Single(a) => CommandArgs::Array2([a, self.into()]),
140            CommandArgs::Array2(a) => {
141                let [item1, item2] = a;
142                CommandArgs::Array3([item1, item2, self.into()])
143            }
144            CommandArgs::Array3(a) => {
145                let [item1, item2, item3] = a;
146                CommandArgs::Array4([item1, item2, item3, self.into()])
147            }
148            CommandArgs::Array4(a) => {
149                let [item1, item2, item3, item4] = a;
150                CommandArgs::Array5([item1, item2, item3, item4, self.into()])
151            }
152            CommandArgs::Array5(a) => {
153                let [item1, item2, item3, item4, item5] = a;
154                CommandArgs::Vec(smallvec![item1, item2, item3, item4, item5, self.into()])
155            }
156            CommandArgs::Vec(mut vec) => {
157                vec.push(self.into());
158                CommandArgs::Vec(vec)
159            }
160        }
161    }
162
163    fn num_args(&self) -> usize {
164        1
165    }
166}
167
168impl<T> IntoArgs for Option<T>
169where
170    T: IntoArgs,
171{
172    #[inline]
173    fn into_args(self, args: CommandArgs) -> CommandArgs {
174        match self {
175            Some(s) => s.into_args(args),
176            None => args,
177        }
178    }
179
180    #[inline]
181    fn num_args(&self) -> usize {
182        match self {
183            Some(t) => t.num_args(),
184            None => 0,
185        }
186    }
187}
188
189impl<T, const N: usize> IntoArgs for [T; N]
190where
191    T: IntoArgs,
192{
193    #[inline]
194    fn into_args(self, args: CommandArgs) -> CommandArgs {
195        let mut args = args;
196        for a in self {
197            args = a.into_args(args);
198        }
199        args
200    }
201
202    #[inline]
203    fn num_args(&self) -> usize {
204        self.iter().fold(0, |acc, t| acc + t.num_args())
205    }
206}
207
208impl<T> IntoArgs for Vec<T>
209where
210    T: IntoArgs,
211{
212    #[inline]
213    fn into_args(self, args: CommandArgs) -> CommandArgs {
214        let mut args = args;
215        for a in self {
216            args = a.into_args(args);
217        }
218        args
219    }
220
221    #[inline]
222    fn num_args(&self) -> usize {
223        self.iter().fold(0, |acc, t| acc + t.num_args())
224    }
225}
226
227impl<T, A> IntoArgs for SmallVec<A>
228where
229    A: smallvec::Array<Item = T>,
230    T: IntoArgs,
231{
232    #[inline]
233    fn into_args(self, args: CommandArgs) -> CommandArgs {
234        let mut args = args;
235        for a in self {
236            args = a.into_args(args);
237        }
238        args
239    }
240
241    #[inline]
242    fn num_args(&self) -> usize {
243        self.iter().fold(0, |acc, t| acc + t.num_args())
244    }
245}
246
247impl<T, S: BuildHasher> IntoArgs for HashSet<T, S>
248where
249    T: IntoArgs,
250{
251    #[inline]
252    fn into_args(self, args: CommandArgs) -> CommandArgs {
253        let mut args = args;
254        for a in self {
255            args = a.into_args(args);
256        }
257        args
258    }
259
260    #[inline]
261    fn num_args(&self) -> usize {
262        self.iter().fold(0, |acc, t| acc + t.num_args())
263    }
264}
265
266impl<T> IntoArgs for BTreeSet<T>
267where
268    T: IntoArgs,
269{
270    #[inline]
271    fn into_args(self, args: CommandArgs) -> CommandArgs {
272        let mut args = args;
273        for a in self {
274            args = a.into_args(args);
275        }
276        args
277    }
278
279    #[inline]
280    fn num_args(&self) -> usize {
281        self.iter().fold(0, |acc, t| acc + t.num_args())
282    }
283}
284
285impl<K, V, S: BuildHasher> IntoArgs for HashMap<K, V, S>
286where
287    K: Into<CommandArg>,
288    V: Into<CommandArg>,
289{
290    #[inline]
291    fn into_args(self, args: CommandArgs) -> CommandArgs {
292        let mut args = args;
293        for (key, value) in self {
294            args = key.into_args(args);
295            args = value.into_args(args);
296        }
297        args
298    }
299
300    #[inline]
301    fn num_args(&self) -> usize {
302        self.iter().fold(0, |acc, (k, v)| acc + k.num_args() + v.num_args())
303    }
304}
305
306impl<K, V> IntoArgs for BTreeMap<K, V>
307where
308    K: Into<CommandArg>,
309    V: Into<CommandArg>,
310{
311    #[inline]
312    fn into_args(self, args: CommandArgs) -> CommandArgs {
313        let mut args = args;
314        for (key, value) in self {
315            args = key.into_args(args);
316            args = value.into_args(args);
317        }
318        args
319    }
320
321    #[inline]
322    fn num_args(&self) -> usize {
323        self.iter().fold(0, |acc, (k, v)| acc + k.num_args() + v.num_args())
324    }
325}
326
327impl<T, U> IntoArgs for (T, U)
328where
329    T: IntoArgs,
330    U: IntoArgs,
331{
332    #[inline]
333    fn into_args(self, args: CommandArgs) -> CommandArgs {
334        let args = self.0.into_args(args);
335        self.1.into_args(args)
336    }
337
338    #[inline]
339    fn num_args(&self) -> usize {
340        self.0.num_args() + self.1.num_args()
341    }
342}
343
344impl<T, U, V> IntoArgs for (T, U, V)
345where
346    T: IntoArgs,
347    U: IntoArgs,
348    V: IntoArgs,
349{
350    #[inline]
351    fn into_args(self, args: CommandArgs) -> CommandArgs {
352        let args = self.0.into_args(args);
353        let args = self.1.into_args(args);
354        self.2.into_args(args)
355    }
356
357    #[inline]
358    fn num_args(&self) -> usize {
359        self.0.num_args() + self.1.num_args() + self.2.num_args()
360    }
361}
362
363/// Allow to merge `CommandArgs` in another `CommandArgs`
364impl IntoArgs for CommandArgs {
365    #[inline]
366    fn into_args(self, args: CommandArgs) -> CommandArgs {
367        match self {
368            CommandArgs::Empty => args,
369            CommandArgs::Single(s) => args.arg(s),
370            CommandArgs::Array2(a) => args.arg(a),
371            CommandArgs::Array3(a) => args.arg(a),
372            CommandArgs::Array4(a) => args.arg(a),
373            CommandArgs::Array5(a) => args.arg(a),
374            CommandArgs::Vec(v) => args.arg(v),
375        }
376    }
377}
378
379/// Generic Marker for Collections of `IntoArgs`
380pub trait ArgsOrCollection<T>: IntoArgs
381where
382    T: IntoArgs,
383{
384}
385
386impl<T, const N: usize> ArgsOrCollection<T> for [T; N] where T: IntoArgs {}
387impl<T> ArgsOrCollection<T> for Vec<T> where T: IntoArgs {}
388impl<T> ArgsOrCollection<T> for T where T: IntoArgs {}
389
390/// Marker for collections of single items (directly convertible to `CommandArg`) of `IntoArgs`
391pub trait SingleArgOrCollection<T>: IntoArgs
392where
393    T: Into<CommandArg>,
394{
395    type IntoIter: Iterator<Item = T>;
396
397    fn into_iter(self) -> Self::IntoIter;
398}
399
400impl<T, const N: usize> SingleArgOrCollection<T> for [T; N]
401where
402    T: Into<CommandArg>,
403{
404    type IntoIter = std::array::IntoIter<T, N>;
405
406    fn into_iter(self) -> Self::IntoIter {
407        IntoIterator::into_iter(self)
408    }
409}
410
411impl<T> SingleArgOrCollection<T> for Vec<T>
412where
413    T: Into<CommandArg>,
414{
415    type IntoIter = std::vec::IntoIter<T>;
416
417    fn into_iter(self) -> Self::IntoIter {
418        IntoIterator::into_iter(self)
419    }
420}
421
422impl<A, T> SingleArgOrCollection<T> for SmallVec<A>
423where
424    A: smallvec::Array<Item = T>,
425    T: Into<CommandArg>,
426{
427    type IntoIter = smallvec::IntoIter<A>;
428
429    fn into_iter(self) -> Self::IntoIter {
430        IntoIterator::into_iter(self)
431    }
432}
433
434impl<T, S: BuildHasher> SingleArgOrCollection<T> for HashSet<T, S>
435where
436    T: Into<CommandArg>,
437{
438    type IntoIter = std::collections::hash_set::IntoIter<T>;
439
440    fn into_iter(self) -> Self::IntoIter {
441        IntoIterator::into_iter(self)
442    }
443}
444
445impl<T> SingleArgOrCollection<T> for BTreeSet<T>
446where
447    T: Into<CommandArg>,
448{
449    type IntoIter = std::collections::btree_set::IntoIter<T>;
450
451    fn into_iter(self) -> Self::IntoIter {
452        IntoIterator::into_iter(self)
453    }
454}
455
456impl<T> SingleArgOrCollection<T> for T
457where
458    T: Into<CommandArg>,
459{
460    type IntoIter = Once<T>;
461
462    fn into_iter(self) -> Self::IntoIter {
463        once(self)
464    }
465}
466
467/// Marker for key/value collections of Args
468pub trait KeyValueArgOrCollection<K, V>: IntoArgs
469where
470    K: Into<CommandArg>,
471    V: Into<CommandArg>,
472{
473}
474
475impl<K, V> KeyValueArgOrCollection<K, V> for Vec<(K, V)>
476where
477    K: Into<CommandArg>,
478    V: Into<CommandArg>,
479{
480}
481
482impl<K, V, const N: usize> KeyValueArgOrCollection<K, V> for [(K, V); N]
483where
484    K: Into<CommandArg>,
485    V: Into<CommandArg>,
486{
487}
488
489impl<K, V> KeyValueArgOrCollection<K, V> for (K, V)
490where
491    K: Into<CommandArg>,
492    V: Into<CommandArg>,
493{
494}
495
496impl<K, V, S: BuildHasher> KeyValueArgOrCollection<K, V> for HashMap<K, V, S>
497where
498    K: Into<CommandArg>,
499    V: Into<CommandArg>,
500{
501}
502
503impl<K, V> KeyValueArgOrCollection<K, V> for BTreeMap<K, V>
504where
505    K: Into<CommandArg>,
506    V: Into<CommandArg>,
507{
508}