Merge pull request #189 from CodaFi/reality-bytes

Switch BitstreamReader Entrypoints to Use ByteString

Author: CodaFi

Sources/TSCUtility/Bits.swift

@@ -9,15 +9,16 @@
 */
 
 import Foundation
+import TSCBasic
 
 struct Bits: RandomAccessCollection {
-  var buffer: Data
+  var buffer: ByteString
 
   var startIndex: Int { return 0 }
   var endIndex: Int { return buffer.count * 8 }
 
   subscript(index: Int) -> UInt8 {
-    let byte = buffer[index / 8]
+    let byte = buffer.contents[index / 8]
     return (byte >> UInt8(index % 8)) & 1
   }
 
@@ -27,7 +28,7 @@ struct Bits: RandomAccessCollection {
     precondition(offset &+ count >= offset)
     precondition(offset &+ count <= self.endIndex)
 
-    return buffer.withUnsafeBytes { (bytes: UnsafeRawBufferPointer) in
+    return buffer.contents.withUnsafeBytes { (bytes: UnsafeRawBufferPointer) in
       let upperBound = offset &+ count
       let topByteIndex = upperBound >> 3
       var result: UInt64 = 0
@@ -56,10 +57,14 @@ struct Bits: RandomAccessCollection {
       self.buffer = buffer
     }
 
-    init(buffer: Data) {
+    init(buffer: ByteString) {
       self.init(buffer: Bits(buffer: buffer))
     }
 
+    var isAtStart: Bool {
+      return offset == buffer.startIndex
+    }
+
     var isAtEnd: Bool {
       return offset == buffer.count
     }
@@ -74,14 +79,14 @@ struct Bits: RandomAccessCollection {
       return try peek(count)
     }
 
-    mutating func read(bytes count: Int) throws -> Data {
+    mutating func read(bytes count: Int) throws -> ArraySlice<UInt8> {
       precondition(count >= 0)
       precondition(offset & 0b111 == 0)
       let newOffset = offset &+ (count << 3)
       precondition(newOffset >= offset)
       if newOffset > buffer.count { throw Error.bufferOverflow }
       defer { offset = newOffset }
-      return buffer.buffer.dropFirst(offset >> 3).prefix((newOffset - offset) >> 3)
+      return buffer.buffer.contents.dropFirst(offset >> 3).prefix((newOffset - offset) >> 3)
     }
 
     mutating func skip(bytes count: Int) throws {

Sources/TSCUtility/BitstreamReader.swift

@@ -9,34 +9,42 @@
 */
 
 import Foundation
+import TSCBasic
 
 extension Bitcode {
   /// Parse a bitstream from data.
+  @available(*, deprecated, message: "Use Bitcode.init(bytes:) instead")
   public init(data: Data) throws {
     precondition(data.count > 4)
-    let signatureValue = UInt32(Bits(buffer: data).readBits(atOffset: 0, count: 32))
-    let bitstreamData = data[4..<data.count]
-
-    var reader = BitstreamReader(buffer: bitstreamData)
+    try self.init(bytes: ByteString(data))
+  }
+  
+  public init(bytes: ByteString) throws {
+    precondition(bytes.count > 4)
+    var reader = BitstreamReader(buffer: bytes)
+    let signature = try reader.readSignature()
     var visitor = CollectingVisitor()
     try reader.readBlock(id: BitstreamReader.fakeTopLevelBlockID,
                          abbrevWidth: 2,
                          abbrevInfo: [],
                          visitor: &visitor)
-    self.init(signature: .init(value: signatureValue),
+    self.init(signature: signature,
               elements: visitor.finalizeTopLevelElements(),
               blockInfo: reader.blockInfo)
   }
 
   /// Traverse a bitstream using the specified `visitor`, which will receive
   /// callbacks when blocks and records are encountered.
+  @available(*, deprecated, message: "Use Bitcode.read(bytes:using:) instead")
   public static func read<Visitor: BitstreamVisitor>(stream data: Data, using visitor: inout Visitor) throws {
     precondition(data.count > 4)
-    let signatureValue = UInt32(Bits(buffer: data).readBits(atOffset: 0, count: 32))
-    try visitor.validate(signature: .init(value: signatureValue))
+    try Self.read(bytes: ByteString(data), using: &visitor)
+  }
 
-    let bitstreamData = data[4..<data.count]
-    var reader = BitstreamReader(buffer: bitstreamData)
+  public static func read<Visitor: BitstreamVisitor>(bytes: ByteString, using visitor: inout Visitor) throws {
+    precondition(bytes.count > 4)
+    var reader = BitstreamReader(buffer: bytes)
+    try visitor.validate(signature: reader.readSignature())
     try reader.readBlock(id: BitstreamReader.fakeTopLevelBlockID,
                          abbrevWidth: 2,
                          abbrevInfo: [],
@@ -113,8 +121,14 @@ private struct BitstreamReader {
   var blockInfo: [UInt64: BlockInfo] = [:]
   var globalAbbrevs: [UInt64: [Bitstream.Abbreviation]] = [:]
 
-  init(buffer: Data) {
-    cursor = Bits.Cursor(buffer: buffer)
+  init(buffer: ByteString) {
+    self.cursor = Bits.Cursor(buffer: buffer)
+  }
+
+  mutating func readSignature() throws -> Bitcode.Signature {
+    precondition(self.cursor.isAtStart)
+    let bits = try UInt32(self.cursor.read(MemoryLayout<UInt32>.size * 8))
+    return Bitcode.Signature(value: bits)
   }
 
   mutating func readAbbrevOp() throws -> Bitstream.Abbreviation.Operand {
@@ -220,7 +234,7 @@ private struct BitstreamReader {
       case .blob:
         let length = Int(try cursor.readVBR(6))
         try cursor.advance(toBitAlignment: 32)
-        payload = .blob(try cursor.read(bytes: length))
+        payload = .blob(try Data(cursor.read(bytes: length)))
         try cursor.advance(toBitAlignment: 32)
       default:
         fatalError()

Sources/TSCUtility/SerializedDiagnostics.swift

@@ -8,6 +8,7 @@
  See http://swift.org/CONTRIBUTORS.txt for Swift project authors
 */
 import Foundation
+import TSCBasic
 
 /// Represents diagnostics serialized in a .dia file by the Swift compiler or Clang.
 public struct SerializedDiagnostics {
@@ -42,13 +43,22 @@ public struct SerializedDiagnostics {
   /// Serialized diagnostics.
   public var diagnostics: [Diagnostic]
 
+  @available(*, deprecated, message: "Use SerializedDiagnostics.init(bytes:) instead")
   public init(data: Data) throws {
     var reader = Reader()
     try Bitcode.read(stream: data, using: &reader)
     guard let version = reader.versionNumber else { throw Error.noMetadataBlock }
     self.versionNumber = version
     self.diagnostics = reader.diagnostics
   }
+
+  public init(bytes: ByteString) throws {
+    var reader = Reader()
+    try Bitcode.read(bytes: bytes, using: &reader)
+    guard let version = reader.versionNumber else { throw Error.noMetadataBlock }
+    self.versionNumber = version
+    self.diagnostics = reader.diagnostics
+  }
 }
 
 extension SerializedDiagnostics {

Tests/TSCUtilityTests/BitstreamTests.swift

@@ -37,7 +37,7 @@ final class BitstreamTests: XCTestCase {
             .appending(components: "Inputs", "serialized.dia")
         let contents = try localFileSystem.readFileContents(bitstreamPath)
         var visitor = LoggingVisitor()
-        try Bitcode.read(stream: Data(contents.contents), using: &visitor)
+        try Bitcode.read(bytes: contents, using: &visitor)
         XCTAssertEqual(visitor.log, [
             "entering block: 8",
             "Record (id: 1, fields: [1], payload: none",
@@ -130,7 +130,7 @@ final class BitstreamTests: XCTestCase {
             .appending(components: "Inputs", "serialized.dia")
         let contents = try localFileSystem.readFileContents(bitstreamPath)
         var visitor = LoggingVisitor()
-        try Bitcode.read(stream: Data(contents.contents), using: &visitor)
+        try Bitcode.read(bytes: contents, using: &visitor)
         XCTAssertEqual(visitor.log, ["skipping block: 8",
                                      "skipping block: 9",
                                      "skipping block: 9",
@@ -155,7 +155,7 @@ final class BitstreamTests: XCTestCase {
         let bitstreamPath = AbsolutePath(#file).parentDirectory
             .appending(components: "Inputs", "serialized.dia")
         let contents = try localFileSystem.readFileContents(bitstreamPath)
-        let bitcode = try Bitcode(data: Data(contents.contents))
+        let bitcode = try Bitcode(bytes: contents)
         XCTAssertEqual(bitcode.signature, .init(string: "DIAG"))
         XCTAssertEqual(bitcode.elements.count, 18)
         guard case .block(let metadataBlock) = bitcode.elements.first else {
@@ -381,7 +381,7 @@ final class BitstreamTests: XCTestCase {
         }
 
         var visitor = RoundTripVisitor()
-        try Bitcode.read(stream: Data(writer.data), using: &visitor)
+        try Bitcode.read(bytes: ByteString(writer.data), using: &visitor)
     }
 
     func testSimpleRecordWrite() {

Tests/TSCUtilityTests/SerializedDiagnosticsTests.swift

@@ -18,7 +18,7 @@ final class SerializedDiagnosticsTests: XCTestCase {
     let serializedDiagnosticsPath = AbsolutePath(#file).parentDirectory
         .appending(components: "Inputs", "serialized.dia")
     let contents = try localFileSystem.readFileContents(serializedDiagnosticsPath)
-    let serializedDiags = try SerializedDiagnostics(data: Data(contents.contents))
+    let serializedDiags = try SerializedDiagnostics(bytes: contents)
 
     XCTAssertEqual(serializedDiags.versionNumber, 1)
     XCTAssertEqual(serializedDiags.diagnostics.count, 17)
@@ -66,7 +66,7 @@ final class SerializedDiagnosticsTests: XCTestCase {
     let serializedDiagnosticsPath = AbsolutePath(#file).parentDirectory
         .appending(components: "Inputs", "clang.dia")
     let contents = try localFileSystem.readFileContents(serializedDiagnosticsPath)
-    let serializedDiags = try SerializedDiagnostics(data: Data(contents.contents))
+    let serializedDiags = try SerializedDiagnostics(bytes: contents)
 
     XCTAssertEqual(serializedDiags.versionNumber, 1)
     XCTAssertEqual(serializedDiags.diagnostics.count, 4)