omgbase_search/
rank.rs

1//! Fusion and boosts (`spec/search` §4) as pure functions over ranked id
2//! lists and per-document facts, plus the `resolve` shaping.
3
4use serde_json::Value as Json;
5
6use crate::fts::split_ws;
7
8/// §4: the reciprocal-rank constant.
9pub const RRF_K: f64 = 60.0;
10/// §4: `hybrid`'s FTS and vector passes both ask for this many hits.
11pub const HYBRID_PASS_LIMIT: usize = 200;
12/// §4: `resolve`'s default limit.
13pub const RESOLVE_DEFAULT_LIMIT: usize = 10;
14/// §4: words in a `resolve` preview.
15pub const PREVIEW_WORDS: usize = 12;
16
17pub const TITLE_BOOST: f64 = 1.25;
18pub const HEADING_BOOST: f64 = 1.15;
19pub const PATH_BOOST: f64 = 1.10;
20
21/// §4 step 4: the `layer` boost; `None` for `proposed` (×1.0, not recorded)
22/// and unknown layers.
23#[must_use]
24pub fn layer_boost(layer: &str) -> Option<f64> {
25    match layer {
26        "canon" => Some(1.3),
27        "working" => Some(1.15),
28        "draft" => Some(0.85),
29        _ => None,
30    }
31}
32
33/// The multiplicative boosts of a hit, each present only when it applies.
34#[derive(Clone, Copy, Debug, Default, PartialEq)]
35pub struct Boosts {
36    pub title: Option<f64>,
37    pub heading: Option<f64>,
38    pub path: Option<f64>,
39    pub layer: Option<f64>,
40    /// Reserved; never set.
41    pub recency: Option<f64>,
42}
43
44impl Boosts {
45    /// The boosts as the `evidence.boosts` object (present keys only).
46    #[must_use]
47    pub fn to_json(&self) -> Json {
48        let mut m = serde_json::Map::new();
49        for (k, v) in [
50            ("title", self.title),
51            ("heading", self.heading),
52            ("path", self.path),
53            ("layer", self.layer),
54            ("recency", self.recency),
55        ] {
56            if let Some(v) = v {
57                m.insert(k.to_owned(), Json::from(v));
58            }
59        }
60        Json::Object(m)
61    }
62}
63
64/// The per-document and per-block facts the boosts read (§4 step 4).
65#[derive(Clone, Debug, Default, PartialEq, Eq)]
66pub struct BoostFacts {
67    /// The merged `title` property rendered with [`property_to_string`]
68    /// (`""` when absent).
69    pub title: String,
70    /// The merged `layer` property rendered likewise.
71    pub layer: String,
72    pub path: String,
73    /// The texts of the section headings whose range contains the block.
74    pub headings: Vec<String>,
75}
76
77/// §4 step 4: `terms` default to `text` split on `\s+`.
78#[must_use]
79pub fn default_terms(text: Option<&str>) -> Vec<String> {
80    text.map(|t| split_ws(t).map(str::to_owned).collect())
81        .unwrap_or_default()
82}
83
84/// §4 step 4: lower-cased for matching, empties dropped.
85#[must_use]
86pub fn lower_terms(terms: &[String]) -> Vec<String> {
87    terms
88        .iter()
89        .map(|t| t.to_lowercase())
90        .filter(|t| !t.is_empty())
91        .collect()
92}
93
94/// JavaScript's `String(value ?? "")` over a merged property value: strings
95/// verbatim, numbers and booleans as JavaScript prints them, `null` → `""`,
96/// arrays joined by `,` (a null element is empty), objects `[object Object]`.
97#[must_use]
98pub fn property_to_string(value: Option<&Json>) -> String {
99    fn element(v: &Json) -> String {
100        match v {
101            Json::Null => String::new(),
102            Json::Array(items) => items.iter().map(element).collect::<Vec<_>>().join(","),
103            other => scalar(other),
104        }
105    }
106    fn scalar(v: &Json) -> String {
107        match v {
108            Json::Null => String::new(),
109            Json::Bool(b) => b.to_string(),
110            Json::Number(n) => n.as_f64().map_or_else(|| n.to_string(), js_number),
111            Json::String(s) => s.clone(),
112            Json::Array(items) => items.iter().map(element).collect::<Vec<_>>().join(","),
113            Json::Object(_) => "[object Object]".to_owned(),
114        }
115    }
116    value.map_or_else(String::new, scalar)
117}
118
119/// JavaScript `Number.prototype.toString()` for the values a property holds:
120/// integers without a fraction, otherwise the shortest round-trip form.
121fn js_number(n: f64) -> String {
122    if n.is_nan() {
123        return "NaN".to_owned();
124    }
125    if n.is_infinite() {
126        return if n > 0.0 { "Infinity" } else { "-Infinity" }.to_owned();
127    }
128    if n == n.trunc() && n.abs() < 1e21 {
129        return format!("{}", n as i128);
130    }
131    format!("{n}")
132}
133
134/// §4 step 4: the boosts for one hit from its facts and the lower-cased terms.
135#[must_use]
136pub fn compute_boosts(facts: &BoostFacts, lower_terms: &[String]) -> Boosts {
137    let mut b = Boosts::default();
138    let title = facts.title.to_lowercase();
139    if !title.is_empty() && lower_terms.iter().any(|t| title.contains(t.as_str())) {
140        b.title = Some(TITLE_BOOST);
141    }
142    if facts.headings.iter().any(|h| {
143        let h = h.to_lowercase();
144        lower_terms.iter().any(|t| h.contains(t.as_str()))
145    }) {
146        b.heading = Some(HEADING_BOOST);
147    }
148    let path = facts.path.to_lowercase();
149    if lower_terms.iter().any(|t| path.contains(t.as_str())) {
150        b.path = Some(PATH_BOOST);
151    }
152    b.layer = layer_boost(&facts.layer);
153    b
154}
155
156/// §4 step 3: `(fts_rank ? 1/(60+fts_rank) : 0) + (vec_rank ? 1/(60+vec_rank) : 0)`.
157#[must_use]
158pub fn rrf_score(fts_rank: Option<usize>, vector_rank: Option<usize>) -> f64 {
159    let term = |r: Option<usize>| r.map_or(0.0, |r| 1.0 / (RRF_K + r as f64));
160    term(fts_rank) + term(vector_rank)
161}
162
163/// §4 step 5: `rrf × Π boosts`, multiplied left to right in the reference's
164/// order (title, heading, path, layer, recency) so the rounding is the same.
165#[must_use]
166pub fn apply_boosts(rrf: f64, boosts: &Boosts) -> f64 {
167    rrf * boosts.title.unwrap_or(1.0)
168        * boosts.heading.unwrap_or(1.0)
169        * boosts.path.unwrap_or(1.0)
170        * boosts.layer.unwrap_or(1.0)
171        * boosts.recency.unwrap_or(1.0)
172}
173
174/// A fused candidate before boosts (§4 steps 1–3).
175#[derive(Clone, Debug, PartialEq)]
176pub struct Candidate {
177    pub block_id: String,
178    /// 1-based rank of the first FTS occurrence.
179    pub fts_rank: Option<usize>,
180    /// 1-based vector rank.
181    pub vector_rank: Option<usize>,
182    pub cosine: Option<f64>,
183    pub rrf: f64,
184}
185
186/// §4 steps 1–3 over the two ranked lists: `fts` in FTS order and `vector`
187/// in cosine order with each cosine; in both, the first occurrence of an id
188/// ranks it. Candidates come out in first-seen order (FTS ids, then new
189/// vector ids).
190#[must_use]
191pub fn fuse(fts: &[String], vector: &[(String, f64)]) -> Vec<Candidate> {
192    let mut out: Vec<Candidate> = Vec::new();
193    for (i, id) in fts.iter().enumerate() {
194        if out.iter().any(|c| &c.block_id == id) {
195            continue;
196        }
197        out.push(Candidate {
198            block_id: id.clone(),
199            fts_rank: Some(i + 1),
200            vector_rank: None,
201            cosine: None,
202            rrf: 0.0,
203        });
204    }
205    for (i, (id, cos)) in vector.iter().enumerate() {
206        match out.iter_mut().find(|c| &c.block_id == id) {
207            Some(c) => {
208                if c.vector_rank.is_none() {
209                    c.vector_rank = Some(i + 1);
210                    c.cosine = Some(*cos);
211                }
212            }
213            None => out.push(Candidate {
214                block_id: id.clone(),
215                fts_rank: None,
216                vector_rank: Some(i + 1),
217                cosine: Some(*cos),
218                rrf: 0.0,
219            }),
220        }
221    }
222    for c in &mut out {
223        c.rrf = rrf_score(c.fts_rank, c.vector_rank);
224    }
225    out
226}
227
228/// The evidence a hybrid hit carries (§4 step 5).
229#[derive(Clone, Debug, PartialEq)]
230pub struct Evidence {
231    pub fts_rank: Option<usize>,
232    pub vector_rank: Option<usize>,
233    pub cosine: Option<f64>,
234    pub rrf: f64,
235    pub boosts: Boosts,
236}
237
238impl Evidence {
239    /// `{ fts_rank?, vector_rank?, cosine?, rrf, boosts }`.
240    #[must_use]
241    pub fn to_json(&self) -> Json {
242        let mut m = serde_json::Map::new();
243        if let Some(r) = self.fts_rank {
244            m.insert("fts_rank".to_owned(), Json::from(r));
245        }
246        if let Some(r) = self.vector_rank {
247            m.insert("vector_rank".to_owned(), Json::from(r));
248        }
249        if let Some(c) = self.cosine {
250            m.insert("cosine".to_owned(), Json::from(c));
251        }
252        m.insert("rrf".to_owned(), Json::from(self.rrf));
253        m.insert("boosts".to_owned(), self.boosts.to_json());
254        Json::Object(m)
255    }
256}
257
258/// §4 step 5's order: score descending, then `block_id` bytewise ascending.
259pub fn sort_by_score<T>(hits: &mut [T], score: impl Fn(&T) -> f64, id: impl Fn(&T) -> &str) {
260    hits.sort_by(|a, b| {
261        score(b)
262            .partial_cmp(&score(a))
263            .unwrap_or(std::cmp::Ordering::Equal)
264            .then_with(|| id(a).as_bytes().cmp(id(b).as_bytes()))
265    });
266}
267
268/// §4 `resolve`: `path + "#" + type + "[" + ordinal + "]"`.
269#[must_use]
270pub fn locator(path: &str, block_type: &str, ordinal: i64) -> String {
271    format!("{path}#{block_type}[{ordinal}]")
272}
273
274/// §4 `resolve`: the first `words` whitespace-separated words of `text`,
275/// `…` appended when cut.
276#[must_use]
277pub fn preview(text: &str, words: usize) -> String {
278    let w: Vec<&str> = split_ws(text).collect();
279    if w.len() <= words {
280        w.join(" ")
281    } else {
282        format!("{}\u{2026}", w[..words].join(" "))
283    }
284}
285
286#[cfg(test)]
287mod tests {
288    use super::*;
289    use serde_json::json;
290
291    fn s(v: &[&str]) -> Vec<String> {
292        v.iter().map(|x| (*x).to_owned()).collect()
293    }
294
295    #[test]
296    fn rrf_arithmetic() {
297        assert_eq!(rrf_score(None, None), 0.0);
298        assert_eq!(rrf_score(Some(1), None), 1.0 / 61.0);
299        assert_eq!(rrf_score(Some(1), Some(1)), 2.0 / 61.0);
300        assert_eq!(rrf_score(Some(3), Some(200)), 1.0 / 63.0 + 1.0 / 260.0);
301    }
302
303    #[test]
304    fn fusion_ranks_first_occurrence() {
305        let fts = s(&["b_1", "b_2", "b_1", "b_3"]);
306        let vec = vec![("b_3".to_owned(), 0.9), ("b_4".to_owned(), 0.5)];
307        let c = fuse(&fts, &vec);
308        assert_eq!(c.len(), 4);
309        assert_eq!(
310            (c[0].block_id.as_str(), c[0].fts_rank, c[0].vector_rank),
311            ("b_1", Some(1), None)
312        );
313        assert_eq!((c[1].block_id.as_str(), c[1].fts_rank), ("b_2", Some(2)));
314        assert_eq!(
315            (
316                c[2].block_id.as_str(),
317                c[2].fts_rank,
318                c[2].vector_rank,
319                c[2].cosine
320            ),
321            ("b_3", Some(4), Some(1), Some(0.9))
322        );
323        assert_eq!(c[2].rrf, 1.0 / 64.0 + 1.0 / 61.0);
324        assert_eq!(
325            (c[3].block_id.as_str(), c[3].fts_rank, c[3].vector_rank),
326            ("b_4", None, Some(2))
327        );
328        assert_eq!(c[3].rrf, 1.0 / 62.0);
329        // A repeated vector id keeps its first rank and cosine.
330        let c = fuse(&[], &[("b_9".to_owned(), 0.9), ("b_9".to_owned(), 0.1)]);
331        assert_eq!(c.len(), 1);
332        assert_eq!((c[0].vector_rank, c[0].cosine), (Some(1), Some(0.9)));
333    }
334
335    #[test]
336    fn boosts_each_and_product() {
337        let facts = BoostFacts {
338            title: "Guides Index".to_owned(),
339            layer: "canon".to_owned(),
340            path: "guides/onboarding.md".to_owned(),
341            headings: s(&["Setup", "First Steps"]),
342        };
343        let b = compute_boosts(&facts, &lower_terms(&s(&["STEPS"])));
344        assert_eq!(
345            b,
346            Boosts {
347                title: None,
348                heading: Some(1.15),
349                path: None,
350                layer: Some(1.3),
351                recency: None
352            }
353        );
354        let b = compute_boosts(&facts, &lower_terms(&s(&["guides"])));
355        assert_eq!((b.title, b.heading, b.path), (Some(1.25), None, Some(1.1)));
356        assert_eq!(apply_boosts(0.5, &b), 0.5 * 1.25 * 1.0 * 1.1 * 1.3 * 1.0);
357        assert_eq!(apply_boosts(0.5, &Boosts::default()), 0.5);
358        assert_eq!(
359            b.to_json(),
360            json!({"title": 1.25, "path": 1.1, "layer": 1.3})
361        );
362        // proposed / unknown layers leave `layer` absent; an empty title never matches.
363        let facts = BoostFacts {
364            layer: "proposed".to_owned(),
365            ..BoostFacts::default()
366        };
367        assert_eq!(
368            compute_boosts(&facts, &lower_terms(&s(&[""]))),
369            Boosts::default()
370        );
371        assert_eq!(layer_boost("draft"), Some(0.85));
372        assert_eq!(layer_boost("working"), Some(1.15));
373        assert_eq!(layer_boost("weird"), None);
374    }
375
376    #[test]
377    fn terms_and_strings() {
378        assert_eq!(
379            default_terms(Some("  Foo  bar\tBAZ ")),
380            s(&["Foo", "bar", "BAZ"])
381        );
382        assert_eq!(default_terms(None), Vec::<String>::new());
383        assert_eq!(lower_terms(&s(&["Foo", "", "É"])), s(&["foo", "é"]));
384        assert_eq!(property_to_string(None), "");
385        assert_eq!(property_to_string(Some(&json!(null))), "");
386        assert_eq!(property_to_string(Some(&json!("T"))), "T");
387        assert_eq!(property_to_string(Some(&json!(3))), "3");
388        assert_eq!(property_to_string(Some(&json!(1.5))), "1.5");
389        assert_eq!(property_to_string(Some(&json!(true))), "true");
390        assert_eq!(
391            property_to_string(Some(&json!(["a", null, 2, ["x", "y"]]))),
392            "a,,2,x,y"
393        );
394        assert_eq!(
395            property_to_string(Some(&json!({"a": 1}))),
396            "[object Object]"
397        );
398    }
399
400    #[test]
401    fn evidence_json_and_sort() {
402        let e = Evidence {
403            fts_rank: Some(2),
404            vector_rank: None,
405            cosine: None,
406            rrf: 1.0 / 62.0,
407            boosts: Boosts::default(),
408        };
409        assert_eq!(
410            e.to_json(),
411            json!({"fts_rank": 2, "rrf": 1.0 / 62.0, "boosts": {}})
412        );
413        let mut hits = vec![("b_2", 0.5), ("b_1", 0.5), ("b_0", 0.7)];
414        sort_by_score(&mut hits, |h| h.1, |h| h.0);
415        assert_eq!(hits, vec![("b_0", 0.7), ("b_1", 0.5), ("b_2", 0.5)]);
416    }
417
418    #[test]
419    fn resolve_shaping() {
420        assert_eq!(locator("a.md", "paragraph", 3), "a.md#paragraph[3]");
421        let twelve = "1 2 3 4 5 6 7 8 9 10 11 12";
422        assert_eq!(preview(twelve, PREVIEW_WORDS), twelve);
423        assert_eq!(
424            preview(&format!("{twelve} 13"), PREVIEW_WORDS),
425            format!("{twelve}…")
426        );
427        assert_eq!(preview("  a \n b ", PREVIEW_WORDS), "a b");
428        assert_eq!(preview("", PREVIEW_WORDS), "");
429    }
430}
omgbase_search/rank.rs

omgbase_search/
rank.rs
