WIP: Keyword completions based on the syntax tree

rdar://98551200

Author: ahoppen

Package.swift

@@ -202,7 +202,7 @@ let package = Package(
     ),
     .testTarget(
       name: "SwiftParserDiagnosticsTest",
-      dependencies: ["SwiftDiagnostics", "SwiftParserDiagnostics"]
+      dependencies: ["SwiftDiagnostics", "SwiftParserDiagnostics", "_SwiftSyntaxTestSupport"]
     ),
     .testTarget(
       name: "SwiftOperatorsTest",

Sources/SwiftParserDiagnostics/KeywordCompletion.swift

@@ -0,0 +1,127 @@
+import SwiftSyntax
+
+extension SyntaxProtocol {
+  static func completions(at keyPath: AnyKeyPath, visitedNodeKinds: inout [SyntaxProtocol.Type]) -> (completions: Set<TokenKind>, hasRequiredToken: Bool) {
+    visitedNodeKinds.append(Self.self)
+    if let keyPath = keyPath as? KeyPath<Self, TokenSyntax> {
+      return (Set(keyPath.tokenChoices), true)
+    } else if let keyPath = keyPath as? KeyPath<Self, TokenSyntax?> {
+      return (Set(keyPath.tokenChoices), false)
+    } else if let value = type(of: keyPath).valueType as? SyntaxOrOptionalProtocol.Type {
+      switch value.syntaxOrOptionalProtocolType {
+      case .optional(let syntaxType):
+        return (syntaxType.completionsAtStartOfNode(visitedNodeKinds: &visitedNodeKinds).completions, false)
+      case .nonOptional(let syntaxType):
+        return (syntaxType.completionsAtStartOfNode(visitedNodeKinds: &visitedNodeKinds).completions, !syntaxType.layout.isCollection)
+      }
+    } else {
+      assertionFailure("Unexpected keypath")
+      return ([], false)
+    }
+  }
+
+
+  static func completionsAtStartOfNode(visitedNodeKinds: inout [SyntaxProtocol.Type]) -> (completions: Set<TokenKind>, hasRequiredToken: Bool) {
+    if visitedNodeKinds.contains(where: { $0 == Self.self }) {
+      return ([], true)
+    }
+    if self == Syntax.self {
+      // An arbitrary syntax node is probably a
+      return ([], true)
+    }
+
+    var hasRequiredToken: Bool
+    var completions: Set<TokenKind> = []
+
+    switch self.layout {
+    case .layout(let keyPaths):
+      hasRequiredToken = false // Only relevant if keyPaths is empty and the loop below isn't traversed
+      for keyPath in keyPaths {
+        let res = self.completions(at: keyPath, visitedNodeKinds: &visitedNodeKinds)
+        completions.formUnion(res.completions)
+        hasRequiredToken = hasRequiredToken || res.hasRequiredToken
+        if hasRequiredToken {
+          break
+        }
+      }
+    case .collection(let collectionElementType):
+      return collectionElementType.completionsAtStartOfNode(visitedNodeKinds: &visitedNodeKinds)
+    case .choices(let choices):
+      hasRequiredToken = true
+      for choice in choices {
+        switch choice {
+        case .node(let nodeType):
+          let res = nodeType.completionsAtStartOfNode(visitedNodeKinds: &visitedNodeKinds)
+          completions.formUnion(res.completions)
+          hasRequiredToken = hasRequiredToken && res.hasRequiredToken
+        case .token(let tokenKind):
+          completions.insert(tokenKind)
+        }
+      }
+    }
+
+    return (completions, hasRequiredToken)
+  }
+}
+
+extension SyntaxProtocol {
+  func completions(afterIndex index: Int) -> Set<TokenKind> {
+    var hasRequiredToken: Bool = false
+
+    var completions: Set<TokenKind> = []
+    var visitedNodeKinds: [SyntaxProtocol.Type] = []
+    if index < childrenKeyPaths.count {
+      for keyPath in childrenKeyPaths[(index + 1)...] {
+        let res = Self.completions(at: keyPath, visitedNodeKinds: &visitedNodeKinds)
+        completions.formUnion(res.completions)
+        hasRequiredToken = res.hasRequiredToken
+        if hasRequiredToken {
+          break
+        }
+      }
+    }
+    if !hasRequiredToken, let parent = parent {
+      completions.formUnion(parent.asProtocol(SyntaxProtocol.self).completions(afterIndex: self.indexInParent))
+    }
+    return completions
+  }
+
+  public func completions(at position: AbsolutePosition) -> Set<TokenKind> {
+    if position <= self.positionAfterSkippingLeadingTrivia {
+      var visitedNodeKinds: [SyntaxProtocol.Type] = []
+      return Self.completionsAtStartOfNode(visitedNodeKinds: &visitedNodeKinds).completions
+    }
+    let finder = TokenFinder(targetPosition: position)
+    finder.walk(self)
+    guard let found = finder.found?.previousToken(viewMode: .sourceAccurate), let parent = found.parent else {
+      return []
+    }
+    return parent.asProtocol(SyntaxProtocol.self).completions(afterIndex: found.indexInParent)
+  }
+}
+
+/// Finds the first token whose text (ignoring trivia) starts after targetPosition.
+class TokenFinder: SyntaxAnyVisitor {
+  var targetPosition: AbsolutePosition
+  var found: TokenSyntax? = nil
+
+  init(targetPosition: AbsolutePosition) {
+    self.targetPosition = targetPosition
+    super.init(viewMode: .sourceAccurate)
+  }
+
+  override func visitAny(_ node: Syntax) -> SyntaxVisitorContinueKind {
+    if found != nil || node.endPosition < targetPosition {
+      return .skipChildren
+    } else {
+      return .visitChildren
+    }
+  }
+
+  override func visit(_ node: TokenSyntax) -> SyntaxVisitorContinueKind {
+    if targetPosition <= node.positionAfterSkippingLeadingTrivia, found == nil {
+      found = node
+    }
+    return .skipChildren
+  }
+}

Sources/SwiftSyntax/Misc.swift.gyb

@@ -56,6 +56,20 @@ extension Syntax {
     }
   }
 
