shardmap 0.3.2

Sharded embedded in-memory map with optional cache, protocol, and server internals
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

#[cfg(feature = "server")]
use bytes::BytesMut;

use crate::commands::redis::{
    array_bulk, bulk, define_redis_command, error, parse_i64, write_frame, wrong_arity,
};
use crate::protocol::Frame;
use crate::storage::{EmbeddedStore, RedisObjectError};

define_redis_command!(ZRandMember, "ZRANDMEMBER", false);

impl crate::commands::redis::RedisCommand for ZRandMember {
    fn execute(store: &EmbeddedStore, args: &[&[u8]]) -> Frame {
        let parsed = match parse_zrandmember_args(args) {
            Ok(parsed) => parsed,
            Err(frame) => return frame,
        };
        let entries = match store.zentries(parsed.key) {
            Ok(entries) => entries,
            Err(RedisObjectError::MissingKey) => Vec::new(),
            Err(RedisObjectError::WrongType) => return crate::commands::redis::wrongtype(),
        };
        zrandmember_frame(entries, parsed.count, parsed.with_scores)
    }

    #[cfg(feature = "server")]
    fn write_resp(store: &EmbeddedStore, args: &[&[u8]], out: &mut BytesMut) {
        write_frame(out, &Self::execute(store, args));
    }
}

#[derive(Clone, Copy)]
struct ZRandMemberArgs<'a> {
    key: &'a [u8],
    count: Option<i64>,
    with_scores: bool,
}

fn parse_zrandmember_args<'a>(
    args: &'a [&'a [u8]],
) -> std::result::Result<ZRandMemberArgs<'a>, Frame> {
    match args {
        [key] => Ok(ZRandMemberArgs {
            key,
            count: None,
            with_scores: false,
        }),
        [key, count] => Ok(ZRandMemberArgs {
            key,
            count: Some(
                parse_i64(count)
                    .map_err(|_| error("ERR value is not an integer or out of range"))?,
            ),
            with_scores: false,
        }),
        [key, count, withscores]
            if crate::commands::redis::eq_ignore_ascii_case(withscores, b"WITHSCORES") =>
        {
            Ok(ZRandMemberArgs {
                key,
                count: Some(
                    parse_i64(count)
                        .map_err(|_| error("ERR value is not an integer or out of range"))?,
                ),
                with_scores: true,
            })
        }
        _ => Err(wrong_arity("ZRANDMEMBER")),
    }
}

fn zrandmember_frame(entries: Vec<(Vec<u8>, f64)>, count: Option<i64>, with_scores: bool) -> Frame {
    if entries.is_empty() {
        return match count {
            None => Frame::Null,
            Some(_) => Frame::Array(Vec::new()),
        };
    }
    match count {
        None => {
            let (member, score) = entries[0].clone();
            if with_scores {
                Frame::Array(vec![bulk(member), bulk(score.to_string().into_bytes())])
            } else {
                bulk(member)
            }
        }
        Some(count) if count >= 0 => {
            let selected = entries.into_iter().take(count as usize).collect::<Vec<_>>();
            zrandmember_array(selected, with_scores)
        }
        Some(count) => {
            let wanted = count.unsigned_abs() as usize;
            let mut selected = Vec::with_capacity(wanted);
            for index in 0..wanted {
                selected.push(entries[index % entries.len()].clone());
            }
            zrandmember_array(selected, with_scores)
        }
    }
}

fn zrandmember_array(entries: Vec<(Vec<u8>, f64)>, with_scores: bool) -> Frame {
    if !with_scores {
        return array_bulk(entries.into_iter().map(|(member, _)| member).collect());
    }
    let mut out = Vec::with_capacity(entries.len().saturating_mul(2));
    for (member, score) in entries {
        out.push(bulk(member));
        out.push(bulk(score.to_string().into_bytes()));
    }
    Frame::Array(out)
}