1use std::{cmp::Ordering, fmt::Debug, num::NonZeroU32};
2
3use idx_binary::{AvltrieeSearch, IdxBinary, IdxFileAvlTriee};
4
5use crate::{Data, FieldName, RowSet};
6
7pub trait CustomSort {
8 fn compare(&self, a: NonZeroU32, b: NonZeroU32) -> Ordering;
9 fn asc(&self) -> Vec<NonZeroU32>;
10 fn desc(&self) -> Vec<NonZeroU32>;
11}
12impl Debug for dyn CustomSort {
13 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
14 writeln!(f, "CustomSort")?;
15 Ok(())
16 }
17}
18
19pub struct NoCustomSort {}
20
21impl CustomSort for NoCustomSort {
22 fn compare(&self, _: NonZeroU32, _: NonZeroU32) -> Ordering {
23 unreachable!()
24 }
25
26 fn asc(&self) -> Vec<NonZeroU32> {
27 unreachable!()
28 }
29
30 fn desc(&self) -> Vec<NonZeroU32> {
31 unreachable!()
32 }
33}
34
35#[derive(Debug)]
36pub enum CustomOrderKey<C: CustomSort> {
37 Serial,
38 Row,
39 TermBegin,
40 TermEnd,
41 LastUpdated,
42 Field(FieldName),
43 Custom(C),
44}
45
46pub type OrderKey = CustomOrderKey<NoCustomSort>;
47
48#[derive(Debug)]
49pub enum Order<C: CustomSort> {
50 Asc(CustomOrderKey<C>),
51 Desc(CustomOrderKey<C>),
52}
53
54impl Data {
55 pub fn sort<C: CustomSort>(&self, rows: &RowSet, orders: &[Order<C>]) -> Vec<NonZeroU32> {
57 let sub_orders = &orders[1..];
58 match &orders[0] {
59 Order::Asc(key) => self.sort_with_key(rows, key, sub_orders),
60 Order::Desc(key) => self.sort_with_key_desc(rows, key, sub_orders),
61 }
62 }
63
64 fn subsort<C: CustomSort>(
65 &self,
66 tmp: Vec<NonZeroU32>,
67 sub_orders: &[Order<C>],
68 ) -> Vec<NonZeroU32> {
69 let mut tmp = tmp;
70 tmp.sort_by(|a, b| {
71 for i in 0..sub_orders.len() {
72 match &sub_orders[i] {
73 Order::Asc(order_key) => match order_key {
74 CustomOrderKey::Serial => {
75 return unsafe {
76 self.serial
77 .value_unchecked(*a)
78 .cmp(self.serial.value_unchecked(*b))
79 };
80 }
81 CustomOrderKey::Row => return a.cmp(b),
82 CustomOrderKey::TermBegin => {
83 if let Some(ref f) = self.term_begin {
84 let ord =
85 unsafe { f.value_unchecked(*a).cmp(f.value_unchecked(*b)) };
86 if ord != Ordering::Equal {
87 return ord;
88 }
89 }
90 }
91 CustomOrderKey::TermEnd => {
92 if let Some(ref f) = self.term_end {
93 let ord =
94 unsafe { f.value_unchecked(*a).cmp(f.value_unchecked(*b)) };
95 if ord != Ordering::Equal {
96 return ord;
97 }
98 }
99 }
100 CustomOrderKey::LastUpdated => {
101 if let Some(ref f) = self.last_updated {
102 let ord =
103 unsafe { f.value_unchecked(*a).cmp(f.value_unchecked(*b)) };
104 if ord != Ordering::Equal {
105 return ord;
106 }
107 }
108 }
109 CustomOrderKey::Field(name) => {
110 if let Some(field) = self.fields.get(name) {
111 let ord = IdxBinary::cmp(
112 field.value(*a).unwrap(),
113 field.value(*b).unwrap(),
114 );
115 if ord != Ordering::Equal {
116 return ord;
117 }
118 }
119 }
120 CustomOrderKey::Custom(custom_order) => {
121 let ord = custom_order.compare(*a, *b);
122 if ord != Ordering::Equal {
123 return ord;
124 }
125 }
126 },
127 Order::Desc(order_key) => match order_key {
128 CustomOrderKey::Serial => {
129 return unsafe {
130 self.serial
131 .value_unchecked(*b)
132 .cmp(self.serial.value_unchecked(*a))
133 };
134 }
135 CustomOrderKey::Row => {
136 return b.cmp(a);
137 }
138 CustomOrderKey::TermBegin => {
139 if let Some(ref f) = self.term_begin {
140 let ord =
141 unsafe { f.value_unchecked(*b).cmp(f.value_unchecked(*a)) };
142 if ord != Ordering::Equal {
143 return ord;
144 }
145 }
146 }
147 CustomOrderKey::TermEnd => {
148 if let Some(ref f) = self.term_end {
149 let ord =
150 unsafe { f.value_unchecked(*b).cmp(f.value_unchecked(*a)) };
151 if ord != Ordering::Equal {
152 return ord;
153 }
154 }
155 }
156 CustomOrderKey::LastUpdated => {
157 if let Some(ref f) = self.last_updated {
158 let ord =
159 unsafe { f.value_unchecked(*b).cmp(f.value_unchecked(*a)) };
160 if ord != Ordering::Equal {
161 return ord;
162 }
163 }
164 }
165 CustomOrderKey::Field(name) => {
166 if let Some(field) = self.fields.get(name) {
167 let ord = IdxBinary::cmp(
168 field.value(*b).unwrap(),
169 field.value(*a).unwrap(),
170 );
171 if ord != Ordering::Equal {
172 return ord;
173 }
174 }
175 }
176 CustomOrderKey::Custom(custom_order) => {
177 let ord = custom_order.compare(*b, *a);
178 if ord != Ordering::Equal {
179 return ord;
180 }
181 }
182 },
183 }
184 }
185 Ordering::Equal
186 });
187 tmp
188 }
189
190 fn sort_with_triee_inner<T: PartialEq, I: ?Sized, C: CustomSort>(
191 &self,
192 rows: &RowSet,
193 triee: &IdxFileAvlTriee<T, I>,
194 iter: impl Iterator<Item = NonZeroU32>,
195 sub_orders: &[Order<C>],
196 ) -> Vec<NonZeroU32> {
197 if sub_orders.len() == 0 {
198 iter.filter_map(|row| rows.contains(&row).then_some(row))
199 .collect()
200 } else {
201 let mut ret = Vec::new();
202
203 let mut before: Option<&T> = None;
204 let mut tmp: Vec<NonZeroU32> = Vec::new();
205 for r in iter {
206 if rows.contains(&r) {
207 let value = unsafe { triee.node_unchecked(r) };
208 if let Some(before) = before {
209 if before.ne(value) {
210 ret.extend(if tmp.len() <= 1 {
211 tmp
212 } else {
213 self.subsort(tmp, sub_orders)
214 });
215 tmp = vec![];
216 }
217 } else {
218 ret.extend(tmp);
219 tmp = vec![];
220 }
221 tmp.push(r);
222 before = Some(value);
223 }
224 }
225 ret.extend(if tmp.len() <= 1 {
226 tmp
227 } else {
228 self.subsort(tmp, sub_orders)
229 });
230 ret
231 }
232 }
233
234 fn sort_with_triee<T: PartialEq, I: ?Sized, C: CustomSort>(
235 &self,
236 rows: &RowSet,
237 triee: &IdxFileAvlTriee<T, I>,
238 sub_orders: &[Order<C>],
239 ) -> Vec<NonZeroU32> {
240 self.sort_with_triee_inner(rows, triee, triee.iter(), sub_orders)
241 }
242
243 fn sort_with_triee_desc<T: PartialEq, I: ?Sized, C: CustomSort>(
244 &self,
245 rows: &RowSet,
246 index: &IdxFileAvlTriee<T, I>,
247 sub_orders: &[Order<C>],
248 ) -> Vec<NonZeroU32> {
249 self.sort_with_triee_inner(rows, index, index.desc_iter(), sub_orders)
250 }
251
252 fn sort_with_key<C: CustomSort>(
253 &self,
254 rows: &RowSet,
255 key: &CustomOrderKey<C>,
256 sub_orders: &[Order<C>],
257 ) -> Vec<NonZeroU32> {
258 match key {
259 CustomOrderKey::Serial => {
260 self.sort_with_triee::<u32, u32, C>(rows, &self.serial, &vec![])
261 }
262 CustomOrderKey::Row => rows.into_iter().cloned().collect(),
263 CustomOrderKey::TermBegin => self.term_begin.as_ref().map_or_else(
264 || rows.into_iter().cloned().collect(),
265 |f| self.sort_with_triee(rows, f, sub_orders),
266 ),
267 CustomOrderKey::TermEnd => self.term_end.as_ref().map_or_else(
268 || rows.into_iter().cloned().collect(),
269 |f| self.sort_with_triee(rows, f, sub_orders),
270 ),
271 CustomOrderKey::LastUpdated => self.term_end.as_ref().map_or_else(
272 || rows.into_iter().cloned().collect(),
273 |f| self.sort_with_triee(rows, f, sub_orders),
274 ),
275 CustomOrderKey::Field(name) => self.fields.get(name).map_or_else(
276 || rows.into_iter().cloned().collect(),
277 |f| self.sort_with_triee(rows, f.as_ref(), sub_orders),
278 ),
279 CustomOrderKey::Custom(custom_order) => custom_order.asc(),
280 }
281 }
282
283 fn sort_with_key_desc<C: CustomSort>(
284 &self,
285 rows: &RowSet,
286 key: &CustomOrderKey<C>,
287 sub_orders: &[Order<C>],
288 ) -> Vec<NonZeroU32> {
289 match key {
290 CustomOrderKey::Serial => {
291 self.sort_with_triee_desc::<u32, u32, C>(rows, &self.serial, &vec![])
292 }
293 CustomOrderKey::Row => rows.into_iter().rev().cloned().collect(),
294 CustomOrderKey::TermBegin => self.term_begin.as_ref().map_or_else(
295 || rows.into_iter().rev().cloned().collect(),
296 |f| self.sort_with_triee_desc(rows, f, sub_orders),
297 ),
298 CustomOrderKey::TermEnd => self.term_end.as_ref().map_or_else(
299 || rows.into_iter().rev().cloned().collect(),
300 |f| self.sort_with_triee_desc(rows, f, sub_orders),
301 ),
302 CustomOrderKey::LastUpdated => self.last_updated.as_ref().map_or_else(
303 || rows.into_iter().rev().cloned().collect(),
304 |f| self.sort_with_triee_desc(rows, f, sub_orders),
305 ),
306 CustomOrderKey::Field(name) => self.fields.get(name).map_or_else(
307 || rows.into_iter().rev().cloned().collect(),
308 |f| self.sort_with_triee_desc(rows, f.as_ref(), sub_orders),
309 ),
310 CustomOrderKey::Custom(custom_order) => custom_order.desc(),
311 }
312 }
313}