src/libstd/io/extensions.rs

   1 // Copyright 2013-2014 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 //! Utility mixins that apply to all Readers and Writers
  12
  13 #![allow(missing_docs)]
  14
  15 // FIXME: Not sure how this should be structured
  16 // FIXME: Iteration should probably be considered separately
  17
  18 use io::{IoError, IoResult, Reader};
  19 use io;
  20 use iter::Iterator;
  21 use num::Int;
  22 use ops::FnOnce;
  23 use option::Option;
  24 use option::Option::{Some, None};
  25 use ptr::PtrExt;
  26 use result::Result::{Ok, Err};
  27 use slice::{SliceExt, AsSlice};
  28
  29 /// An iterator that reads a single byte on each iteration,
  30 /// until `.read_byte()` returns `EndOfFile`.
  31 ///
  32 /// # Notes about the Iteration Protocol
  33 ///
  34 /// The `Bytes` may yield `None` and thus terminate
  35 /// an iteration, but continue to yield elements if iteration
  36 /// is attempted again.
  37 ///
  38 /// # Error
  39 ///
  40 /// Any error other than `EndOfFile` that is produced by the underlying Reader
  41 /// is returned by the iterator and should be handled by the caller.
  42 pub struct Bytes<'r, T:'r> {
  43     reader: &'r mut T,
  44 }
  45
  46 impl<'r, R: Reader> Bytes<'r, R> {
  47     /// Constructs a new byte iterator from the given Reader instance.
  48     pub fn new(r: &'r mut R) -> Bytes<'r, R> {
  49         Bytes {
  50             reader: r,
  51         }
  52     }
  53 }
  54
  55 impl<'r, R: Reader> Iterator for Bytes<'r, R> {
  56     type Item = IoResult<u8>;
  57
  58     #[inline]
  59     fn next(&mut self) -> Option<IoResult<u8>> {
  60         match self.reader.read_byte() {
  61             Ok(x) => Some(Ok(x)),
  62             Err(IoError { kind: io::EndOfFile, .. }) => None,
  63             Err(e) => Some(Err(e))
  64         }
  65     }
  66 }
  67
  68 /// Converts an 8-bit to 64-bit unsigned value to a little-endian byte
  69 /// representation of the given size. If the size is not big enough to
  70 /// represent the value, then the high-order bytes are truncated.
  71 ///
  72 /// Arguments:
  73 ///
  74 /// * `n`: The value to convert.
  75 /// * `size`: The size of the value, in bytes. This must be 8 or less, or task
  76 ///           panic occurs. If this is less than 8, then a value of that
  77 ///           many bytes is produced. For example, if `size` is 4, then a
  78 ///           32-bit byte representation is produced.
  79 /// * `f`: A callback that receives the value.
  80 ///
  81 /// This function returns the value returned by the callback, for convenience.
  82 pub fn u64_to_le_bytes<T, F>(n: u64, size: uint, f: F) -> T where
  83     F: FnOnce(&[u8]) -> T,
  84 {
  85     use mem::transmute;
  86
  87     // LLVM fails to properly optimize this when using shifts instead of the to_le* intrinsics
  88     assert!(size <= 8u);
  89     match size {
  90       1u => f(&[n as u8]),
  91       2u => f(unsafe { & transmute::<_, [u8; 2]>((n as u16).to_le()) }),
  92       4u => f(unsafe { & transmute::<_, [u8; 4]>((n as u32).to_le()) }),
  93       8u => f(unsafe { & transmute::<_, [u8; 8]>(n.to_le()) }),
  94       _ => {
  95
  96         let mut bytes = vec!();
  97         let mut i = size;
  98         let mut n = n;
  99         while i > 0u {
 100             bytes.push((n & 255_u64) as u8);
 101             n >>= 8;
 102             i -= 1u;
 103         }
 104         f(bytes.as_slice())
 105       }
 106     }
 107 }
 108
 109 /// Converts an 8-bit to 64-bit unsigned value to a big-endian byte
 110 /// representation of the given size. If the size is not big enough to
 111 /// represent the value, then the high-order bytes are truncated.
 112 ///
 113 /// Arguments:
 114 ///
 115 /// * `n`: The value to convert.
 116 /// * `size`: The size of the value, in bytes. This must be 8 or less, or task
 117 ///           panic occurs. If this is less than 8, then a value of that
 118 ///           many bytes is produced. For example, if `size` is 4, then a
 119 ///           32-bit byte representation is produced.
 120 /// * `f`: A callback that receives the value.
 121 ///
 122 /// This function returns the value returned by the callback, for convenience.
 123 pub fn u64_to_be_bytes<T, F>(n: u64, size: uint, f: F) -> T where
 124     F: FnOnce(&[u8]) -> T,
 125 {
 126     use mem::transmute;
 127
 128     // LLVM fails to properly optimize this when using shifts instead of the to_be* intrinsics
 129     assert!(size <= 8u);
 130     match size {
 131       1u => f(&[n as u8]),
 132       2u => f(unsafe { & transmute::<_, [u8; 2]>((n as u16).to_be()) }),
 133       4u => f(unsafe { & transmute::<_, [u8; 4]>((n as u32).to_be()) }),
 134       8u => f(unsafe { & transmute::<_, [u8; 8]>(n.to_be()) }),
 135       _ => {
 136         let mut bytes = vec!();
 137         let mut i = size;
 138         while i > 0u {
 139             let shift = (i - 1u) * 8u;
 140             bytes.push((n >> shift) as u8);
 141             i -= 1u;
 142         }
 143         f(bytes.as_slice())
 144       }
 145     }
 146 }
 147
 148 /// Extracts an 8-bit to 64-bit unsigned big-endian value from the given byte
 149 /// buffer and returns it as a 64-bit value.
 150 ///
 151 /// Arguments:
 152 ///
 153 /// * `data`: The buffer in which to extract the value.
 154 /// * `start`: The offset at which to extract the value.
 155 /// * `size`: The size of the value in bytes to extract. This must be 8 or
 156 ///           less, or task panic occurs. If this is less than 8, then only
 157 ///           that many bytes are parsed. For example, if `size` is 4, then a
 158 ///           32-bit value is parsed.
 159 pub fn u64_from_be_bytes(data: &[u8], start: uint, size: uint) -> u64 {
 160     use ptr::{copy_nonoverlapping_memory};
 161     use slice::SliceExt;
 162
 163     assert!(size <= 8u);
 164
 165     if data.len() - start < size {
 166         panic!("index out of bounds");
 167     }
 168
 169     let mut buf = [0u8; 8];
 170     unsafe {
 171         let ptr = data.as_ptr().offset(start as int);
 172         let out = buf.as_mut_ptr();
 173         copy_nonoverlapping_memory(out.offset((8 - size) as int), ptr, size);
 174         (*(out as *const u64)).to_be()
 175     }
 176 }
 177
 178 #[cfg(test)]
 179 mod test {
 180     use prelude::v1::*;
 181     use io;
 182     use io::{MemReader, BytesReader};
 183
 184     struct InitialZeroByteReader {
 185         count: int,
 186     }
 187
 188     impl Reader for InitialZeroByteReader {
 189         fn read(&mut self, buf: &mut [u8]) -> io::IoResult<uint> {
 190             if self.count == 0 {
 191                 self.count = 1;
 192                 Ok(0)
 193             } else {
 194                 buf[0] = 10;
 195                 Ok(1)
 196             }
 197         }
 198     }
 199
 200     struct EofReader;
 201
 202     impl Reader for EofReader {
 203         fn read(&mut self, _: &mut [u8]) -> io::IoResult<uint> {
 204             Err(io::standard_error(io::EndOfFile))
 205         }
 206     }
 207
 208     struct ErroringReader;
 209
 210     impl Reader for ErroringReader {
 211         fn read(&mut self, _: &mut [u8]) -> io::IoResult<uint> {
 212             Err(io::standard_error(io::InvalidInput))
 213         }
 214     }
 215
 216     struct PartialReader {
 217         count: int,
 218     }
 219
 220     impl Reader for PartialReader {
 221         fn read(&mut self, buf: &mut [u8]) -> io::IoResult<uint> {
 222             if self.count == 0 {
 223                 self.count = 1;
 224                 buf[0] = 10;
 225                 buf[1] = 11;
 226                 Ok(2)
 227             } else {
 228                 buf[0] = 12;
 229                 buf[1] = 13;
 230                 Ok(2)
 231             }
 232         }
 233     }
 234
 235     struct ErroringLaterReader {
 236         count: int,
 237     }
 238
 239     impl Reader for ErroringLaterReader {
 240         fn read(&mut self, buf: &mut [u8]) -> io::IoResult<uint> {
 241             if self.count == 0 {
 242                 self.count = 1;
 243                 buf[0] = 10;
 244                 Ok(1)
 245             } else {
 246                 Err(io::standard_error(io::InvalidInput))
 247             }
 248         }
 249     }
 250
 251     struct ThreeChunkReader {
 252         count: int,
 253     }
 254
 255     impl Reader for ThreeChunkReader {
 256         fn read(&mut self, buf: &mut [u8]) -> io::IoResult<uint> {
 257             if self.count == 0 {
 258                 self.count = 1;
 259                 buf[0] = 10;
 260                 buf[1] = 11;
 261                 Ok(2)
 262             } else if self.count == 1 {
 263                 self.count = 2;
 264                 buf[0] = 12;
 265                 buf[1] = 13;
 266                 Ok(2)
 267             } else {
 268                 Err(io::standard_error(io::EndOfFile))
 269             }
 270         }
 271     }
 272
 273     #[test]
 274     fn read_byte() {
 275         let mut reader = MemReader::new(vec!(10));
 276         let byte = reader.read_byte();
 277         assert!(byte == Ok(10));
 278     }
 279
 280     #[test]
 281     fn read_byte_0_bytes() {
 282         let mut reader = InitialZeroByteReader {
 283             count: 0,
 284         };
 285         let byte = reader.read_byte();
 286         assert!(byte == Ok(10));
 287     }
 288
 289     #[test]
 290     fn read_byte_eof() {
 291         let mut reader = EofReader;
 292         let byte = reader.read_byte();
 293         assert!(byte.is_err());
 294     }
 295
 296     #[test]
 297     fn read_byte_error() {
 298         let mut reader = ErroringReader;
 299         let byte = reader.read_byte();
 300         assert!(byte.is_err());
 301     }
 302
 303     #[test]
 304     fn bytes_0_bytes() {
 305         let mut reader = InitialZeroByteReader {
 306             count: 0,
 307         };
 308         let byte = reader.bytes().next();
 309         assert!(byte == Some(Ok(10)));
 310     }
 311
 312     #[test]
 313     fn bytes_eof() {
 314         let mut reader = EofReader;
 315         let byte = reader.bytes().next();
 316         assert!(byte.is_none());
 317     }
 318
 319     #[test]
 320     fn bytes_error() {
 321         let mut reader = ErroringReader;
 322         let mut it = reader.bytes();
 323         let byte = it.next();
 324         assert!(byte.unwrap().is_err());
 325     }
 326
 327     #[test]
 328     fn read_bytes() {
 329         let mut reader = MemReader::new(vec!(10, 11, 12, 13));
 330         let bytes = reader.read_exact(4).unwrap();
 331         assert!(bytes == vec!(10, 11, 12, 13));
 332     }
 333
 334     #[test]
 335     fn read_bytes_partial() {
 336         let mut reader = PartialReader {
 337             count: 0,
 338         };
 339         let bytes = reader.read_exact(4).unwrap();
 340         assert!(bytes == vec!(10, 11, 12, 13));
 341     }
 342
 343     #[test]
 344     fn read_bytes_eof() {
 345         let mut reader = MemReader::new(vec!(10, 11));
 346         assert!(reader.read_exact(4).is_err());
 347     }
 348
 349     #[test]
 350     fn push_at_least() {
 351         let mut reader = MemReader::new(vec![10, 11, 12, 13]);
 352         let mut buf = vec![8, 9];
 353         assert!(reader.push_at_least(4, 4, &mut buf).is_ok());
 354         assert!(buf == vec![8, 9, 10, 11, 12, 13]);
 355     }
 356
 357     #[test]
 358     fn push_at_least_partial() {
 359         let mut reader = PartialReader {
 360             count: 0,
 361         };
 362         let mut buf = vec![8, 9];
 363         assert!(reader.push_at_least(4, 4, &mut buf).is_ok());
 364         assert!(buf == vec![8, 9, 10, 11, 12, 13]);
 365     }
 366
 367     #[test]
 368     fn push_at_least_eof() {
 369         let mut reader = MemReader::new(vec![10, 11]);
 370         let mut buf = vec![8, 9];
 371         assert!(reader.push_at_least(4, 4, &mut buf).is_err());
 372         assert!(buf == vec![8, 9, 10, 11]);
 373     }
 374
 375     #[test]
 376     fn push_at_least_error() {
 377         let mut reader = ErroringLaterReader {
 378             count: 0,
 379         };
 380         let mut buf = vec![8, 9];
 381         assert!(reader.push_at_least(4, 4, &mut buf).is_err());
 382         assert!(buf == vec![8, 9, 10]);
 383     }
 384
 385     #[test]
 386     fn read_to_end() {
 387         let mut reader = ThreeChunkReader {
 388             count: 0,
 389         };
 390         let buf = reader.read_to_end().unwrap();
 391         assert!(buf == vec!(10, 11, 12, 13));
 392     }
 393
 394     #[test]
 395     #[should_fail]
 396     fn read_to_end_error() {
 397         let mut reader = ThreeChunkReader {
 398             count: 0,
 399         };
 400         let buf = reader.read_to_end().unwrap();
 401         assert!(buf == vec!(10, 11));
 402     }
 403
 404     #[test]
 405     fn test_read_write_le_mem() {
 406         let uints = [0, 1, 2, 42, 10_123, 100_123_456, ::u64::MAX];
 407
 408         let mut writer = Vec::new();
 409         for i in uints.iter() {
 410             writer.write_le_u64(*i).unwrap();
 411         }
 412
 413         let mut reader = MemReader::new(writer);
 414         for i in uints.iter() {
 415             assert!(reader.read_le_u64().unwrap() == *i);
 416         }
 417     }
 418
 419
 420     #[test]
 421     fn test_read_write_be() {
 422         let uints = [0, 1, 2, 42, 10_123, 100_123_456, ::u64::MAX];
 423
 424         let mut writer = Vec::new();
 425         for i in uints.iter() {
 426             writer.write_be_u64(*i).unwrap();
 427         }
 428
 429         let mut reader = MemReader::new(writer);
 430         for i in uints.iter() {
 431             assert!(reader.read_be_u64().unwrap() == *i);
 432         }
 433     }
 434
 435     #[test]
 436     fn test_read_be_int_n() {
 437         let ints = [::i32::MIN, -123456, -42, -5, 0, 1, ::i32::MAX];
 438
 439         let mut writer = Vec::new();
 440         for i in ints.iter() {
 441             writer.write_be_i32(*i).unwrap();
 442         }
 443
 444         let mut reader = MemReader::new(writer);
 445         for i in ints.iter() {
 446             // this tests that the sign extension is working
 447             // (comparing the values as i32 would not test this)
 448             assert!(reader.read_be_int_n(4).unwrap() == *i as i64);
 449         }
 450     }
 451
 452     #[test]
 453     fn test_read_f32() {
 454         //big-endian floating-point 8.1250
 455         let buf = vec![0x41, 0x02, 0x00, 0x00];
 456
 457         let mut writer = Vec::new();
 458         writer.write(buf.as_slice()).unwrap();
 459
 460         let mut reader = MemReader::new(writer);
 461         let f = reader.read_be_f32().unwrap();
 462         assert!(f == 8.1250);
 463     }
 464
 465     #[test]
 466     fn test_read_write_f32() {
 467         let f:f32 = 8.1250;
 468
 469         let mut writer = Vec::new();
 470         writer.write_be_f32(f).unwrap();
 471         writer.write_le_f32(f).unwrap();
 472
 473         let mut reader = MemReader::new(writer);
 474         assert!(reader.read_be_f32().unwrap() == 8.1250);
 475         assert!(reader.read_le_f32().unwrap() == 8.1250);
 476     }
 477
 478     #[test]
 479     fn test_u64_from_be_bytes() {
 480         use super::u64_from_be_bytes;
 481
 482         let buf = [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09];
 483
 484         // Aligned access
 485         assert_eq!(u64_from_be_bytes(&buf, 0, 0), 0);
 486         assert_eq!(u64_from_be_bytes(&buf, 0, 1), 0x01);
 487         assert_eq!(u64_from_be_bytes(&buf, 0, 2), 0x0102);
 488         assert_eq!(u64_from_be_bytes(&buf, 0, 3), 0x010203);
 489         assert_eq!(u64_from_be_bytes(&buf, 0, 4), 0x01020304);
 490         assert_eq!(u64_from_be_bytes(&buf, 0, 5), 0x0102030405);
 491         assert_eq!(u64_from_be_bytes(&buf, 0, 6), 0x010203040506);
 492         assert_eq!(u64_from_be_bytes(&buf, 0, 7), 0x01020304050607);
 493         assert_eq!(u64_from_be_bytes(&buf, 0, 8), 0x0102030405060708);
 494
 495         // Unaligned access
 496         assert_eq!(u64_from_be_bytes(&buf, 1, 0), 0);
 497         assert_eq!(u64_from_be_bytes(&buf, 1, 1), 0x02);
 498         assert_eq!(u64_from_be_bytes(&buf, 1, 2), 0x0203);
 499         assert_eq!(u64_from_be_bytes(&buf, 1, 3), 0x020304);
 500         assert_eq!(u64_from_be_bytes(&buf, 1, 4), 0x02030405);
 501         assert_eq!(u64_from_be_bytes(&buf, 1, 5), 0x0203040506);
 502         assert_eq!(u64_from_be_bytes(&buf, 1, 6), 0x020304050607);
 503         assert_eq!(u64_from_be_bytes(&buf, 1, 7), 0x02030405060708);
 504         assert_eq!(u64_from_be_bytes(&buf, 1, 8), 0x0203040506070809);
 505     }
 506 }
 507
 508 #[cfg(test)]
 509 mod bench {
 510     extern crate test;
 511
 512     use prelude::v1::*;
 513     use self::test::Bencher;
 514
 515     // why is this a macro? wouldn't an inlined function work just as well?
 516     macro_rules! u64_from_be_bytes_bench_impl {
 517         ($b:expr, $size:expr, $stride:expr, $start_index:expr) =>
 518         ({
 519             use super::u64_from_be_bytes;
 520
 521             let data = Vec::from_fn($stride*100+$start_index, |i| i as u8);
 522             let mut sum = 0u64;
 523             $b.iter(|| {
 524                 let mut i = $start_index;
 525                 while i < data.len() {
 526                     sum += u64_from_be_bytes(data.as_slice(), i, $size);
 527                     i += $stride;
 528                 }
 529             });
 530         })
 531     }
 532
 533     #[bench]
 534     fn u64_from_be_bytes_4_aligned(b: &mut Bencher) {
 535         u64_from_be_bytes_bench_impl!(b, 4, 4, 0);
 536     }
 537
 538     #[bench]
 539     fn u64_from_be_bytes_4_unaligned(b: &mut Bencher) {
 540         u64_from_be_bytes_bench_impl!(b, 4, 4, 1);
 541     }
 542
 543     #[bench]
 544     fn u64_from_be_bytes_7_aligned(b: &mut Bencher) {
 545         u64_from_be_bytes_bench_impl!(b, 7, 8, 0);
 546     }
 547
 548     #[bench]
 549     fn u64_from_be_bytes_7_unaligned(b: &mut Bencher) {
 550         u64_from_be_bytes_bench_impl!(b, 7, 8, 1);
 551     }
 552
 553     #[bench]
 554     fn u64_from_be_bytes_8_aligned(b: &mut Bencher) {
 555         u64_from_be_bytes_bench_impl!(b, 8, 8, 0);
 556     }
 557
 558     #[bench]
 559     fn u64_from_be_bytes_8_unaligned(b: &mut Bencher) {
 560         u64_from_be_bytes_bench_impl!(b, 8, 8, 1);
 561     }
 562 }