1use alloc::collections::{BTreeMap, BTreeSet};
2use alloc::string::{String, ToString};
3use alloc::sync::Arc;
4use alloc::vec::Vec;
5use core::str::from_utf8;
6
7use crate::dsl::{VALUE_IS_PLACEHOLDER_MASK, VALUE_WORD_ID_MASK};
8use crate::index::Index;
9use crate::provided::{Context as ContextTrait, ContextError, StoreError, LoadError, Tree};
10use crate::debug_log;
11use crate::required::{Stores, SetOutcome};
12
13pub struct Context<'r> {
16 index: Arc<Index>,
17 stores: &'r dyn Stores,
18 cache_keys: Vec<u16>,
19 cache_vals: Vec<Tree>,
20 called_paths: BTreeSet<u16>,
21}
22
23impl<'r> Context<'r> {
24 pub fn new(index: Arc<Index>, stores: &'r dyn Stores) -> Self {
25 Self {
26 index,
27 stores,
28 cache_keys: Vec::new(),
29 cache_vals: Vec::new(),
30 called_paths: BTreeSet::new(),
31 }
32 }
33
34 fn cache_get(&self, path_id: u16) -> Option<&Tree> {
35 self.cache_keys.iter()
36 .position(|&k| k == path_id)
37 .and_then(|i| self.cache_vals.get(i))
38 }
39
40 fn cache_set(&mut self, path_id: u16, value: Tree) {
41 if let Some(i) = self.cache_keys.iter().position(|&k| k == path_id) {
42 self.cache_vals[i] = value;
43 } else {
44 self.cache_keys.push(path_id);
45 self.cache_vals.push(value);
46 }
47 }
48
49 fn cache_remove(&mut self, path_id: u16) {
50 if let Some(i) = self.cache_keys.iter().position(|&k| k == path_id) {
51 self.cache_keys[i] = u16::MAX;
52 self.cache_vals[i] = Tree::Null;
53 }
54 }
55
56 fn guard_recursion(&self, path_id: u16) -> Result<(), ContextError> {
57 if self.called_paths.contains(&path_id) {
58 return Err(ContextError::RecursionLimitExceeded);
59 }
60 Ok(())
61 }
62}
63
64impl<'r> ContextTrait for Context<'r> {
67 fn get(&mut self, key: &str) -> Result<Option<Tree>, ContextError> {
68 debug_log!("Context", "get", key);
69 let leaves = self.index.traverse(key);
70 if leaves.is_empty() {
71 return Err(ContextError::KeyNotFound(key.to_string()));
72 }
73
74 if leaves.len() == 1 {
75 let leaf = &leaves[0];
76 self.guard_recursion(leaf.path_id)?;
77 self.called_paths.insert(leaf.path_id);
78 let result = self.resolve_leaf(leaf.path_id, leaf.leaf_id, leaf.value_id);
79 self.called_paths.remove(&leaf.path_id);
80 result
81 } else {
82 let mut pairs: Vec<(Vec<u8>, Tree)> = Vec::new();
83 for leaf in leaves.iter() {
84 self.guard_recursion(leaf.path_id)?;
85 self.called_paths.insert(leaf.path_id);
86 let value = self.resolve_leaf(leaf.path_id, leaf.leaf_id, leaf.value_id)?;
87 self.called_paths.remove(&leaf.path_id);
88 if let Some(v) = value {
89 let keyword = self.index.keyword_of(leaf.path_id).to_vec();
90 pairs.push((keyword, v));
91 }
92 }
93 Ok(if pairs.is_empty() { None } else { Some(Tree::Mapping(pairs)) })
94 }
95 }
96
97 fn set(&mut self, key: &str, value: Tree) -> Result<bool, ContextError> {
98 debug_log!("Context", "set", key, &crate::debug_log::format_arg(&value));
99 let leaves = self.index.traverse(key);
100 if leaves.is_empty() {
101 return Err(ContextError::KeyNotFound(key.to_string()));
102 }
103 let leaf = &leaves[0];
104
105 let (store_id, key_frags, _mk, _mv, args_keys, args_vals) = self.index.set_meta(leaf);
106 let (key_frags, args_keys, args_vals) = (key_frags.to_vec(), args_keys.to_vec(), args_vals.to_vec());
107 let store = self.stores.store_for(store_id)
108 .ok_or_else(|| ContextError::StoreFailed(
109 StoreError::ClientNotFound(store_id.to_string())
110 ))?;
111
112 let id_str = ToString::to_string(&leaf.path_id);
113 let store_key = self.resolve_key_frags(&key_frags)?
114 .unwrap_or_else(|| id_str.clone());
115
116 let mut owned_args = self.resolve_args(&args_keys, &args_vals)?;
117 owned_args.insert("value".to_string(), value.clone());
118 let store_args: BTreeMap<&str, Tree> = owned_args.iter().map(|(k, v)| (k.as_str(), v.clone())).collect();
119
120 match store.set(store_key.as_bytes(), &store_args) {
121 Some(SetOutcome::Created(_)) | Some(SetOutcome::Updated) => {
122 self.cache_set(leaf.path_id, value);
123 Ok(true)
124 }
125 None => Ok(false),
126 }
127 }
128
129 fn delete(&mut self, key: &str) -> Result<bool, ContextError> {
130 debug_log!("Context", "delete", key);
131 let leaves = self.index.traverse(key);
132 if leaves.is_empty() {
133 return Err(ContextError::KeyNotFound(key.to_string()));
134 }
135 let leaf = &leaves[0];
136
137 let (store_id, key_frags, _mk, _mv, args_keys, args_vals) = self.index.set_meta(leaf);
138 let (key_frags, args_keys, args_vals) = (key_frags.to_vec(), args_keys.to_vec(), args_vals.to_vec());
139 let store = self.stores.store_for(store_id)
140 .ok_or_else(|| ContextError::StoreFailed(
141 StoreError::ClientNotFound(store_id.to_string())
142 ))?;
143
144 let id_str = ToString::to_string(&leaf.path_id);
145 let store_key = self.resolve_key_frags(&key_frags)?
146 .unwrap_or_else(|| id_str.clone());
147
148 let owned_args = self.resolve_args(&args_keys, &args_vals)?;
149 let store_args: BTreeMap<&str, Tree> = owned_args.iter().map(|(k, v)| (k.as_str(), v.clone())).collect();
150 let ok = store.delete(store_key.as_bytes(), &store_args);
151 if ok {
152 self.cache_remove(leaf.path_id);
153 }
154 Ok(ok)
155 }
156
157 fn exists(&mut self, key: &str) -> Result<bool, ContextError> {
158 debug_log!("Context", "exists", key);
159 let leaves = self.index.traverse(key);
160 if leaves.is_empty() {
161 return Err(ContextError::KeyNotFound(key.to_string()));
162 }
163 let leaf = &leaves[0];
164
165 if let Some(v) = self.cache_get(leaf.path_id) {
166 debug_log!("Context", "exists", key, "-> cache hit");
167 return Ok(!matches!(v, Tree::Null));
168 }
169
170 let (store_id, key_frags, _mk, _mv, args_keys, args_vals) = self.index.set_meta(leaf);
171 let (key_frags, args_keys, args_vals) = (key_frags.to_vec(), args_keys.to_vec(), args_vals.to_vec());
172 let Some(store) = self.stores.store_for(store_id) else {
173 return Ok(false);
174 };
175
176 let id_str = ToString::to_string(&leaf.path_id);
177 let store_key = self.resolve_key_frags(&key_frags)?
178 .unwrap_or_else(|| id_str.clone());
179
180 let owned_args = self.resolve_args(&args_keys, &args_vals)?;
181 let store_args: BTreeMap<&str, Tree> = owned_args.iter().map(|(k, v)| (k.as_str(), v.clone())).collect();
182 Ok(store.get(store_key.as_bytes(), &store_args).is_some())
183 }
184}
185
186impl<'r> Context<'r> {
189 fn resolve_key_frags(&mut self, key_frags: &[u16]) -> Result<Option<String>, ContextError> {
190 if key_frags.is_empty() {
191 return Ok(None);
192 }
193 let mut buf = String::new();
194 for &f in key_frags {
195 let is_ph = (f & VALUE_IS_PLACEHOLDER_MASK) != 0;
196 let word_id = (f & VALUE_WORD_ID_MASK) as usize;
197 let segment = from_utf8(self.index.word_bytes(word_id)).unwrap_or("").to_string();
198 if is_ph {
199 match self.get(&segment)? {
200 Some(Tree::Scalar(b)) => buf.push_str(from_utf8(&b).unwrap_or("")),
201 Some(_) => {}
202 None => return Err(ContextError::LoadFailed(LoadError::NotFound(segment))),
203 }
204 } else {
205 buf.push_str(&segment);
206 }
207 }
208 Ok(Some(buf))
209 }
210
211 fn resolve_args(
212 &mut self,
213 args_keys: &[u16],
214 args_vals: &[u16],
215 ) -> Result<BTreeMap<String, Tree>, ContextError> {
216 let pairs: Vec<(String, u16)> = args_keys.iter().zip(args_vals.iter())
217 .filter_map(|(&key_word_id, &val_values_id)| {
218 let k = from_utf8(self.index.word_bytes(key_word_id as usize)).unwrap_or("").to_string();
219 if k.is_empty() { return None; }
220 Some((k, val_values_id))
221 })
222 .collect();
223 let mut map: BTreeMap<String, Tree> = BTreeMap::new();
224 for (k, val_values_id) in pairs {
225 let v = if val_values_id == 0 {
226 Tree::Null
227 } else {
228 let frags: Vec<u16> = self.index.values_slice(val_values_id as usize)
229 .unwrap_or(&[]).to_vec();
230 match self.resolve_key_frags(&frags) {
231 Ok(Some(s)) => Tree::Scalar(s.into_bytes()),
232 _ => Tree::Null,
233 }
234 };
235 map.insert(k, v);
236 }
237 Ok(map)
238 }
239
240 fn resolve_leaf(&mut self, path_id: u16, leaf_id: u16, value_id: u16) -> Result<Option<Tree>, ContextError> {
241 if let Some(v) = self.cache_get(path_id) {
242 debug_log!("Context", "resolve_leaf", &alloc::format!("path_id={path_id}"), "-> cache hit");
243 return Ok(Some(v.clone()));
244 }
245
246 let leaf_ref = crate::index::LeafRef { path_id, leaf_id, value_id };
247
248 let (set_store_id, set_key_frags, _set_map_keys, _set_map_vals, set_args_keys, set_args_vals) = self.index.set_meta(&leaf_ref);
249 let set_key_frags: Vec<u16> = set_key_frags.to_vec();
250 let set_args_keys: Vec<u16> = set_args_keys.to_vec();
251 let set_args_vals: Vec<u16> = set_args_vals.to_vec();
252 if set_store_id != 0 {
253 if let Some(store) = self.stores.store_for(set_store_id) {
254 let id_str = ToString::to_string(&path_id);
255 let store_key = self.resolve_key_frags(&set_key_frags)?
256 .unwrap_or_else(|| id_str.clone());
257 let owned_args = self.resolve_args(&set_args_keys, &set_args_vals)?;
258 let store_args: BTreeMap<&str, Tree> = owned_args.iter().map(|(k, v)| (k.as_str(), v.clone())).collect();
259 if let Some(value) = store.get(store_key.as_bytes(), &store_args) {
260 debug_log!("Context", "resolve_leaf", &alloc::format!("path_id={path_id}"), "-> _set hit");
261 self.cache_set(path_id, value.clone());
262 return Ok(Some(value));
263 }
264 }
265 }
266
267 let frags: Vec<(bool, Vec<u8>)> = self.index.leaf_fragments(&leaf_ref)
268 .into_iter()
269 .map(|(is_ph, b)| (is_ph, b.to_vec()))
270 .collect();
271 if !frags.is_empty() {
272 let value = if frags.len() == 1 && frags[0].0 {
273 let path_str = from_utf8(&frags[0].1)
274 .map_err(|_| ContextError::LoadFailed(
275 LoadError::ConfigMissing("placeholder utf8".to_string())
276 ))?
277 .to_string();
278 self.get(&path_str)?
279 .ok_or_else(|| ContextError::LoadFailed(
280 LoadError::NotFound(path_str.clone())
281 ))?
282 } else {
283 let mut buf = String::new();
284 for (is_ph, bytes) in frags {
285 if is_ph {
286 let path_str = from_utf8(&bytes)
287 .map_err(|_| ContextError::LoadFailed(
288 LoadError::ConfigMissing("placeholder utf8".to_string())
289 ))?;
290 match self.get(path_str)? {
291 Some(Tree::Scalar(b)) => {
292 buf.push_str(from_utf8(&b).unwrap_or(""));
293 }
294 Some(_) => {}
295 None => return Err(ContextError::LoadFailed(
296 LoadError::NotFound(path_str.to_string())
297 )),
298 }
299 } else {
300 buf.push_str(from_utf8(&bytes).unwrap_or(""));
301 }
302 }
303 Tree::Scalar(buf.into_bytes())
304 };
305 self.cache_set(path_id, value.clone());
306 return Ok(Some(value));
307 }
308
309 let (get_store_id, get_key_frags, get_map_keys, get_map_vals, get_args_keys, get_args_vals) = self.index.get_meta(&leaf_ref);
310 if get_store_id == 0 {
311 return Ok(None);
312 }
313 let get_key_frags: Vec<u16> = get_key_frags.to_vec();
314 let get_map_keys: Vec<u16> = get_map_keys.to_vec();
315 let get_map_vals: Vec<u16> = get_map_vals.to_vec();
316 let get_args_keys: Vec<u16> = get_args_keys.to_vec();
317 let get_args_vals: Vec<u16> = get_args_vals.to_vec();
318
319 let store = self.stores.store_for(get_store_id)
320 .ok_or_else(|| ContextError::LoadFailed(
321 LoadError::ClientNotFound(get_store_id.to_string())
322 ))?;
323
324 let id_str = ToString::to_string(&path_id);
325 let store_key = self.resolve_key_frags(&get_key_frags)?
326 .unwrap_or_else(|| id_str.clone());
327
328 let owned_get_args = self.resolve_args(&get_args_keys, &get_args_vals)?;
329 let store_args: BTreeMap<&str, Tree> = owned_get_args.iter().map(|(k, v)| (k.as_str(), v.clone())).collect();
330 let fetched = store.get(store_key.as_bytes(), &store_args)
331 .ok_or_else(|| ContextError::LoadFailed(
332 LoadError::NotFound(store_key.clone())
333 ))?;
334 debug_log!("Context", "resolve_leaf", &alloc::format!("path_id={path_id}"), "-> _get hit");
335
336 let value = if !get_map_keys.is_empty() {
337 if let Tree::Mapping(fetched_pairs) = &fetched {
338 for (&dst_path_id, &src_word_id) in get_map_keys.iter().zip(get_map_vals.iter()) {
339 let src_key = self.index.word_bytes(src_word_id as usize);
340 let fetched_val = fetched_pairs.iter()
341 .find(|(k, _)| k.as_slice() == src_key)
342 .map(|(_, v)| v.clone())
343 .unwrap_or(Tree::Null);
344 self.cache_set(dst_path_id, fetched_val);
345 }
346 }
347 debug_log!("Context", "resolve_leaf", &alloc::format!("path_id={path_id}"), "-> map expanded");
348 match self.cache_get(path_id) {
349 Some(v) => v.clone(),
350 None => return Ok(None),
351 }
352 } else {
353 fetched
354 };
355
356 if set_store_id != 0 {
357 if let Some(set_store) = self.stores.store_for(set_store_id) {
358 let set_store_key = self.resolve_key_frags(&set_key_frags)?
359 .unwrap_or_else(|| id_str.clone());
360 let mut sargs_owned = self.resolve_args(&set_args_keys, &set_args_vals)?;
361 sargs_owned.insert("value".to_string(), value.clone());
362 let sargs: BTreeMap<&str, Tree> = sargs_owned.iter().map(|(k, v)| (k.as_str(), v.clone())).collect();
363 debug_log!("Context", "resolve_leaf", &alloc::format!("path_id={path_id}"), "-> write-through to _set");
364 set_store.set(set_store_key.as_bytes(), &sargs);
365 let _ = id_str;
366 }
367 }
368
369 self.cache_set(path_id, value.clone());
370 Ok(Some(value))
371 }
372}