Skip to content

Commit 9d0a6d5

Browse files
intitniintitni
committed
Update AST parser
1 parent b03be81 commit 9d0a6d5

11 files changed

Lines changed: 451 additions & 9 deletions

File tree

‎Core/Sources/ChatContextCollectors/ActiveDocumentChatContextCollector/ActiveDocumentChatContextCollector.swift‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
import ASTParser
2+
import ChatContextCollector
3+
import Foundation
4+
import OpenAIService
5+
import Preferences
6+
import SuggestionModel
7+
import XcodeInspector
8+
9+
public final class ActiveDocumentChatContextCollector: ChatContextCollector {
10+
public init() {}
11+
12+
public func generateContext(
13+
history: [ChatMessage],
14+
scopes: Set<String>,
15+
content: String
16+
) -> ChatContext? {
17+
guard scopes.contains("file") else { return nil }
18+
let info = getEditorInformation()
19+
20+
return .init(
21+
systemPrompt: extractSystemPrompt(info),
22+
functions: []
23+
)
24+
}
25+
26+
func extractSystemPrompt(_ info: EditorInformation) -> String {
27+
let relativePath = info.documentURL.path
28+
.replacingOccurrences(of: info.projectURL.path, with: "")
29+
let selectionRange = info.editorContent?.selections.first ?? .outOfScope
30+
let lineAnnotations = info.editorContent?.lineAnnotations ?? []
31+
32+
var result = """
33+
Active Document Context:###
34+
Document Relative Path: \(relativePath)
35+
Language: \(info.language.rawValue)
36+
Selection Range [line, character]: \
37+
[\(selectionRange.start.line), \(selectionRange.start.character)] - \
38+
[\(selectionRange.end.line), \(selectionRange.end.character)]
39+
###
40+
"""
41+
42+
if !lineAnnotations.isEmpty {
43+
result += """
44+
Line Annotations:
45+
\(lineAnnotations.map { " - \($0)" }.joined(separator: "\n"))
46+
"""
47+
}
48+
49+
return result
50+
}
51+
}
52+
53+

