Skip to content

Commit ecfdfc0

Browse files
committed
Merge branch 'feature/objective-c-code-finder' into develop
2 parents 221a034 + e8c3ba4 commit ecfdfc0

22 files changed

Lines changed: 2087 additions & 904 deletions

Pro

Submodule Pro updated from 14b2333 to c0ffa88

Tool/Package.swift

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -63,11 +63,7 @@ let package = Package(
6363
.package(url: "https://github.com/GottaGetSwifty/CodableWrappers", from: "2.0.7"),
6464

6565
// TreeSitter
66-
.package(url: "https://github.com/ChimeHQ/SwiftTreeSitter", from: "0.7.1"),
67-
.package(
68-
url: "https://github.com/alex-pinkus/tree-sitter-swift",
69-
branch: "with-generated-files"
70-
),
66+
.package(url: "https://github.com/intitni/SwiftTreeSitter.git", branch: "main"),
7167
.package(url: "https://github.com/lukepistrol/tree-sitter-objc", branch: "feature/spm"),
7268
],
7369
targets: [
@@ -197,7 +193,6 @@ let package = Package(
197193
"SuggestionModel",
198194
.product(name: "SwiftTreeSitter", package: "SwiftTreeSitter"),
199195
.product(name: "TreeSitterObjC", package: "tree-sitter-objc"),
200-
.product(name: "TreeSitterSwift", package: "tree-sitter-swift"),
201196
]),
202197

203198
.testTarget(name: "ASTParserTests", dependencies: ["ASTParser"]),
@@ -228,10 +223,15 @@ let package = Package(
228223
dependencies: [
229224
"Preferences",
230225
"ASTParser",
226+
"SuggestionModel",
231227
.product(name: "SwiftSyntax", package: "swift-syntax"),
232228
.product(name: "SwiftParser", package: "swift-syntax"),
233229
]
234230
),
231+
.testTarget(
232+
name: "FocusedCodeFinderTests",
233+
dependencies: ["FocusedCodeFinder"]
234+
),
235235

236236
// MARK: - Services
237237

Tool/Sources/ASTParser/ASTParser.swift

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,12 @@ import SuggestionModel
22
import SwiftTreeSitter
33
import tree_sitter
44
import TreeSitterObjC
5-
import TreeSitterSwift
65

