Skip to main content

rupg_common/
sort.rs

1//! `pg_qsort` of `sort_template.h`. Code that must give equal items in the same order as PostgreSQL uses it in place of the sort of the standard library.
2
3use std::cmp::Ordering;
4
5/// `med3`: the index of the median of three items.
6fn med3<T, E, F>(data: &[T], a: usize, b: usize, c: usize, cmp: &mut F) -> Result<usize, E>
7where
8    F: FnMut(&T, &T) -> Result<Ordering, E>,
9{
10    Ok(if cmp(&data[a], &data[b])?.is_lt() {
11        if cmp(&data[b], &data[c])?.is_lt() {
12            b
13        } else if cmp(&data[a], &data[c])?.is_lt() {
14            c
15        } else {
16            a
17        }
18    } else if cmp(&data[b], &data[c])?.is_gt() {
19        b
20    } else if cmp(&data[a], &data[c])?.is_lt() {
21        a
22    } else {
23        c
24    })
25}
26
27/// `swapn`: swaps `n` items at `a` with `n` items at `b`.
28fn swapn<T>(data: &mut [T], a: usize, b: usize, n: usize) {
29    for i in 0..n {
30        data.swap(a + i, b + i);
31    }
32}
33
34/// `pg_qsort`: a quicksort with the median of three or of nine, a partition in three parts and an insertion sort for fewer than 7 items. It is not stable. The sort stops at the first error of `cmp`.
35///
36/// # Errors
37///
38/// The first error of `cmp`.
39pub fn qsort<T, E, F>(data: &mut [T], cmp: &mut F) -> Result<(), E>
40where
41    F: FnMut(&T, &T) -> Result<Ordering, E>,
42{
43    let mut a = 0;
44    let mut n = data.len();
45    loop {
46        if n < 7 {
47            for pm in a + 1..a + n {
48                let mut pl = pm;
49                while pl > a && cmp(&data[pl - 1], &data[pl])?.is_gt() {
50                    data.swap(pl, pl - 1);
51                    pl -= 1;
52                }
53            }
54            return Ok(());
55        }
56        let mut presorted = true;
57        for pm in a + 1..a + n {
58            if cmp(&data[pm - 1], &data[pm])?.is_gt() {
59                presorted = false;
60                break;
61            }
62        }
63        if presorted {
64            return Ok(());
65        }
66        let mut pm = a + n / 2;
67        if n > 7 {
68            let mut pl = a;
69            let mut last = a + n - 1;
70            if n > 40 {
71                let d = n / 8;
72                pl = med3(data, pl, pl + d, pl + 2 * d, cmp)?;
73                pm = med3(data, pm - d, pm, pm + d, cmp)?;
74                last = med3(data, last - 2 * d, last - d, last, cmp)?;
75            }
76            pm = med3(data, pl, pm, last, cmp)?;
77        }
78        data.swap(a, pm);
79        let (mut pa, mut pb) = (a + 1, a + 1);
80        let (mut pc, mut pd) = (a + n - 1, a + n - 1);
81        loop {
82            while pb <= pc {
83                let r = cmp(&data[pb], &data[a])?;
84                if r.is_gt() {
85                    break;
86                }
87                if r.is_eq() {
88                    data.swap(pa, pb);
89                    pa += 1;
90                }
91                pb += 1;
92            }
93            while pb <= pc {
94                let r = cmp(&data[pc], &data[a])?;
95                if r.is_lt() {
96                    break;
97                }
98                if r.is_eq() {
99                    data.swap(pc, pd);
100                    pd -= 1;
101                }
102                pc -= 1;
103            }
104            if pb > pc {
105                break;
106            }
107            data.swap(pb, pc);
108            pb += 1;
109            pc -= 1;
110        }
111        let end = a + n;
112        let d1 = (pa - a).min(pb - pa);
113        swapn(data, a, pb - d1, d1);
114        let d1 = (pd - pc).min(end - pd - 1);
115        swapn(data, pb, end - d1, d1);
116        let (d1, d2) = (pb - pa, pd - pc);
117        if d1 <= d2 {
118            qsort(&mut data[a..a + d1], cmp)?;
119            a = end - d2;
120            n = d2;
121        } else {
122            qsort(&mut data[end - d2..end], cmp)?;
123            n = d1;
124        }
125    }
126}
127
128#[cfg(test)]
129mod tests {
130    use super::*;
131
132    /// The order of the keys only, so that a test can see which of two equal items comes first.
133    fn by_key(a: &(i32, usize), b: &(i32, usize)) -> Result<Ordering, ()> {
134        Ok(a.0.cmp(&b.0))
135    }
136
137    #[test]
138    fn qsort_sorts() {
139        for n in [0, 1, 5, 7, 8, 40, 41, 200, 1000] {
140            let mut items: Vec<(i32, usize)> =
141                (0..n).map(|i| (i32::try_from(i).unwrap() * 7919 % 13, i)).collect();
142            qsort(&mut items, &mut by_key).unwrap();
143            assert!(items.windows(2).all(|w| w[0].0 <= w[1].0), "n = {n}");
144            let mut seen: Vec<usize> = items.iter().map(|i| i.1).collect();
145            seen.sort_unstable();
146            assert_eq!(seen, (0..n).collect::<Vec<_>>());
147        }
148    }
149
150    #[test]
151    fn same_order_as_postgres() {
152        // The order comes from the qsort of sort_template.h, compiled in C.
153        let mut items: Vec<(i32, usize)> =
154            (0..10).map(|i| (i % 2, usize::try_from(i).unwrap())).collect();
155        qsort(&mut items, &mut by_key).unwrap();
156        assert_eq!(items.iter().map(|i| i.1).collect::<Vec<_>>(), [2, 6, 4, 8, 0, 9, 1, 3, 5, 7]);
157    }
158}