Support setting max focused code line count

Author: intitni

Core/Sources/ChatContextCollectors/ActiveDocumentChatContextCollector/ActiveDocumentChatContextCollector.swift

@@ -229,7 +229,7 @@ struct ActiveDocumentContext {
             case .builtIn(.swift):
                 return SwiftFocusedCodeFinder()
             default:
-                return UnknownLanguageFocusedCodeFinder()
+                return UnknownLanguageFocusedCodeFinder(proposedSearchRange: 5)
             }
         }()
 

Core/Sources/ChatContextCollectors/ActiveDocumentChatContextCollector/FocusedCodeFinder/FocusedCodeFinder.swift

@@ -14,10 +14,11 @@ struct CodeContext: Equatable {
             return []
         case .top:
             return ["Top level of the file"]
-        case .scope(let signature):
+        case let .scope(signature):
             return signature
         }
     }
+
     var scope: Scope
     var contextRange: CursorRange
     var focusedRange: CursorRange
@@ -37,6 +38,12 @@ protocol FocusedCodeFinder {
 }
 
 struct UnknownLanguageFocusedCodeFinder: FocusedCodeFinder {
+    let proposedSearchRange: Int
+
+    init(proposedSearchRange: Int) {
+        self.proposedSearchRange = proposedSearchRange
+    }
+
     func findFocusedCode(
         containingRange: CursorRange,
         activeDocumentContext: ActiveDocumentContext
@@ -45,25 +52,32 @@ struct UnknownLanguageFocusedCodeFinder: FocusedCodeFinder {
 
         // when user is not selecting any code.
         if containingRange.start == containingRange.end {
-            // search up and down for up to 7 lines.
+            // search up and down for up to `proposedSearchRange * 2 + 1` lines.
             let lines = activeDocumentContext.lines
-            var startLineIndex = max(containingRange.start.line - 3, 0)
-            let endLineIndex = min(containingRange.start.line + 3, lines.count - 1)
-            if endLineIndex - startLineIndex <= 6, startLineIndex > 0 {
-                startLineIndex = max(startLineIndex - (6 - (endLineIndex - startLineIndex)), 0)
+            let proposedLineCount = proposedSearchRange * 2 + 1
+            var startLineIndex = max(containingRange.start.line - proposedSearchRange, 0)
+            let endLineIndex = min(containingRange.start.line + proposedSearchRange, lines.count - 1)
+            if endLineIndex - startLineIndex < proposedLineCount, startLineIndex > 0 {
+                startLineIndex = max(
+                    startLineIndex - ((proposedSearchRange * 2) - (endLineIndex - startLineIndex)),
+                    0
+                )
             }
             let focusedLines = lines[startLineIndex...endLineIndex]
 
-            let contextStartLine = max(startLineIndex - 3, 0)
-            let contextEndLine = min(endLineIndex + 3, lines.count - 1)
+            let contextStartLine = max(startLineIndex - 5, 0)
+            let contextEndLine = min(endLineIndex + 5, lines.count - 1)
 
             return .init(
                 scope: .top,
                 contextRange: .init(
                     start: .init(line: contextStartLine, character: 0),
-                    end: .init(line: contextEndLine, character: 0)
+                    end: .init(line: contextEndLine, character: lines[contextEndLine].count)
+                ),
+                focusedRange: .init(
+                    start: .init(line: startLineIndex, character: 0),
+                    end: .init(line: endLineIndex, character: lines[endLineIndex].count)
                 ),
-                focusedRange: containingRange,
                 focusedCode: focusedLines.joined(separator: "\n"),
                 imports: []
             )
@@ -82,7 +96,10 @@ struct UnknownLanguageFocusedCodeFinder: FocusedCodeFinder {
             scope: .top,
             contextRange: .init(
                 start: .init(line: contextStartLine, character: 0),
-                end: .init(line: contextEndLine, character: 0)
+                end: .init(
+                    line: contextEndLine,
+                    character: activeDocumentContext.lines[contextEndLine].count
+                )
             ),
             focusedRange: containingRange,
             focusedCode: focusedLines.joined(separator: "\n"),

Core/Sources/ChatContextCollectors/ActiveDocumentChatContextCollector/FocusedCodeFinder/SwiftFocusedCodeFinder.swift

@@ -5,6 +5,12 @@ import SwiftParser
 import SwiftSyntax
 
 struct SwiftFocusedCodeFinder: FocusedCodeFinder {
+    let maxFocusedCodeLineCount: Int
+
+    init(maxFocusedCodeLineCount: Int = UserDefaults.shared.value(for: \.maxFocusedCodeLineCount)) {
+        self.maxFocusedCodeLineCount = maxFocusedCodeLineCount
+    }
+
     func findFocusedCode(
         containingRange range: CursorRange,
         activeDocumentContext: ActiveDocumentContext
@@ -27,8 +33,7 @@ struct SwiftFocusedCodeFinder: FocusedCodeFinder {
 
         var nodes = visitor.findScopeHierarchy()
 
-        let code: String
-        let codeRange: CursorRange
+        var codeRange: CursorRange
 
         func convertRange(_ node: SyntaxProtocol) -> CursorRange {
             .init(sourceRange: node.sourceRange(converter: locationConverter))
@@ -52,10 +57,11 @@ struct SwiftFocusedCodeFinder: FocusedCodeFinder {
                 }
             }
             guard let focusedNode else {
-                var result = UnknownLanguageFocusedCodeFinder().findFocusedCode(
-                    containingRange: range,
-                    activeDocumentContext: activeDocumentContext
-                )
+                var result = UnknownLanguageFocusedCodeFinder(proposedSearchRange: 5)
+                    .findFocusedCode(
+                        containingRange: range,
+                        activeDocumentContext: activeDocumentContext
+                    )
                 result.imports = visitor.imports
                 return result
             }
@@ -64,8 +70,26 @@ struct SwiftFocusedCodeFinder: FocusedCodeFinder {
             codeRange = range
         }
 
-        code = EditorInformation
-            .code(in: activeDocumentContext.lines, inside: codeRange, ignoreColumns: true).code
+        let result = EditorInformation
+            .code(in: activeDocumentContext.lines, inside: codeRange, ignoreColumns: true)
+
+        var code = result.code
+
+        if range.isEmpty, result.lines.count > maxFocusedCodeLineCount {
+            // if the focused code is too long, truncate it to be shorter
+            let centerLine = range.start.line
+            let relativeCenterLine = centerLine - codeRange.start.line
+            let startLine = max(0, relativeCenterLine - maxFocusedCodeLineCount / 2)
+            let endLine = min(result.lines.count - 1, startLine + maxFocusedCodeLineCount)
+            code = result.lines[startLine...endLine].joined(separator: "\n")
+            codeRange = .init(
+                start: .init(line: startLine + codeRange.start.line, character: 0),
+                end: .init(
+                    line: endLine + codeRange.start.line,
+                    character: result.lines[endLine].count
+                )
+            )
+        }
 
         var contextRange = CursorRange.zero
         var signature = [String]()
@@ -219,7 +243,7 @@ extension SwiftFocusedCodeFinder {
             let signature = node.bindings.first?.typeAnnotation?.trimmedDescription ?? ""
 
             return (.init(
-                signature: "\(type) \(name)\(signature.isEmpty ? "" : ": \(signature)")"
+                signature: "\(type) \(name)\(signature.isEmpty ? "" : "\(signature)")"
                     .prefixedModifiers(node.modifierAndAttributeText(extractText))
                     .replacingOccurrences(of: "\n", with: " "),
                 contextRange: convertRange(node)
@@ -471,10 +495,6 @@ extension SwiftFocusedCodeFinder {
 
         // skip if possible
 
-        override func visit(_ node: CodeBlockItemSyntax) -> SyntaxVisitorContinueKind {
-            skipChildrenIfPossible(node)
-        }
-
         override func visit(_ node: MemberDeclBlockSyntax) -> SyntaxVisitorContinueKind {
             skipChildrenIfPossible(node)
         }
@@ -486,7 +506,7 @@ extension SwiftFocusedCodeFinder {
         // capture if possible
 
         override func visit(_ node: ImportDeclSyntax) -> SyntaxVisitorContinueKind {
-            imports.append(node.trimmedDescription)
+            imports.append(node.path.trimmedDescription)
             return .skipChildren
         }
 

Tool/Sources/SuggestionModel/EditorInformation.swift

@@ -11,7 +11,7 @@ public struct EditorInformation {
     public struct SourceEditorContent {
         /// The content of the source editor.
         public var content: String
-        /// The content of the source editor in lines.
+        /// The content of the source editor in lines. Every line should ends with `\n`.
         public var lines: [String]
         /// The selection ranges of the source editor.
         public var selections: [CursorRange]