‎Core/Sources/ChatContextCollectors/ActiveDocumentChatContextCollector/GetCodeFunction.swift‎

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
1+
import ASTParser
2+
import Foundation
3+
import OpenAIService
4+
import SuggestionModel
5+
6+
struct GetCodeFunction: ChatGPTFunction {
7+
enum CodeType: String, Codable {
8+
case selected
9+
case focused
10+
}
11+
12+
struct Arguments: Codable {
13+
var codeType: CodeType
14+
}
15+
16+
struct Result: ChatGPTFunctionResult {
17+
struct Context {
18+
var parentName: String
19+
var parentType: String
20+
}
21+
22+
var relativePath: String
23+
var code: String
24+
var range: CursorRange
25+
var context: Context
26+
var type: CodeType
27+
var language: String
28+
29+
var botReadableContent: String {
30+
"""
31+
The \(type.rawValue) code is a part of `\(context.parentType) \(context.parentName)` \
32+
in file \(relativePath).
33+
Range [\(range.start.line), \(range.start.character)] - \
34+
[\(range.end.line), \(range.end.character)]
35+
```\(language)
36+
\(code)
37+
```
38+
"""
39+
}
40+
}
41+
42+
var reportProgress: (String) async -> Void = { _ in }
43+
44+
var name: String {
45+
"getCode"
46+
}
47+
48+
var description: String {
49+
"Get selected or focused code from the active document."
50+
}
51+
52+
var argumentSchema: JSONSchemaValue { [
53+
.type: "object",
54+
.properties: [:],
55+
] }
56+
57+
func prepare() async {
58+
await reportProgress("Reading code..")
59+
}
60+
61+
func call(arguments: Arguments) async throws -> Result {
62+
await reportProgress("Reading code..")
63+
let content = getEditorInformation()
64+
let selectionRange = content.editorContent?.selections.first ?? .outOfScope
65+
let editorContent = {
66+
if selectionRange.start == selectionRange.end {
67+
return content.editorContent?.content ?? ""
68+
} else {
69+
return content.selectedContent
70+
}
71+
}()
72+
73+
let language = content.language.rawValue
74+
let type = CodeType.selected
75+
let relativePath = content.documentURL.path
76+
.replacingOccurrences(of: content.projectURL.path, with: "")
77+
let context = Result.Context(
78+
parentName: content.documentURL.lastPathComponent,
79+
parentType: "File"
80+
)
81+
let range = selectionRange
82+
83+
await reportProgress("Finish reading code..")
84+
return .init(
85+
relativePath: relativePath,
86+
code: editorContent,
87+
range: range,
88+
context: context,
89+
type: type,
90+
language: language
91+
)
92+
}
93+
}
94+
95+
struct GetCodeResultParser {
96+
let editorInformation: EditorInformation
97+
98+
func parse() -> GetCodeFunction.Result {
99+
let language = editorInformation.language.rawValue
100+
let relativePath = editorInformation.relativePath
101+
let selectionRange = editorInformation.editorContent?.selections.first
102+
103+
if let selectionRange, let node = findSmallestScopeContainingRange(selectionRange) {
104+
let code = {
105+
if editorInformation.selectedContent.isEmpty {
106+
return editorInformation.selectedLines.first ?? ""
107+
}
108+
return editorInformation.selectedContent
109+
}()
110+
111+
return .init(
112+
relativePath: relativePath,
113+
code: code,
114+
range: selectionRange,
115+
context: .init(parentName: "", parentType: ""),
116+
type: .selected,
117+
language: language
118+
)
119+
}
120+
121+
return .init(
122+
relativePath: relativePath,
123+
code: "",
124+
range: selectionRange ?? .zero,
125+
context: .init(parentName: "", parentType: ""),
126+
type: .focused,
127+
language: language
128+
)
129+
}
130+
131+
func findSmallestScopeContainingRange(_ range: CursorRange) -> ASTNode? {
132+
guard let language = {
133+
switch editorInformation.language {
134+
case .builtIn(.swift):
135+
return ParsableLanguage.swift
136+
case .builtIn(.objc), .builtIn(.objcpp):
137+
return ParsableLanguage.objectiveC
138+
default:
139+
return nil
140+
}
141+
}() else { return nil }
142+
143+
let parser = ASTParser(language: language)
144+
guard let tree = parser.parse(editorInformation.editorContent?.content ?? "")
145+
else { return nil }
146+
147+
return tree.smallestNodeContainingRange(range) { node in
148+
ScopeType.allCases.map { $0.rawValue }.contains(node.nodeType)
149+
}
150+
}
151+
}
152+
153+
enum ScopeType: String, CaseIterable {
154+
case protocolDeclaration = "protocol_declaration"
155+
case classDeclaration = "class_declaration"
156+
case functionDeclaration = "function_declaration"
157+
case propertyDeclaration = "property_declaration"
158+
case computedProperty = "computed_property"
159+
}
160+

‎Core/Sources/ChatContextCollectors/ActiveDocumentChatContextCollector/GetEditorInfo.swift‎

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,19 +2,23 @@ import Foundation
22
import SuggestionModel
33
import XcodeInspector
44

