nihav-core/src/io/byteio.rs

   1 use std::io::SeekFrom;
   2 use std::fs::File;
   3 use std::io::prelude::*;
   4
   5 #[derive(Debug)]
   6 pub enum ByteIOError {
   7     EOF,
   8     WrongRange,
   9     WrongIOMode,
  10     NotImplemented,
  11     ReadError,
  12     WriteError,
  13     SeekError,
  14 }
  15
  16 pub type ByteIOResult<T> = Result<T, ByteIOError>;
  17
  18 pub trait ByteIO {
  19     fn read_buf(&mut self, buf: &mut [u8]) -> ByteIOResult<usize>;
  20     fn read_buf_some(&mut self, buf: &mut [u8]) -> ByteIOResult<usize>;
  21     fn peek_buf(&mut self, buf: &mut [u8]) -> ByteIOResult<usize>;
  22     fn read_byte(&mut self) -> ByteIOResult<u8>;
  23     fn peek_byte(&mut self) -> ByteIOResult<u8>;
  24     fn write_buf(&mut self, buf: &[u8]) -> ByteIOResult<()>;
  25     fn tell(&mut self) -> u64;
  26     fn seek(&mut self, pos: SeekFrom) -> ByteIOResult<u64>;
  27     fn is_eof(&mut self) -> bool;
  28     fn is_seekable(&mut self) -> bool;
  29     fn size(&mut self) -> i64;
  30 }
  31
  32 #[allow(dead_code)]
  33 pub struct ByteReader<'a> {
  34     io: &'a mut ByteIO,
  35 }
  36
  37 pub struct MemoryReader<'a> {
  38     buf:      &'a [u8],
  39     size:     usize,
  40     pos:      usize,
  41 }
  42
  43 pub struct FileReader<'a> {
  44     file:     &'a File,
  45     eof:      bool,
  46 }
  47
  48 macro_rules! read_int {
  49     ($s: ident, $inttype: ty, $size: expr, $which: ident) => ({
  50         let mut buf = [0; $size];
  51         $s.read_buf(&mut buf)?;
  52         unsafe {
  53             Ok((*(buf.as_ptr() as *const $inttype)).$which())
  54         }
  55     })
  56 }
  57
  58 macro_rules! peek_int {
  59     ($s: ident, $inttype: ty, $size: expr, $which: ident) => ({
  60         let mut buf = [0; $size];
  61         $s.peek_buf(&mut buf)?;
  62         unsafe {
  63             Ok((*(buf.as_ptr() as *const $inttype)).$which())
  64         }
  65     })
  66 }
  67
  68 impl<'a> ByteReader<'a> {
  69     pub fn new(io: &'a mut ByteIO) -> Self { ByteReader { io: io } }
  70
  71     pub fn read_buf(&mut self, buf: &mut [u8])  -> ByteIOResult<usize> {
  72         self.io.read_buf(buf)
  73     }
  74
  75     pub fn read_buf_some(&mut self, buf: &mut [u8])  -> ByteIOResult<usize> {
  76         self.io.read_buf_some(buf)
  77     }
  78
  79     pub fn peek_buf(&mut self, buf: &mut [u8])  -> ByteIOResult<usize> {
  80         self.io.peek_buf(buf)
  81     }
  82
  83     pub fn read_byte(&mut self) -> ByteIOResult<u8> {
  84         self.io.read_byte()
  85     }
  86
  87     pub fn peek_byte(&mut self) -> ByteIOResult<u8> {
  88         self.io.peek_byte()
  89     }
  90
  91     pub fn read_u16be(&mut self) -> ByteIOResult<u16> {
  92         read_int!(self, u16, 2, to_be)
  93     }
  94
  95     pub fn peek_u16be(&mut self) -> ByteIOResult<u16> {
  96         peek_int!(self, u16, 2, to_be)
  97     }
  98
  99     pub fn read_u24be(&mut self) -> ByteIOResult<u32> {
 100         let p16 = self.read_u16be()?;
 101         let p8 = self.read_byte()?;
 102         Ok(((p16 as u32) << 8) | (p8 as u32))
 103     }
 104
 105     pub fn peek_u24be(&mut self) -> ByteIOResult<u32> {
 106         let mut src: [u8; 3] = [0; 3];
 107         self.peek_buf(&mut src)?;
 108         Ok(((src[0] as u32) << 16) | ((src[1] as u32) << 8) | (src[2] as u32))
 109     }
 110
 111     pub fn read_u32be(&mut self) -> ByteIOResult<u32> {
 112         read_int!(self, u32, 4, to_be)
 113     }
 114
 115     pub fn peek_u32be(&mut self) -> ByteIOResult<u32> {
 116         peek_int!(self, u32, 4, to_be)
 117     }
 118
 119     pub fn read_u64be(&mut self) -> ByteIOResult<u64> {
 120         read_int!(self, u64, 8, to_be)
 121     }
 122
 123     pub fn peek_u64be(&mut self) -> ByteIOResult<u64> {
 124         peek_int!(self, u64, 8, to_be)
 125     }
 126
 127     pub fn read_u16le(&mut self) -> ByteIOResult<u16> {
 128         read_int!(self, u16, 2, to_le)
 129     }
 130
 131     pub fn peek_u16le(&mut self) -> ByteIOResult<u16> {
 132         peek_int!(self, u16, 2, to_le)
 133     }
 134
 135     pub fn read_u24le(&mut self) -> ByteIOResult<u32> {
 136         let p8 = self.read_byte()?;
 137         let p16 = self.read_u16le()?;
 138         Ok(((p16 as u32) << 8) | (p8 as u32))
 139     }
 140
 141     pub fn peek_u24le(&mut self) -> ByteIOResult<u32> {
 142         let mut src: [u8; 3] = [0; 3];
 143         self.peek_buf(&mut src)?;
 144         Ok((src[0] as u32) | ((src[1] as u32) << 8) | ((src[2] as u32) << 16))
 145     }
 146
 147     pub fn read_u32le(&mut self) -> ByteIOResult<u32> {
 148         read_int!(self, u32, 4, to_le)
 149     }
 150
 151     pub fn peek_u32le(&mut self) -> ByteIOResult<u32> {
 152         peek_int!(self, u32, 4, to_le)
 153     }
 154
 155     pub fn read_u64le(&mut self) -> ByteIOResult<u64> {
 156         read_int!(self, u64, 8, to_le)
 157     }
 158
 159     pub fn peek_u64le(&mut self) -> ByteIOResult<u64> {
 160         peek_int!(self, u64, 8, to_le)
 161     }
 162
 163     pub fn read_skip(&mut self, len: usize) -> ByteIOResult<()> {
 164         if self.io.is_seekable() {
 165             self.io.seek(SeekFrom::Current(len as i64))?;
 166         } else {
 167             let mut ssize = len;
 168             let mut buf : [u8; 16] = [0; 16];
 169             let bref = &mut buf;
 170             while ssize > bref.len() {
 171                 self.io.read_buf(bref)?;
 172                 ssize -= bref.len();
 173             }
 174             while ssize > 0 {
 175                 self.io.read_byte()?;
 176                 ssize = ssize - 1;
 177             }
 178         }
 179         Ok(())
 180     }
 181
 182     pub fn tell(&mut self) -> u64 {
 183         self.io.tell()
 184     }
 185
 186     pub fn seek(&mut self, pos: SeekFrom) -> ByteIOResult<u64> {
 187         self.io.seek(pos)
 188     }
 189
 190     pub fn is_eof(&mut self) -> bool {
 191         self.io.is_eof()
 192     }
 193
 194     pub fn size(&mut self) -> i64 {
 195         self.io.size()
 196     }
 197
 198     pub fn left(&mut self) -> i64 {
 199         let size = self.io.size();
 200         if size == -1 { return -1; }
 201         return size - (self.io.tell() as i64)
 202     }
 203 }
 204
 205 impl<'a> MemoryReader<'a> {
 206
 207     pub fn new_read(buf: &'a [u8]) -> Self {
 208         MemoryReader { buf: buf, size: buf.len(), pos: 0 }
 209     }
 210
 211     fn real_seek(&mut self, pos: i64) -> ByteIOResult<u64> {
 212         if pos < 0 || (pos as usize) > self.size {
 213             return Err(ByteIOError::WrongRange)
 214         }
 215         self.pos = pos as usize;
 216         Ok(pos as u64)
 217     }
 218 }
 219
 220 impl<'a> ByteIO for MemoryReader<'a> {
 221     fn read_byte(&mut self) -> ByteIOResult<u8> {
 222         if self.is_eof() { return Err(ByteIOError::EOF); }
 223         let res = self.buf[self.pos];
 224         self.pos = self.pos + 1;
 225         Ok(res)
 226     }
 227
 228     fn peek_byte(&mut self) -> ByteIOResult<u8> {
 229         if self.is_eof() { return Err(ByteIOError::EOF); }
 230         Ok(self.buf[self.pos])
 231     }
 232
 233     fn peek_buf(&mut self, buf: &mut [u8]) -> ByteIOResult<usize> {
 234         let copy_size = if self.size - self.pos < buf.len() { self.size } else { buf.len() };
 235         if copy_size == 0 { return Err(ByteIOError::EOF); }
 236         for i in 0..copy_size {
 237             buf[i] = self.buf[self.pos + i];
 238         }
 239         Ok(copy_size)
 240     }
 241
 242     fn read_buf(&mut self, buf: &mut [u8]) -> ByteIOResult<usize> {
 243         let read_size = self.peek_buf(buf)?;
 244         if read_size < buf.len() { return Err(ByteIOError::EOF); }
 245         self.pos += read_size;
 246         Ok(read_size)
 247     }
 248
 249     fn read_buf_some(&mut self, buf: &mut [u8]) -> ByteIOResult<usize> {
 250         let read_size = self.peek_buf(buf)?;
 251         self.pos += read_size;
 252         Ok(read_size)
 253     }
 254
 255     #[allow(unused_variables)]
 256     fn write_buf(&mut self, buf: &[u8]) -> ByteIOResult<()> {
 257         Err(ByteIOError::NotImplemented)
 258     }
 259
 260     fn tell(&mut self) -> u64 {
 261         self.pos as u64
 262     }
 263
 264     fn seek(&mut self, pos: SeekFrom) -> ByteIOResult<u64> {
 265         let cur_pos  = self.pos  as i64;
 266         let cur_size = self.size as i64;
 267         match pos {
 268             SeekFrom::Start(x)   => self.real_seek(x as i64),
 269             SeekFrom::Current(x) => self.real_seek(cur_pos + x),
 270             SeekFrom::End(x)     => self.real_seek(cur_size + x),
 271         }
 272     }
 273
 274     fn is_eof(&mut self) -> bool {
 275         self.pos >= self.size
 276     }
 277
 278     fn is_seekable(&mut self) -> bool {
 279         true
 280     }
 281
 282     fn size(&mut self) -> i64 {
 283         self.buf.len() as i64
 284     }
 285 }
 286
 287 impl<'a> FileReader<'a> {
 288
 289     pub fn new_read(file: &'a mut File) -> Self {
 290         FileReader { file: file, eof : false }
 291     }
 292 }
 293
 294 impl<'a> ByteIO for FileReader<'a> {
 295     fn read_byte(&mut self) -> ByteIOResult<u8> {
 296         let mut byte : [u8; 1] = [0];
 297         let err = self.file.read(&mut byte);
 298         if let Err(_) = err { return Err(ByteIOError::ReadError); }
 299         let sz = err.unwrap();
 300         if sz == 0 { self.eof = true; return Err(ByteIOError::EOF); }
 301         Ok (byte[0])
 302     }
 303
 304     fn peek_byte(&mut self) -> ByteIOResult<u8> {
 305         let b = self.read_byte()?;
 306         self.seek(SeekFrom::Current(-1))?;
 307         Ok(b)
 308     }
 309
 310     fn read_buf(&mut self, buf: &mut [u8]) -> ByteIOResult<usize> {
 311         let res = self.file.read(buf);
 312         if let Err(_) = res { return Err(ByteIOError::ReadError); }
 313         let sz = res.unwrap();
 314         if sz < buf.len() { self.eof = true; return Err(ByteIOError::EOF); }
 315         Ok(sz)
 316     }
 317
 318     fn read_buf_some(&mut self, buf: &mut [u8]) -> ByteIOResult<usize> {
 319         let res = self.file.read(buf);
 320         if let Err(_) = res { return Err(ByteIOError::ReadError); }
 321         let sz = res.unwrap();
 322         if sz < buf.len() { self.eof = true; }
 323         Ok(sz)
 324     }
 325
 326     fn peek_buf(&mut self, buf: &mut [u8]) -> ByteIOResult<usize> {
 327         let size = self.read_buf(buf)?;
 328         self.seek(SeekFrom::Current(-(size as i64)))?;
 329         Ok(size)
 330     }
 331
 332     #[allow(unused_variables)]
 333     fn write_buf(&mut self, buf: &[u8]) -> ByteIOResult<()> {
 334         Err(ByteIOError::NotImplemented)
 335     }
 336
 337     fn tell(&mut self) -> u64 {
 338         self.file.seek(SeekFrom::Current(0)).unwrap()
 339     }
 340
 341     fn seek(&mut self, pos: SeekFrom) -> ByteIOResult<u64> {
 342         let res = self.file.seek(pos);
 343         match res {
 344             Ok(r) => Ok(r),
 345             Err(_) => Err(ByteIOError::SeekError),
 346         }
 347     }
 348
 349     fn is_eof(&mut self) -> bool {
 350         self.eof
 351     }
 352
 353     fn is_seekable(&mut self) -> bool {
 354         true
 355     }
 356
 357     fn size(&mut self) -> i64 {
 358         -1
 359     }
 360 }
 361
 362 #[allow(dead_code)]
 363 pub struct ByteWriter<'a> {
 364     io: &'a mut ByteIO,
 365 }
 366
 367 pub struct MemoryWriter<'a> {
 368     buf:      &'a mut [u8],
 369     size:     usize,
 370     pos:      usize,
 371 }
 372
 373 pub struct FileWriter {
 374     file:     File,
 375 }
 376
 377 impl<'a> ByteWriter<'a> {
 378     pub fn new(io: &'a mut ByteIO) -> Self { ByteWriter { io: io } }
 379
 380     pub fn write_buf(&mut self, buf: &[u8])  -> ByteIOResult<()> {
 381         self.io.write_buf(buf)
 382     }
 383
 384     pub fn write_byte(&mut self, val: u8) -> ByteIOResult<()> {
 385         let buf: [u8; 1] = [val];
 386         self.io.write_buf(&buf)
 387     }
 388
 389     pub fn write_u16be(&mut self, val: u16) -> ByteIOResult<()> {
 390         let buf: [u8; 2] = [((val >> 8) & 0xFF) as u8, (val & 0xFF) as u8];
 391         self.io.write_buf(&buf)
 392     }
 393
 394     pub fn write_u16le(&mut self, val: u16) -> ByteIOResult<()> {
 395         let buf: [u8; 2] = [(val & 0xFF) as u8, ((val >> 8) & 0xFF) as u8];
 396         self.io.write_buf(&buf)
 397     }
 398
 399     pub fn write_u24be(&mut self, val: u32) -> ByteIOResult<()> {
 400         let buf: [u8; 3] = [((val >> 16) & 0xFF) as u8, ((val >> 8) & 0xFF) as u8, (val & 0xFF) as u8];
 401         self.write_buf(&buf)
 402     }
 403
 404     pub fn write_u24le(&mut self, val: u32) -> ByteIOResult<()> {
 405         let buf: [u8; 3] = [(val & 0xFF) as u8, ((val >> 8) & 0xFF) as u8, ((val >> 16) & 0xFF) as u8];
 406         self.write_buf(&buf)
 407     }
 408
 409     pub fn write_u32be(&mut self, val: u32) -> ByteIOResult<()> {
 410         self.write_u16be(((val >> 16) & 0xFFFF) as u16)?;
 411         self.write_u16be((val & 0xFFFF) as u16)
 412     }
 413
 414     pub fn write_u32le(&mut self, val: u32) -> ByteIOResult<()> {
 415         self.write_u16le((val & 0xFFFF) as u16)?;
 416         self.write_u16le(((val >> 16) & 0xFFFF) as u16)
 417     }
 418
 419     pub fn write_u64be(&mut self, val: u64) -> ByteIOResult<()> {
 420         self.write_u32be(((val >> 32) & 0xFFFFFFFF) as u32)?;
 421         self.write_u32be((val & 0xFFFFFFFF) as u32)
 422     }
 423
 424     pub fn write_u64le(&mut self, val: u64) -> ByteIOResult<()> {
 425         self.write_u32le((val & 0xFFFFFFFF) as u32)?;
 426         self.write_u32le(((val >> 32) & 0xFFFFFFFF) as u32)
 427     }
 428
 429     pub fn tell(&mut self) -> u64 {
 430         self.io.tell()
 431     }
 432
 433     pub fn seek(&mut self, pos: SeekFrom) -> ByteIOResult<u64> {
 434         self.io.seek(pos)
 435     }
 436
 437     pub fn size_left(&mut self) -> i64 {
 438         let sz = self.io.size();
 439         if sz == -1 { return -1; }
 440         sz - (self.tell() as i64)
 441     }
 442 }
 443
 444 impl<'a> MemoryWriter<'a> {
 445
 446     pub fn new_write(buf: &'a mut [u8]) -> Self {
 447         let len = buf.len();
 448         MemoryWriter { buf: buf, size: len, pos: 0 }
 449     }
 450
 451     fn real_seek(&mut self, pos: i64) -> ByteIOResult<u64> {
 452         if pos < 0 || (pos as usize) > self.size {
 453             return Err(ByteIOError::WrongRange)
 454         }
 455         self.pos = pos as usize;
 456         Ok(pos as u64)
 457     }
 458 }
 459
 460 impl<'a> ByteIO for MemoryWriter<'a> {
 461     #[allow(unused_variables)]
 462     fn read_byte(&mut self) -> ByteIOResult<u8> {
 463         Err(ByteIOError::NotImplemented)
 464     }
 465
 466     #[allow(unused_variables)]
 467     fn peek_byte(&mut self) -> ByteIOResult<u8> {
 468         Err(ByteIOError::NotImplemented)
 469     }
 470
 471     #[allow(unused_variables)]
 472     fn read_buf(&mut self, buf: &mut [u8]) -> ByteIOResult<usize> {
 473         Err(ByteIOError::NotImplemented)
 474     }
 475
 476     #[allow(unused_variables)]
 477     fn read_buf_some(&mut self, buf: &mut [u8]) -> ByteIOResult<usize> {
 478         Err(ByteIOError::NotImplemented)
 479     }
 480
 481     #[allow(unused_variables)]
 482     fn peek_buf(&mut self, buf: &mut [u8]) -> ByteIOResult<usize> {
 483         Err(ByteIOError::NotImplemented)
 484     }
 485
 486     fn write_buf(&mut self, buf: &[u8]) -> ByteIOResult<()> {
 487         if self.pos + buf.len() > self.size { return Err(ByteIOError::WriteError); }
 488         for i in 0..buf.len() {
 489             self.buf[self.pos + i] = buf[i];
 490         }
 491         self.pos += buf.len();
 492         Ok(())
 493     }
 494
 495     fn tell(&mut self) -> u64 {
 496         self.pos as u64
 497     }
 498
 499     fn seek(&mut self, pos: SeekFrom) -> ByteIOResult<u64> {
 500         let cur_pos  = self.pos  as i64;
 501         let cur_size = self.size as i64;
 502         match pos {
 503             SeekFrom::Start(x)   => self.real_seek(x as i64),
 504             SeekFrom::Current(x) => self.real_seek(cur_pos + x),
 505             SeekFrom::End(x)     => self.real_seek(cur_size + x),
 506         }
 507     }
 508
 509     fn is_eof(&mut self) -> bool {
 510         self.pos >= self.size
 511     }
 512
 513     fn is_seekable(&mut self) -> bool {
 514         true
 515     }
 516
 517     fn size(&mut self) -> i64 {
 518         self.buf.len() as i64
 519     }
 520 }
 521
 522 impl FileWriter {
 523     pub fn new_write(file: File) -> Self {
 524         FileWriter { file: file }
 525     }
 526 }
 527
 528 impl ByteIO for FileWriter {
 529     #[allow(unused_variables)]
 530     fn read_byte(&mut self) -> ByteIOResult<u8> {
 531         Err(ByteIOError::NotImplemented)
 532     }
 533
 534     #[allow(unused_variables)]
 535     fn peek_byte(&mut self) -> ByteIOResult<u8> {
 536         Err(ByteIOError::NotImplemented)
 537     }
 538
 539     #[allow(unused_variables)]
 540     fn read_buf(&mut self, buf: &mut [u8]) -> ByteIOResult<usize> {
 541         Err(ByteIOError::NotImplemented)
 542     }
 543
 544     #[allow(unused_variables)]
 545     fn read_buf_some(&mut self, buf: &mut [u8]) -> ByteIOResult<usize> {
 546         Err(ByteIOError::NotImplemented)
 547     }
 548
 549     #[allow(unused_variables)]
 550     fn peek_buf(&mut self, buf: &mut [u8]) -> ByteIOResult<usize> {
 551         Err(ByteIOError::NotImplemented)
 552     }
 553
 554     fn write_buf(&mut self, buf: &[u8]) -> ByteIOResult<()> {
 555         match self.file.write_all(buf) {
 556             Ok(()) => Ok(()),
 557             Err(_) => Err(ByteIOError::WriteError),
 558         }
 559     }
 560
 561     fn tell(&mut self) -> u64 {
 562         self.file.seek(SeekFrom::Current(0)).unwrap()
 563     }
 564
 565     fn seek(&mut self, pos: SeekFrom) -> ByteIOResult<u64> {
 566         let res = self.file.seek(pos);
 567         match res {
 568             Ok(r) => Ok(r),
 569             Err(_) => Err(ByteIOError::SeekError),
 570         }
 571     }
 572
 573     fn is_eof(&mut self) -> bool {
 574         false
 575     }
 576
 577     fn is_seekable(&mut self) -> bool {
 578         true
 579     }
 580
 581     fn size(&mut self) -> i64 {
 582         -1
 583     }
 584 }
 585
 586 #[cfg(test)]
 587 mod test {
 588     use super::*;
 589     use std::fs::File;
 590
 591     #[test]
 592     fn test_read() {
 593         //const DATA : &'static [u8] = include_bytes!("../../assets/file");
 594         let buf: [u8; 64] = [1; 64];
 595         let mut mr = MemoryReader::new_read(&buf);
 596         let mut reader = ByteReader::new(&mut mr);
 597         assert_eq!(reader.read_byte().unwrap(),  0x01u8);
 598         assert_eq!(reader.read_u16le().unwrap(), 0x0101u16);
 599         assert_eq!(reader.read_u24le().unwrap(), 0x010101u32);
 600         assert_eq!(reader.read_u32le().unwrap(), 0x01010101u32);
 601         assert_eq!(reader.read_u64le().unwrap(), 0x0101010101010101u64);
 602         let mut file = File::open("assets/MaoMacha.asx").unwrap();
 603         let mut fr = FileReader::new_read(&mut file);
 604         let mut br2 = ByteReader::new(&mut fr);
 605         assert_eq!(br2.read_byte().unwrap(), 0x30);
 606         assert_eq!(br2.read_u24be().unwrap(), 0x26B275);
 607         assert_eq!(br2.read_u24le().unwrap(), 0xCF668E);
 608         assert_eq!(br2.read_u32be().unwrap(), 0x11A6D900);
 609         assert_eq!(br2.read_u32le().unwrap(), 0xCE6200AA);
 610     }
 611     #[test]
 612     fn test_write() {
 613         let mut buf: [u8; 64] = [0; 64];
 614         {
 615             let mut mw = MemoryWriter::new_write(&mut buf);
 616             let mut bw = ByteWriter::new(&mut mw);
 617             bw.write_byte(0x00).unwrap();
 618             bw.write_u16be(0x0102).unwrap();
 619             bw.write_u24be(0x030405).unwrap();
 620             bw.write_u32be(0x06070809).unwrap();
 621             bw.write_u64be(0x0A0B0C0D0E0F1011).unwrap();
 622             bw.write_byte(0x00).unwrap();
 623             bw.write_u16le(0x0201).unwrap();
 624             bw.write_u24le(0x050403).unwrap();
 625             bw.write_u32le(0x09080706).unwrap();
 626             bw.write_u64le(0x11100F0E0D0C0B0A).unwrap();
 627             assert_eq!(bw.size_left(), 28);
 628         }
 629         for i in 0..0x12 {
 630             assert_eq!(buf[(i + 0x00) as usize], i);
 631             assert_eq!(buf[(i + 0x12) as usize], i);
 632         }
 633     }
 634 }