76
public enum ParsableLanguage {
8-
case swift
97
case objectiveC
108

119
var tsLanguage: UnsafeMutablePointer<TSLanguage> {
1210
switch self {
13-
case .swift:
14-
return tree_sitter_swift()
1511
case .objectiveC:
1612
return tree_sitter_objc()
1713
}
Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
import Foundation
2+
import SwiftTreeSitter
3+
4+
public enum ASTTreeVisitorContinueKind {
5+
/// The visitor should visit the descendants of the current node.
6+
case visitChildren
7+
/// The visitor should avoid visiting the descendants of the current node.
8+
case skipChildren
9+
}
10+
11+
// A SwiftSyntax style tree visitor.
12+
open class ASTTreeVisitor {
13+
public let tree: ASTTree
14+
15+
public init(tree: ASTTree) {
16+
self.tree = tree
17+
}
18+
19+
public func walk() {
20+
guard let cursor = tree.rootNode?.treeCursor else { return }
21+
visit(cursor)
22+
}
23+
24+
public func walk(_ node: ASTNode) {
25+
let cursor = node.treeCursor
26+
visit(cursor)
27+
}
28+
29+
open func visit(_: ASTNode) -> ASTTreeVisitorContinueKind {
30+
// do nothing
31+
return .skipChildren
32+
}
33+
34+
open func visitPost(_: ASTNode) {
35+
// do nothing
36+
}
37+
38+
private func visit(_ cursor: TreeCursor) {
39+
guard let currentNode = cursor.currentNode else { return }
40+
let continueKind = visit(currentNode)
41+
42+
switch continueKind {
43+
case .skipChildren:
44+
visitPost(currentNode)
45+
case .visitChildren:
46+
visitChildren(cursor)
47+
visitPost(currentNode)
48+
}
49+
}
50+
51+
private func visitChildren(_ cursor: TreeCursor) {
52+
let hasChild = cursor.goToFirstChild()
53+
guard hasChild else { return }
54+
visit(cursor)
55+
while cursor.goToNextSibling() {
56+
visit(cursor)
57+
}
58+
_ = cursor.gotoParent()
59+
}
60+
}
61+
Lines changed: 35 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -1,76 +1,55 @@
11
import SwiftTreeSitter
2+
import SwiftUI
23

34
public extension ASTTree {
45
/// Dumps the syntax tree as a string, for debugging purposes.
5-
func dump() -> String {
6+
func dump() -> AttributedString {
67
guard let tree, let root = tree.rootNode else { return "" }
7-
var result = ""
8+
var result: AttributedString = ""
89

9-
let appendNode: (_ level: Int, _ node: Node) -> Void = { level, node in
10+
let appendNode: (_ level: Int, _ node: Node, _ name: String) -> Void = {
11+
level, node, name in
1012
let range = node.pointRange
1113
let lowerBoundL = range.lowerBound.row
1214
let lowerBoundC = range.lowerBound.column / 2
1315
let upperBoundL = range.upperBound.row
1416
let upperBoundC = range.upperBound.column / 2
15-
let line =
16-
"\(String(repeating: " ", count: level))\(node.nodeType ?? "N/A") [\(lowerBoundL), \(lowerBoundC)] - [\(upperBoundL), \(upperBoundC)]"
17-
result += line + "\n"
17+
let indentation = AttributedString(String(repeating: " ", count: level))
18+
let nodeInfo = {
19+
if name.isEmpty {
20+
return AttributedString(node.nodeType ?? "N/A", attributes: .init([
21+
.foregroundColor: NSColor.blue,
22+
]))
23+
} else {
24+
var string = AttributedString("\(name): ", attributes: .init([
25+
.foregroundColor: NSColor.brown,
26+
]))
27+
string.append(AttributedString(node.nodeType ?? "N/A", attributes: .init([
28+
.foregroundColor: NSColor.blue,
29+
])))
30+
return string
31+
}
32+
}()
33+
let rangeText = "[\(lowerBoundL), \(lowerBoundC)] - [\(upperBoundL), \(upperBoundC)]"
34+
35+
var line: AttributedString = ""
36+
line.append(indentation)
37+
line.append(nodeInfo)
38+
line.append(AttributedString(" \(rangeText)\n"))
39+
40+
result.append(line)
1841
}
1942

20-
guard let node = root.descendant(in: root.byteRange) else { return result }
21-
22-
appendNode(0, node)
23-
24-
let cursor = node.treeCursor
25-
let level = 0
26-
27-
if cursor.goToFirstChild(for: node.byteRange.lowerBound) == false {
28-
return result
29-
}
30-
31-
cursor.enumerateCurrentAndDescendents(level: level + 1) { level, node in
32-
appendNode(level, node)
33-
}
34-
35-
while cursor.goToNextSibling() {
36-
guard let node = cursor.currentNode else {
37-
assertionFailure("no current node when gotoNextSibling succeeded")
38-
break
39-
}
40-
41-
// once we are past the interesting range, stop
42-
if node.byteRange.lowerBound > root.byteRange.upperBound {
43-
break
44-
}
45-
46-
cursor.enumerateCurrentAndDescendents(level: level + 1) { level, node in
47-
appendNode(level, node)
43+
func enumerate(_ node: Node, level: Int, name: String) {
44+
appendNode(level, node, name)
45+
for i in 0..<node.childCount {
46+
let n = node.child(at: i)!
47+
enumerate(n, level: level + 1, name: node.fieldNameForChild(at: i) ?? "")
4848
}
4949
}
5050

51+
enumerate(root, level: 0, name: "root")
5152
return result
5253
}
5354
}
5455

55-
private extension TreeCursor {
56-
func enumerateCurrentAndDescendents(level: Int, block: (Int, Node) throws -> Void) rethrows {
57-
if let node = currentNode {
58-
try block(level, node)
59-
}
60-
61-
if goToFirstChild() == false {
62-
return
63-
}
64-
65-
try enumerateCurrentAndDescendents(level: level + 1, block: block)
66-
67-
while goToNextSibling() {
68-
try enumerateCurrentAndDescendents(level: level + 1, block: block)
69-
}
70-
71-
let success = gotoParent()
72-
73-
assert(success)
74-
}
75-
}
76-

Tool/Sources/ASTParser/TreeCursor.swift

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import Foundation
22
import SwiftTreeSitter
33

4-
extension TreeCursor {
4+
public extension TreeCursor {
55
/// Deep first search nodes.
66
/// - Parameter skipChildren: Check if children of a `Node` should be skipped.
77
func deepFirstSearch(
@@ -13,7 +13,7 @@ extension TreeCursor {
1313

1414
// MARK: - Search
1515

16-
protocol Cursor {
16+
public protocol Cursor {
1717
associatedtype Node
1818
var currentNode: Node? { get }
1919
func goToFirstChild() -> Bool
@@ -22,32 +22,32 @@ protocol Cursor {
2222
}
2323

2424
extension TreeCursor: Cursor {
25-
func goToNextSibling() -> Bool {
25+
public func goToNextSibling() -> Bool {
2626
gotoNextSibling()
2727
}
28-
29-
func goToParent() -> Bool {
28+
29+
public func goToParent() -> Bool {
3030
gotoParent()
3131
}
3232
}
3333

34-
struct CursorDeepFirstSearchSequence<C: Cursor>: Sequence {
34+
public struct CursorDeepFirstSearchSequence<C: Cursor>: Sequence {
3535
let cursor: C
3636
let skipChildren: (C.Node) -> Bool
3737

38-
func makeIterator() -> CursorDeepFirstSearchIterator<C> {
38+
public func makeIterator() -> CursorDeepFirstSearchIterator {
3939
return CursorDeepFirstSearchIterator(
4040
cursor: cursor,
4141
skipChildren: skipChildren
4242
)
4343
}
4444

45-
struct CursorDeepFirstSearchIterator<C: Cursor>: IteratorProtocol {
45+
public struct CursorDeepFirstSearchIterator: IteratorProtocol {
4646
let cursor: C
4747
let skipChildren: (C.Node) -> Bool
4848
var isEnded = false
4949

50-
mutating func next() -> C.Node? {
50+
public mutating func next() -> C.Node? {
5151
guard !isEnded else { return nil }
5252
let currentNode = cursor.currentNode
5353
let hasChild = {

Tool/Sources/ChatContextCollectors/ActiveDocumentChatContextCollector/ActiveDocumentChatContextCollector.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ public final class ActiveDocumentChatContextCollector: ChatContextCollector {
7474

7575
func getActiveDocumentContext(_ info: EditorInformation) -> ActiveDocumentContext {
7676
var activeDocumentContext = activeDocumentContext ?? .init(
77-
filePath: "",
77+
documentURL: .init(fileURLWithPath: "/"),
7878
relativePath: "",
7979
language: .builtIn(.swift),
8080
fileContent: "",

Tool/Sources/FocusedCodeFinder/ActiveDocumentContext.swift

Lines changed: 8 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import Foundation
22
import SuggestionModel
33

44
public struct ActiveDocumentContext {
5-
public var filePath: String
5+
public var documentURL: URL
66
public var relativePath: String
77
public var language: CodeLanguage
88
public var fileContent: String
@@ -52,7 +52,7 @@ public struct ActiveDocumentContext {
5252
public var focusedContext: FocusedContext?
5353

5454
public init(
55-
filePath: String,
55+
documentURL: URL,
5656
relativePath: String,
5757
language: CodeLanguage,
5858
fileContent: String,
@@ -63,7 +63,7 @@ public struct ActiveDocumentContext {
6363
imports: [String],
6464
focusedContext: FocusedContext? = nil
6565
) {
66-
self.filePath = filePath
66+
self.documentURL = documentURL
6767
self.relativePath = relativePath
6868
self.language = language
6969
self.fileContent = fileContent
@@ -92,18 +92,12 @@ public struct ActiveDocumentContext {
9292
}
9393

9494
public mutating func moveToCodeContainingRange(_ range: CursorRange) {
95-
let finder: FocusedCodeFinder = {
96-
switch language {
97-
case .builtIn(.swift):
98-
return SwiftFocusedCodeFinder()
99-
default:
100-
return UnknownLanguageFocusedCodeFinder(proposedSearchRange: 5)
101-
}
102-
}()
103-
95+
let finder = FocusedCodeFinder()
96+
10497
let codeContext = finder.findFocusedCode(
98+
in: .init(documentURL: documentURL, content: fileContent, lines: lines),
10599
containingRange: range,
106-
activeDocumentContext: self
100+
language: language
107101
)
108102

109103
imports = codeContext.imports
@@ -141,7 +135,7 @@ public struct ActiveDocumentContext {
141135
return false
142136
}()
143137

144-
filePath = info.documentURL.path
138+
documentURL = info.documentURL
145139
relativePath = info.relativePath
146140
language = info.language
147141
fileContent = info.editorContent?.content ?? ""

0 commit comments

Comments
 (0)