Skip to main content

weighted_gss/
linear_prefix.rs

1use crate::Weight;
2use crate::gss::WeightedGss;
3use crate::nodes::{UKind, URef, WKind, u_has_empty, u_segment, w_shared};
4use crate::segment::Segment;
5use smallvec::SmallVec;
6use std::hash::Hash;
7use std::sync::Arc;
8
9/// Try to expose a mutable linear top prefix over an unchanged hidden floor.
10///
11/// Returns `None` when the current representation does not have one homogeneous
12/// weight and a directly accessible linear prefix.
13#[must_use]
14pub fn linear_prefix<S, W>(gss: &WeightedGss<S, W>) -> Option<LinearPrefix<S, W>>
15where
16    S: Clone + Eq + Hash,
17    W: Weight,
18{
19    LinearPrefix::from_gss(gss)
20}
21
22/// Mutable view of a linear top prefix over an unchanged hidden floor.
23///
24/// The hidden floor may still be branched. Use [`Self::floor_is_empty`] to test
25/// whether the prefix represents one complete concrete stack.
26#[derive(Clone)]
27pub struct LinearPrefix<S, W> {
28    values: Option<Segment<S>>,
29    next: URef<S>,
30    weight: Arc<W>,
31    pending: SmallVec<[S; 2]>,
32}
33
34impl<S, W> LinearPrefix<S, W>
35where
36    S: Clone + Eq + Hash,
37    W: Weight,
38{
39    pub(crate) fn from_gss(gss: &WeightedGss<S, W>) -> Option<Self> {
40        let WKind::Shared { weight, stacks } = &gss.root.kind else {
41            return None;
42        };
43        let UKind::Segment { values, next } = &stacks.kind else {
44            return None;
45        };
46        Some(Self {
47            values: Some(values.clone()),
48            next: next.clone(),
49            weight: weight.clone(),
50            pending: SmallVec::new(),
51        })
52    }
53
54    /// Return the number of values in the accessible linear prefix.
55    #[must_use]
56    pub fn len(&self) -> usize {
57        let current = self.values.as_ref().map_or(0, Segment::len);
58        self.pending
59            .len()
60            .saturating_add(current)
61            .saturating_add(segment_chain_len(&self.next))
62    }
63
64    /// Return whether the accessible prefix is empty.
65    #[must_use]
66    pub fn is_empty(&self) -> bool {
67        self.len() == 0
68    }
69
70    /// Return a visible value by depth from the top.
71    #[must_use]
72    pub fn get(&self, mut depth_from_top: usize) -> Option<&S> {
73        if depth_from_top < self.pending.len() {
74            return self.pending.iter().rev().nth(depth_from_top);
75        }
76        depth_from_top -= self.pending.len();
77
78        if let Some(values) = &self.values {
79            if depth_from_top < values.len() {
80                return values.get(depth_from_top);
81            }
82            depth_from_top -= values.len();
83        }
84
85        let mut next = &self.next;
86        loop {
87            match &next.kind {
88                UKind::Segment {
89                    values,
90                    next: following,
91                } => {
92                    if depth_from_top < values.len() {
93                        return values.get(depth_from_top);
94                    }
95                    depth_from_top -= values.len();
96                    next = following;
97                }
98                UKind::Branch { .. } => return None,
99            }
100        }
101    }
102
103    /// Return whether the hidden floor is exactly the empty stack.
104    #[must_use]
105    pub fn floor_is_empty(&self) -> bool {
106        let floor = segment_chain_floor(&self.next);
107        floor.paths == 1 && u_has_empty(floor)
108    }
109
110    /// Push one value onto the prefix.
111    pub fn push(&mut self, value: S) {
112        self.pending.push(value);
113    }
114
115    /// Pop values from the prefix and return the number that reached its floor.
116    pub fn popn(&mut self, mut count: usize) -> usize {
117        while count > 0 && !self.pending.is_empty() {
118            self.pending.pop();
119            count -= 1;
120        }
121
122        while count > 0 {
123            let Some(values) = self.values.take() else {
124                break;
125            };
126            if count < values.len() {
127                self.values = values.drop_front(count);
128                return 0;
129            }
130            count -= values.len();
131            match &self.next.kind {
132                UKind::Segment { values, next } => {
133                    self.values = Some(values.clone());
134                    self.next = next.clone();
135                }
136                UKind::Branch { .. } => {
137                    self.values = None;
138                    break;
139                }
140            }
141        }
142        count
143    }
144
145    /// Convert the view back into a weighted GSS.
146    #[must_use]
147    pub fn into_gss(self) -> WeightedGss<S, W> {
148        let mut stacks = match self.values {
149            Some(values) => u_segment(values, self.next),
150            None => self.next,
151        };
152        if !self.pending.is_empty() {
153            stacks = u_segment(
154                Segment::from_top_first(self.pending.into_iter().rev().collect()),
155                stacks,
156            );
157        }
158        WeightedGss {
159            root: w_shared(self.weight, stacks),
160        }
161    }
162}
163
164fn segment_chain_len<S>(mut node: &URef<S>) -> usize {
165    let mut len = 0usize;
166    while let UKind::Segment { values, next } = &node.kind {
167        len = len.saturating_add(values.len());
168        node = next;
169    }
170    len
171}
172
173fn segment_chain_floor<S>(mut node: &URef<S>) -> &URef<S> {
174    while let UKind::Segment { next, .. } = &node.kind {
175        node = next;
176    }
177    node
178}