lean_ctx/core/
attention_placement.rs

1//! U-shaped placement (primacy/recency): alternate placing ranked items at context head vs tail.
2
3/// Sort by descending importance, then assign alternating fronts (`front`) vs backs (`back`).
4pub fn reorder_for_attention(items: &mut [(String, f64)]) {
5    if items.is_empty() {
6        return;
7    }
8    let mut sorted: Vec<(String, f64)> = items.iter().map(|(s, r)| (s.clone(), *r)).collect();
9    sorted.sort_by(|a, b| {
10        b.1.partial_cmp(&a.1)
11            .unwrap_or(std::cmp::Ordering::Equal)
12            .then_with(|| a.0.cmp(&b.0))
13    });
14
15    let n = sorted.len();
16    let mut front: Vec<(String, f64)> = Vec::with_capacity(n.div_ceil(2));
17    let mut back: Vec<(String, f64)> = Vec::with_capacity(n / 2);
18
19    for (idx, pair) in sorted.into_iter().enumerate() {
20        if idx % 2 == 0 {
21            front.push(pair);
22        } else {
23            back.push(pair);
24        }
25    }
26
27    let mut out = Vec::with_capacity(n);
28    out.extend(front);
29    out.extend(back.into_iter().rev());
30    for (i, pair) in out.into_iter().enumerate() {
31        items[i] = pair;
32    }
33}
34
35#[cfg(test)]
36mod tests {
37    use super::*;
38
39    #[test]
40    fn empty_stable() {
41        let mut v: Vec<(String, f64)> = vec![];
42        reorder_for_attention(&mut v);
43        assert!(v.is_empty());
44    }
45
46    #[test]
47    fn highest_front_second_back_next_near_front() {
48        let mut v = vec![
49            ("c".into(), 3.0),
50            ("a".into(), 1.0),
51            ("d".into(), 4.0),
52            ("b".into(), 2.0),
53        ];
54        reorder_for_attention(&mut v);
55        assert_eq!(v[0].0, "d");
56        assert_eq!(v[v.len() - 1].0, "c");
57        assert_eq!(v[1].0, "b");
58        assert_eq!(v[v.len() - 2].0, "a");
59    }
60
61    #[test]
62    fn ties_lexicographic_stable_middle_behavior() {
63        let mut v = vec![
64            ("z_last".into(), 1.0),
65            ("a_first".into(), 1.0),
66            ("m_mid".into(), 1.0),
67        ];
68        reorder_for_attention(&mut v);
69        assert_eq!(v[0].0, "a_first");
70        assert_eq!(v[v.len() - 1].0, "m_mid");
71        assert_eq!(v[1].0, "z_last");
72    }
73}
lean_ctx/core/attention_placement.rs

lean_ctx/core/
attention_placement.rs
