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