crates/stdx/src/lib.rs

   1 //! Missing batteries for standard libraries.
   2 use std::{
   3     sync::atomic::{AtomicUsize, Ordering},
   4     time::Instant,
   5 };
   6
   7 mod macros;
   8
   9 #[inline(always)]
  10 pub fn is_ci() -> bool {
  11     option_env!("CI").is_some()
  12 }
  13
  14 #[must_use]
  15 pub fn timeit(label: &'static str) -> impl Drop {
  16     struct Guard {
  17         label: &'static str,
  18         start: Instant,
  19     }
  20
  21     impl Drop for Guard {
  22         fn drop(&mut self) {
  23             eprintln!("{}: {:.2?}", self.label, self.start.elapsed())
  24         }
  25     }
  26
  27     Guard { label, start: Instant::now() }
  28 }
  29
  30 pub fn to_lower_snake_case(s: &str) -> String {
  31     let mut buf = String::with_capacity(s.len());
  32     let mut prev = false;
  33     for c in s.chars() {
  34         if c.is_ascii_uppercase() && prev {
  35             buf.push('_')
  36         }
  37         prev = true;
  38
  39         buf.push(c.to_ascii_lowercase());
  40     }
  41     buf
  42 }
  43
  44 pub fn replace(buf: &mut String, from: char, to: &str) {
  45     if !buf.contains(from) {
  46         return;
  47     }
  48     // FIXME: do this in place.
  49     *buf = buf.replace(from, to)
  50 }
  51
  52 // https://github.com/rust-lang/rust/issues/74773
  53 pub fn split_once(haystack: &str, delim: char) -> Option<(&str, &str)> {
  54     let mut split = haystack.splitn(2, delim);
  55     let prefix = split.next()?;
  56     let suffix = split.next()?;
  57     Some((prefix, suffix))
  58 }
  59 pub fn rsplit_once(haystack: &str, delim: char) -> Option<(&str, &str)> {
  60     let mut split = haystack.rsplitn(2, delim);
  61     let suffix = split.next()?;
  62     let prefix = split.next()?;
  63     Some((prefix, suffix))
  64 }
  65
  66 pub fn trim_indent(mut text: &str) -> String {
  67     if text.starts_with('\n') {
  68         text = &text[1..];
  69     }
  70     let indent = text
  71         .lines()
  72         .filter(|it| !it.trim().is_empty())
  73         .map(|it| it.len() - it.trim_start().len())
  74         .min()
  75         .unwrap_or(0);
  76     lines_with_ends(text)
  77         .map(
  78             |line| {
  79                 if line.len() <= indent {
  80                     line.trim_start_matches(' ')
  81                 } else {
  82                     &line[indent..]
  83                 }
  84             },
  85         )
  86         .collect()
  87 }
  88
  89 pub fn lines_with_ends(text: &str) -> LinesWithEnds {
  90     LinesWithEnds { text }
  91 }
  92
  93 pub struct LinesWithEnds<'a> {
  94     text: &'a str,
  95 }
  96
  97 impl<'a> Iterator for LinesWithEnds<'a> {
  98     type Item = &'a str;
  99     fn next(&mut self) -> Option<&'a str> {
 100         if self.text.is_empty() {
 101             return None;
 102         }
 103         let idx = self.text.find('\n').map_or(self.text.len(), |it| it + 1);
 104         let (res, next) = self.text.split_at(idx);
 105         self.text = next;
 106         Some(res)
 107     }
 108 }
 109
 110 // https://github.com/rust-lang/rust/issues/73831
 111 pub fn partition_point<T, P>(slice: &[T], mut pred: P) -> usize
 112 where
 113     P: FnMut(&T) -> bool,
 114 {
 115     let mut left = 0;
 116     let mut right = slice.len();
 117
 118     while left != right {
 119         let mid = left + (right - left) / 2;
 120         // SAFETY:
 121         // When left < right, left <= mid < right.
 122         // Therefore left always increases and right always decreases,
 123         // and either of them is selected.
 124         // In both cases left <= right is satisfied.
 125         // Therefore if left < right in a step,
 126         // left <= right is satisfied in the next step.
 127         // Therefore as long as left != right, 0 <= left < right <= len is satisfied
 128         // and if this case 0 <= mid < len is satisfied too.
 129         let value = unsafe { slice.get_unchecked(mid) };
 130         if pred(value) {
 131             left = mid + 1;
 132         } else {
 133             right = mid;
 134         }
 135     }
 136
 137     left
 138 }
 139
 140 pub struct RacyFlag(AtomicUsize);
 141
 142 impl RacyFlag {
 143     pub const fn new() -> RacyFlag {
 144         RacyFlag(AtomicUsize::new(!0))
 145     }
 146
 147     pub fn get(&self, init: impl FnMut() -> bool) -> bool {
 148         let mut init = Some(init);
 149         self.get_impl(&mut || init.take().map_or(false, |mut f| f()))
 150     }
 151
 152     fn get_impl(&self, init: &mut dyn FnMut() -> bool) -> bool {
 153         match self.0.load(Ordering::Relaxed) {
 154             0 => false,
 155             1 => true,
 156             _ => {
 157                 let res = init();
 158                 self.0.store(if res { 1 } else { 0 }, Ordering::Relaxed);
 159                 res
 160             }
 161         }
 162     }
 163 }
 164
 165 #[cfg(test)]
 166 mod tests {
 167     use super::*;
 168
 169     #[test]
 170     fn test_trim_indent() {
 171         assert_eq!(trim_indent(""), "");
 172         assert_eq!(
 173             trim_indent(
 174                 "
 175             hello
 176             world
 177 "
 178             ),
 179             "hello\nworld\n"
 180         );
 181         assert_eq!(
 182             trim_indent(
 183                 "
 184             hello
 185             world"
 186             ),
 187             "hello\nworld"
 188         );
 189         assert_eq!(trim_indent("    hello\n    world\n"), "hello\nworld\n");
 190         assert_eq!(
 191             trim_indent(
 192                 "
 193             fn main() {
 194                 return 92;
 195             }
 196         "
 197             ),
 198             "fn main() {\n    return 92;\n}\n"
 199         );
 200     }
 201 }