Skip to main content

kevy_store/
string_rmw.rs

1//! Read-modify-write string ops split out of `string.rs` (500-LOC
2//! rule): `APPEND` / `GETSET` / `GETDEL` / `INCRBYFLOAT`. All four are
3//! encoding-aware across `Value::Str` / `Value::Int` / `Value::ArcBulk`
4//! — earlier `Str`-only arms replied WRONGTYPE where Redis succeeds;
5//! guard: `tests_string_encoding.rs`.
6
7#[cfg(not(feature = "std"))]
8use crate::nostd_prelude::*;
9use crate::string_set::pick_value_for_set_owned;
10use crate::util::{fmt_num, format_i64_into, itoa_i64_stack, parse_f64};
11use crate::value::{SmallBytes, Value};
12use crate::{Entry, Store, StoreError};
13
14impl Store {
15    /// Append to a string, creating it if absent, and return the new
16    /// length. An integer-encoded value is formatted back to bytes first —
17    /// APPEND leaves nothing integer-encoded.
18    pub fn append(&mut self, key: &[u8], data: &[u8]) -> Result<usize, StoreError> {
19        self.tier_resolve(key, crate::value::COLD_TAG_STRING)?; // cold string pages in
20        let outcome = match self.live_entry_mut(key) {
21            Some(e) => match &mut e.value {
22                Value::Str(v) => {
23                    // SmallBytes is immutable; pop out, grow via Vec, re-wrap.
24                    let mut owned = core::mem::take(v).into_vec();
25                    owned.extend_from_slice(data);
26                    let new_len = owned.len();
27                    *v = SmallBytes::from_vec(owned);
28                    AppendOutcome::Reweigh(new_len)
29                }
30                // L2: Int — materialise digits, append, re-pick the
31                // encoding (canonical results go straight back to Int).
32                Value::Int(n) => {
33                    let mut buf = itoa_i64_stack();
34                    let mut owned = format_i64_into(*n, &mut buf).to_vec();
35                    owned.extend_from_slice(data);
36                    let new_len = owned.len();
37                    e.value = pick_value_for_set_owned(owned);
38                    AppendOutcome::Reweigh(new_len)
39                }
40                // L1: APPEND on Arc-backed bulk → materialise to a fresh
41                // Vec (no other reader has refs to the old Arc post-replace),
42                // append, then pick the new encoding via SET routing rules.
43                Value::ArcBulk(a) => {
44                    let mut owned: Vec<u8> = a.as_ref().to_vec();
45                    owned.extend_from_slice(data);
46                    let new_len = owned.len();
47                    e.value = pick_value_for_set_owned(owned);
48                    AppendOutcome::Reweigh(new_len)
49                }
50                _ => return Err(StoreError::WrongType),
51            },
52            None => AppendOutcome::Insert,
53        };
54        match outcome {
55            AppendOutcome::Reweigh(new_len) => {
56                self.reweigh_entry(key);
57                Ok(new_len)
58            }
59            AppendOutcome::Insert => {
60                self.insert_entry(
61                    SmallBytes::from_slice(key),
62                    Entry::new(Value::Str(SmallBytes::from_slice(data)), None),
63                );
64                Ok(data.len())
65            }
66        }
67    }
68    /// `GETSET` — set to `val`, return the previous string (WRONGTYPE if the old
69    /// value isn't a string). Clears any TTL, like SET.
70    pub fn getset(&mut self, key: &[u8], val: Vec<u8>) -> Result<Option<Vec<u8>>, StoreError> {
71        self.tier_resolve(key, crate::value::COLD_TAG_STRING)?;
72        let old = match self.live_entry(key) {
73            Some(e) => match &e.value {
74                Value::Str(v) => Some(v.to_vec()),
75                Value::Int(n) => {
76                    let mut buf = itoa_i64_stack();
77                    Some(format_i64_into(*n, &mut buf).to_vec())
78                }
79                Value::ArcBulk(a) => Some(a.as_ref().to_vec()),
80                _ => return Err(StoreError::WrongType),
81            },
82            None => None,
83        };
84        // Route through the SET encoding rules so a canonical-int new
85        // value takes the Int fast path, same as a plain SET would.
86        self.insert_entry(
87            SmallBytes::from_slice(key),
88            Entry::new(pick_value_for_set_owned(val), None),
89        );
90        Ok(old)
91    }
92
93    /// `GETDEL` — get then delete (WRONGTYPE if non-string).
94    pub fn getdel(&mut self, key: &[u8]) -> Result<Option<Vec<u8>>, StoreError> {
95        self.tier_resolve(key, crate::value::COLD_TAG_STRING)?;
96        match self.live_entry(key) {
97            None => return Ok(None),
98            Some(e) => match &e.value {
99                Value::Str(_) | Value::Int(_) | Value::ArcBulk(_) => {}
100                _ => return Err(StoreError::WrongType),
101            },
102        }
103        match self.remove_entry(key) {
104            Some(Entry { value: Value::Str(v), .. }) => Ok(Some(v.into_vec())),
105            Some(Entry { value: Value::Int(n), .. }) => {
106                let mut buf = itoa_i64_stack();
107                Ok(Some(format_i64_into(n, &mut buf).to_vec()))
108            }
109            Some(Entry { value: Value::ArcBulk(a), .. }) => Ok(Some(a.as_ref().to_vec())),
110            _ => Ok(None),
111        }
112    }
113
114    /// `INCRBYFLOAT` — returns the new value formatted as Redis would. Preserves TTL.
115    pub fn incr_by_float(&mut self, key: &[u8], delta: f64) -> Result<Vec<u8>, StoreError> {
116        self.tier_resolve(key, crate::value::COLD_TAG_STRING)?;
117        let outcome = if let Some(e) = self.live_entry_mut(key) {
118            match &mut e.value {
119                Value::Str(v) => {
120                    let cur = parse_f64(v.as_slice()).ok_or(StoreError::NotFloat)?;
121                    let bytes = float_incr_bytes(cur, delta)?;
122                    *v = SmallBytes::from_slice(&bytes);
123                    FloatOutcome::Reweigh(bytes)
124                }
125                Value::Int(n) => {
126                    let bytes = float_incr_bytes(*n as f64, delta)?;
127                    e.value = Value::Str(SmallBytes::from_slice(&bytes));
128                    FloatOutcome::Reweigh(bytes)
129                }
130                Value::ArcBulk(a) => {
131                    let cur = parse_f64(a.as_ref()).ok_or(StoreError::NotFloat)?;
132                    let bytes = float_incr_bytes(cur, delta)?;
133                    e.value = Value::Str(SmallBytes::from_slice(&bytes));
134                    FloatOutcome::Reweigh(bytes)
135                }
136                _ => return Err(StoreError::WrongType),
137            }
138        } else {
139            // Absent/expired ⇒ start from 0.0.
140            if !delta.is_finite() {
141                return Err(StoreError::NotFloat);
142            }
143            FloatOutcome::Insert(fmt_num(delta))
144        };
145        match outcome {
146            FloatOutcome::Reweigh(bytes) => {
147                self.reweigh_entry(key);
148                Ok(bytes)
149            }
150            FloatOutcome::Insert(bytes) => {
151                self.insert_entry(
152                    SmallBytes::from_slice(key),
153                    Entry::new(Value::Str(SmallBytes::from_slice(&bytes)), None),
154                );
155                Ok(bytes)
156            }
157        }
158    }
159}
160
161enum AppendOutcome {
162    Reweigh(usize),
163    Insert,
164}
165
166enum FloatOutcome {
167    Reweigh(Vec<u8>),
168    Insert(Vec<u8>),
169}
170
171/// Shared tail of the three INCRBYFLOAT arms: add `delta`, reject a
172/// non-finite result (Redis `NaN`/`inf` guard), format Redis-style.
173#[inline]
174fn float_incr_bytes(cur: f64, delta: f64) -> Result<Vec<u8>, StoreError> {
175    let next = cur + delta;
176    if !next.is_finite() {
177        return Err(StoreError::NotFloat);
178    }
179    Ok(fmt_num(next))
180}