Skip to main content

context_engine/
context.rs

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
13// ── Context ───────────────────────────────────────────────────────────────────
14
15pub 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
64// ── Context trait impl ────────────────────────────────────────────────────────
65
66impl<'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
186// ── private helpers ───────────────────────────────────────────────────────────
187
188impl<'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}