+  public static var layout: SyntaxNodeLayoutDescription {
+    return .choices([
+      .node(UnknownSyntax.self),
+      .node(TokenSyntax.self),
+% for node in NON_BASE_SYNTAX_NODES:
+%   if node.is_base():
+      .node(Unknown${node.name}.self),
+%   else:
+      .node(${node.name}.self),
+%   end
+% end
+    ])
+  }
+
   public func childNameForDiagnostics(_ index: SyntaxChildrenIndex) -> String? {
     switch self.as(SyntaxEnum.self) {
     case .token(let node):

Sources/SwiftSyntax/Syntax.swift

@@ -10,6 +10,38 @@
 //
 //===----------------------------------------------------------------------===//
 
+public enum SyntaxChoice {
+  case node(SyntaxProtocol.Type)
+  case token(TokenKind)
+}
+
+public enum SyntaxNodeLayoutDescription {
+  case layout([AnyKeyPath])
+  case collection(SyntaxProtocol.Type)
+  case choices([SyntaxChoice])
+
+  public var isLayout: Bool {
+    switch self {
+    case .layout: return true
+    default: return false
+    }
+  }
+
+  public var isCollection: Bool {
+    switch self {
+    case .collection: return true
+    default: return false
+    }
+  }
+
+  public var isChoices: Bool {
+    switch self {
+    case .choices: return true
+    default: return false
+    }
+  }
+}
+
 /// A Syntax node represents a tree of nodes with tokens at the leaves.
 /// Each node has accessors for its known children, and allows efficient
 /// iteration over the children through its `children` property.
@@ -49,10 +81,6 @@ public struct Syntax: SyntaxProtocol, SyntaxHashable {
     self = syntax._syntaxNode
   }
 
-  public static var childrenKeyPaths: [AnyKeyPath] {
-    fatalError("Syntax is a base type and thus doesn't have statically known children")
-  }
-
   public func hash(into hasher: inout Hasher) {
     return data.nodeId.hash(into: &hasher)
   }
@@ -126,10 +154,25 @@ public extension SyntaxHashable {
   }
 }
 
+public enum SyntaxOrOptionalProtocolType {
+  case optional(SyntaxProtocol.Type)
+  case nonOptional(SyntaxProtocol.Type)
+}
+
+public protocol SyntaxOrOptionalProtocol {
+  static var syntaxOrOptionalProtocolType: SyntaxOrOptionalProtocolType { get }
+}
+
+extension Optional: SyntaxOrOptionalProtocol where Wrapped: SyntaxProtocol {
+  public static var syntaxOrOptionalProtocolType: SyntaxOrOptionalProtocolType {
+    return .optional(Wrapped.self)
+  }
+}
+
 /// Provide common functionality for specialized syntax nodes. Extend this
 /// protocol to provide common functionality for all syntax nodes.
 /// DO NOT CONFORM TO THIS PROTOCOL YOURSELF!
-public protocol SyntaxProtocol: CustomStringConvertible,
+public protocol SyntaxProtocol: SyntaxOrOptionalProtocol, CustomStringConvertible,
     CustomDebugStringConvertible, TextOutputStreamable {
 
   /// Retrieve the generic syntax node that is represented by this node.
@@ -148,7 +191,7 @@ public protocol SyntaxProtocol: CustomStringConvertible,
   ///
   /// Technically, this should have type `PartialKeyPath<Self>` but we cannot do
   /// that because then `SyntaxProtocol` would have a Self type requirement.
-  static var childrenKeyPaths: [AnyKeyPath] { get }
+  static var layout: SyntaxNodeLayoutDescription { get }
 
   /// Return a name with which the child at the given `index` can be referred to
   /// in diagnostics.
@@ -157,9 +200,20 @@ public protocol SyntaxProtocol: CustomStringConvertible,
   func childNameForDiagnostics(_ index: SyntaxChildrenIndex) -> String?
 }
 
+public extension SyntaxProtocol {
+  static var syntaxOrOptionalProtocolType: SyntaxOrOptionalProtocolType {
+    return .nonOptional(Self.self)
+  }
+}
+
 public extension SyntaxProtocol {
   /// Returns the value of the static `childrenKeyPaths` for this node's type.
-  var childrenKeyPaths: [AnyKeyPath] { return Self.childrenKeyPaths }
+  var childrenKeyPaths: [AnyKeyPath] {
+    switch Self.layout {
+    case .layout(let keyPaths): return keyPaths
+    default: return []
+    }
+  }
 }
 
 public extension SyntaxProtocol {

Sources/SwiftSyntax/SyntaxBaseNodes.swift.gyb

@@ -142,11 +142,14 @@ public struct ${node.name}: ${node.name}Protocol, SyntaxHashable {
     fatalError("${node.name} is a base type and thus doesn't have statically known children")
   }
 
-  public var childrenKeyPaths: [AnyKeyPath] {
-    return Syntax(self).childrenKeyPaths
+  public static var layout: SyntaxNodeLayoutDescription {
+    return .choices([
+%     for child_node in [child_node for child_node in SYNTAX_NODES if child_node.base_kind == node.syntax_kind]:
+      .node(${child_node.name}.self),
+%     end
+    ])
   }
 
-
   public func childNameForDiagnostics(_ index: SyntaxChildrenIndex) -> String? {
     return Syntax(self).childNameForDiagnostics(index)
   }

Sources/SwiftSyntax/SyntaxCollections.swift.gyb

@@ -18,13 +18,17 @@
 //
 //===----------------------------------------------------------------------===//
 
-public protocol SyntaxCollection: SyntaxProtocol, Sequence {
+public protocol SyntaxCollection: SyntaxProtocol, Sequence where Element: SyntaxProtocol {
   /// The number of elements, `present` or `missing`, in this collection.
   var count: Int { get }
 }
 
 public extension SyntaxCollection {
   static var childrenKeyPaths: [AnyKeyPath] { [] }
+
+  static var layout: SyntaxNodeLayoutDescription {
+    return .collection(Element.self)
+  }
 }
 
 % for node in SYNTAX_NODES:
@@ -79,6 +83,16 @@ public struct ${node.name}: SyntaxCollection, SyntaxHashable {
 %   end
       return nil
     }
+
+
+    public static var layout: SyntaxNodeLayoutDescription {
+      return .choices([
+%         for choice_name in node.collection_element_choices:
+%           choice = NODE_MAP[choice_name]
+        .node(${choice.name}.self),
+%         end
+      ])
+    }
   }
 % else:
   public typealias Element = ${node.collection_element_type}

Sources/SwiftSyntax/SyntaxNodes.swift.gyb.template

@@ -107,6 +107,20 @@ public struct ${node.name}: ${base_type}Protocol, SyntaxHashable {
 %         end
       return nil
     }
+
+    public static var layout: SyntaxNodeLayoutDescription {
+      return .choices([
+%         for choice in child.node_choices:
+%           if choice.is_token() and choice.token.text:
+        .token(.${choice.token.swift_kind()}),
+%           elif choice.is_token() and not choice.token.text:
+        .token(.${choice.token.swift_kind()}("")),
+%           else:
+        .node(${choice.type_name}.self),
+%           end
+%         end
+      ])
+    }
   }
 
 %       end
@@ -236,12 +250,12 @@ public struct ${node.name}: ${base_type}Protocol, SyntaxHashable {
   }
 %     end
 
-  public static var childrenKeyPaths: [AnyKeyPath] {
-    [
+  public static var layout: SyntaxNodeLayoutDescription {
+    return .layout([
 %     for (index, child) in enumerate(node.children):
       \Self.${child.swift_name},
 %     end
-    ]
+    ])
   }
 
   public func childNameForDiagnostics(_ index: SyntaxChildrenIndex) -> String? {

Sources/SwiftSyntax/SyntaxOtherNodes.swift

@@ -38,7 +38,9 @@ public struct UnknownSyntax: SyntaxProtocol, SyntaxHashable {
     self.init(data)
   }
 
-  public static var childrenKeyPaths: [AnyKeyPath] { [] }
+  public static var layout: SyntaxNodeLayoutDescription {
+    return .layout([])
+  }
 
   public func childNameForDiagnostics(_ index: SyntaxChildrenIndex) -> String? {
     return nil
@@ -202,7 +204,9 @@ public struct TokenSyntax: SyntaxProtocol, SyntaxHashable {
     return raw.totalLength
   }
 
-  public static var childrenKeyPaths: [AnyKeyPath] { [] }
+  public static var layout: SyntaxNodeLayoutDescription {
+    return .layout([])
+  }
 
   public func childNameForDiagnostics(_ index: SyntaxChildrenIndex) -> String? {
     return nil