metrale_model_layers/speculative/
tree_shape.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Spine-plus-hedge draft trees for tree speculative decoding: the
4//! tree shape, its verify-row layout and the accepted path.
5//!
6//! Owner: model-layers (speculative).
7//! Invariants:
8//! - `TreeShape::parse` returns only shapes that pass `validate`: every depth
9//!   holds 1..=`MAX_RANK` nodes, and the tree holds at most `MAX_NODES`.
10//! - `TreeDraft::rows` puts the root at row 0, the spine at rows `1..=L` in
11//!   depth order, then the hedges in the order of `hedges`.
12//!
13//! A tree is a spine (the drafter's top-1 chain) plus hedge leaves: lower-ranked
14//! siblings of spine nodes, with no children. So every node's parent is the
15//! root or a spine node. A shape is a comma-separated list of per-depth node
16//! counts, e.g. `"1,2,2,2"`: depth d holds one spine node and `count_d - 1`
17//! hedges, and the verify width is the node count plus the root row.
18
19use anyhow::{Result, bail};
20
21/// 2026-09-25: Largest node count `validate` accepts at one depth: the spine
22/// node plus up to three hedges.
23pub const MAX_RANK: usize = 4;
24/// 2026-09-25: Largest node count `validate` accepts for the whole tree,
25/// excluding the root row.
26pub const MAX_NODES: usize = 7;
27
28#[derive(Debug, Clone, PartialEq, Eq)]
29pub struct TreeShape {
30    /// 2026-09-25: `counts[d-1]` is the number of nodes at depth d; 1 means
31    /// the spine node only.
32    pub counts: Vec<u8>,
33}
34
35impl TreeShape {
36    pub fn parse(s: &str) -> Result<Self> {
37        let counts: Vec<u8> = s
38            .split(',')
39            .map(|t| t.trim().parse::<u8>())
40            .collect::<std::result::Result<_, _>>()
41            .map_err(|e| anyhow::anyhow!("METRALE_TREE_SHAPE parse: {e}"))?;
42        let shape = Self { counts };
43        shape.validate()?;
44        Ok(shape)
45    }
46
47    pub fn validate(&self) -> Result<()> {
48        if self.counts.is_empty() {
49            bail!("tree shape: empty");
50        }
51        if self.counts.iter().any(|&c| c == 0 || c as usize > MAX_RANK) {
52            bail!("tree shape: per-depth count must be 1..={MAX_RANK}");
53        }
54        if self.nodes() > MAX_NODES {
55            bail!("tree shape: {} nodes > max {MAX_NODES}", self.nodes());
56        }
57        Ok(())
58    }
59
60    pub fn spine_len(&self) -> usize {
61        self.counts.len()
62    }
63
64    /// 2026-09-25: Total tree nodes (spine and hedges), excluding the root row.
65    pub fn nodes(&self) -> usize {
66        self.counts.iter().map(|&c| c as usize).sum()
67    }
68
69    /// 2026-09-25: Verify width: the root row plus the nodes.
70    pub fn verify_width(&self) -> usize {
71        self.nodes() + 1
72    }
73
74    /// 2026-09-25: The counts packed 4 bits per depth. Distinct for distinct
75    /// valid shapes: every count is in 1..=4, so no nibble is zero.
76    pub fn shape_id(&self) -> u64 {
77        self.counts
78            .iter()
79            .fold(0u64, |acc, &c| (acc << 4) | (c as u64))
80    }
81
82    /// 2026-09-25: True when no depth holds a hedge.
83    pub fn is_chain(&self) -> bool {
84        self.counts.iter().all(|&c| c == 1)
85    }
86}
87
88/// 2026-09-25: One hedge leaf: the drafter's rank-`rank` candidate at `depth`.
89#[derive(Debug, Clone, Copy, PartialEq, Eq)]
90pub struct HedgeNode {
91    pub depth: usize,
92    /// 2026-09-25: Candidate rank. The spine holds rank 1, so a hedge's rank
93    /// is 2 or more.
94    pub rank: usize,
95    pub token: u32,
96}
97
98/// 2026-09-25: A proposed draft tree for one verify step.
99#[derive(Debug, Clone)]
100pub struct TreeDraft {
101    pub shape: TreeShape,
102    /// 2026-09-25: `spine[d-1]` is the top-1 draft at depth d.
103    pub spine: Vec<u32>,
104    /// 2026-09-25: Hedge leaves; `rows` lays them out in this order, and
105    /// callers keep it sorted by (depth, rank).
106    pub hedges: Vec<HedgeNode>,
107}
108
109/// 2026-09-25: One verify row. Row 0 is the root.
110#[derive(Debug, Clone, Copy, PartialEq, Eq)]
111pub struct TreeRow {
112    pub token: u32,
113    pub depth: usize,
114    pub parent_row: usize,
115}
116
117impl TreeDraft {
118    /// 2026-09-25: Verify rows `[root, spine..., hedges...]`. `root_token` is
119    /// the last committed token. A depth-d node's parent is row d-1, the
120    /// spine node (or the root) one depth up.
121    pub fn rows(&self, root_token: u32) -> Vec<TreeRow> {
122        let l = self.spine.len();
123        let mut rows = Vec::with_capacity(1 + l + self.hedges.len());
124        rows.push(TreeRow {
125            token: root_token,
126            depth: 0,
127            parent_row: 0,
128        });
129        for (i, &t) in self.spine.iter().enumerate() {
130            rows.push(TreeRow {
131                token: t,
132                depth: i + 1,
133                parent_row: i,
134            });
135        }
136        for h in &self.hedges {
137            rows.push(TreeRow {
138                token: h.token,
139                depth: h.depth,
140                parent_row: h.depth - 1,
141            });
142        }
143        rows
144    }
145
146    /// 2026-09-25: The accepted rows, root to leaf, given `v[row]`, the
147    /// target's next token after each row. Returns them with the bonus token
148    /// `v` at the last accepted row (row 0 when none is accepted).
149    ///
150    /// Each step accepts the child whose token equals `v[parent]`, so the
151    /// accepted tokens are the ones greedy decoding of the target emits. The
152    /// spine child is checked first, then the hedges in order, and the first
153    /// match wins. Accepting a hedge ends the path, since a hedge is a leaf.
154    ///
155    /// Panics when `rows` or `v` is shorter than the row layout of this draft.
156    pub fn accept_path(&self, rows: &[TreeRow], v: &[u32]) -> (Vec<usize>, u32) {
157        let l = self.spine.len();
158        let mut path = Vec::with_capacity(l);
159        let mut cur = 0usize;
160        for d in 1..=l {
161            let want = v[cur];
162            // 2026-09-25: `cur` is a spine row at depth d-1 (or the root); its
163            // children are the spine row at depth d and every hedge at depth d.
164            let spine_row = d;
165            let mut next = None;
166            if rows[spine_row].token == want {
167                next = Some(spine_row);
168            } else {
169                for (i, h) in self.hedges.iter().enumerate() {
170                    if h.depth == d && h.token == want {
171                        next = Some(1 + l + i);
172                        break;
173                    }
174                }
175            }
176            match next {
177                Some(r) => {
178                    path.push(r);
179                    cur = r;
180                    if r > l {
181                        break;
182                    }
183                }
184                None => break,
185            }
186        }
187        (path, v[cur])
188    }
189}
190
191#[cfg(test)]
192mod tests {
193    use super::*;
194
195    fn draft(shape: &str, spine: &[u32], hedges: &[(usize, usize, u32)]) -> TreeDraft {
196        TreeDraft {
197            shape: TreeShape::parse(shape).unwrap(),
198            spine: spine.to_vec(),
199            hedges: hedges
200                .iter()
201                .map(|&(depth, rank, token)| HedgeNode { depth, rank, token })
202                .collect(),
203        }
204    }
205
206    #[test]
207    fn parse_and_validate() {
208        let s = TreeShape::parse("1,2,2,2").unwrap();
209        assert_eq!(s.spine_len(), 4);
210        assert_eq!(s.nodes(), 7);
211        assert_eq!(s.verify_width(), 8);
212        assert!(!s.is_chain());
213        assert!(TreeShape::parse("1,1").unwrap().is_chain());
214        assert!(TreeShape::parse("").is_err());
215        assert!(TreeShape::parse("1,0").is_err());
216        assert!(TreeShape::parse("5").is_err());
217        assert!(TreeShape::parse("2,2,2,2").is_err());
218        assert_ne!(
219            TreeShape::parse("1,2").unwrap().shape_id(),
220            TreeShape::parse("2,1").unwrap().shape_id()
221        );
222    }
223
224    #[test]
225    fn row_layout_spine_prefix() {
226        let t = draft("2,2", &[10, 20], &[(1, 2, 11), (2, 2, 21)]);
227        let rows = t.rows(5);
228        assert_eq!(rows.len(), 5);
229        assert_eq!(
230            rows[0],
231            TreeRow {
232                token: 5,
233                depth: 0,
234                parent_row: 0
235            }
236        );
237        assert_eq!(
238            rows[1],
239            TreeRow {
240                token: 10,
241                depth: 1,
242                parent_row: 0
243            }
244        );
245        assert_eq!(
246            rows[2],
247            TreeRow {
248                token: 20,
249                depth: 2,
250                parent_row: 1
251            }
252        );
253        assert_eq!(
254            rows[3],
255            TreeRow {
256                token: 11,
257                depth: 1,
258                parent_row: 0
259            }
260        );
261        assert_eq!(
262            rows[4],
263            TreeRow {
264                token: 21,
265                depth: 2,
266                parent_row: 1
267            }
268        );
269    }
270
271    #[test]
272    fn accept_full_spine() {
273        let t = draft("1,1", &[10, 20], &[]);
274        let rows = t.rows(5);
275        let (path, bonus) = t.accept_path(&rows, &[10, 20, 99]);
276        assert_eq!(path, vec![1, 2]);
277        assert_eq!(bonus, 99);
278    }
279
280    #[test]
281    fn accept_hedge_rescue_is_terminal() {
282        // 2026-09-25: The hedge's bonus equals the depth-2 spine token, so a
283        // path that continued below the hedge would accept row 2 as well.
284        let t = draft("2,1", &[10, 20], &[(1, 2, 11)]);
285        let rows = t.rows(5);
286        let (path, bonus) = t.accept_path(&rows, &[11, 55, 66, 20]);
287        assert_eq!(path, vec![3]);
288        assert_eq!(bonus, 20);
289    }
290
291    #[test]
292    fn accept_hedge_after_spine_prefix_is_terminal() {
293        // 2026-09-25: The first spine token is accepted before a depth-2 hedge
294        // matches, so the hedge lookup runs after `cur` has advanced.
295        let t = draft("1,2,1", &[10, 20, 30], &[(2, 2, 21)]);
296        let rows = t.rows(5);
297        let (path, bonus) = t.accept_path(&rows, &[10, 21, 66, 77, 88]);
298        assert_eq!(path, vec![1, 4]);
299        assert_eq!(bonus, 88);
300    }
301
302    #[test]
303    fn accept_reject_all_gives_root_bonus() {
304        let t = draft("2,1", &[10, 20], &[(1, 2, 11)]);
305        let rows = t.rows(5);
306        let (path, bonus) = t.accept_path(&rows, &[42, 0, 0, 0]);
307        assert!(path.is_empty());
308        assert_eq!(bonus, 42);
309    }
310}
metrale_model_layers/speculative/tree_shape.rs

metrale_model_layers/speculative/
tree_shape.rs
