1use std::convert::TryInto;
2use std::ffi::CString;
3use std::fs::File;
4use std::io::Read;
5use std::os::unix::ffi::OsStrExt;
6use std::os::unix::prelude::AsRawFd;
7use std::path::Path;
8use std::{fmt, io, mem};
9
10use byte_parser::{ParseIterator, StrParser};
11
12use libc::c_int;
13
14const DEF_PRECISION: usize = 2;
15
16#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct DataSize {
19 bytes: u128,
20}
21
22impl DataSize {
23 pub(crate) fn from_str(s: &str) -> Option<Self> {
25 let mut iter = StrParser::new(s);
26 let float = parse_f64(&mut iter)?;
27 let unit = iter.record().consume_to_str().trim();
29
30 let unit = DataSizeUnit::from_str(unit)?;
31 Some(Self {
32 bytes: unit.to_byte(float),
33 })
34 }
35
36 pub(crate) fn from_size_bytes(bytes: impl TryInto<u128>) -> Option<Self> {
37 bytes.try_into().ok().map(|bytes| Self { bytes })
38 }
39
40 pub fn to(self, unit: &DataSizeUnit) -> f64 {
42 DataSizeUnit::convert(self.bytes, unit)
43 }
44}
45
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47pub enum DataSizeUnit {
48 B,
50 Kb, Mb, Gb, Tb, }
59
60impl DataSizeUnit {
61 const fn val(&self) -> u128 {
62 match self {
63 Self::B => 1,
64 Self::Kb => 1_024,
65 Self::Mb => 1_024 * 1_024,
66 Self::Gb => 1_024 * 1_024 * 1_024,
67 Self::Tb => 1_024 * 1_024 * 1_024 * 1_024,
68 }
69 }
70
71 fn from_str(s: &str) -> Option<Self> {
72 Some(match s {
73 "" => Self::B,
74 s if eqs(s, "b") => Self::B,
75 s if eqs(s, "kb") => Self::Kb,
76 s if eqs(s, "mb") => Self::Mb,
77 s if eqs(s, "gb") => Self::Gb,
78 s if eqs(s, "tb") => Self::Tb,
79 _ => return None,
80 })
81 }
82
83 fn to_byte(&self, val: f64) -> u128 {
84 (val * self.val() as f64) as u128
86 }
87
88 fn adjust_to(byte: u128) -> Self {
89 match byte {
90 b if b < Self::Kb.val() => Self::B,
91 b if b < Self::Mb.val() => Self::Kb,
92 b if b < Self::Gb.val() => Self::Mb,
93 b if b < Self::Tb.val() => Self::Gb,
94 _ => Self::Tb,
95 }
96 }
97
98 fn convert(byte: u128, to: &Self) -> f64 {
99 byte as f64 / to.val() as f64
100 }
101
102 const fn as_str(&self) -> &'static str {
103 match self {
104 Self::B => "b",
105 Self::Kb => "kb",
106 Self::Mb => "mb",
107 Self::Gb => "gb",
108 Self::Tb => "tb",
109 }
110 }
111
112 fn fmt_val(&self, val: f64, f: &mut fmt::Formatter<'_>) -> fmt::Result {
114 if val == 0f64 {
115 return write!(f, "{}", val);
116 }
117
118 let fract = val.fract();
120 let precision = if fract == 0f64 {
121 0
122 } else {
123 let max = f.precision().unwrap_or(DEF_PRECISION);
124
125 if max <= 1 {
126 max
127 } else {
128 calculate_precision(fract, max)
129 }
130 };
131
132 write!(f, "{:.*} {}", precision, val, self.as_str())
133 }
134}
135
136fn calculate_precision(mut fract: f64, max: usize) -> usize {
137 let max_number = 10usize.pow(max as u32) as f64;
138
139 fract *= max_number;
141 fract = fract.round();
142
143 for m in (1..=max).rev() {
145 fract /= 10f64;
146 if fract.fract() != 0f64 {
147 return m;
148 }
149 }
150
151 0
152}
153
154#[inline(always)]
155fn eqs(a: &str, b: &str) -> bool {
156 a.eq_ignore_ascii_case(b)
157}
158
159impl fmt::Display for DataSize {
160 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
161 let unit = DataSizeUnit::adjust_to(self.bytes);
162 let val = DataSizeUnit::convert(self.bytes, &unit);
163 unit.fmt_val(val, f)
164 }
165}
166
167fn parse_f64<'s, I>(iter: &mut I) -> Option<f64>
170where
171 I: ParseIterator<'s>,
172{
173 let mut iter = iter.record();
174
175 iter.while_byte_fn(u8::is_ascii_digit)
177 .consume_at_least(1)
178 .ok()?;
179
180 let has_dot = iter.next_if(|&b| b == b'.').is_some();
182
183 if has_dot {
184 iter.consume_while_byte_fn(u8::is_ascii_digit);
186 }
187
188 iter.to_str().parse().ok()
189}
190
191pub fn read_to_string_mut(
194 path: impl AsRef<Path>,
195 s: &mut String,
196) -> io::Result<()> {
197 s.clear();
198 let mut file = File::open(path)?;
199 file.read_to_string(s).map(|_| ())
200}
201
202fn cstr(path: impl AsRef<Path>) -> io::Result<CString> {
203 CString::new(path.as_ref().as_os_str().as_bytes()).map_err(From::from)
204}
205
206pub fn statfs(path: impl AsRef<Path>) -> io::Result<libc::statfs> {
208 unsafe {
209 let mut stat = mem::MaybeUninit::<libc::statfs>::uninit();
210 let c = cstr(path)?;
211 let r = libc::statfs(c.as_ptr(), stat.as_mut_ptr());
212 match r {
213 0 => Ok(stat.assume_init()),
214 -1 => Err(io::Error::last_os_error()),
215 _ => panic!("unexpected return value from statfs {:?}", r),
216 }
217 }
218}
219
220pub fn blkdev_sector_size(fd: impl AsRawFd) -> io::Result<u64> {
223 let s = unsafe {
224 let mut size: c_int = 0; match blksszget(fd.as_raw_fd(), &mut size) {
226 -1 => return Err(io::Error::last_os_error()),
227 _ => size,
228 }
229 };
230
231 s.try_into()
232 .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
233}
234
235#[cfg(any(
236 target_arch = "x86",
237 target_arch = "arm",
238 target_arch = "x86_64",
239 target_arch = "aarch64"
240))]
241unsafe fn blksszget(fd: c_int, data: *mut c_int) -> c_int {
242 let nr = (0x12 << 8) | (104 << 0);
243 libc::ioctl(fd, nr, data)
244}
245
246#[cfg(test)]
247mod tests {
248 use super::*;
249
250 #[test]
251 fn test_size() {
252 let size = DataSize::from_str("24576 kB").unwrap();
253 assert_eq!(size.to(&DataSizeUnit::Kb), 24576.0);
254 }
255
256 #[test]
257 fn size_str() {
258 let s = DataSize::from_str("1024").unwrap();
260 assert_eq!(s.to_string(), "1 kb");
261 let s = DataSize::from_str("10 kb").unwrap();
262 assert_eq!(s.to_string(), "10 kb");
263 let s = DataSize::from_str("42.1 mB").unwrap();
264 assert_eq!(s.to_string(), "42.1 mb");
265 let s = DataSize::from_str("4.22 Gb").unwrap();
266 assert_eq!(s.to_string(), "4.22 gb");
267 let s = DataSize::from_str("2000 Tb").unwrap();
268 assert_eq!(s.to_string(), "2000 tb");
269 assert_eq!(
271 format!("{:.0}", DataSize::from_str("1.2 kb").unwrap()),
272 "1 kb"
273 );
274 }
275
276 #[test]
277 fn test_precision() {
278 assert_eq!(calculate_precision(0.00005, 4), 4);
279 assert_eq!(calculate_precision(0.00001, 4), 0);
280 assert_eq!(calculate_precision(0.0001, 4), 4);
281 assert_eq!(calculate_precision(0.001, 4), 3);
282 assert_eq!(calculate_precision(0.01, 4), 2);
283 assert_eq!(calculate_precision(0.1, 4), 1);
284 assert_eq!(calculate_precision(0.0, 4), 0);
285 }
286
287 #[test]
288 fn run_statfs() {
289 statfs("/").unwrap();
290 }
291}