Merge pull request #19294 from nkcsgexi/cherry-pick-byte-tree-deserialization

[swift-5.0-branch] cherry pick Alex's byte tree deserialization patches

Author: nkcsgexi

tools/SwiftSyntax/ByteTreeDeserialization.swift

@@ -14,6 +14,14 @@ import Foundation
 
 // MARK: - ByteTree decoder protocols
 
+struct ByteTreeUserInfoKey: Hashable {
+  let rawValue: String
+
+  init(rawValue: String) {
+    self.rawValue = rawValue
+  }
+}
+
 /// A type that can be deserialized from ByteTree into a scalar value that
 /// doesn't have any child nodes
 protocol ByteTreeScalarDecodable {
@@ -23,7 +31,8 @@ protocol ByteTreeScalarDecodable {
   ///   - pointer: The pointer pointing to the start of the serialized data
   ///   - size: The length of the serialized data in bytes
   /// - Returns: The deserialized value
-  static func read(from pointer: UnsafeRawPointer, size: Int) -> Self
+  static func read(from pointer: UnsafeRawPointer, size: Int,
+                   userInfo: [ByteTreeUserInfoKey: Any]) -> Self
 }
 
 /// A type that can be deserialized from ByteTree into an object with child
@@ -39,7 +48,8 @@ protocol ByteTreeObjectDecodable {
   ///   - numFields: The number of fields that are present in the serialized
   ///                object
   /// - Returns: The deserialized object
-  static func read(from reader: ByteTreeObjectReader, numFields: Int) -> Self
+  static func read(from reader: ByteTreeObjectReader, numFields: Int,
+                   userInfo: [ByteTreeUserInfoKey: Any]) -> Self
 }
 
 // MARK: - Reader objects
@@ -119,14 +129,30 @@ class ByteTreeObjectReader {
 
 /// Reader for reading the ByteTree format into Swift objects
 class ByteTreeReader {
+  enum DeserializationError: Error, CustomStringConvertible {
+    case versionValidationFailed(ByteTreeReader.ProtocolVersion)
+
+    public var description: String {
+      switch self {
+      case .versionValidationFailed(let version):
+        return "The serialized ByteTree version \(version) cannot be parsed " +
+               "by this version of swiftSyntax"
+      }
+    }
+  }
+
   /// The type as which the protocol version is encoded in ByteTree
   typealias ProtocolVersion = UInt32
 
   /// A pointer pointing to the next byte of serialized data to be read
   private var pointer: UnsafeRawPointer
 
-  private init(pointer: UnsafeRawPointer) {
+  private var userInfo: [ByteTreeUserInfoKey: Any]
+
+  private init(pointer: UnsafeRawPointer,
+               userInfo: [ByteTreeUserInfoKey: Any]) {
     self.pointer = pointer
+    self.userInfo = userInfo
   }
 
   // MARK: Public entrance function
@@ -137,22 +163,45 @@ class ByteTreeReader {
   /// - Parameters:
   ///   - rootObjectType: The type of the root object in the deserialized tree
   ///   - pointer: The memory location at which the serialized data resides
-  ///   - protocolVerisonValidation: A callback to determine if the data can be
+  ///   - protocolVersionValidation: A callback to determine if the data can be
   ///       read, based on the format's protocol version. If the callback
-  ///       returns `false`, `nil` will be returned and reading aborded.
+  ///       returns `false` an error will be thrown
   /// - Returns: The deserialized tree or `nil` if protocol version validation
   ///            failed
   static func read<T: ByteTreeObjectDecodable>(
     _ rootObjectType: T.Type, from pointer: UnsafeRawPointer,
-    protocolVerisonValidation: (ProtocolVersion) -> Bool
-  ) -> T? {
-    let reader = ByteTreeReader(pointer: pointer)
-    if !reader.readAndValidateProtocolVersion(protocolVerisonValidation) {
-      return nil
-    }
+    userInfo: [ByteTreeUserInfoKey: Any],
+    protocolVersionValidation: (ProtocolVersion) -> Bool
+  ) throws -> T {
+    let reader = ByteTreeReader(pointer: pointer, userInfo: userInfo)
+    try reader.readAndValidateProtocolVersion(protocolVersionValidation)
     return reader.read(rootObjectType)
   }
 
+  /// Deserialize an object tree from the ByteTree data at the given memory
+  /// location.
+  ///
+  /// - Parameters:
+  ///   - rootObjectType: The type of the root object in the deserialized tree
+  ///   - data: The data to deserialize
+  ///   - protocolVersionValidation: A callback to determine if the data can be
+  ///       read, based on the format's protocol version. If the callback
+  ///       returns `false` an error will be thrown
+  /// - Returns: The deserialized tree
+  static func read<T: ByteTreeObjectDecodable>(
+    _ rootObjectType: T.Type, from data: Data,
+    userInfo: [ByteTreeUserInfoKey: Any],
+    protocolVersionValidation versionValidate: (ProtocolVersion) -> Bool
+  ) throws -> T {
+    return try data.withUnsafeBytes { (pointer: UnsafePointer<UInt8>) in
+      let rawPointer = UnsafeRawPointer(pointer)
+      return try ByteTreeReader.read(rootObjectType, from: rawPointer,
+                                     userInfo: userInfo,
+                                     protocolVersionValidation: versionValidate)
+    }
+  }
+
+
   // MARK: Internal read functions
 
   /// Cast the current pointer location to the given type and advance `pointer`
@@ -181,12 +230,13 @@ class ByteTreeReader {
   ///             protocol version can be read
   private func readAndValidateProtocolVersion(
     _ validationCallback: (ProtocolVersion) -> Bool
-  ) -> Bool {
+  ) throws {
     let protocolVersion = ProtocolVersion(littleEndian: 
       readRaw(ProtocolVersion.self))
     let result = validationCallback(protocolVersion)
-    pointer = pointer.advanced(by: MemoryLayout<ProtocolVersion>.size)
-    return result
+    if !result {
+      throw DeserializationError.versionValidationFailed(protocolVersion)
+    }
   }
 
   /// Read the next field in the tree as an object of the specified type.
@@ -199,7 +249,7 @@ class ByteTreeReader {
     let numFields = readFieldLength()
     let objectReader = ByteTreeObjectReader(reader: self,
                                             numFields: numFields)
-    return T.read(from: objectReader, numFields: numFields)
+    return T.read(from: objectReader, numFields: numFields, userInfo: userInfo)
   }
 
   /// Read the next field in the tree as a scalar of the specified type.
@@ -213,7 +263,7 @@ class ByteTreeReader {
     defer {
       pointer = pointer.advanced(by: fieldSize)
     }
-    return T.read(from: pointer, size: fieldSize)
+    return T.read(from: pointer, size: fieldSize, userInfo: userInfo)
   }
 
   /// Discard the next scalar field, advancing the pointer to the next field
@@ -230,7 +280,9 @@ class ByteTreeReader {
 // Implemenation for reading an integer from memory to be shared between
 // multiple types
 extension ByteTreeScalarDecodable where Self : FixedWidthInteger {
-  static func read(from pointer: UnsafeRawPointer, size: Int) -> Self {
+  static func read(from pointer: UnsafeRawPointer, size: Int,
+                   userInfo: [ByteTreeUserInfoKey: Any]
+  ) -> Self {
     assert(size == MemoryLayout<Self>.size)
     return pointer.bindMemory(to: Self.self, capacity: 1).pointee
   }
@@ -241,7 +293,8 @@ extension UInt16: ByteTreeScalarDecodable {}
 extension UInt32: ByteTreeScalarDecodable {}
 
 extension String: ByteTreeScalarDecodable {
-  static func read(from pointer: UnsafeRawPointer, size: Int) -> String {
+  static func read(from pointer: UnsafeRawPointer, size: Int,
+                   userInfo: [ByteTreeUserInfoKey: Any]) -> String {
     let data = Data(bytes: pointer, count: size)
     return String(data: data, encoding: .utf8)!
   }
@@ -250,21 +303,24 @@ extension String: ByteTreeScalarDecodable {
 extension Optional: ByteTreeObjectDecodable
   where
   Wrapped: ByteTreeObjectDecodable {
-  static func read(from reader: ByteTreeObjectReader, numFields: Int) ->
-    Optional<Wrapped> {
+  static func read(from reader: ByteTreeObjectReader, numFields: Int,
+                   userInfo: [ByteTreeUserInfoKey: Any]
+  ) -> Optional<Wrapped> {
     if numFields == 0 {
       return nil
     } else {
-      return Wrapped.read(from: reader, numFields: numFields)
+      return Wrapped.read(from: reader, numFields: numFields,
+                          userInfo: userInfo)
     }
   }
 }
 
 extension Array: ByteTreeObjectDecodable
   where
   Element: ByteTreeObjectDecodable {
-  static func read(from reader: ByteTreeObjectReader, numFields: Int) ->
-    Array<Element> {
+  static func read(from reader: ByteTreeObjectReader, numFields: Int,
+                   userInfo: [ByteTreeUserInfoKey: Any]
+  ) -> Array<Element> {
     return (0..<numFields).map {
       return reader.readField(Element.self, index: $0)
     }

tools/SwiftSyntax/RawSyntax.swift

@@ -24,6 +24,27 @@ extension CodingUserInfoKey {
     CodingUserInfoKey(rawValue: "SwiftSyntax.RawSyntax.OmittedNodeLookup")!
 }
 
+extension ByteTreeUserInfoKey {
+  /// Callback that will be called whenever a `RawSyntax` node is decoded
+  /// Value must have signature `(RawSyntax) -> Void`
+  static let rawSyntaxDecodedCallback =
+    ByteTreeUserInfoKey(rawValue: "SwiftSyntax.RawSyntax.DecodedCallback")
+  /// Function that shall be used to look up nodes that were omitted in the
+  /// syntax tree transfer.
+  /// Value must have signature `(SyntaxNodeId) -> RawSyntax`
+  static let omittedNodeLookupFunction =
+    ByteTreeUserInfoKey(rawValue: "SwiftSyntax.RawSyntax.OmittedNodeLookup")
+}
+
+/// Box a value type into a reference type
+class Box<T> {
+  let value: T
+
+  init(_ value: T) {
+    self.value = value
+  }
+}
+
 /// A ID that uniquely identifies a syntax node and stays stable across multiple
 /// incremental parses
 public struct SyntaxNodeId: Hashable, Codable {
@@ -54,7 +75,7 @@ public struct SyntaxNodeId: Hashable, Codable {
 }
 
 /// The data that is specific to a tree or token node
-fileprivate indirect enum RawSyntaxData {
+fileprivate enum RawSyntaxData {
   /// A tree node with a kind and an array of children
   case node(kind: SyntaxKind, layout: [RawSyntax?])
   /// A token with a token kind, leading trivia, and trailing trivia
@@ -72,11 +93,11 @@ struct RawSyntax: Codable {
   /// incremental parses
   let id: SyntaxNodeId
 
-  var _contentLength = AtomicCache<SourceLength>()
+  var _contentLength = AtomicCache<Box<SourceLength>>()
 
   /// The length of this node excluding its leading and trailing trivia
   var contentLength: SourceLength {
-    return _contentLength.value() {
+    return _contentLength.value({
       switch data {
       case .node(kind: _, layout: let layout):
         let firstElementIndex = layout.firstIndex(where: { $0 != nil })
@@ -97,11 +118,11 @@ struct RawSyntax: Codable {
             contentLength += element.trailingTriviaLength
           }
         }
-        return contentLength
+        return Box(contentLength)
       case .token(kind: let kind, leadingTrivia: _, trailingTrivia: _):
-        return SourceLength(of: kind.text)
+        return Box(SourceLength(of: kind.text))
       }
-    }
+    }).value
   }
 
   init(kind: SyntaxKind, layout: [RawSyntax?], presence: SourcePresence,
@@ -372,43 +393,64 @@ extension RawSyntax: ByteTreeObjectDecodable {
   enum SyntaxType: UInt8, ByteTreeScalarDecodable {
     case token = 0
     case layout = 1
+    case omitted = 2
 
-    static func read(from pointer: UnsafeRawPointer, size: Int) ->
-      SyntaxType {
-      let rawValue = UInt8.read(from: pointer, size: size)
+    static func read(from pointer: UnsafeRawPointer, size: Int,
+                     userInfo: [ByteTreeUserInfoKey: Any]
+    ) -> SyntaxType {
+      let rawValue = UInt8.read(from: pointer, size: size, userInfo: userInfo)
       guard let type = SyntaxType(rawValue: rawValue) else {
         fatalError("Unknown RawSyntax node type \(rawValue)")
       }
       return type
     }
   }
 
-  static func read(from reader: ByteTreeObjectReader, numFields: Int) ->
-    RawSyntax {
+  static func read(from reader: ByteTreeObjectReader, numFields: Int,
+                   userInfo: [ByteTreeUserInfoKey: Any]
+  ) -> RawSyntax {
+    let syntaxNode: RawSyntax
     let type = reader.readField(SyntaxType.self, index: 0)
+    let id = reader.readField(SyntaxNodeId.self, index: 1)
     switch type {
     case .token:
-      let presence = reader.readField(SourcePresence.self, index: 1)
-      let id = reader.readField(SyntaxNodeId.self, index: 2)
+      let presence = reader.readField(SourcePresence.self, index: 2)
       let kind = reader.readField(TokenKind.self, index: 3)
       let leadingTrivia = reader.readField(Trivia.self, index: 4)
       let trailingTrivia = reader.readField(Trivia.self, index: 5)
-      return RawSyntax(kind: kind, leadingTrivia: leadingTrivia,
-                       trailingTrivia: trailingTrivia,
-                       presence: presence, id: id)
+      syntaxNode = RawSyntax(kind: kind, leadingTrivia: leadingTrivia,
+                             trailingTrivia: trailingTrivia,
+                             presence: presence, id: id)
     case .layout:
-      let presence = reader.readField(SourcePresence.self, index: 1)
-      let id = reader.readField(SyntaxNodeId.self, index: 2)
+      let presence = reader.readField(SourcePresence.self, index: 2)
       let kind = reader.readField(SyntaxKind.self, index: 3)
       let layout = reader.readField([RawSyntax?].self, index: 4)
-      return RawSyntax(kind: kind, layout: layout, presence: presence, id: id)
+      syntaxNode = RawSyntax(kind: kind, layout: layout, presence: presence,
+                             id: id)
+    case .omitted:
+      guard let lookupFunc = userInfo[.omittedNodeLookupFunction] as?
+        (SyntaxNodeId) -> RawSyntax? else {
+          fatalError("omittedNodeLookupFunction is required when decoding an " +
+                     "incrementally transferred syntax tree")
+      }
+      guard let lookupNode = lookupFunc(id) else {
+        fatalError("Node lookup for id \(id) failed")
+      }
+      syntaxNode = lookupNode
+    }
+    if let callback = userInfo[.rawSyntaxDecodedCallback] as?
+      (RawSyntax) -> Void {
+      callback(syntaxNode)
     }
+    return syntaxNode
   }
 }
 
 extension SyntaxNodeId: ByteTreeScalarDecodable {
-  static func read(from pointer: UnsafeRawPointer, size: Int) -> SyntaxNodeId {
-    let rawValue = UInt8.read(from: pointer, size: size)
+  static func read(from pointer: UnsafeRawPointer, size: Int,
+                   userInfo: [ByteTreeUserInfoKey: Any]
+  ) -> SyntaxNodeId {
+    let rawValue = UInt32.read(from: pointer, size: size, userInfo: userInfo)
     return SyntaxNodeId(rawValue: UInt(rawValue))
   }
 }

tools/SwiftSyntax/SourceLength.swift

@@ -13,7 +13,7 @@
 /// The length a syntax node spans in the source code. From any AbsolutePosition
 /// you reach a node's end location by either adding its UTF-8 length or by
 /// inserting `lines` newlines and then moving `columns` columns to the right.
-public final class SourceLength {
+public struct SourceLength {
   public let newlines: Int
   public let columnsAtLastLine: Int
   public let utf8Length: Int
@@ -90,4 +90,4 @@ extension AbsolutePosition {
   public static func +=(lhs: inout AbsolutePosition, rhs: SourceLength) {
     lhs = lhs + rhs
   }
-}
+}

tools/SwiftSyntax/SourcePresence.swift

@@ -25,7 +25,9 @@ public enum SourcePresence: String, Codable {
 }
 
 extension SourcePresence: ByteTreeScalarDecodable {
-  static func read(from pointer: UnsafeRawPointer, size: Int) -> SourcePresence {
+  static func read(from pointer: UnsafeRawPointer, size: Int,
+                   userInfo: [ByteTreeUserInfoKey: Any]
+  ) -> SourcePresence {
     let rawValue = pointer.bindMemory(to: UInt8.self, capacity: 1).pointee
     switch rawValue {
     case 0: return .missing