5-
struct Information {
5+
struct EditorInformation {
66
let editorContent: SourceEditor.Content?
77
let selectedContent: String
8+
let selectedLines: [String]
89
let documentURL: URL
910
let projectURL: URL
11+
let relativePath: String
1012
let language: CodeLanguage
1113
}
1214

13-
func getEditorInformation() -> Information {
15+
func getEditorInformation() -> EditorInformation {
1416
let editorContent = XcodeInspector.shared.focusedEditor?.content
1517
let documentURL = XcodeInspector.shared.activeDocumentURL
1618
let projectURL = XcodeInspector.shared.activeProjectURL
1719
let language = languageIdentifierFromFileURL(documentURL)
20+
let relativePath = documentURL.path
21+
.replacingOccurrences(of: projectURL.path, with: "")
1822

1923
if let editorContent, let range = editorContent.selections.first {
2024
let startIndex = min(
@@ -25,21 +29,34 @@ func getEditorInformation() -> Information {
2529
max(startIndex, range.end.line),
2630
editorContent.lines.endIndex - 1
2731
)
28-
let selectedContent = editorContent.lines[startIndex...endIndex]
32+
let selectedLines = editorContent.lines[startIndex...endIndex]
33+
var selectedContent = selectedLines
34+
if selectedContent.count > 0 {
35+
selectedContent[0] = String(selectedContent[0].dropFirst(range.start.character))
36+
selectedContent[selectedContent.endIndex - 1] = String(
37+
selectedContent[selectedContent.endIndex - 1].dropLast(
38+
selectedContent[selectedContent.endIndex - 1].count - range.end.character
39+
)
40+
)
41+
}
2942
return .init(
3043
editorContent: editorContent,
3144
selectedContent: selectedContent.joined(),
45+
selectedLines: Array(selectedLines),
3246
documentURL: documentURL,
3347
projectURL: projectURL,
48+
relativePath: relativePath,
3449
language: language
3550
)
3651
}
3752

3853
return .init(
3954
editorContent: editorContent,
4055
selectedContent: "",
56+
selectedLines: [],
4157
documentURL: documentURL,
4258
projectURL: projectURL,
59+
relativePath: relativePath,
4360
language: language
4461
)
4562
}

‎Playground.playground/Pages/ASTParsing.xcplaygroundpage/Contents.swift‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ struct ParsingForm: View {
2626
Section("Result") {
2727
Text(result)
2828
.fontDesign(.monospaced)
29+
.textSelection(.enabled)
2930
}
3031
}
3132
.formStyle(.grouped)
@@ -35,3 +36,5 @@ struct ParsingForm: View {
3536

3637
PlaygroundPage.current.needsIndefiniteExecution = true
3738
PlaygroundPage.current.setLiveView(NSHostingController(rootView: ParsingForm()))
39+
// protocol_declaration, class_declaration, function_declaration, property_declaration, computed_property
40+
// type_identifier, simple_identifier (for variables and funcs)

‎Playground.playground/Pages/ASTParsing.xcplaygroundpage/timeline.xctimeline‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,4 +5,6 @@
55
</TimelineItems>
66
<TimelineItems>
77
</TimelineItems>
8+
<TimelineItems>
9+
</TimelineItems>
810
</Timeline>

‎Tool/Package.swift‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -157,6 +157,8 @@ let package = Package(
157157
.product(name: "TreeSitterSwift", package: "tree-sitter-swift"),
158158
]),
159159

160+
.testTarget(name: "ASTParserTests", dependencies: ["ASTParser"]),
161+
160162
// MARK: - Services
161163

162164
.target(

‎Tool/Sources/ASTParser/ASTParser.swift‎

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import SuggestionModel
12
import SwiftTreeSitter
23
import tree_sitter
34
import TreeSitterObjC
@@ -32,11 +33,52 @@ public struct ASTParser {
3233
}
3334
}
3435

36+
public typealias ASTNode = Node
37+
38+
public typealias ASTPoint = Point
39+
3540
public struct ASTTree {
3641
public let tree: Tree?
3742

38-
public var rootNode: Node? {
43+
public var rootNode: ASTNode? {
3944
return tree?.rootNode
4045
}
46+
47+
public func smallestNodeContainingRange(
48+
_ range: CursorRange,
49+
filter: (ASTNode) -> Bool = { _ in true }
50+
) -> ASTNode? {
51+
guard var targetNode = rootNode else { return nil }
52+
53+
func rangeContains(_ range: Range<Point>, _ another: Range<Point>) -> Bool {
54+
return range.lowerBound <= another.lowerBound && range.upperBound >= another.upperBound
55+
}
56+
57+
for node in targetNode.treeCursor.deepFirstSearch(skipChildren: { node in
58+
!rangeContains(node.pointRange, range.pointRange)
59+
}) {
60+
guard filter(node) else { continue }
61+
if rangeContains(node.pointRange, range.pointRange) {
62+
targetNode = node
63+
}
64+
}
65+
66+
return targetNode
67+
}
68+
}
69+
70+
extension CursorRange {
71+
var pointRange: Range<Point> {
72+
let bytePerCharacter = 2 // tree sitter uses UTF-16
73+
let startPoint = Point(row: start.line, column: start.character * bytePerCharacter)
74+
let endPoint = Point(row: end.line, column: end.character * bytePerCharacter)
75+
guard endPoint > startPoint else {
76+
return startPoint..<Point(
77+
row: start.line,
78+
column: (start.character + 1) * bytePerCharacter
79+
)
80+
}
81+
return startPoint..<endPoint
82+
}
4183
}
4284

0 commit comments

Comments
 (0)