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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179

#![no_std]

use core::str::Utf8Error;

use adapters::StoreAdapter;
use modular_bitfield::prelude::*;

mod alloc;
mod grasshopper;
mod store;

pub mod adapters;

pub use alloc::*;
pub use grasshopper::*;
pub use store::*;

pub const MAX_KEY_LEN: usize = 256;
pub const MAX_VALUE_LEN: usize = 64 * 1024;

pub type Address = usize;

#[derive(Debug)]
pub struct Bucket {
    index: usize,
    raw: RawBucket,
}

impl Bucket {
    pub fn val_address(&self) -> Address {
        let addr = self.raw.address() + self.raw.key_len() as u32;
        addr as Address
    }

    pub(crate) fn index(&self) -> usize {
        self.index
    }

    pub(crate) fn address(&self) -> Address {
        self.raw.address() as Address
    }

    pub fn key_len(&self) -> usize {
        self.raw.key_len() as usize
    }

    pub fn val_len(&self) -> usize {
        self.raw.val_len() as usize
    }

    pub fn record_len(&self) -> usize {
        self.key_len() + self.val_len()
    }
}

#[derive(Debug, PartialEq)]
pub enum Error<E> {
    AdapterError(E),
    IndexOverflow,
    InvalidCapacity,
    InvalidNonce,
    InvalidPatchOffset,
    KeyNotFound,
    ReadOnlyStore,
    StoreNotFound,
    StoreOverflow,
    ValueOverflow,
    KeyOverflow,
    Utf8Error(Utf8Error),
    #[cfg(feature = "serde")]
    SerializationError(postcard::Error),
}

#[bitfield]
pub(crate) struct StoreHeader {
    magic: B32,
    nonce: B16,
    buckets: B16,
}

#[bitfield]
#[derive(Default, Debug, Clone)]
pub(crate) struct RawBucket {
    val_len: B16,
    key_len: B8,
    address: B24,
    hash: B16,
}

pub struct KeysIterator<'a, 'b, A, const BUCKETS: usize, const SLOTS: usize>
where
    A: StoreAdapter,
{
    store: &'a mut KVStore<A, BUCKETS, SLOTS>,
    prefix: Option<&'b [u8]>,
    cursor: usize,
}

impl<'a, 'b, A, const BUCKETS: usize, const SLOTS: usize> KeysIterator<'a, 'b, A, BUCKETS, SLOTS>
where
    A: StoreAdapter,
{
    pub fn new(store: &'a mut KVStore<A, BUCKETS, SLOTS>) -> Self {
        Self {
            store,
            cursor: 0,
            prefix: None,
        }
    }

    pub fn with_prefix(store: &'a mut KVStore<A, BUCKETS, SLOTS>, prefix: &'b [u8]) -> Self {
        Self {
            store,
            cursor: 0,
            prefix: Some(prefix),
        }
    }
}

impl<'a, 'b, A, const BUCKETS: usize, const SLOTS: usize> Iterator
    for KeysIterator<'a, 'b, A, BUCKETS, SLOTS>
where
    A: StoreAdapter,
{
    type Item = KeyReference;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if self.cursor >= BUCKETS {
                return None;
            }

            let raw = self.store.load_bucket(self.cursor).unwrap_or_default();
            let index = self.cursor;
            self.cursor += 1;

            let key_len = raw.key_len() as usize;
            let prefix_len = self.prefix.map_or(0, |prefix| prefix.len());

            if key_len > prefix_len {
                let val_len = raw.val_len() as usize;
                let bucket = Bucket { index, raw };
                let address = bucket.raw.address() as Address;
                let mut scratch = [0; MAX_KEY_LEN];

                self.store
                    .adapter()
                    .read(address, &mut scratch[..key_len])
                    .ok();

                if matches!(self.prefix, Some(prefix) if &scratch[..prefix_len] != prefix) {
                    continue;
                }

                return Some(KeyReference {
                    key_len,
                    val_len,
                    scratch,
                });
            }
        }
    }
}

pub struct KeyReference {
    key_len: usize,
    val_len: usize,
    scratch: [u8; MAX_KEY_LEN],
}

impl KeyReference {
    pub fn key(&self) -> &[u8] {
        &self.scratch[..self.key_len]
    }

    pub fn val_len(&self) -> usize {
        self.val_len
    }
}