weighted_gss/
linear_prefix.rs1use 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#[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#[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 #[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 #[must_use]
66 pub fn is_empty(&self) -> bool {
67 self.len() == 0
68 }
69
70 #[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 #[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 pub fn push(&mut self, value: S) {
112 self.pending.push(value);
113 }
114
115 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 #[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}