Cleanup

Author: marcoroth

javascript/packages/language-server/src/folding_range_service.ts

@@ -1,27 +1,50 @@
 import { FoldingRange, FoldingRangeKind } from "vscode-languageserver/node"
+import { TextDocument } from "vscode-languageserver-textdocument"
+
 import { Visitor } from "@herb-tools/core"
+import { ParserService } from "./parser_service"
+
+import { isERBIfNode } from "@herb-tools/core"
+import { lspLine } from "./range_utils"
 
 import type {
   Node,
   ERBNode,
+  ERBContentNode,
   HTMLElementNode,
   HTMLOpenTagNode,
   HTMLAttributeValueNode,
   HTMLCommentNode,
+  HTMLConditionalElementNode,
+  CDATANode,
   ERBIfNode,
-  DocumentNode,
+  ERBUnlessNode,
+  ERBCaseNode,
+  ERBCaseMatchNode,
+  ERBBeginNode,
+  SerializedPosition,
 } from "@herb-tools/core"
 
 export class FoldingRangeService {
-  getFoldingRanges(document: DocumentNode): FoldingRange[] {
+  private parserService: ParserService
+
+  constructor(parserService: ParserService) {
+    this.parserService = parserService
+  }
+
+  getFoldingRanges(textDocument: TextDocument): FoldingRange[] {
+    const parseResult = this.parserService.parseDocument(textDocument)
     const collector = new FoldingRangeCollector()
-    collector.visit(document)
+
+    collector.visit(parseResult.document)
+
     return collector.ranges
   }
 }
 
-class FoldingRangeCollector extends Visitor {
-  ranges: FoldingRange[] = []
+export class FoldingRangeCollector extends Visitor {
+  public ranges: FoldingRange[] = []
+  private processedIfNodes: Set<ERBIfNode> = new Set()
 
   visitHTMLElementNode(node: HTMLElementNode): void {
     this.addRangeForNode(node, node.body)
@@ -34,9 +57,7 @@ class FoldingRangeCollector extends Visitor {
   }
 
   visitHTMLCommentNode(node: HTMLCommentNode): void {
-    const startLine = this.toZeroBased(node.location.start.line)
-    const endLine = this.toZeroBased(node.location.end.line)
-    this.addRange(startLine, endLine, FoldingRangeKind.Comment)
+    this.addRange(node.location.start, node.location.end, FoldingRangeKind.Comment)
     this.visitChildNodes(node)
   }
 
@@ -45,43 +66,107 @@ class FoldingRangeCollector extends Visitor {
     this.visitChildNodes(node)
   }
 
+  visitCDATANode(node: CDATANode): void {
+    this.addRange(node.location.start, node.location.end)
+    this.visitChildNodes(node)
+  }
+
+  visitHTMLConditionalElementNode(node: HTMLConditionalElementNode): void {
+    this.addRange(node.location.start, node.location.end)
+    this.visitChildNodes(node)
+  }
+
   visitERBNode(node: ERBNode): void {
-    const startLine = this.toZeroBased(node.location.start.line)
-    const endLine = this.toZeroBased(node.location.end.line)
-    this.addRange(startLine, endLine)
+    this.addRange(node.location.start, node.location.end)
+  }
+
+  visitERBContentNode(node: ERBContentNode): void {
+    this.addRange(node.location.start, node.location.end)
+    this.visitChildNodes(node)
   }
 
   visitERBIfNode(node: ERBIfNode): void {
-    const startLine = this.toZeroBased(node.location.start.line)
-    const endLine = this.toZeroBased(node.location.end.line)
-    this.addRange(startLine, endLine)
+    if (this.processedIfNodes.has(node)) {
+      this.visitChildNodes(node)
+      return
+    }
+
+    this.markIfChainAsProcessed(node)
+    this.addRange(node.location.start, node.location.end)
 
     if (node.statements.length > 0) {
       const firstStatement = node.statements[0]
       const lastStatement = node.statements[node.statements.length - 1]
-      const statementsStartLine = this.toZeroBased(firstStatement.location.start.line)
-      const statementsEndLine = this.toZeroBased(lastStatement.location.end.line)
-      this.addRange(statementsStartLine, statementsEndLine)
+
+      this.addRange(firstStatement.location.start, lastStatement.location.end)
     }
 
+    let current: Node | null = node.subsequent
+
+    while (current) {
+      if (isERBIfNode(current)) {
+        if (current.statements.length > 0) {
+          const firstStatement = current.statements[0]
+          const lastStatement = current.statements[current.statements.length - 1]
+          this.addRange(firstStatement.location.start, lastStatement.location.end)
+        }
+
+        current = current.subsequent
+      } else {
+        break
+      }
+    }
+
+    this.visitChildNodes(node)
+  }
+
+  visitERBUnlessNode(node: ERBUnlessNode): void {
+    this.addRange(node.location.start, node.location.end)
     this.visitChildNodes(node)
   }
 
-  private addRange(startLine: number, endLine: number, kind?: FoldingRangeKind): void {
+  visitERBCaseNode(node: ERBCaseNode): void {
+    this.addRange(node.location.start, node.location.end)
+    this.visitChildNodes(node)
+  }
+
+  visitERBCaseMatchNode(node: ERBCaseMatchNode): void {
+    this.addRange(node.location.start, node.location.end)
+    this.visitChildNodes(node)
+  }
+
+  visitERBBeginNode(node: ERBBeginNode): void {
+    this.addRange(node.location.start, node.location.end)
+    this.visitChildNodes(node)
+  }
+
+  private markIfChainAsProcessed(node: ERBIfNode): void {
+    this.processedIfNodes.add(node)
+
+    let current: Node | null = node.subsequent
+
+    while (current) {
+      if (isERBIfNode(current)) {
+        this.processedIfNodes.add(current)
+        current = current.subsequent
+      } else {
+        break
+      }
+    }
+  }
+
+  private addRange(start: SerializedPosition, end: SerializedPosition, kind?: FoldingRangeKind): void {
+    const startLine = lspLine(start)
+    const endLine = lspLine(end)
+
     if (endLine > startLine) {
       this.ranges.push({ startLine, endLine, kind })
     }
   }
 
   private addRangeForNode(node: Node, children: Node[]) {
     if (children.length > 0) {
-      const startLine = this.toZeroBased(node.location.start.line)
-      const endLine = this.toZeroBased(node.location.end.line)
-      this.addRange(startLine, endLine)
+      this.addRange(node.location.start, node.location.end)
     }
   }
-
-  private toZeroBased(line: number): number {
-    return line - 1
-  }
 }

javascript/packages/language-server/src/range_utils.ts

@@ -6,6 +6,10 @@ export function lspPosition(herbPosition: SerializedPosition): Position {
   return Position.create(herbPosition.line - 1, herbPosition.column)
 }
 
+export function lspLine(herbPosition: SerializedPosition): number {
+  return herbPosition.line - 1
+}
+
 export function lspRangeFromLocation(herbLocation: SerializedLocation): Range {
   return Range.create(lspPosition(herbLocation.start), lspPosition(herbLocation.end))
 }

javascript/packages/language-server/src/server.ts

@@ -193,9 +193,7 @@ export class Server {
 
       if (!document) return []
 
-      const parseResult = this.service.parserService.parseDocument(document)
-
-      return this.service.foldingRangeService.getFoldingRanges(parseResult.document)
+      return this.service.foldingRangeService.getFoldingRanges(document)
     })
   }
 

javascript/packages/language-server/src/service.ts

@@ -48,7 +48,7 @@ export class Service {
     this.codeActionService = new CodeActionService(this.project, this.config)
     this.diagnostics = new Diagnostics(this.connection, this.documentService, this.parserService, this.linterService, this.configService)
     this.documentSaveService = new DocumentSaveService(this.connection, this.settings, this.autofixService, this.formattingService)
-    this.foldingRangeService = new FoldingRangeService()
+    this.foldingRangeService = new FoldingRangeService(this.parserService)
     this.documentHighlightService = new DocumentHighlightService(this.parserService)
 
     if (params.initializationOptions) {