weighted_gss/
stack_visit.rs

1use crate::Weight;
2use crate::gss::WeightedGss;
3use crate::nodes::*;
4use smallvec::SmallVec;
5use std::fmt;
6use std::hash::Hash;
7
8/// Error returned when a bounded operation would materialise too many stacks.
9///
10/// The error is intentionally opaque: the caller already supplied the limit.
11#[derive(Clone, Copy, Debug, PartialEq, Eq)]
12pub struct StackLimitExceeded(());
13
14impl StackLimitExceeded {
15    pub(crate) const fn new() -> Self {
16        Self(())
17    }
18}
19
20impl fmt::Display for StackLimitExceeded {
21    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
22        formatter.write_str("the weighted GSS exceeds the configured distinct-stack limit")
23    }
24}
25
26impl std::error::Error for StackLimitExceeded {}
27
28/// Visit distinct concrete stacks in top-first order without unbounded expansion.
29///
30/// At most `max_stacks` callbacks are made. Equal concrete stacks are coalesced
31/// and their weights joined before the callback. Returns [`StackLimitExceeded`]
32/// without invoking the callback when the complete result would exceed the
33/// supplied limit.
34pub fn for_each_stack_top_first<S, W>(
35    gss: &WeightedGss<S, W>,
36    max_stacks: usize,
37    mut visit: impl FnMut(&[S], &W),
38) -> Result<(), StackLimitExceeded>
39where
40    S: Clone + Eq + Hash,
41    W: Weight,
42{
43    if gss.root.paths == 1 {
44        let mut stack = SmallVec::<[S; 16]>::new();
45        if let Some(weight) = single_weighted_path(&gss.root, &mut stack) {
46            if max_stacks == 0 {
47                return Err(StackLimitExceeded::new());
48            }
49            visit(&stack, weight);
50            return Ok(());
51        }
52    }
53
54    let mut stacks = crate::materialize::materialize_stacks(&gss.root, max_stacks)?;
55    for (stack, weight) in &mut stacks {
56        stack.reverse();
57        visit(stack, weight);
58    }
59    Ok(())
60}
61
62pub(crate) fn collect_stacks_top_first<S, W>(
63    gss: &WeightedGss<S, W>,
64    max_stacks: usize,
65) -> Result<Vec<(Vec<S>, W)>, StackLimitExceeded>
66where
67    S: Clone + Eq + Hash,
68    W: Weight,
69{
70    if gss.root.paths == 1 {
71        let mut stack = SmallVec::<[S; 16]>::new();
72        if let Some(weight) = single_weighted_path(&gss.root, &mut stack) {
73            if max_stacks == 0 {
74                return Err(StackLimitExceeded::new());
75            }
76            return Ok(vec![(stack.into_vec(), weight.clone())]);
77        }
78    }
79
80    let mut stacks = crate::materialize::materialize_stacks(&gss.root, max_stacks)?;
81    for (stack, _) in &mut stacks {
82        stack.reverse();
83    }
84    Ok(stacks)
85}
86
87fn single_weighted_path<'a, S, W>(
88    mut node: &'a WRef<S, W>,
89    output: &mut SmallVec<[S; 16]>,
90) -> Option<&'a W>
91where
92    S: Clone,
93{
94    loop {
95        match &node.kind {
96            WKind::Shared { weight, stacks } => {
97                single_unweighted_path(stacks, output)?;
98                return Some(weight.as_ref());
99            }
100            WKind::Segment { values, next } => {
101                output.extend(values.iter().cloned());
102                node = next;
103            }
104            WKind::Branch { empty, children } => {
105                if empty.len() == 1 && children.is_empty() {
106                    return Some(empty[0].as_ref());
107                }
108                let mut entries = children
109                    .iter()
110                    .flat_map(|(top, values)| values.iter().map(move |child| (top, child)));
111                let (top, child) = entries.next()?;
112                if entries.next().is_some() || !empty.is_empty() {
113                    return None;
114                }
115                output.push(top.clone());
116                node = child;
117            }
118        }
119    }
120}
121
122fn single_unweighted_path<S>(mut node: &URef<S>, output: &mut SmallVec<[S; 16]>) -> Option<()>
123where
124    S: Clone,
125{
126    loop {
127        match &node.kind {
128            UKind::Segment { values, next } => {
129                output.extend(values.iter().cloned());
130                node = next;
131            }
132            UKind::Branch { empty, children } => {
133                if *empty && children.is_empty() {
134                    return Some(());
135                }
136                let mut entries = children
137                    .iter()
138                    .flat_map(|(top, values)| values.iter().map(move |child| (top, child)));
139                let (top, child) = entries.next()?;
140                if entries.next().is_some() || *empty {
141                    return None;
142                }
143                output.push(top.clone());
144                node = child;
145            }
146        }
147    }
148}
weighted_gss/stack_visit.rs

weighted_gss/
stack_visit.rs
