use std::collections::{HashMap, HashSet};
pub type NodeId = u32;
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct Location {
pub page_start: u32,
pub page_end: u32,
pub offset_start: u32,
pub offset_end: u32,
}
#[derive(Debug, Clone)]
pub struct PixNode {
pub id: NodeId,
pub title: String,
pub summary: String,
pub location: Location,
pub children: Vec<NodeId>,
pub content: Option<String>,
}
impl PixNode {
pub fn is_leaf(&self) -> bool {
self.children.is_empty()
}
}
#[derive(Debug, Clone)]
pub struct PixTree {
nodes: HashMap<NodeId, PixNode>,
root: NodeId,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TreeError {
MissingRoot(NodeId),
DanglingChild(NodeId),
MultipleParents(NodeId),
NotConnected(NodeId),
Cycle(NodeId),
}
impl PixTree {
pub fn new(nodes: Vec<PixNode>, root: NodeId) -> Result<Self, TreeError> {
let mut map: HashMap<NodeId, PixNode> = HashMap::with_capacity(nodes.len());
for n in nodes {
map.insert(n.id, n);
}
if !map.contains_key(&root) {
return Err(TreeError::MissingRoot(root));
}
let mut indegree: HashMap<NodeId, u32> = HashMap::new();
for node in map.values() {
for &c in &node.children {
if !map.contains_key(&c) {
return Err(TreeError::DanglingChild(c));
}
let e = indegree.entry(c).or_insert(0);
*e += 1;
if *e > 1 {
return Err(TreeError::MultipleParents(c));
}
}
}
if indegree.get(&root).copied().unwrap_or(0) != 0 {
return Err(TreeError::MultipleParents(root));
}
let mut seen: HashSet<NodeId> = HashSet::new();
let mut frontier = vec![root];
seen.insert(root);
while let Some(id) = frontier.pop() {
for &c in &map[&id].children {
if !seen.insert(c) {
return Err(TreeError::Cycle(c));
}
frontier.push(c);
}
}
if seen.len() != map.len() {
let orphan = map.keys().find(|k| !seen.contains(k)).copied().unwrap_or(root);
return Err(TreeError::NotConnected(orphan));
}
Ok(PixTree { nodes: map, root })
}
pub fn root(&self) -> &PixNode {
&self.nodes[&self.root]
}
pub fn node(&self, id: NodeId) -> Option<&PixNode> {
self.nodes.get(&id)
}
pub fn children_of(&self, id: NodeId) -> Vec<&PixNode> {
self.nodes
.get(&id)
.map(|n| n.children.iter().filter_map(|c| self.nodes.get(c)).collect())
.unwrap_or_default()
}
pub fn height(&self) -> usize {
fn depth(t: &PixTree, id: NodeId) -> usize {
let node = &t.nodes[&id];
if node.is_leaf() {
0
} else {
1 + node.children.iter().map(|&c| depth(t, c)).max().unwrap_or(0)
}
}
depth(self, self.root)
}
pub fn len(&self) -> usize {
self.nodes.len()
}
pub fn is_empty(&self) -> bool {
self.nodes.len() <= 1
}
}
pub trait RelevanceScorer {
fn score(&self, query: &str, node: &PixNode, path: &[NodeId]) -> f64;
}
#[derive(Debug, Clone)]
pub struct NavConfig {
pub b_max: usize,
pub d_max: usize,
pub theta_ratio: f64,
}
impl Default for NavConfig {
fn default() -> Self {
NavConfig { b_max: 3, d_max: 4, theta_ratio: 0.5 }
}
}
#[derive(Debug, Clone)]
pub struct NavStep {
pub depth: usize,
pub from: NodeId,
pub scored: Vec<(NodeId, f64)>,
pub threshold: f64,
pub selected: Vec<NodeId>,
}
#[derive(Debug, Clone)]
pub struct RetrievedLeaf {
pub id: NodeId,
pub path: Vec<NodeId>,
pub content: String,
pub path_gain: f64,
}
#[derive(Debug, Clone)]
pub struct NavResult {
pub leaves: Vec<RetrievedLeaf>,
pub trail: Vec<NavStep>,
pub total_gain: f64,
}
pub fn pix_navigate(
tree: &PixTree,
query: &str,
cfg: &NavConfig,
scorer: &dyn RelevanceScorer,
) -> NavResult {
let mut trail: Vec<NavStep> = Vec::new();
let mut leaves: Vec<RetrievedLeaf> = Vec::new();
let mut total_gain = 0.0_f64;
struct Frontier {
id: NodeId,
path: Vec<NodeId>,
gain: f64,
}
let root = tree.root();
if root.is_leaf() {
leaves.push(RetrievedLeaf {
id: root.id,
path: vec![root.id],
content: root.content.clone().unwrap_or_default(),
path_gain: 0.0,
});
return NavResult { leaves, trail, total_gain };
}
let mut frontier = vec![Frontier { id: root.id, path: vec![root.id], gain: 0.0 }];
for depth in 0..cfg.d_max {
if frontier.is_empty() {
break;
}
let mut next: Vec<Frontier> = Vec::new();
for f in &frontier {
let node = &tree.nodes[&f.id];
if node.is_leaf() {
leaves.push(RetrievedLeaf {
id: node.id,
path: f.path.clone(),
content: node.content.clone().unwrap_or_default(),
path_gain: f.gain,
});
continue;
}
let scored: Vec<(NodeId, f64)> = node
.children
.iter()
.map(|&c| {
let s = scorer.score(query, &tree.nodes[&c], &f.path).clamp(0.0, 1.0);
(c, s)
})
.collect();
let max_score = scored.iter().map(|(_, s)| *s).fold(0.0_f64, f64::max);
let threshold = cfg.theta_ratio * max_score;
let mut survivors: Vec<(NodeId, f64)> =
scored.iter().copied().filter(|(_, s)| *s >= threshold).collect();
survivors.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
survivors.truncate(cfg.b_max);
let selected: Vec<NodeId> = survivors.iter().map(|(c, _)| *c).collect();
trail.push(NavStep {
depth,
from: node.id,
scored: scored.clone(),
threshold,
selected: selected.clone(),
});
for (c, s) in survivors {
total_gain += s;
let mut path = f.path.clone();
path.push(c);
next.push(Frontier { id: c, path, gain: f.gain + s });
}
}
frontier = next;
}
for f in frontier {
let node = &tree.nodes[&f.id];
leaves.push(RetrievedLeaf {
id: node.id,
path: f.path,
content: node
.content
.clone()
.unwrap_or_else(|| node.summary.clone()),
path_gain: f.gain,
});
}
NavResult { leaves, trail, total_gain }
}
pub fn pix_drill(
tree: &PixTree,
subtree_root: NodeId,
query: &str,
cfg: &NavConfig,
scorer: &dyn RelevanceScorer,
) -> Option<NavResult> {
if !tree.nodes.contains_key(&subtree_root) {
return None;
}
let subtree = PixTree { nodes: tree.nodes.clone(), root: subtree_root };
Some(pix_navigate(&subtree, query, cfg, scorer))
}
pub fn pix_trail(tree: &PixTree, result: &NavResult) -> Vec<String> {
result
.trail
.iter()
.map(|step| {
let from_title = tree
.nodes
.get(&step.from)
.map(|n| n.title.as_str())
.unwrap_or("?");
let picks: Vec<String> = step
.selected
.iter()
.map(|id| {
let title = tree.nodes.get(id).map(|n| n.title.as_str()).unwrap_or("?");
let score = step
.scored
.iter()
.find(|(c, _)| c == id)
.map(|(_, s)| *s)
.unwrap_or(0.0);
format!("{title} (I={score:.2})")
})
.collect();
format!("@{} {from_title} → [{}]", step.depth, picks.join(", "))
})
.collect()
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum IndexError {
Empty,
NoHeadings,
}
pub fn index_markdown(text: &str) -> Result<PixTree, IndexError> {
if text.trim().is_empty() {
return Err(IndexError::Empty);
}
struct Section {
level: usize,
title: String,
body: String,
}
let mut sections: Vec<Section> = Vec::new();
for line in text.lines() {
let trimmed = line.trim_start();
let hashes = trimmed.chars().take_while(|&c| c == '#').count();
if hashes > 0 && trimmed[hashes..].starts_with(' ') {
sections.push(Section {
level: hashes,
title: trimmed[hashes..].trim().to_string(),
body: String::new(),
});
} else if let Some(s) = sections.last_mut() {
if !line.trim().is_empty() {
if !s.body.is_empty() {
s.body.push(' ');
}
s.body.push_str(line.trim());
}
}
}
if sections.is_empty() {
return Err(IndexError::NoHeadings);
}
let mut nodes: Vec<PixNode> = vec![PixNode {
id: 0,
title: "root".to_string(),
summary: "document root".to_string(),
location: Location::default(),
children: vec![],
content: None,
}];
let mut stack: Vec<(usize, usize)> = vec![(0, 0)];
for (i, sec) in sections.iter().enumerate() {
let id = (i + 1) as NodeId;
while let Some(&(lvl, _)) = stack.last() {
if lvl >= sec.level && stack.len() > 1 {
stack.pop();
} else {
break;
}
}
let parent_idx = stack.last().map(|&(_, idx)| idx).unwrap_or(0);
nodes[parent_idx].children.push(id);
let snippet: String = sec.body.chars().take(160).collect();
nodes.push(PixNode {
id,
title: sec.title.clone(),
summary: if snippet.is_empty() {
sec.title.clone()
} else {
format!("{} — {}", sec.title, snippet)
},
location: Location::default(),
children: vec![],
content: Some(sec.body.clone()),
});
stack.push((sec.level, nodes.len() - 1));
}
let child_bearers: HashSet<NodeId> = nodes
.iter()
.filter(|n| !n.children.is_empty())
.map(|n| n.id)
.collect();
for n in nodes.iter_mut() {
if child_bearers.contains(&n.id) {
n.content = None;
}
}
PixTree::new(nodes, 0).map_err(|_| IndexError::NoHeadings)
}
pub struct LexicalScorer {
pub epsilon: f64,
}
impl Default for LexicalScorer {
fn default() -> Self {
LexicalScorer { epsilon: 0.05 }
}
}
fn tokenize(s: &str) -> HashSet<String> {
s.split(|c: char| !c.is_alphanumeric())
.filter(|w| w.len() >= 2)
.map(|w| w.to_lowercase())
.collect()
}
impl RelevanceScorer for LexicalScorer {
fn score(&self, query: &str, node: &PixNode, _path: &[NodeId]) -> f64 {
let q = tokenize(query);
if q.is_empty() {
return self.epsilon;
}
let mut text = tokenize(&node.title);
text.extend(tokenize(&node.summary));
let hits = q.iter().filter(|t| text.contains(*t)).count();
let coverage = hits as f64 / q.len() as f64;
coverage.max(self.epsilon).min(1.0)
}
}
pub fn build_score_prompt(query: &str, node: &PixNode) -> String {
format!(
"You are navigating a document by relevance, reading ONLY section \
summaries (not full text). Rate from 0 to 100 how likely this section — \
or one of its subsections — CONTAINS THE ANSWER to the query. Reply with \
ONLY the integer.\n\nQuery: {query}\nSection title: {}\nSection summary: \
{}\n\nScore (0-100):",
node.title, node.summary
)
}
pub fn parse_score(response: &str) -> f64 {
let mut num = String::new();
for ch in response.chars() {
if ch.is_ascii_digit() || (ch == '.' && !num.contains('.')) {
num.push(ch);
} else if !num.is_empty() {
break;
}
}
let n: f64 = num.trim_matches('.').parse().unwrap_or(0.0);
let frac = if n > 1.0 { n / 100.0 } else { n };
frac.clamp(0.0, 1.0)
}
pub struct LlmRelevanceScorer<F>
where
F: Fn(&str) -> String,
{
pub complete: F,
}
impl<F> RelevanceScorer for LlmRelevanceScorer<F>
where
F: Fn(&str) -> String,
{
fn score(&self, query: &str, node: &PixNode, _path: &[NodeId]) -> f64 {
parse_score(&(self.complete)(&build_score_prompt(query, node)))
}
}
pub fn find_by_title_path(tree: &PixTree, titles: &[&str]) -> Option<NodeId> {
let mut current = tree.root().id;
for want in titles {
let want_lc = want.trim().to_lowercase();
if want_lc.is_empty() {
continue;
}
let next = tree
.children_of(current)
.into_iter()
.find(|n| n.title.to_lowercase() == want_lc)?;
current = next.id;
}
if current == tree.root().id {
None
} else {
Some(current)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn leaf(id: NodeId, title: &str, content: &str) -> PixNode {
PixNode {
id,
title: title.into(),
summary: format!("summary of {title}"),
location: Location::default(),
children: vec![],
content: Some(content.into()),
}
}
fn internal(id: NodeId, title: &str, children: Vec<NodeId>) -> PixNode {
PixNode {
id,
title: title.into(),
summary: format!("summary of {title}"),
location: Location::default(),
children,
content: None,
}
}
fn sample_tree() -> PixTree {
PixTree::new(
vec![
internal(0, "root", vec![1, 2]),
internal(1, "left", vec![3, 4]),
internal(2, "right", vec![5, 6]),
leaf(3, "L3", "content-3"),
leaf(4, "L4", "content-4"),
leaf(5, "L5", "content-5"),
leaf(6, "L6", "content-6"),
],
0,
)
.expect("valid tree")
}
struct KeywordScorer;
impl RelevanceScorer for KeywordScorer {
fn score(&self, query: &str, node: &PixNode, _path: &[NodeId]) -> f64 {
if query.contains(&node.title) {
1.0
} else if node.title == "left" && query.contains("L3") {
1.0
} else if node.title == "left" && query.contains("L4") {
1.0
} else if node.title == "right" && (query.contains("L5") || query.contains("L6")) {
1.0
} else {
0.1
}
}
}
#[test]
fn valid_tree_builds() {
let t = sample_tree();
assert_eq!(t.len(), 7);
assert_eq!(t.height(), 2);
assert_eq!(t.root().title, "root");
}
#[test]
fn missing_root_rejected() {
let e = PixTree::new(vec![leaf(1, "a", "x")], 99).unwrap_err();
assert_eq!(e, TreeError::MissingRoot(99));
}
#[test]
fn dangling_child_rejected() {
let e = PixTree::new(vec![internal(0, "r", vec![7])], 0).unwrap_err();
assert_eq!(e, TreeError::DanglingChild(7));
}
#[test]
fn shared_child_rejected_t2() {
let e = PixTree::new(
vec![
internal(0, "r", vec![1, 2]),
internal(1, "a", vec![3]),
internal(2, "b", vec![3]),
leaf(3, "c", "x"),
],
0,
)
.unwrap_err();
assert_eq!(e, TreeError::MultipleParents(3));
}
#[test]
fn cycle_rejected_t3() {
let e = PixTree::new(
vec![internal(0, "r", vec![1]), internal(1, "a", vec![0])],
0,
)
.unwrap_err();
assert!(matches!(e, TreeError::Cycle(_) | TreeError::MultipleParents(0)));
}
#[test]
fn disconnected_rejected_t1() {
let e = PixTree::new(
vec![internal(0, "r", vec![1]), leaf(1, "a", "x"), leaf(9, "orphan", "y")],
0,
)
.unwrap_err();
assert_eq!(e, TreeError::NotConnected(9));
}
#[test]
fn navigates_to_the_targeted_leaf() {
let t = sample_tree();
let r = pix_navigate(&t, "find L5", &NavConfig::default(), &KeywordScorer);
assert!(r.leaves.iter().any(|l| l.id == 5), "expected to retrieve L5: {:?}", r.leaves);
let l5 = r.leaves.iter().find(|l| l.id == 5).unwrap();
assert_eq!(l5.content, "content-5");
assert_eq!(l5.path, vec![0, 2, 5], "reasoning path root→right→L5");
}
#[test]
fn navigation_terminates_within_d_max() {
let t = sample_tree();
let cfg = NavConfig { b_max: 3, d_max: 2, theta_ratio: 0.5 };
let r = pix_navigate(&t, "find L6", &cfg, &KeywordScorer);
assert!(r.trail.iter().all(|s| s.depth < cfg.d_max));
assert!(r.leaves.iter().all(|l| l.path.len() <= cfg.d_max + 1));
}
#[test]
fn every_leaf_has_a_path() {
let t = sample_tree();
let r = pix_navigate(&t, "find L3", &NavConfig::default(), &KeywordScorer);
assert!(!r.leaves.is_empty());
for l in &r.leaves {
assert_eq!(l.path.first(), Some(&0), "path starts at root");
assert_eq!(l.path.last(), Some(&l.id), "path ends at the leaf");
}
}
#[test]
fn path_gain_is_monotone_nondecreasing() {
let t = sample_tree();
let r = pix_navigate(&t, "find L4", &NavConfig::default(), &KeywordScorer);
for l in &r.leaves {
let mut running = 0.0_f64;
for w in l.path.windows(2) {
let (from, to) = (w[0], w[1]);
if let Some(step) = r.trail.iter().find(|s| s.from == from) {
let s = step.scored.iter().find(|(c, _)| *c == to).map(|(_, s)| *s).unwrap_or(0.0);
assert!(s >= 0.0, "every information score is non-negative");
let prev = running;
running += s;
assert!(running >= prev, "cumulative gain is monotone non-decreasing");
}
}
assert!((l.path_gain - running).abs() < 1e-9, "path_gain matches the trail");
}
assert!(r.total_gain >= 0.0);
}
#[test]
fn branching_is_capped_at_b_max() {
let t = PixTree::new(
vec![
internal(0, "root", vec![1, 2, 3, 4]),
leaf(1, "a", "x"),
leaf(2, "b", "x"),
leaf(3, "c", "x"),
leaf(4, "d", "x"),
],
0,
)
.unwrap();
let cfg = NavConfig { b_max: 2, d_max: 4, theta_ratio: 0.0 };
let r = pix_navigate(&t, "anything", &cfg, &KeywordScorer);
for step in &r.trail {
assert!(step.selected.len() <= cfg.b_max, "b_max respected");
}
}
#[test]
fn drill_navigates_a_subtree() {
let t = sample_tree();
let r = pix_drill(&t, 2, "find L6", &NavConfig::default(), &KeywordScorer).unwrap();
assert!(r.leaves.iter().any(|l| l.id == 6));
let l6 = r.leaves.iter().find(|l| l.id == 6).unwrap();
assert_eq!(l6.path.first(), Some(&2));
}
#[test]
fn drill_unknown_subtree_is_none() {
let t = sample_tree();
assert!(pix_drill(&t, 999, "q", &NavConfig::default(), &KeywordScorer).is_none());
}
#[test]
fn trail_renders_reasoning_path() {
let t = sample_tree();
let r = pix_navigate(&t, "find L5", &NavConfig::default(), &KeywordScorer);
let trail = pix_trail(&t, &r);
assert!(!trail.is_empty());
assert!(trail[0].contains("root"), "trail starts at root: {trail:?}");
assert!(trail.iter().any(|s| s.contains("I=")));
}
#[test]
fn root_leaf_is_returned_directly() {
let t = PixTree::new(vec![leaf(0, "only", "the-answer")], 0).unwrap();
let r = pix_navigate(&t, "q", &NavConfig::default(), &KeywordScorer);
assert_eq!(r.leaves.len(), 1);
assert_eq!(r.leaves[0].content, "the-answer");
assert!(r.trail.is_empty());
}
const DOC: &str = r#"
# Liability
General liability terms.
## Indemnification
The seller indemnifies the buyer against third-party claims.
## Limitation
Liability is capped at the contract value.
# Termination
## Notice
Either party may terminate with thirty days written notice.
"#;
#[test]
fn index_markdown_builds_a_valid_tree() {
let t = index_markdown(DOC).expect("indexable");
assert_eq!(t.len(), 6);
let liability = t
.children_of(t.root().id)
.into_iter()
.find(|n| n.title == "Liability")
.expect("Liability under root");
assert_eq!(t.children_of(liability.id).len(), 2);
assert!(liability.content.is_none(), "internal node has no content");
let indemn = t
.children_of(liability.id)
.into_iter()
.find(|n| n.title == "Indemnification")
.unwrap();
assert!(indemn.content.as_deref().unwrap().contains("indemnifies"));
}
#[test]
fn index_rejects_empty_and_headingless() {
assert_eq!(index_markdown(" ").unwrap_err(), IndexError::Empty);
assert_eq!(
index_markdown("just prose, no headings").unwrap_err(),
IndexError::NoHeadings
);
}
#[test]
fn lexical_scorer_is_embeddings_free_and_floored() {
let s = LexicalScorer::default();
let n = leaf(1, "Indemnification", "the seller indemnifies the buyer");
assert!(s.score("indemnification clause", &n, &[]) > 0.4);
assert!((s.score("quantum chromodynamics", &n, &[]) - s.epsilon).abs() < 1e-9);
}
#[test]
fn end_to_end_index_then_navigate_retrieves_the_right_section() {
let tree = index_markdown(DOC).unwrap();
let r = pix_navigate(
&tree,
"what is the cap on liability limitation",
&NavConfig::default(),
&LexicalScorer::default(),
);
let got: Vec<&str> = r.leaves.iter().map(|l| l.content.as_str()).collect();
assert!(
got.iter().any(|c| c.contains("capped at the contract value")),
"expected the Limitation section, got {got:?}"
);
let trail = pix_trail(&tree, &r);
assert!(trail.iter().any(|s| s.contains("Liability")));
}
#[test]
fn find_by_title_path_locates_a_subtree() {
let tree = index_markdown(DOC).unwrap();
let id = find_by_title_path(&tree, &["liability", "limitation"]).unwrap();
assert_eq!(tree.node(id).unwrap().title, "Limitation");
let id2 = find_by_title_path(&tree, &["termination"]).unwrap();
assert_eq!(tree.node(id2).unwrap().title, "Termination");
assert!(find_by_title_path(&tree, &["nonexistent"]).is_none());
}
#[test]
fn parse_score_normalises_to_unit_interval() {
assert!((parse_score("85") - 0.85).abs() < 1e-9); assert!((parse_score("100") - 1.0).abs() < 1e-9);
assert!((parse_score("0") - 0.0).abs() < 1e-9);
assert!((parse_score("0.7") - 0.7).abs() < 1e-9); assert!((parse_score("Score: 42 / 100") - 0.42).abs() < 1e-9); assert!((parse_score("the answer is likely here") - 0.0).abs() < 1e-9); assert!((parse_score("250") - 1.0).abs() < 1e-9); }
#[test]
fn build_score_prompt_carries_query_and_summary() {
let n = PixNode {
id: 1,
title: "Limitation".into(),
summary: "liability is capped at contract value".into(),
location: Location::default(),
children: vec![],
content: Some("full text".into()),
};
let p = build_score_prompt("what is the cap", &n);
assert!(p.contains("what is the cap"), "carries the query");
assert!(p.contains("Limitation"), "carries the title");
assert!(p.contains("liability is capped at contract value"), "carries the summary");
assert!(p.contains("ONLY the integer"), "instructs a bare-integer reply");
}
#[test]
fn llm_scorer_uses_the_injected_completion() {
let scorer = LlmRelevanceScorer {
complete: |prompt: &str| {
if prompt.contains("Limitation") { "90".to_string() } else { "5".to_string() }
},
};
let hit = leaf(1, "Limitation", "capped");
let miss = leaf(2, "Preamble", "intro");
assert!((scorer.score("cap", &hit, &[]) - 0.90).abs() < 1e-9);
assert!((scorer.score("cap", &miss, &[]) - 0.05).abs() < 1e-9);
let tree = index_markdown(DOC).unwrap();
let r = pix_navigate(&tree, "limitation cap", &NavConfig::default(), &scorer);
assert!(r.leaves.iter().any(|l| l.content.contains("capped at the contract value")));
}
}