Add utilities for reading the LLVM bitstream container format

Originally authored by Jordan Rose
Modified to drop some printing and logging code that's not needed for the driver's use cases

Co-authored-by: Jordan Rose <[email protected]>

Author: owenv

Sources/SwiftDriver/CMakeLists.txt

@@ -59,6 +59,8 @@ add_library(SwiftDriver
   Toolchains/GenericUnixToolchain.swift
   Toolchains/Toolchain.swift
 
+  Utilities/Bits.swift
+  Utilities/Bitstream.swift
   Utilities/DOTJobGraphSerializer.swift
   Utilities/DateAdditions.swift
   Utilities/Diagnostics.swift

Sources/SwiftDriver/Incremental Compilation/SourceFileDependencyGraph.swift

@@ -111,3 +111,109 @@ extension SourceFileDependencyGraph {
     }
   }
 }
+
+extension SourceFileDependencyGraph {
+  private static let recordBlockId = 8
+
+  private enum RecordKind: UInt64 {
+    case metadata = 1
+    case sourceFileDepGraphNode
+    case fingerprintNode
+    case dependsOnDefinitionNode
+    case identifierNode
+  }
+
+  private enum ReadError: Error {
+    case badMagic
+    case noRecordBlock
+    case malformedMetadataRecord
+    case unexpectedMetadataRecord
+    case malformedFingerprintRecord
+    case malformedDependsOnDefinitionRecord
+    case malformedIdentifierRecord
+    case malformedSourceFileDepGraphNodeRecord
+    case unknownRecord
+    case unexpectedSubblock
+  }
+
+  public init(data: Data) throws {
+    let bitcode = try Bitcode(data: data)
+    guard bitcode.signature == 0x53_50_45_44 else { throw ReadError.badMagic } // DEPS
+
+    guard bitcode.elements.count == 1,
+          case .block(let recordBlock) = bitcode.elements.first,
+          recordBlock.id == Self.recordBlockId else { throw ReadError.noRecordBlock }
+
+    guard case .record(let metadataRecord) = recordBlock.elements.first,
+          RecordKind(rawValue: metadataRecord.id) == .metadata,
+          metadataRecord.fields.count == 2,
+          case .blob(let compilerVersionBlob) = metadataRecord.payload,
+          let compilerVersionString = String(data: compilerVersionBlob, encoding: .utf8)
+    else { throw ReadError.malformedMetadataRecord }
+
+    self.majorVersion = metadataRecord.fields[0]
+    self.minorVersion = metadataRecord.fields[1]
+    self.compilerVersionString = compilerVersionString
+
+    var nodes: [Node] = []
+    var node: Node? = nil
+    var identifiers: [String] = [""] // The empty string is hardcoded as identifiers[0]
+    var sequenceNumber = 0
+    for element in recordBlock.elements.dropFirst() {
+      guard case .record(let record) = element else { throw ReadError.unexpectedSubblock }
+      guard let kind = RecordKind(rawValue: record.id) else { throw ReadError.unknownRecord }
+      switch kind {
+      case .metadata:
+        throw ReadError.unexpectedMetadataRecord
+      case .sourceFileDepGraphNode:
+        if let node = node {
+          nodes.append(node)
+        }
+        guard record.fields.count == 5,
+              let nodeKind = DependencyKey.Kind(rawValue: record.fields[0]),
+              let declAspect = DeclAspect(rawValue: record.fields[1]),
+              record.fields[2] < identifiers.count,
+              record.fields[3] < identifiers.count else {
+          throw ReadError.malformedSourceFileDepGraphNodeRecord
+        }
+        let context = identifiers[Int(record.fields[2])]
+        let identifier = identifiers[Int(record.fields[3])]
+        let isProvides = record.fields[4] != 0
+        node = Node(key: .init(kind: nodeKind,
+                               aspect: declAspect,
+                               context: context,
+                               name: identifier),
+                    fingerprint: nil,
+                    sequenceNumber: sequenceNumber,
+                    defsIDependUpon: [],
+                    isProvides: isProvides)
+        sequenceNumber += 1
+      case .fingerprintNode:
+        guard node != nil,
+              record.fields.count == 0,
+              case .blob(let fingerprintBlob) = record.payload,
+              let fingerprint = String(data: fingerprintBlob, encoding: .utf8) else {
+          throw ReadError.malformedFingerprintRecord
+        }
+        node?.fingerprint = fingerprint
+      case .dependsOnDefinitionNode:
+        guard node != nil,
+              record.fields.count == 1 else { throw ReadError.malformedDependsOnDefinitionRecord }
+        node?.defsIDependUpon.append(Int(record.fields[0]))
+      case .identifierNode:
+        guard record.fields.count == 0,
+              case .blob(let identifierBlob) = record.payload,
+              let identifier = String(data: identifierBlob, encoding: .utf8) else {
+          throw ReadError.malformedIdentifierRecord
+        }
+        identifiers.append(identifier)
+      }
+    }
+
+    if let node = node {
+      nodes.append(node)
+    }
+
+    self.allNodes = nodes
+  }
+}

