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