util.rs

  1use futures::Future;
  2use rand::prelude::*;
  3use std::{
  4    cmp::Ordering,
  5    pin::Pin,
  6    task::{Context, Poll},
  7};
  8pub use sum_tree::Bias;
  9
 10pub fn post_inc(value: &mut usize) -> usize {
 11    let prev = *value;
 12    *value += 1;
 13    prev
 14}
 15
 16/// Extend a sorted vector with a sorted sequence of items, maintaining the vector's sort order and
 17/// enforcing a maximum length. Sort the items according to the given callback. Before calling this,
 18/// both `vec` and `new_items` should already be sorted according to the `cmp` comparator.
 19pub fn extend_sorted<T, I, F>(vec: &mut Vec<T>, new_items: I, limit: usize, mut cmp: F)
 20where
 21    I: IntoIterator<Item = T>,
 22    F: FnMut(&T, &T) -> Ordering,
 23{
 24    let mut start_index = 0;
 25    for new_item in new_items {
 26        if let Err(i) = vec[start_index..].binary_search_by(|m| cmp(m, &new_item)) {
 27            let index = start_index + i;
 28            if vec.len() < limit {
 29                vec.insert(index, new_item);
 30            } else if index < vec.len() {
 31                vec.pop();
 32                vec.insert(index, new_item);
 33            }
 34            start_index = index;
 35        }
 36    }
 37}
 38
 39pub struct RandomCharIter<T: Rng>(T);
 40
 41impl<T: Rng> RandomCharIter<T> {
 42    #[cfg(test)]
 43    pub fn new(rng: T) -> Self {
 44        Self(rng)
 45    }
 46}
 47
 48impl<T: Rng> Iterator for RandomCharIter<T> {
 49    type Item = char;
 50
 51    fn next(&mut self) -> Option<Self::Item> {
 52        match self.0.gen_range(0..100) {
 53            // whitespace
 54            0..=19 => [' ', '\n', '\t'].choose(&mut self.0).copied(),
 55            // two-byte greek letters
 56            20..=32 => char::from_u32(self.0.gen_range(('α' as u32)..('ω' as u32 + 1))),
 57            // three-byte characters
 58            33..=45 => ['✋', '✅', '❌', '❎', '⭐'].choose(&mut self.0).copied(),
 59            // four-byte characters
 60            46..=58 => ['🍐', '🏀', '🍗', '🎉'].choose(&mut self.0).copied(),
 61            // ascii letters
 62            _ => Some(self.0.gen_range(b'a'..b'z' + 1).into()),
 63        }
 64    }
 65}
 66
 67pub trait ResultExt {
 68    type Ok;
 69
 70    fn log_err(self) -> Option<Self::Ok>;
 71}
 72
 73impl<T> ResultExt for anyhow::Result<T> {
 74    type Ok = T;
 75
 76    fn log_err(self) -> Option<T> {
 77        match self {
 78            Ok(value) => Some(value),
 79            Err(error) => {
 80                log::error!("{:?}", error);
 81                None
 82            }
 83        }
 84    }
 85}
 86
 87pub trait TryFutureExt {
 88    fn log_err(self) -> LogErrorFuture<Self>
 89    where
 90        Self: Sized;
 91}
 92
 93impl<F, T> TryFutureExt for F
 94where
 95    F: Future<Output = anyhow::Result<T>>,
 96{
 97    fn log_err(self) -> LogErrorFuture<Self>
 98    where
 99        Self: Sized,
100    {
101        LogErrorFuture(self)
102    }
103}
104
105pub struct LogErrorFuture<F>(F);
106
107impl<F, T> Future for LogErrorFuture<F>
108where
109    F: Future<Output = anyhow::Result<T>>,
110{
111    type Output = Option<T>;
112
113    fn poll(self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
114        let inner = unsafe { Pin::new_unchecked(&mut self.get_unchecked_mut().0) };
115        match inner.poll(cx) {
116            Poll::Ready(output) => Poll::Ready(output.log_err()),
117            Poll::Pending => Poll::Pending,
118        }
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125
126    #[test]
127    fn test_extend_sorted() {
128        let mut vec = vec![];
129
130        extend_sorted(&mut vec, vec![21, 17, 13, 8, 1, 0], 5, |a, b| b.cmp(a));
131        assert_eq!(vec, &[21, 17, 13, 8, 1]);
132
133        extend_sorted(&mut vec, vec![101, 19, 17, 8, 2], 8, |a, b| b.cmp(a));
134        assert_eq!(vec, &[101, 21, 19, 17, 13, 8, 2, 1]);
135
136        extend_sorted(&mut vec, vec![1000, 19, 17, 9, 5], 8, |a, b| b.cmp(a));
137        assert_eq!(vec, &[1000, 101, 21, 19, 17, 13, 9, 8]);
138    }
139}
140
141// Allow surf Results to accept context like other Results do when
142// using anyhow.
143pub trait SurfResultExt {
144    fn context<C>(self, cx: C) -> Self
145    where
146        C: std::fmt::Display + Send + Sync + 'static;
147
148    fn with_context<C, F>(self, f: F) -> Self
149    where
150        C: std::fmt::Display + Send + Sync + 'static,
151        F: FnOnce() -> C;
152}
153
154impl<T> SurfResultExt for surf::Result<T> {
155    fn context<C>(self, cx: C) -> Self
156    where
157        C: std::fmt::Display + Send + Sync + 'static,
158    {
159        self.map_err(|e| surf::Error::new(e.status(), e.into_inner().context(cx)))
160    }
161
162    fn with_context<C, F>(self, f: F) -> Self
163    where
164        C: std::fmt::Display + Send + Sync + 'static,
165        F: FnOnce() -> C,
166    {
167        self.map_err(|e| surf::Error::new(e.status(), e.into_inner().context(f())))
168    }
169}