mq_db/
query.rs

1use crate::{
2    block::{Block, BlockType},
3    document::Document,
4    store::DocumentStore,
5};
6
7enum SectionAnchor {
8    /// Directly supply a known (pre, post) interval.
9    Interval { pre: u32, post: u32 },
10    /// Find the first heading matching the given content and optional depth.
11    Heading { content: String, depth: Option<u8> },
12}
13
14/// A query result pairing a matched [`Block`] with its parent [`Document`].
15pub struct QueryResult<'a> {
16    pub block: &'a Block,
17    pub document: &'a Document,
18}
19
20/// Chainable, lazy query builder over a [`DocumentStore`].
21///
22/// Filters are applied in evaluation order (cheapest first):
23/// 1. Document-level zone-map skip (applied per document before scanning blocks)
24/// 2. Section `UNDER` constraint (interval check)
25/// 3. Block-level predicates
26///
27/// # Example – RAG chunk extraction
28///
29/// ```rust
30/// use mq_db::{DocumentStore, block::BlockType};
31///
32/// let mut store = DocumentStore::new();
33/// store.add_str("# Doc\n\n## Architecture\n\nExplanation\n\n```rust\ncode\n```\n").unwrap();
34///
35/// let results = store.query()
36///     .under_heading("Architecture", Some(2))
37///     .filter(|b| matches!(b.block_type, BlockType::Paragraph | BlockType::Code))
38///     .blocks();
39///
40/// assert_eq!(results.len(), 2);
41/// ```
42type DocPredicate<'store> = Box<dyn Fn(&Document) -> bool + 'store>;
43type BlockPredicate<'store> = Box<dyn Fn(&Block) -> bool + 'store>;
44
45pub struct Query<'store> {
46    store: &'store DocumentStore,
47    doc_predicate: Option<DocPredicate<'store>>,
48    block_predicates: Vec<BlockPredicate<'store>>,
49    anchor: Option<SectionAnchor>,
50    limit: Option<usize>,
51}
52
53impl<'store> Query<'store> {
54    pub(crate) fn new(store: &'store DocumentStore) -> Self {
55        Self {
56            store,
57            doc_predicate: None,
58            block_predicates: Vec::new(),
59            anchor: None,
60            limit: None,
61        }
62    }
63
64    /// Skip documents for which `predicate` returns `false`.
65    ///
66    /// Use this to leverage zone-map statistics before scanning blocks:
67    /// ```rust
68    /// # use mq_db::DocumentStore;
69    /// # let store = DocumentStore::new();
70    /// store.query()
71    ///     .documents(|doc| doc.zone_maps.code_languages.contains("python"));
72    /// ```
73    pub fn documents<F>(mut self, f: F) -> Self
74    where
75        F: Fn(&Document) -> bool + 'store,
76    {
77        self.doc_predicate = Some(Box::new(f));
78        self
79    }
80
81    /// Restrict results to blocks that fall within the heading section
82    /// identified by `content` and optional `depth`.
83    ///
84    /// Equivalent to the SQL `WHERE b UNDER (SELECT id FROM blocks WHERE ...)`.
85    ///
86    /// For best performance, chain a `.documents(|d| d.zone_maps.heading_contents.contains("..."))`
87    /// filter before this to skip irrelevant documents via zone maps.
88    pub fn under_heading(mut self, content: impl Into<String>, depth: Option<u8>) -> Self {
89        self.anchor = Some(SectionAnchor::Heading {
90            content: content.into(),
91            depth,
92        });
93        self
94    }
95
96    /// Restrict results to blocks that fall within the interval `(pre, post)`.
97    ///
98    /// Use this when you already know the ancestor block's interval values.
99    pub fn under_interval(mut self, pre: u32, post: u32) -> Self {
100        self.anchor = Some(SectionAnchor::Interval { pre, post });
101        self
102    }
103
104    /// Keep only blocks for which `predicate` returns `true`.
105    pub fn filter<F>(mut self, f: F) -> Self
106    where
107        F: Fn(&Block) -> bool + 'store,
108    {
109        self.block_predicates.push(Box::new(f));
110        self
111    }
112
113    /// Keep only blocks with the given [`BlockType`].
114    pub fn block_type(self, ty: BlockType) -> Self {
115        self.filter(move |b| b.block_type == ty)
116    }
117
118    /// Keep only heading blocks at the given depth.
119    pub fn heading_depth(self, depth: u8) -> Self {
120        self.filter(move |b| b.block_type == BlockType::Heading && b.heading_depth() == Some(depth))
121    }
122
123    /// Keep only code blocks with the given language tag.
124    pub fn code_lang(self, lang: impl Into<String>) -> Self {
125        let lang = lang.into();
126        self.filter(move |b| b.code_lang() == Some(lang.as_str()))
127    }
128
129    /// Keep only blocks whose content contains `substring` (case-sensitive).
130    pub fn content_contains(self, substring: impl Into<String>) -> Self {
131        let s = substring.into();
132        self.filter(move |b| b.content.contains(s.as_str()))
133    }
134
135    /// Keep only blocks whose content matches `pattern` (case-insensitive).
136    pub fn content_contains_ci(self, pattern: impl Into<String>) -> Self {
137        let pat = pattern.into().to_lowercase();
138        self.filter(move |b| b.content.to_lowercase().contains(pat.as_str()))
139    }
140
141    /// Stop collecting after `n` results.
142    pub fn limit(mut self, n: usize) -> Self {
143        self.limit = Some(n);
144        self
145    }
146
147    /// Execute the query, returning matched (block, document) pairs in
148    /// document order.
149    pub fn collect(&self) -> Vec<QueryResult<'_>> {
150        let mut results: Vec<QueryResult<'_>> = Vec::new();
151
152        'doc: for doc in self.store.documents() {
153            // Zone-map skip
154            if let Some(dp) = &self.doc_predicate
155                && !dp(doc)
156            {
157                continue 'doc;
158            }
159
160            // Resolve the section anchor for this document
161            let interval: Option<(u32, u32)> = match &self.anchor {
162                None => None,
163                Some(SectionAnchor::Interval { pre, post }) => Some((*pre, *post)),
164                Some(SectionAnchor::Heading { content, depth }) => doc
165                    .blocks
166                    .iter()
167                    .find(|b| {
168                        b.block_type == BlockType::Heading
169                            && b.content == *content
170                            && depth.is_none_or(|d| b.heading_depth() == Some(d))
171                    })
172                    .map(|h| (h.pre, h.post)),
173            };
174
175            for block in &doc.blocks {
176                // UNDER interval check
177                if let Some((anc_pre, anc_post)) = interval
178                    && !block.is_under_interval(anc_pre, anc_post)
179                {
180                    continue;
181                }
182
183                // Block predicates
184                if self.block_predicates.iter().all(|f| f(block)) {
185                    results.push(QueryResult {
186                        block,
187                        document: doc,
188                    });
189
190                    if let Some(limit) = self.limit
191                        && results.len() >= limit
192                    {
193                        return results;
194                    }
195                }
196            }
197        }
198
199        results
200    }
201
202    /// Like [`collect`], but returns cloned blocks (discarding document context).
203    ///
204    /// Returns owned `Vec<Block>` so it can be used on temporaries:
205    /// ```rust
206    /// # use mq_db::{DocumentStore, block::BlockType};
207    /// # let mut store = DocumentStore::new();
208    /// # store.add_str("# Hello\n").unwrap();
209    /// let blocks = store.query().heading_depth(1).blocks();
210    /// assert_eq!(blocks.len(), 1);
211    /// ```
212    pub fn blocks(&self) -> Vec<Block> {
213        self.collect()
214            .into_iter()
215            .map(|r| r.block.clone())
216            .collect()
217    }
218
219    /// Returns the number of matching blocks without materialising them.
220    pub fn count(&self) -> usize {
221        self.collect().len()
222    }
223
224    /// Find all (heading, next_sibling) pairs where the heading matches the
225    /// given depth and the immediately following sibling has one of the
226    /// `forbidden_types`.
227    ///
228    /// This is the foundation for structural lint rules such as:
229    /// > "An H2 heading must not be immediately followed by a list."
230    ///
231    /// # Example
232    ///
233    /// ```rust
234    /// use mq_db::{DocumentStore, block::BlockType};
235    ///
236    /// let mut store = DocumentStore::new();
237    /// store.add_str("## Section\n\n- item\n").unwrap();
238    ///
239    /// let q = store.query();
240    /// let violations = q.lint_heading_followed_by(2, &[BlockType::List]);
241    ///
242    /// assert_eq!(violations.len(), 1);
243    /// ```
244    pub fn lint_heading_followed_by(
245        &self,
246        heading_depth: u8,
247        forbidden_types: &[BlockType],
248    ) -> Vec<LintViolation<'_>> {
249        let mut violations = Vec::new();
250
251        for doc in self.store.documents() {
252            if let Some(dp) = &self.doc_predicate
253                && !dp(doc)
254            {
255                continue;
256            }
257
258            for block in &doc.blocks {
259                if block.block_type != BlockType::Heading {
260                    continue;
261                }
262                if block.heading_depth() != Some(heading_depth) {
263                    continue;
264                }
265
266                if let Some(next) = doc.first_child(block)
267                    && forbidden_types.contains(&next.block_type)
268                {
269                    violations.push(LintViolation {
270                        heading: block,
271                        offending: next,
272                        document: doc,
273                    });
274                }
275            }
276        }
277
278        violations
279    }
280}
281
282/// A structural lint violation: a heading followed by a forbidden block type.
283pub struct LintViolation<'a> {
284    pub heading: &'a Block,
285    pub offending: &'a Block,
286    pub document: &'a Document,
287}
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292    use crate::DocumentStore;
293    use rstest::rstest;
294
295    fn store_with(sources: &[&str]) -> DocumentStore {
296        let mut s = DocumentStore::new();
297        for src in sources {
298            s.add_str(src).unwrap();
299        }
300        s
301    }
302
303    #[test]
304    fn test_block_type_filter() {
305        let store = store_with(&["# H1\n\n## H2\n\nParagraph\n\n```rust\ncode\n```\n"]);
306        let q = store.query().block_type(BlockType::Heading);
307        let headings = q.blocks();
308        assert_eq!(headings.len(), 2);
309    }
310
311    #[test]
312    fn test_heading_depth_filter() {
313        let store = store_with(&["# H1\n\n## H2\n\n### H3\n"]);
314        let q = store.query().heading_depth(2);
315        let h2s = q.blocks();
316        assert_eq!(h2s.len(), 1);
317        assert_eq!(h2s[0].content, "H2");
318    }
319
320    #[test]
321    fn test_under_heading_filter() {
322        let store = store_with(&[
323            "# Doc\n\n## Architecture\n\nExplanation\n\n```rust\ncode\n```\n\n## Other\n\nOther para\n",
324        ]);
325
326        let q = store
327            .query()
328            .under_heading("Architecture", Some(2))
329            .filter(|b| matches!(b.block_type, BlockType::Paragraph | BlockType::Code));
330        let results = q.blocks();
331
332        // Should contain "Explanation" and the code block, but NOT "Other para"
333        assert_eq!(results.len(), 2);
334        assert!(results.iter().any(|b| b.content.contains("Explanation")));
335        assert!(results.iter().any(|b| b.block_type == BlockType::Code));
336        assert!(!results.iter().any(|b| b.content.contains("Other para")));
337    }
338
339    #[test]
340    fn test_limit() {
341        let store = store_with(&["# A\n\n# B\n\n# C\n\n# D\n"]);
342        let q = store.query().heading_depth(1).limit(2);
343        let results = q.blocks();
344        assert_eq!(results.len(), 2);
345    }
346
347    #[test]
348    fn test_content_contains() {
349        let store = store_with(&["# Hello World\n\n## Goodbye\n"]);
350        let q = store.query().content_contains("World");
351        let results = q.blocks();
352        assert_eq!(results.len(), 1);
353        assert_eq!(results[0].content, "Hello World");
354    }
355
356    #[test]
357    fn test_code_lang_filter() {
358        let store = store_with(&["```rust\nfn x(){}\n```\n\n```python\nx=1\n```\n"]);
359        let q = store.query().code_lang("rust");
360        let results = q.blocks();
361        assert_eq!(results.len(), 1);
362        assert_eq!(results[0].code_lang(), Some("rust"));
363    }
364
365    #[test]
366    fn test_zone_map_document_skip() {
367        let store = store_with(&["```python\nx=1\n```\n", "```rust\nfn x(){}\n```\n"]);
368
369        let q = store
370            .query()
371            .documents(|doc| doc.zone_maps.code_languages.contains("rust"))
372            .block_type(BlockType::Code);
373        let results = q.blocks();
374
375        assert_eq!(results.len(), 1);
376        assert_eq!(results[0].code_lang(), Some("rust"));
377    }
378
379    #[test]
380    fn test_lint_heading_followed_by_list() {
381        let store = store_with(&[
382            "## Good\n\nParagraph intro\n\n- item\n",
383            "## Bad\n\n- item without intro\n",
384        ]);
385
386        let q = store.query();
387        let violations = q.lint_heading_followed_by(2, &[BlockType::List]);
388        // Only the second doc has a violation
389        assert_eq!(violations.len(), 1);
390        assert_eq!(violations[0].heading.content, "Bad");
391        assert_eq!(violations[0].offending.block_type, BlockType::List);
392    }
393
394    #[test]
395    fn test_multi_document_query() {
396        let store = store_with(&[
397            "# Doc1\n\n```rust\nfn a(){}\n```\n",
398            "# Doc2\n\n```python\nx=1\n```\n",
399            "# Doc3\n\n```rust\nfn b(){}\n```\n",
400        ]);
401
402        let q = store.query().code_lang("rust");
403        let results = q.blocks();
404        assert_eq!(results.len(), 2);
405    }
406
407    #[rstest]
408    #[case(BlockType::Heading, 3)]
409    #[case(BlockType::Paragraph, 1)]
410    #[case(BlockType::Code, 1)]
411    #[case(BlockType::List, 1)]
412    fn test_block_type_filter_count_param(#[case] block_type: BlockType, #[case] expected: usize) {
413        let store =
414            store_with(&["# H1\n\n## H2\n\n### H3\n\nParagraph\n\n```rust\ncode\n```\n\n- item\n"]);
415        assert_eq!(
416            store.query().block_type(block_type).blocks().len(),
417            expected
418        );
419    }
420
421    #[rstest]
422    #[case(1, 1, "H1")]
423    #[case(2, 1, "H2")]
424    #[case(3, 1, "H3")]
425    fn test_heading_depth_count_param(
426        #[case] depth: u8,
427        #[case] expected: usize,
428        #[case] content: &str,
429    ) {
430        let store = store_with(&["# H1\n\n## H2\n\n### H3\n"]);
431        let results = store.query().heading_depth(depth).blocks();
432        assert_eq!(results.len(), expected);
433        assert_eq!(results[0].content, content);
434    }
435
436    #[rstest]
437    #[case("Hello", 1)]
438    #[case("World", 1)]
439    #[case("Goodbye", 1)]
440    #[case("nonexistent_xyz", 0)]
441    fn test_content_contains_count_param(#[case] needle: &str, #[case] expected: usize) {
442        let store = store_with(&["# Hello World\n\n## Goodbye\n"]);
443        assert_eq!(
444            store.query().content_contains(needle).blocks().len(),
445            expected
446        );
447    }
448
449    #[rstest]
450    #[case("rust", 1)]
451    #[case("python", 1)]
452    #[case("go", 0)]
453    fn test_code_lang_count_param(#[case] lang: &str, #[case] expected: usize) {
454        let store = store_with(&["```rust\nfn x(){}\n```\n\n```python\nx=1\n```\n"]);
455        assert_eq!(store.query().code_lang(lang).blocks().len(), expected);
456    }
457
458    #[rstest]
459    #[case(1, 1)]
460    #[case(2, 2)]
461    #[case(3, 3)]
462    #[case(100, 4)]
463    fn test_limit_count_param(#[case] limit: usize, #[case] expected: usize) {
464        let store = store_with(&["# A\n\n# B\n\n# C\n\n# D\n"]);
465        let results = store.query().heading_depth(1).limit(limit).blocks();
466        assert_eq!(results.len(), expected);
467    }
468}
mq_db/query.rs

mq_db/
query.rs
