Fix that context retriever was incorrect when the context is too long

Author: intitni

Pro

@@ -82,7 +82,8 @@ public final class ActiveDocumentChatContextCollector: ChatContextCollector {
             selectedCode: "",
             selectionRange: .outOfScope,
             lineAnnotations: [],
-            imports: []
+            imports: [],
+            includes: []
         )
 
         activeDocumentContext.update(info)

Tool/Sources/ChatContextCollectors/ActiveDocumentChatContextCollector/ActiveDocumentChatContextCollector.swift

@@ -11,22 +11,24 @@ public struct ActiveDocumentContext {
     public var selectionRange: CursorRange
     public var lineAnnotations: [EditorInformation.LineAnnotation]
     public var imports: [String]
+    public var includes: [String]
 
     public struct FocusedContext {
         public struct Context: Equatable {
             public var signature: String
             public var name: String
             public var range: CursorRange
-            
+
             public init(signature: String, name: String, range: CursorRange) {
                 self.signature = signature
                 self.name = name
                 self.range = range
             }
         }
-        
+
         public var context: [Context]
         public var contextRange: CursorRange
+        public var smallestContextRange: CursorRange
         public var codeRange: CursorRange
         public var code: String
         public var lineAnnotations: [EditorInformation.LineAnnotation]
@@ -35,13 +37,15 @@ public struct ActiveDocumentContext {
         public init(
             context: [Context],
             contextRange: CursorRange,
+            smallestContextRange: CursorRange,
             codeRange: CursorRange,
             code: String,
             lineAnnotations: [EditorInformation.LineAnnotation],
             otherLineAnnotations: [EditorInformation.LineAnnotation]
         ) {
             self.context = context
             self.contextRange = contextRange
+            self.smallestContextRange = smallestContextRange
             self.codeRange = codeRange
             self.code = code
             self.lineAnnotations = lineAnnotations
@@ -61,6 +65,7 @@ public struct ActiveDocumentContext {
         selectionRange: CursorRange,
         lineAnnotations: [EditorInformation.LineAnnotation],
         imports: [String],
+        includes: [String],
         focusedContext: FocusedContext? = nil
     ) {
         self.documentURL = documentURL
@@ -72,6 +77,7 @@ public struct ActiveDocumentContext {
         self.selectionRange = selectionRange
         self.lineAnnotations = lineAnnotations
         self.imports = imports
+        self.includes = includes
         self.focusedContext = focusedContext
     }
 
@@ -92,15 +98,18 @@ public struct ActiveDocumentContext {
     }
 
     public mutating func moveToCodeContainingRange(_ range: CursorRange) {
-        let finder = FocusedCodeFinder()
-        
+        let finder = FocusedCodeFinder(
+            maxFocusedCodeLineCount: UserDefaults.shared.value(for: \.maxFocusedCodeLineCount)
+        )
+
         let codeContext = finder.findFocusedCode(
             in: .init(documentURL: documentURL, content: fileContent, lines: lines),
             containingRange: range,
             language: language
         )
 
         imports = codeContext.imports
+        includes = codeContext.includes
 
         let startLine = codeContext.focusedRange.start.line
         let endLine = codeContext.focusedRange.end.line
@@ -117,6 +126,7 @@ public struct ActiveDocumentContext {
         focusedContext = .init(
             context: codeContext.scopeContexts,
             contextRange: codeContext.contextRange,
+            smallestContextRange: codeContext.smallestContextRange,
             codeRange: codeContext.focusedRange,
             code: codeContext.focusedCode,
             lineAnnotations: matchedAnnotations,
@@ -144,6 +154,7 @@ public struct ActiveDocumentContext {
         selectionRange = info.editorContent?.selections.first ?? .zero
         lineAnnotations = info.editorContent?.lineAnnotations ?? []
         imports = []
+        includes = []
 
         if changed {
             moveToFocusedCode()

Tool/Sources/FocusedCodeFinder/ActiveDocumentContext.swift

@@ -23,6 +23,7 @@ public struct CodeContext: Equatable {
 
     public var scope: Scope
     public var contextRange: CursorRange
+    public var smallestContextRange: CursorRange
     public var focusedRange: CursorRange
     public var focusedCode: String
     public var imports: [String]
@@ -32,6 +33,7 @@ public struct CodeContext: Equatable {
         .init(
             scope: .file,
             contextRange: .zero,
+            smallestContextRange: .zero,
             focusedRange: .zero,
             focusedCode: "",
             imports: [],
@@ -42,13 +44,15 @@ public struct CodeContext: Equatable {
     public init(
         scope: Scope,
         contextRange: CursorRange,
+        smallestContextRange: CursorRange,
         focusedRange: CursorRange,
         focusedCode: String,
         imports: [String],
         includes: [String]
     ) {
         self.scope = scope
         self.contextRange = contextRange
+        self.smallestContextRange = smallestContextRange
         self.focusedRange = focusedRange
         self.focusedCode = focusedCode
         self.imports = imports
@@ -57,7 +61,11 @@ public struct CodeContext: Equatable {
 }
 
 public struct FocusedCodeFinder {
-    public init() {}
+    public let maxFocusedCodeLineCount: Int
+
+    public init(maxFocusedCodeLineCount: Int) {
+        self.maxFocusedCodeLineCount = maxFocusedCodeLineCount
+    }
 
     public struct Document {
         var documentURL: URL
@@ -79,10 +87,12 @@ public struct FocusedCodeFinder {
         let finder: FocusedCodeFinderType = {
             switch language {
             case .builtIn(.swift):
-                return SwiftFocusedCodeFinder()
+                return SwiftFocusedCodeFinder(maxFocusedCodeLineCount: maxFocusedCodeLineCount)
             case .builtIn(.objc), .builtIn(.objcpp), .builtIn(.c):
-                #warning("TODO: Implement C++ focused code finder, use it for C and metal shading language")
-                return ObjectiveCFocusedCodeFinder()
+                #warning(
+                    "TODO: Implement C++ focused code finder, use it for C and metal shading language"
+                )
+                return ObjectiveCFocusedCodeFinder(maxFocusedCodeLineCount: maxFocusedCodeLineCount)
             default:
                 return UnknownLanguageFocusedCodeFinder(proposedSearchRange: 5)
             }

Tool/Sources/FocusedCodeFinder/FocusedCodeFinder.swift

@@ -126,7 +126,8 @@ public extension KnownLanguageFocusedCodeFinderType {
 
         return .init(
             scope: scopeContexts.isEmpty ? .file : .scope(signature: scopeContexts),
-            contextRange: contextRange,
+            contextRange: contextRange, 
+            smallestContextRange: codeRange,
             focusedRange: focusedRange,
             focusedCode: code,
             imports: contextInfo.imports,

Tool/Sources/FocusedCodeFinder/KnownLanguageFocusedCodeFinder.swift

@@ -14,9 +14,7 @@ public class ObjectiveCFocusedCodeFinder: KnownLanguageFocusedCodeFinder<
     ASTNode,
     TreeSitterTextPosition
 > {
-    override public init(
-        maxFocusedCodeLineCount: Int = UserDefaults.shared.value(for: \.maxFocusedCodeLineCount)
-    ) {
+    override public init(maxFocusedCodeLineCount: Int) {
         super.init(maxFocusedCodeLineCount: maxFocusedCodeLineCount)
     }
 
@@ -99,7 +97,7 @@ public class ObjectiveCFocusedCodeFinder: KnownLanguageFocusedCodeFinder<
         var prefix = ""
         /// Generics, super class, etc.
         var extra = ""
-        
+
         if let nameNode = node.child(byFieldName: "name") {
             name = textProvider(.node(nameNode))
             prefix = textProvider(.range(
@@ -135,7 +133,7 @@ public class ObjectiveCFocusedCodeFinder: KnownLanguageFocusedCodeFinder<
         prefix = prefix.split(separator: "\n")
             .joined(separator: " ")
             .trimmingCharacters(in: .whitespacesAndNewlines)
-        
+
         extra = extra.split(separator: "\n")
             .joined(separator: " ")
             .trimmingCharacters(in: .whitespacesAndNewlines)
@@ -152,7 +150,7 @@ public class ObjectiveCFocusedCodeFinder: KnownLanguageFocusedCodeFinder<
             canBeUsedAsCodeRange: true
         )
     }
-    
+
     func parseMethodDefinitionNode(
         _ node: ASTNode,
         textProvider: @escaping TextProvider

Tool/Sources/FocusedCodeFinder/ObjectiveC/ObjectiveCCodeFinder.swift

@@ -10,9 +10,7 @@ public class SwiftFocusedCodeFinder: KnownLanguageFocusedCodeFinder<
     SyntaxProtocol,
     SyntaxProtocol
 > {
-    override public init(
-        maxFocusedCodeLineCount: Int = UserDefaults.shared.value(for: \.maxFocusedCodeLineCount)
-    ) {
+    override public init(maxFocusedCodeLineCount: Int) {
         super.init(maxFocusedCodeLineCount: maxFocusedCodeLineCount)
     }
 

Tool/Sources/FocusedCodeFinder/Swift/SwiftFocusedCodeFinder.swift

@@ -28,17 +28,22 @@ public struct UnknownLanguageFocusedCodeFinder: FocusedCodeFinderType {
                 min(startLineIndex + proposedLineCount - 1, lines.count - 1)
             )
 
+            if lines.endIndex <= endLineIndex { return .empty }
+            
             let focusedLines = lines[startLineIndex...endLineIndex]
 
             let contextStartLine = max(startLineIndex - 5, 0)
             let contextEndLine = min(endLineIndex + 5, lines.count - 1)
+            
+            let contextRange = CursorRange(
+                start: .init(line: contextStartLine, character: 0),
+                end: .init(line: contextEndLine, character: lines[contextEndLine].count)
+            )
 
             return .init(
                 scope: .top,
-                contextRange: .init(
-                    start: .init(line: contextStartLine, character: 0),
-                    end: .init(line: contextEndLine, character: lines[contextEndLine].count)
-                ),
+                contextRange: contextRange,
+                smallestContextRange: contextRange,
                 focusedRange: .init(
                     start: .init(line: startLineIndex, character: 0),
                     end: .init(line: endLineIndex, character: lines[endLineIndex].count)
@@ -57,16 +62,19 @@ public struct UnknownLanguageFocusedCodeFinder: FocusedCodeFinderType {
         let focusedLines = document.lines[startLine...endLine]
         let contextStartLine = max(startLine - 3, 0)
         let contextEndLine = min(endLine + 3, document.lines.count - 1)
+        
+        let contextRange = CursorRange(
+            start: .init(line: contextStartLine, character: 0),
+            end: .init(
+                line: contextEndLine,
+                character: document.lines[contextEndLine].count
+            )
+        )
 
         return CodeContext(
             scope: .top,
-            contextRange: .init(
-                start: .init(line: contextStartLine, character: 0),
-                end: .init(
-                    line: contextEndLine,
-                    character: document.lines[contextEndLine].count
-                )
-            ),
+            contextRange: contextRange,
+            smallestContextRange: contextRange,
             focusedRange: containingRange,
             focusedCode: focusedLines.joined(),
             imports: [],

Tool/Sources/FocusedCodeFinder/UnknownLanguageFocusCodeFinder.swift

@@ -36,7 +36,8 @@ final class ObjectiveCFocusedCodeFinder_Selection_Tests: XCTestCase {
                     range: .init(startPair: (1, 0), endPair: (5, 1))
                 ),
             ]),
-            contextRange: .init(startPair: (0, 0), endPair: (6, 4)),
+            contextRange: .init(startPair: (0, 0), endPair: (6, 4)), 
+            smallestContextRange: range,
             focusedRange: range,
             focusedCode: """
                 NSInteger foo = 0;
@@ -69,6 +70,7 @@ final class ObjectiveCFocusedCodeFinder_Selection_Tests: XCTestCase {
                 ),
             ]),
             contextRange: .init(startPair: (0, 0), endPair: (4, 1)),
+            smallestContextRange: range,
             focusedRange: range,
             focusedCode: """
                 NSLog(@"Hello");
@@ -107,6 +109,7 @@ final class ObjectiveCFocusedCodeFinder_Selection_Tests: XCTestCase {
                 ),
             ]),
             contextRange: .init(startPair: (0, 0), endPair: (7, 4)),
+            smallestContextRange: range,
             focusedRange: range,
             focusedCode: """
             - (void)fooWith:(NSInteger)foo {
@@ -146,6 +149,7 @@ final class ObjectiveCFocusedCodeFinder_Selection_Tests: XCTestCase {
                 ),
             ]),
             contextRange: .init(startPair: (0, 0), endPair: (4, 4)),
+            smallestContextRange: range,
             focusedRange: range,
             focusedCode: """
             - (void)fooWith:(NSInteger)foo;
@@ -183,6 +187,7 @@ final class ObjectiveCFocusedCodeFinder_Selection_Tests: XCTestCase {
                 ),
             ]),
             contextRange: .init(startPair: (0, 0), endPair: (4, 4)),
+            smallestContextRange: range,
             focusedRange: range,
             focusedCode: """
             - (void)fooWith:(NSInteger)foo;
@@ -220,6 +225,7 @@ final class ObjectiveCFocusedCodeFinder_Selection_Tests: XCTestCase {
                 ),
             ]),
             contextRange: .init(startPair: (0, 0), endPair: (4, 4)),
+            smallestContextRange: range,
             focusedRange: range,
             focusedCode: """
             - (void)fooWith:(NSInteger)foo;
@@ -257,6 +263,7 @@ final class ObjectiveCFocusedCodeFinder_Selection_Tests: XCTestCase {
                 ),
             ]),
             contextRange: .init(startPair: (0, 0), endPair: (4, 1)),
+            smallestContextRange: range,
             focusedRange: range,
             focusedCode: """
                 NSInteger foo;
@@ -294,6 +301,7 @@ final class ObjectiveCFocusedCodeFinder_Selection_Tests: XCTestCase {
                 ),
             ]),
             contextRange: .init(startPair: (0, 0), endPair: (4, 1)),
+            smallestContextRange: range,
             focusedRange: range,
             focusedCode: """
                 foo,
@@ -331,6 +339,7 @@ final class ObjectiveCFocusedCodeFinder_Selection_Tests: XCTestCase {
                 ),
             ]),
             contextRange: .init(startPair: (0, 0), endPair: (4, 2)),
+            smallestContextRange: range,
             focusedRange: range,
             focusedCode: """
                 foo,
@@ -372,6 +381,7 @@ final class ObjectiveCFocusedCodeFinder_Focus_Tests: XCTestCase {
                 ),
             ]),
             contextRange: .init(startPair: (0, 0), endPair: (6, 4)),
+            smallestContextRange: range,
             focusedRange: .init(startPair: (1, 0), endPair: (5, 1)),
             focusedCode: """
             - (void)fooWith:(NSInteger)foo {
@@ -408,6 +418,7 @@ final class ObjectiveCFocusedCodeFinder_Focus_Tests: XCTestCase {
         XCTAssertEqual(context, .init(
             scope: .file,
             contextRange: .init(startPair: (0, 0), endPair: (0, 0)),
+            smallestContextRange: range,
             focusedRange: .init(startPair: (0, 0), endPair: (4, 4)),
             focusedCode: """
             @interface __GENERICS(NSArray, ObjectType) (BlocksKit)