Merge pull request #4 from koliyo:feature/improved-document-symbols

Improve document symbol handler

Sources/hylo-lsp-client/hylo-lsp-client.swift

@@ -94,6 +94,7 @@ struct HyloLspCommand: AsyncParsableCommand {
           SemanticToken.self,
           Diagnostics.self,
           Definition.self,
+          Symbols.self,
         ],
         defaultSubcommand: nil)
 }
@@ -126,6 +127,46 @@ func initServer(workspace: String? = nil, documents: [String]) async throws -> R
 }
 
 extension HyloLspCommand {
+  struct Symbols : AsyncParsableCommand {
+    @OptionGroup var options: Options
+    func run() async throws {
+      logger.logLevel = options.log
+      let (doc, line, char) = try options.parseDocument()
+      let docURL = URL.init(fileURLWithPath: doc)
+
+      let server = try await initServer(documents: [doc])
+      let params = DocumentSymbolParams(textDocument: TextDocumentIdentifier(uri: docURL.absoluteString))
+
+      let symbols = try await server.documentSymbol(params: params)
+
+      switch symbols {
+        case nil:
+          print("No symbols")
+        case let .optionA(s):
+          if s.isEmpty {
+            print("No symbols")
+          }
+          for s in s {
+            printSymbol(symbolInformation(s, in: doc))
+          }
+        case let .optionB(s):
+          if s.isEmpty {
+            print("No symbols")
+          }
+          for s in s {
+            printSymbol(s)
+          }
+      }
+    }
+
+    func symbolInformation(_ s: DocumentSymbol, in uri: DocumentUri) -> SymbolInformation {
+      SymbolInformation(name: s.name, kind: s.kind, location: Location(uri: uri, range: s.range))
+    }
+
+    func printSymbol(_ s: SymbolInformation) {
+      print("\(cliLink(uri: s.location.uri, range: s.location.range)) name: \(s.name), kind: \(s.kind)")
+    }
+  }
 
   struct Definition : AsyncParsableCommand {
     @OptionGroup var options: Options

Sources/hylo-lsp/AST+FindTranslationUnit.swift

@@ -8,7 +8,7 @@ extension AST {
   private struct TranslationUnitFinder: ASTWalkObserver {
     // var outermostFunctions: [FunctionDecl.ID] = []
     let query: DocumentUri
-    private(set) var match: TranslationUnit?
+    private(set) var match: TranslationUnit.ID?
 
 
     public init(_ query: DocumentUri) {
@@ -19,9 +19,9 @@ extension AST {
       let node = ast[n]
       let site = node.site
 
-      if let t = node as? TranslationUnit {
+      if node is TranslationUnit {
         if site.file.url.absoluteString == query {
-          match = t
+          match = TranslationUnit.ID(n)
         }
         return false
       }
@@ -30,7 +30,7 @@ extension AST {
     }
   }
 
-  public func findTranslationUnit(_ url: DocumentUri) -> TranslationUnit? {
+  public func findTranslationUnit(_ url: DocumentUri) -> TranslationUnit.ID? {
     var finder = TranslationUnitFinder(url)
     for m in modules {
       walk(m, notifying: &finder)

Sources/hylo-lsp/AST+GetSemanticTokens.swift

@@ -41,9 +41,8 @@ struct SemanticTokensWalker {
   mutating func addDecl(_ node: Node) {
 
     switch node {
-    case _ as NamespaceDecl:
-      break
-
+    case let d as NamespaceDecl:
+      addMembers(d.members)
     case let d as BindingDecl:
       addBinding(d)
     case let d as InitializerDecl:
@@ -650,14 +649,14 @@ struct SemanticTokensWalker {
 extension AST {
 
   public func getSematicTokens(_ document: DocumentUri, _ program: TypedProgram) -> [SemanticToken] {
-    logger.debug("List symbols in document: \(document)")
+    logger.debug("List semantic tokens in document: \(document)")
 
     guard let translationUnit = findTranslationUnit(document) else {
       logger.error("Failed to locate translation unit: \(document)")
       return []
     }
 
-    var finder = SemanticTokensWalker(document: document, translationUnit: translationUnit, program: program, ast: self)
-    return finder.walk()
+    var walker = SemanticTokensWalker(document: document, translationUnit: self[translationUnit], program: program, ast: self)
+    return walker.walk()
   }
 }

Sources/hylo-lsp/AST+ListDocumentSymbols.swift

@@ -0,0 +1,110 @@
+import Core
+import FrontEnd
+import IR
+import LanguageServerProtocol
+
+struct DocumentSymbolWalker {
+  public let document: DocumentUri
+  public let translationUnit: TranslationUnit
+  public let program: TypedProgram
+  public let ast: AST
+  private(set) var symbols: [AnyDeclID]
+
+  public init(document: DocumentUri, translationUnit: TranslationUnit, program: TypedProgram, ast: AST) {
+    self.document = document
+    self.translationUnit = translationUnit
+    self.program = program
+    self.ast = ast
+    self.symbols = []
+  }
+
+
+  public mutating func walk() -> [AnyDeclID] {
+    precondition(symbols.isEmpty)
+    addMembers(translationUnit.decls)
+    return symbols
+  }
+
+  mutating func addMembers(_ members: [AnyDeclID]) {
+    for m in members {
+      addDecl(m)
+    }
+  }
+
+
+  mutating func addDecl(_ id: AnyDeclID) {
+    let node = ast[id]
+    logger.debug("Found symbol node: \(id), site: \(node.site)")
+    symbols.append(id)
+    addMembers(id: id, node: node)
+  }
+
+  mutating func addMembers(id: AnyDeclID, node: Node) {
+
+    switch node {
+    case let d as NamespaceDecl:
+      addMembers(d.members)
+    case let d as ProductTypeDecl:
+      addMembers(d.members)
+    case let d as ExtensionDecl:
+      addMembers(d.members)
+    case let d as ConformanceDecl:
+      addMembers(d.members)
+    case let d as TraitDecl:
+      addMembers(d.members)
+    default:
+      // print("Ignored declaration node: \(node)")
+      break
+    }
+  }
+
+
+  mutating func willEnter(_ n: AnyNodeID, in ast: AST) -> Bool {
+    let node = ast[n]
+    let site = node.site
+
+    if let scheme = site.file.url.scheme {
+      if scheme == "synthesized" {
+        return true
+      }
+      else if n.kind == TranslationUnit.self && scheme == "file" {
+        if site.file.url.absoluteString != document {
+          // logger.debug("Ignore file: \(site.file.url)")
+          return false
+        }
+        // logger.debug("Enter file: \(site.file.url)")
+      }
+    }
+
+    if n.kind == NamespaceDecl.self || n.kind == TranslationUnit.self {
+      return true
+    }
+
+
+    if n.kind == BindingDecl.self {
+      return true
+    }
+
+    // if n.kind == FunctionDecl.self || n.kind == VarDecl.self {
+    if let d = AnyDeclID(n) {
+      logger.debug("Found symbol node: \(d), site: \(site)")
+      symbols.append(d)
+      return false
+    }
+
+    // logger.debug("Ignore node: \(n)")
+    return true
+  }
+}
+
+extension AST {
+  public func listDocumentSymbols(_ document: DocumentUri, _ program: TypedProgram) -> [AnyDeclID] {
+    logger.debug("List symbols in document: \(document)")
+    guard let translationUnit = findTranslationUnit(document) else {
+      logger.error("Failed to locate translation unit: \(document)")
+      return []
+    }
+    var walker = DocumentSymbolWalker(document: document, translationUnit: self[translationUnit], program: program, ast: self)
+    return walker.walk()
+  }
+}

Sources/hylo-lsp/DocumentCache.swift

@@ -4,14 +4,14 @@ import FrontEnd
 
 public struct Document {
   public let uri: DocumentUri
-  public var task: Task<TypedProgram, Error>
+  public var task: Task<(AST, TypedProgram), Error>
 
-  public init(uri: DocumentUri, program: TypedProgram) {
+  public init(uri: DocumentUri, ast: AST, program: TypedProgram) {
     self.uri = uri
-    self.task = Task { program }
+    self.task = Task { (ast, program) }
   }
 
-  public init(uri: DocumentUri, task: Task<TypedProgram, Error>) {
+  public init(uri: DocumentUri, task: Task<(AST, TypedProgram), Error>) {
     self.uri = uri
     self.task = task
   }