Sources/SwiftDriver/Utilities/Bits.swift

@@ -0,0 +1,107 @@
+//===----------------- Bits.swift - Collection of bits --------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2020 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import Foundation
+
+// FIXME: Should be RandomAccessCollection
+struct Bits: BidirectionalCollection {
+  typealias Index = Int
+
+  var buffer: Data
+
+  var startIndex: Int { return 0 }
+  var endIndex: Int { return buffer.count * 8 }
+
+  subscript(index: Int) -> UInt8 {
+    let byte = buffer[index / 8]
+    return (byte >> UInt8(index % 8)) & 1
+  }
+
+  func index(after i: Int) -> Int {
+    return i + 1
+  }
+  func index(before i: Int) -> Int {
+    return i - 1
+  }
+
+  subscript(extractAt offset: Int, count count: Int) -> UInt64 {
+    precondition(count >= 0 && count <= 64)
+    precondition(offset >= 0)
+    precondition(offset &+ count >= offset)
+    precondition(offset &+ count <= self.endIndex)
+
+    return buffer.withUnsafeBytes { (bytes: UnsafeRawBufferPointer) in
+      let upperBound = offset &+ count
+      let topByteIndex = upperBound >> 3
+      var result: UInt64 = 0
+      if upperBound & 7 != 0 {
+        let mask: UInt8 = (1 << UInt8(upperBound & 7)) &- 1
+        result = UInt64(bytes[topByteIndex] & mask)
+      }
+      for i in ((offset >> 3)..<(upperBound >> 3)).reversed() {
+        result <<= 8
+        result |= UInt64(bytes[i])
+      }
+      if offset & 7 != 0 {
+        result >>= UInt64(offset & 7)
+      }
+      return result
+    }
+  }
+
+  struct Cursor {
+    enum Error: Swift.Error { case bufferOverflow }
+
+    let buffer: Bits
+    private var offset: Int = 0
+
+    init(buffer: Bits) {
+      self.buffer = buffer
+    }
+    
+    init(buffer: Data) {
+      self.init(buffer: Bits(buffer: buffer))
+    }
+
+    var isAtEnd: Bool {
+      return offset == buffer.count
+    }
+
+    func peek(_ count: Int) throws -> UInt64 {
+      if buffer.count - offset < count { throw Error.bufferOverflow }
+      return buffer[extractAt: offset, count: count]
+    }
+
+    mutating func read(_ count: Int) throws -> UInt64 {
+      defer { offset += count }
+      return try peek(count)
+    }
+
+    mutating func read(bytes count: Int) throws -> Data {
+      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.subdata(in: (offset >> 3)..<(newOffset >> 3))
+    }
+
+    mutating func advance(toBitAlignment align: Int) throws {
+      precondition(align > 0)
+      precondition(offset &+ (align&-1) >= offset)
+      if offset % align == 0 { return }
+      offset = ((offset / align) &+ 1) &* align
+      if offset > buffer.count { throw Error.bufferOverflow }
+    }
+  }
+}