From 167579afaab4d2ee0a637036a12227dbe5429e72 Mon Sep 17 00:00:00 2001
From: Ben Davis <imben123@gmail.com>
Date: Sun, 13 Mar 2016 11:34:53 +0000
Subject: [PATCH] Can advance byte index in byte stream

---
 BitTorrent/Utilities/ByteStream.swift         |  2 +
 BitTorrent/Utilities/NSDataByteStream.swift   | 14 ++++++
 .../Utilities/NSDataByteStreamTests.swift     | 49 +++++++++++++++++++
 3 files changed, 65 insertions(+)

diff --git a/BitTorrent/Utilities/ByteStream.swift b/BitTorrent/Utilities/ByteStream.swift
index fc1bd44..c8587fa 100644
--- a/BitTorrent/Utilities/ByteStream.swift
+++ b/BitTorrent/Utilities/ByteStream.swift
@@ -17,5 +17,7 @@ public protocol ByteStream {
     func nextBytes(numberOfBytes: Int) -> NSData?
 
     func indexIsValid(index: Int) -> Bool
+    
+    func advanceBy(numberOfBytes: Int)
 
 }
\ No newline at end of file
diff --git a/BitTorrent/Utilities/NSDataByteStream.swift b/BitTorrent/Utilities/NSDataByteStream.swift
index 5c88843..cdb4978 100644
--- a/BitTorrent/Utilities/NSDataByteStream.swift
+++ b/BitTorrent/Utilities/NSDataByteStream.swift
@@ -48,4 +48,18 @@ class NSDataByteStream: ByteStream {
         return index >= 0 && index <= self.length
     }
     
+    func advanceBy(numberOfBytes: Int) {
+        
+        let finalIndex = self.currentIndex + numberOfBytes
+        
+        if finalIndex > self.length {
+            self.advancePointer(self.length - self.currentIndex)
+        } else if finalIndex < 0 {
+            self.advancePointer(-self.currentIndex)
+        } else {
+            self.advancePointer(numberOfBytes)
+        }
+        
+    }
+    
 }
\ No newline at end of file
diff --git a/BitTorrent/Utilities/NSDataByteStreamTests.swift b/BitTorrent/Utilities/NSDataByteStreamTests.swift
index d925e12..d490b57 100644
--- a/BitTorrent/Utilities/NSDataByteStreamTests.swift
+++ b/BitTorrent/Utilities/NSDataByteStreamTests.swift
@@ -84,4 +84,53 @@ class NSDataByteStreamTests: XCTestCase {
         XCTAssertEqual(bytes, NSData(byteArray: [5, 6]))
         XCTAssertEqual(byte, UInt8(7))
     }
+    
+    func testCanAdvanceForward1Byte() {
+        let byteStream = NSDataByteStream(data: NSData(byteArray: [5, 6, 7, 8]))
+        byteStream.advanceBy(1)
+        let byte = byteStream.nextByte()
+        XCTAssertEqual(byte, 6)
+    }
+    
+    func testCanAdvanceForwardMultipleBytes() {
+        let byteStream = NSDataByteStream(data: NSData(byteArray: [5, 6, 7, 8]))
+        byteStream.advanceBy(3)
+        let byte = byteStream.nextByte()
+        XCTAssertEqual(byte, 8)
+    }
+    
+    func testCanAdvanceBackward1Byte() {
+        let byteStream = NSDataByteStream(data: NSData(byteArray: [5, 6, 7, 8]))
+        
+        var byte = byteStream.nextByte()
+        XCTAssertEqual(byte, 5)
+        
+        byteStream.advanceBy(-1)
+        byte = byteStream.nextByte()
+        XCTAssertEqual(byte, 5)
+    }
+    
+    func testCanAdvanceBackwardMultipleBytes() {
+        let byteStream = NSDataByteStream(data: NSData(byteArray: [5, 6, 7, 8]))
+        
+        var byte = byteStream.nextByte()
+        byte = byteStream.nextByte()
+        XCTAssertEqual(byte, 6)
+        
+        byteStream.advanceBy(-2)
+        byte = byteStream.nextByte()
+        XCTAssertEqual(byte, 5)
+    }
+    
+    func testAdvanceForwardTooFarStopsAtEnd() {
+        let byteStream = NSDataByteStream(data: NSData(byteArray: [5, 6, 7, 8]))
+        byteStream.advanceBy(999)
+        XCTAssertEqual(byteStream.currentIndex, 4)
+    }
+    
+    func testAdvanceBackTooFarStopsAtBeggining() {
+        let byteStream = NSDataByteStream(data: NSData(byteArray: [5, 6, 7, 8]))
+        byteStream.advanceBy(-999)
+        XCTAssertEqual(byteStream.currentIndex, 0)
+    }
 }