Skip to main content

context_engine/
index.rs

1use alloc::vec::Vec;
2use crate::dsl::{
3    PATH_IS_LEAF_MASK,
4    PATH_KEYWORD_ID_SHIFT, PATH_KEYWORD_ID_MASK,
5    PATH_PARENT_ID_SHIFT,  PATH_PARENT_ID_MASK,
6    PATH_CHILD_ID_SHIFT,   PATH_CHILD_ID_MASK,
7    PATH_VALUE_ID_MASK,
8    VALUE_IS_PLACEHOLDER_MASK, VALUE_WORD_ID_MASK,
9    LEAF_WIDTH,
10    LEAF_GET_STORE_ID,    LEAF_GET_KEY_ID,
11    LEAF_SET_STORE_ID,    LEAF_SET_KEY_ID,
12    LEAF_GET_MAP_KEY_ID,  LEAF_GET_MAP_VAL_ID,
13    LEAF_GET_ARGS_KEY_ID, LEAF_GET_ARGS_VAL_ID,
14    LEAF_SET_MAP_KEY_ID,  LEAF_SET_MAP_VAL_ID,
15    LEAF_SET_ARGS_KEY_ID, LEAF_SET_ARGS_VAL_ID,
16};
17use crate::list::{List, VariableList};
18use crate::debug_log;
19
20const EMPTY_SLICE: &[u16] = &[];
21
22// ── LeafRef ───────────────────────────────────────────────────────────────────
23
24pub struct LeafRef {
25    pub path_id: u16,
26    pub leaf_id:  u16,
27    pub value_id: u16,
28}
29
30// ── Index ─────────────────────────────────────────────────────────────────────
31
32pub struct Index {
33    paths:     List<u64>,
34    children:  VariableList<u16>,
35    leaves:    VariableList<u16>,
36    values:    VariableList<u16>,
37    words:     VariableList<u8>,
38    map_keys:  VariableList<u16>,
39    map_vals:  VariableList<u16>,
40    args_keys: VariableList<u16>,
41    args_vals: VariableList<u16>,
42}
43
44impl Index {
45    pub fn new(
46        paths:     List<u64>,
47        children:  VariableList<u16>,
48        leaves:    VariableList<u16>,
49        values:    VariableList<u16>,
50        words:     VariableList<u8>,
51        map_keys:  VariableList<u16>,
52        map_vals:  VariableList<u16>,
53        args_keys: VariableList<u16>,
54        args_vals: VariableList<u16>,
55    ) -> Self {
56        Self { paths, children, leaves, values, words, map_keys, map_vals, args_keys, args_vals }
57    }
58
59    /// Traverse to the path node matching `path` (dot-separated keywords),
60    /// then collect all leaf descendants into a flat list.
61    ///
62    /// ```
63    /// # extern crate alloc;
64    /// use context_engine::{Tree, Index, dsl::Dsl};
65    /// let tree = Tree::Mapping(alloc::vec![(b"id".to_vec(), Tree::Null)]);
66    /// let (paths, children, leaves, values, words, map_keys, map_vals, args_keys, args_vals) = Dsl::compile(&tree, &[]).unwrap();
67    /// let index = Index::new(paths, children, leaves, values, words, map_keys, map_vals, args_keys, args_vals);
68    /// assert_eq!(index.traverse("id").len(), 1);
69    /// assert!(index.traverse("missing").is_empty());
70    /// ```
71    pub fn traverse(&self, path: &str) -> alloc::boxed::Box<[LeafRef]> {
72        let mut result = Vec::new();
73        let Some(path_id) = self.find(path) else {
74            debug_log!("Index", "traverse", path, "-> not found");
75            return result.into_boxed_slice();
76        };
77        self.collect_leaves(path_id, &mut result);
78        debug_log!("Index", "traverse", path, &alloc::format!("-> {} leaf(s)", result.len()));
79        result.into_boxed_slice()
80    }
81
82    /// Return the keyword bytes for a path node.
83    ///
84    /// ```
85    /// # extern crate alloc;
86    /// use context_engine::{Tree, Index, dsl::Dsl};
87    /// let tree = Tree::Mapping(alloc::vec![(b"name".to_vec(), Tree::Null)]);
88    /// let (paths, children, leaves, values, words, map_keys, map_vals, args_keys, args_vals) = Dsl::compile(&tree, &[]).unwrap();
89    /// let index = Index::new(paths, children, leaves, values, words, map_keys, map_vals, args_keys, args_vals);
90    /// assert_eq!(index.keyword_of(1), b"name");
91    /// ```
92    pub fn keyword_of(&self, path_id: u16) -> &[u8] {
93        let path = self.paths.data[path_id as usize];
94        let word_id = ((path & PATH_KEYWORD_ID_MASK) >> PATH_KEYWORD_ID_SHIFT) as usize;
95        let kw = self.word_bytes(word_id);
96        debug_log!("Index", "keyword_of", &alloc::format!("path_id={path_id}"), &alloc::format!("-> {:?}", core::str::from_utf8(kw).unwrap_or("?")));
97        kw
98    }
99
100    /// Extract `_get` meta for the given leaf.
101    /// Returns (store_id, key_frags, map_keys, map_vals, args_keys, args_vals).
102    /// store_id=0 means no `_get` configured; all slices are empty in that case.
103    /// key_frags: VALUE_IS_PLACEHOLDER_MASK-flagged u16 fragments (same encoding as leaf values).
104    /// map_keys: dst path_id列, map_vals: src word_id列.
105    /// args_keys: arg name word_id列, args_vals: arg value values_id列 (each indexes into values).
106    pub fn get_meta(&self, leaf: &LeafRef) -> (u8, &[u16], &[u16], &[u16], &[u16], &[u16]) {
107        let r = self.decode_meta(leaf, MetaKind::Get);
108        debug_log!("Index", "get_meta", &alloc::format!("path_id={}", leaf.path_id), &alloc::format!("-> store_id={}", r.0));
109        r
110    }
111
112    /// Extract `_set` meta for the given leaf.
113    /// Same layout as `get_meta`.
114    pub fn set_meta(&self, leaf: &LeafRef) -> (u8, &[u16], &[u16], &[u16], &[u16], &[u16]) {
115        let r = self.decode_meta(leaf, MetaKind::Set);
116        debug_log!("Index", "set_meta", &alloc::format!("path_id={}", leaf.path_id), &alloc::format!("-> store_id={}", r.0));
117        r
118    }
119
120    /// Return the value fragments encoded in the leaf.
121    ///
122    /// Each element is `(is_placeholder, bytes)`:
123    /// - `false` — static string literal
124    /// - `true`  — `${path}` reference whose bytes are the path string
125    ///
126    /// An empty slice means the leaf has no value (`Null`).
127    ///
128    /// ```
129    /// # extern crate alloc;
130    /// use context_engine::{Tree, Index, dsl::Dsl};
131    /// let tree = Tree::Mapping(alloc::vec![
132    ///     (b"driver".to_vec(), Tree::Scalar(b"postgres".to_vec())),
133    /// ]);
134    /// let (paths, children, leaves, values, words, map_keys, map_vals, args_keys, args_vals) = Dsl::compile(&tree, &[]).unwrap();
135    /// let index = Index::new(paths, children, leaves, values, words, map_keys, map_vals, args_keys, args_vals);
136    /// let leaves = index.traverse("driver");
137    /// let frags = index.leaf_fragments(&leaves[0]);
138    /// assert_eq!(frags.len(), 1);
139    /// assert_eq!(frags[0].0, false);
140    /// assert_eq!(frags[0].1, b"postgres");
141    /// ```
142    pub fn leaf_fragments(&self, leaf: &LeafRef) -> Vec<(bool, &[u8])> {
143        let value_id = leaf.value_id as usize;
144        if value_id == 0 {
145            debug_log!("Index", "leaf_fragments", &alloc::format!("path_id={}", leaf.path_id), "-> null");
146            return Vec::new();
147        }
148        let Some(frags) = self.values_slice(value_id) else { return Vec::new(); };
149        let result: Vec<(bool, &[u8])> = frags.iter().map(|&f| {
150            let is_ph  = (f & VALUE_IS_PLACEHOLDER_MASK) != 0;
151            let word_id = (f & VALUE_WORD_ID_MASK) as usize;
152            (is_ph, self.word_bytes(word_id))
153        }).collect();
154        debug_log!("Index", "leaf_fragments", &alloc::format!("path_id={}", leaf.path_id), &alloc::format!("-> {} frag(s)", result.len()));
155        result
156    }
157
158    /// leaf の path_id と同じ親を持つ sibling を keyword で探して path_id を返す。
159    /// map dst の cache_set 先を特定するために使う。
160    pub fn path_id_of_keyword(&self, leaf_path_id: u16, keyword: &[u8]) -> Option<u16> {
161        let path = self.paths.data[leaf_path_id as usize];
162        let parent_id = ((path & PATH_PARENT_ID_MASK) >> PATH_PARENT_ID_SHIFT) as u16;
163        let mut current = parent_id;
164        for segment in keyword.split(|&b| b == b'.') {
165            current = self.find_child(current, segment)?;
166        }
167        Some(current)
168    }
169}
170
171// ── private ───────────────────────────────────────────────────────────────────
172
173enum MetaKind { Get, Set }
174
175impl Index {
176    fn find(&self, path: &str) -> Option<u16> {
177        let mut current: u16 = 0;
178        for segment in path.split('.') {
179            current = self.find_child(current, segment.as_bytes())?;
180        }
181        Some(current)
182    }
183
184    fn find_child(&self, path_id: u16, keyword: &[u8]) -> Option<u16> {
185        let path = self.paths.data[path_id as usize];
186        if path & PATH_IS_LEAF_MASK != 0 { return None; }
187        let children_id = ((path & PATH_CHILD_ID_MASK) >> PATH_CHILD_ID_SHIFT) as usize;
188        if children_id == 0 { return None; }
189        let Some(children) = self.children_slice(children_id) else { return None; };
190        children.iter().copied().find(|&child_id| self.keyword_of(child_id) == keyword)
191    }
192
193    fn collect_leaves(&self, path_id: u16, out: &mut Vec<LeafRef>) {
194        let path = self.paths.data[path_id as usize];
195        if path & PATH_IS_LEAF_MASK != 0 {
196            let leaf_id  = ((path & PATH_CHILD_ID_MASK) >> PATH_CHILD_ID_SHIFT) as u16;
197            let value_id = (path & PATH_VALUE_ID_MASK) as u16;
198            out.push(LeafRef { path_id, leaf_id, value_id });
199            return;
200        }
201        let children_id = ((path & PATH_CHILD_ID_MASK) >> PATH_CHILD_ID_SHIFT) as usize;
202        if children_id == 0 { return; }
203        let children: Vec<u16> = match self.children_slice(children_id) {
204            Some(s) => s.to_vec(),
205            None    => return,
206        };
207        for child_id in children {
208            self.collect_leaves(child_id, out);
209        }
210    }
211
212    fn decode_meta(&self, leaf: &LeafRef, kind: MetaKind) -> (u8, &[u16], &[u16], &[u16], &[u16], &[u16]) {
213        let leaf_id = leaf.leaf_id as usize;
214        if leaf_id == 0 { return (0, EMPTY_SLICE, EMPTY_SLICE, EMPTY_SLICE, EMPTY_SLICE, EMPTY_SLICE); }
215        let Some(leaf_data) = self.leaves_slice(leaf_id) else {
216            return (0, EMPTY_SLICE, EMPTY_SLICE, EMPTY_SLICE, EMPTY_SLICE, EMPTY_SLICE);
217        };
218        if leaf_data.len() < LEAF_WIDTH {
219            return (0, EMPTY_SLICE, EMPTY_SLICE, EMPTY_SLICE, EMPTY_SLICE, EMPTY_SLICE);
220        }
221
222        let (store_id_raw, key_id, map_key_id, map_val_id, args_key_id, args_val_id) = match kind {
223            MetaKind::Get => (
224                leaf_data[LEAF_GET_STORE_ID],
225                leaf_data[LEAF_GET_KEY_ID],
226                leaf_data[LEAF_GET_MAP_KEY_ID],
227                leaf_data[LEAF_GET_MAP_VAL_ID],
228                leaf_data[LEAF_GET_ARGS_KEY_ID],
229                leaf_data[LEAF_GET_ARGS_VAL_ID],
230            ),
231            MetaKind::Set => (
232                leaf_data[LEAF_SET_STORE_ID],
233                leaf_data[LEAF_SET_KEY_ID],
234                leaf_data[LEAF_SET_MAP_KEY_ID],
235                leaf_data[LEAF_SET_MAP_VAL_ID],
236                leaf_data[LEAF_SET_ARGS_KEY_ID],
237                leaf_data[LEAF_SET_ARGS_VAL_ID],
238            ),
239        };
240
241        let store_id = store_id_raw as u8;
242        if store_id == 0 {
243            return (0, EMPTY_SLICE, EMPTY_SLICE, EMPTY_SLICE, EMPTY_SLICE, EMPTY_SLICE);
244        }
245
246        let key_frags  = if key_id != 0 { self.values_slice(key_id as usize).unwrap_or(EMPTY_SLICE) } else { EMPTY_SLICE };
247        let map_keys   = if map_key_id != 0 { self.map_keys_slice(map_key_id as usize).unwrap_or(EMPTY_SLICE) } else { EMPTY_SLICE };
248        let map_vals   = if map_val_id != 0 { self.map_vals_slice(map_val_id as usize).unwrap_or(EMPTY_SLICE) } else { EMPTY_SLICE };
249        let args_keys  = if args_key_id != 0 { self.args_keys_slice(args_key_id as usize).unwrap_or(EMPTY_SLICE) } else { EMPTY_SLICE };
250        let args_vals  = if args_val_id != 0 { self.args_vals_slice(args_val_id as usize).unwrap_or(EMPTY_SLICE) } else { EMPTY_SLICE };
251
252        (store_id, key_frags, map_keys, map_vals, args_keys, args_vals)
253    }
254
255    // ── slice helpers ─────────────────────────────────────────────────────────
256
257    fn children_slice(&self, id: usize) -> Option<&[u16]> {
258        self.vl_slice_u16(&self.children, id)
259    }
260
261    fn leaves_slice(&self, id: usize) -> Option<&[u16]> {
262        self.vl_slice_u16(&self.leaves, id)
263    }
264
265    pub fn values_slice(&self, id: usize) -> Option<&[u16]> {
266        self.vl_slice_u16(&self.values, id)
267    }
268
269    fn map_keys_slice(&self, id: usize) -> Option<&[u16]> {
270        self.vl_slice_u16(&self.map_keys, id)
271    }
272
273    fn map_vals_slice(&self, id: usize) -> Option<&[u16]> {
274        self.vl_slice_u16(&self.map_vals, id)
275    }
276
277    fn args_keys_slice(&self, id: usize) -> Option<&[u16]> {
278        self.vl_slice_u16(&self.args_keys, id)
279    }
280
281    fn args_vals_slice(&self, id: usize) -> Option<&[u16]> {
282        self.vl_slice_u16(&self.args_vals, id)
283    }
284
285    fn vl_slice_u16<'a>(&'a self, vl: &'a VariableList<u16>, id: usize) -> Option<&'a [u16]> {
286        let identity_start = id * 2;
287        let start = *vl.identity.get(identity_start)?;
288        let end   = *vl.identity.get(identity_start + 1)?;
289        vl.data.get(start..end)
290    }
291
292    pub fn word_bytes(&self, id: usize) -> &[u8] {
293        let identity_start = id * 2;
294        let start = match self.words.identity.get(identity_start) {
295            Some(&s) => s,
296            None     => return b"",
297        };
298        let end = match self.words.identity.get(identity_start + 1) {
299            Some(&e) => e,
300            None     => return b"",
301        };
302        self.words.data.get(start..end).unwrap_or(b"")
303    }
304}
305
306// ── tests ─────────────────────────────────────────────────────────────────────
307
308#[cfg(test)]
309mod tests {
310    use super::*;
311    use alloc::vec;
312    use crate::dsl::Dsl;
313    use crate::provided::Tree;
314
315    fn scalar(s: &str) -> Tree { Tree::Scalar(s.as_bytes().to_vec()) }
316    fn mapping(pairs: Vec<(&str, Tree)>) -> Tree {
317        Tree::Mapping(pairs.into_iter().map(|(k, v)| (k.as_bytes().to_vec(), v)).collect())
318    }
319
320    fn make_index(tree: &Tree) -> Index {
321        let (paths, children, leaves, values, words, map_keys, map_vals, args_keys, args_vals)
322            = Dsl::compile(tree, &[]).unwrap();
323        Index::new(paths, children, leaves, values, words, map_keys, map_vals, args_keys, args_vals)
324    }
325
326    fn make_index_with_stores<'a>(tree: &Tree, store_ids: &[&'a str]) -> Index {
327        let (paths, children, leaves, values, words, map_keys, map_vals, args_keys, args_vals)
328            = Dsl::compile(tree, store_ids).unwrap();
329        Index::new(paths, children, leaves, values, words, map_keys, map_vals, args_keys, args_vals)
330    }
331
332    // --- traverse ---
333
334    #[test]
335    fn traverse_leaf_path() {
336        let index = make_index(&mapping(vec![
337            ("session", mapping(vec![
338                ("user", mapping(vec![
339                    ("id", Tree::Null),
340                ])),
341            ])),
342        ]));
343        let leaves = index.traverse("session.user.id");
344        assert_eq!(leaves.len(), 1);
345    }
346
347    #[test]
348    fn traverse_intermediate_collects_all_leaves() {
349        let index = make_index(&mapping(vec![
350            ("session", mapping(vec![
351                ("user", mapping(vec![
352                    ("id",   Tree::Null),
353                    ("name", Tree::Null),
354                ])),
355            ])),
356        ]));
357        let leaves = index.traverse("session.user");
358        assert_eq!(leaves.len(), 2);
359    }
360
361    #[test]
362    fn traverse_nonexistent_returns_empty() {
363        let index = make_index(&mapping(vec![
364            ("session", mapping(vec![
365                ("user", mapping(vec![
366                    ("id", Tree::Null),
367                ])),
368            ])),
369        ]));
370        let leaves = index.traverse("session.user.missing");
371        assert!(leaves.is_empty());
372    }
373
374    // --- keyword_of ---
375
376    #[test]
377    fn keyword_of_leaf() {
378        let index = make_index(&mapping(vec![
379            ("user", mapping(vec![
380                ("id", Tree::Null),
381            ])),
382        ]));
383        assert_eq!(index.keyword_of(2), b"id");
384    }
385
386    #[test]
387    fn keyword_of_intermediate() {
388        let index = make_index(&mapping(vec![
389            ("user", mapping(vec![
390                ("id", Tree::Null),
391            ])),
392        ]));
393        assert_eq!(index.keyword_of(1), b"user");
394    }
395
396    // --- get_meta ---
397
398    #[test]
399    fn get_meta_store_id() {
400        let index = make_index_with_stores(&mapping(vec![
401            ("session", mapping(vec![
402                ("_get", mapping(vec![
403                    ("store", scalar("Memory")),
404                    ("key",    scalar("session:1")),
405                ])),
406                ("user", mapping(vec![
407                    ("id", Tree::Null),
408                ])),
409            ])),
410        ]), &["Memory"]);
411        let leaves = index.traverse("session.user.id");
412        let (store_id, _, _, _, _, _) = index.get_meta(&leaves[0]);
413        assert_eq!(store_id, 1); // "Memory" is store_ids[0] → id=1
414    }
415
416    #[test]
417    fn get_meta_key_frags() {
418        let index = make_index_with_stores(&mapping(vec![
419            ("session", mapping(vec![
420                ("_get", mapping(vec![
421                    ("store", scalar("Memory")),
422                    ("key",    scalar("session:1")),
423                ])),
424                ("user", mapping(vec![
425                    ("id", Tree::Null),
426                ])),
427            ])),
428        ]), &["Memory"]);
429        let leaves = index.traverse("session.user.id");
430        let (_, key_frags, _, _, _, _) = index.get_meta(&leaves[0]);
431        assert_eq!(key_frags.len(), 1);
432        let word_id = (key_frags[0] & crate::dsl::VALUE_WORD_ID_MASK) as usize;
433        assert_eq!(index.word_bytes(word_id), b"session:1");
434    }
435
436    #[test]
437    fn get_meta_no_get_returns_empty() {
438        let index = make_index(&mapping(vec![
439            ("user", mapping(vec![
440                ("id", Tree::Null),
441            ])),
442        ]));
443        let leaves = index.traverse("user.id");
444        let (store_id, key_frags, map_keys, map_vals, args_keys, args_vals) = index.get_meta(&leaves[0]);
445        assert_eq!(store_id, 0);
446        assert!(key_frags.is_empty());
447        assert!(map_keys.is_empty());
448        assert!(map_vals.is_empty());
449        assert!(args_keys.is_empty());
450        assert!(args_vals.is_empty());
451    }
452
453    // --- leaf_fragments ---
454
455    #[test]
456    fn leaf_fragments_static_value() {
457        let index = make_index(&mapping(vec![
458            ("driver", scalar("postgres")),
459        ]));
460        let leaves = index.traverse("driver");
461        let frags = index.leaf_fragments(&leaves[0]);
462        assert_eq!(frags.len(), 1);
463        assert_eq!(frags[0], (false, b"postgres" as &[u8]));
464    }
465
466    #[test]
467    fn leaf_fragments_null_value_is_empty() {
468        let index = make_index(&mapping(vec![
469            ("id", Tree::Null),
470        ]));
471        let leaves = index.traverse("id");
472        let frags = index.leaf_fragments(&leaves[0]);
473        assert!(frags.is_empty());
474    }
475
476    #[test]
477    fn leaf_fragments_single_placeholder() {
478        let index = make_index(&mapping(vec![
479            ("copy", scalar("${session.user.name}")),
480        ]));
481        let leaves = index.traverse("copy");
482        let frags = index.leaf_fragments(&leaves[0]);
483        assert_eq!(frags.len(), 1);
484        assert_eq!(frags[0], (true, b"session.user.name" as &[u8]));
485    }
486
487    #[test]
488    fn leaf_fragments_template() {
489        let index = make_index(&mapping(vec![
490            ("key", scalar("prefix.${some.path}.suffix")),
491        ]));
492        let leaves = index.traverse("key");
493        let frags = index.leaf_fragments(&leaves[0]);
494        assert_eq!(frags.len(), 3);
495        assert_eq!(frags[0], (false, b"prefix." as &[u8]));
496        assert_eq!(frags[1], (true,  b"some.path" as &[u8]));
497        assert_eq!(frags[2], (false, b".suffix" as &[u8]));
498    }
499
500    // --- set_meta ---
501
502    #[test]
503    fn set_meta_store_id() {
504        let index = make_index_with_stores(&mapping(vec![
505            ("session", mapping(vec![
506                ("_set", mapping(vec![
507                    ("store", scalar("Kvs")),
508                    ("key",    scalar("session:1")),
509                ])),
510                ("user", mapping(vec![
511                    ("id", Tree::Null),
512                ])),
513            ])),
514        ]), &["Memory", "Kvs"]);
515        let leaves = index.traverse("session.user.id");
516        let (store_id, _, _, _, _, _) = index.set_meta(&leaves[0]);
517        assert_eq!(store_id, 2); // "Kvs" is store_ids[1] → id=2
518    }
519
520    #[test]
521    fn set_meta_no_set_returns_empty() {
522        let index = make_index(&mapping(vec![
523            ("user", mapping(vec![
524                ("id", Tree::Null),
525            ])),
526        ]));
527        let leaves = index.traverse("user.id");
528        let (store_id, key_frags, map_keys, map_vals, args_keys, args_vals) = index.set_meta(&leaves[0]);
529        assert_eq!(store_id, 0);
530        assert!(key_frags.is_empty());
531        assert!(map_keys.is_empty());
532        assert!(map_vals.is_empty());
533        assert!(args_keys.is_empty());
534        assert!(args_vals.is_empty());
535    }
536}