Skip to content

Commit 5c4b432

Browse files
intitniintitni
committed
Update SuggestionInjector to support range replacement, fix tests
1 parent 74449f8 commit 5c4b432

File tree

2 files changed

+173
-37
lines changed

2 files changed

+173
-37
lines changed

‎Core/Sources/SuggestionInjector/SuggestionInjector.swift‎

Lines changed: 60 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@ let suggestionEnd = "*///======== End of Copilot Suggestion"
66

77
// NOTE: Every lines from Xcode Extension has a line break at its end, even the last line.
88
// NOTE: Copilot's completion always start at character 0, no matter where the cursor is.
9-
// NOTE: range.end and position in Copilot's completion are useless, don't bother looking at them.
109

1110
public struct SuggestionInjector {
1211
public init() {}
@@ -139,40 +138,69 @@ public struct SuggestionInjector {
139138
extraInfo.didChangeCursorPosition = true
140139
extraInfo.suggestionRange = nil
141140
let start = completion.range.start
141+
let end = completion.range.end
142142
let suggestionContent = completion.text
143143

144144
let existedLine = start.line < content.endIndex ? content[start.line] : nil
145145
let commonPrefix = longestCommonPrefix(of: suggestionContent, and: existedLine ?? "")
146146

147-
if let existedLine, existedLine.count > 1, !commonPrefix.isEmpty {
148-
extraInfo.modifications.append(.deleted(start.line...start.line))
149-
content.remove(at: start.line)
150-
} else if content.count > start.line,
151-
content[start.line].isEmpty || content[start.line] == "\n"
147+
let firstRemovedLine = content[safe: start.line]
148+
let lastRemovedLine = content[safe: end.line]
149+
let startLine = max(0, start.line)
150+
let endLine = max(start.line, min(end.line, content.endIndex - 1))
151+
extraInfo.modifications.append(.deleted(startLine...endLine))
152+
content.removeSubrange(startLine...endLine)
153+
154+
var toBeInserted = suggestionContent.breakLines(appendLineBreakToLastLine: true)
155+
156+
if let firstRemovedLine,
157+
!firstRemovedLine.isEmptyOrNewLine,
158+
start.character > 0,
159+
start.character < firstRemovedLine.count,
160+
!toBeInserted.isEmpty
152161
{
153-
extraInfo.modifications.append(.deleted(start.line...start.line))
154-
content.remove(at: start.line)
162+
let leftoverRange = firstRemovedLine.startIndex..<(firstRemovedLine.index(
163+
firstRemovedLine.startIndex,
164+
offsetBy: start.character,
165+
limitedBy: firstRemovedLine.endIndex
166+
) ?? firstRemovedLine.endIndex)
167+
var leftover = firstRemovedLine[leftoverRange]
168+
if leftover.hasSuffix("\n") {
169+
leftover.removeLast(1)
170+
}
171+
toBeInserted[0].insert(
172+
contentsOf: leftover,
173+
at: toBeInserted[0].startIndex
174+
)
155175
}
156176

157-
let toBeInserted = suggestionContent.breakLines(appendLineBreakToLastLine: true)
158-
if content.endIndex < start.line {
159-
extraInfo.modifications.append(.inserted(content.endIndex, toBeInserted))
160-
content.append(contentsOf: toBeInserted)
161-
cursorPosition = .init(
162-
line: toBeInserted.endIndex,
163-
character: (toBeInserted.last?.count ?? 1) - 1
164-
)
165-
} else {
166-
extraInfo.modifications.append(.inserted(start.line, toBeInserted))
167-
content.insert(
168-
contentsOf: toBeInserted,
169-
at: start.line
170-
)
171-
cursorPosition = .init(
172-
line: start.line + toBeInserted.count - 1,
173-
character: (toBeInserted.last?.count ?? 1) - 1
174-
)
177+
let cursorCol = toBeInserted[toBeInserted.endIndex - 1].count - 1
178+
if let lastRemovedLine,
179+
!lastRemovedLine.isEmptyOrNewLine,
180+
end.character >= 0,
181+
end.character - 1 < lastRemovedLine.count,
182+
!toBeInserted.isEmpty
183+
{
184+
let leftoverRange = (lastRemovedLine.index(
185+
lastRemovedLine.startIndex,
186+
offsetBy: end.character,
187+
limitedBy: lastRemovedLine.endIndex
188+
) ?? lastRemovedLine.endIndex)..<lastRemovedLine.endIndex
189+
if toBeInserted[toBeInserted.endIndex - 1].hasSuffix("\n") {
190+
toBeInserted[toBeInserted.endIndex - 1].removeLast(1)
191+
}
192+
let leftover = lastRemovedLine[leftoverRange]
193+
toBeInserted[toBeInserted.endIndex - 1]
194+
.append(contentsOf: leftover)
175195
}
196+
197+
let insertingIndex = min(start.line, content.endIndex)
198+
content.insert(contentsOf: toBeInserted, at: insertingIndex)
199+
extraInfo.modifications.append(.inserted(insertingIndex, toBeInserted))
200+
cursorPosition = .init(
201+
line: startLine + toBeInserted.count - 1,
202+
character: max(0, cursorCol)
203+
)
176204
}
177205
}
178206

@@ -225,3 +253,9 @@ func longestCommonPrefix(of a: String, and b: String) -> String {
225253

226254
return prefix
227255
}
256+
257+
extension Array {
258+
subscript(safe index: Index) -> Element? {
259+
indices.contains(index) ? self[index] : nil
260+
}
261+
}

‎Core/Tests/SuggestionInjectorTests/AcceptSuggestionTests.swift‎

Lines changed: 113 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,11 @@ final class AcceptSuggestionTests: XCTestCase {
1616
"""
1717
let suggestion = CopilotCompletion(
1818
text: text,
19-
position: .init(line: 2, character: 19),
19+
position: .init(line: 0, character: 1),
2020
uuid: "",
2121
range: .init(
2222
start: .init(line: 1, character: 0),
23-
end: .init(line: 2, character: 18)
23+
end: .init(line: 1, character: 0)
2424
),
2525
displayText: ""
2626
)
@@ -52,16 +52,17 @@ final class AcceptSuggestionTests: XCTestCase {
5252
}
5353
"""
5454
let text = """
55+
struct Cat {
5556
var name: String
5657
var age: String
5758
"""
5859
let suggestion = CopilotCompletion(
5960
text: text,
60-
position: .init(line: 2, character: 19),
61+
position: .init(line: 0, character: 12),
6162
uuid: "",
6263
range: .init(
63-
start: .init(line: 1, character: 0),
64-
end: .init(line: 2, character: 18)
64+
start: .init(line: 0, character: 0),
65+
end: .init(line: 0, character: 12)
6566
),
6667
displayText: ""
6768
)
@@ -100,11 +101,11 @@ final class AcceptSuggestionTests: XCTestCase {
100101
"""
101102
let suggestion = CopilotCompletion(
102103
text: text,
103-
position: .init(line: 2, character: 19),
104+
position: .init(line: 1, character: 12),
104105
uuid: "",
105106
range: .init(
106107
start: .init(line: 1, character: 0),
107-
end: .init(line: 2, character: 18)
108+
end: .init(line: 1, character: 12)
108109
),
109110
displayText: ""
110111
)
@@ -144,11 +145,11 @@ final class AcceptSuggestionTests: XCTestCase {
144145
"""
145146
let suggestion = CopilotCompletion(
146147
text: text,
147-
position: .init(line: 0, character: 0),
148+
position: .init(line: 0, character: 18),
148149
uuid: "",
149150
range: .init(
150151
start: .init(line: 0, character: 0),
151-
end: .init(line: 5, character: 15)
152+
end: .init(line: 0, character: 20)
152153
),
153154
displayText: ""
154155
)
@@ -191,11 +192,11 @@ final class AcceptSuggestionTests: XCTestCase {
191192
"""
192193
let suggestion = CopilotCompletion(
193194
text: text,
194-
position: .init(line: 0, character: 0),
195+
position: .init(line: 0, character: 18),
195196
uuid: "",
196197
range: .init(
197198
start: .init(line: 1, character: 0),
198-
end: .init(line: 5, character: 15)
199+
end: .init(line: 1, character: 0)
199200
),
200201
displayText: ""
201202
)
@@ -225,4 +226,105 @@ final class AcceptSuggestionTests: XCTestCase {
225226
226227
""")
227228
}
229+
230+
func test_replacing_multiple_lines() async throws {
231+
let content = """
232+
struct Cat {
233+
func speak() { print("meow") }
234+
}
235+
"""
236+
let text = """
237+
struct Dog {
238+
func speak() {
239+
print("woof")
240+
}
241+
}
242+
"""
243+
let suggestion = CopilotCompletion(
244+
text: text,
245+
position: .init(line: 0, character: 7),
246+
uuid: "",
247+
range: .init(
248+
start: .init(line: 0, character: 0),
249+
end: .init(line: 2, character: 1)
250+
),
251+
displayText: ""
252+
)
253+
254+
var extraInfo = SuggestionInjector.ExtraInfo()
255+
var lines = content.breakLines()
256+
var cursor = CursorPosition(line: 0, character: 0)
257+
SuggestionInjector().acceptSuggestion(
258+
intoContentWithoutSuggestion: &lines,
259+
cursorPosition: &cursor,
260+
completion: suggestion,
261+
extraInfo: &extraInfo
262+
)
263+
264+
XCTAssertTrue(extraInfo.didChangeContent)
265+
XCTAssertTrue(extraInfo.didChangeCursorPosition)
266+
XCTAssertNil(extraInfo.suggestionRange)
267+
XCTAssertEqual(lines, content.breakLines().applying(extraInfo.modifications))
268+
XCTAssertEqual(cursor, .init(line: 4, character: 1))
269+
XCTAssertEqual(lines.joined(separator: ""), text)
270+
}
271+
272+
func test_replacing_multiple_lines_in_the_middle() async throws {
273+
let content = """
274+
protocol Animal {
275+
func speak()
276+
}
277+
278+
struct Cat: Animal {
279+
func speak() { print("meow") }
280+
}
281+
282+
func foo() {}
283+
"""
284+
let text = """
285+
Dog {
286+
func speak() {
287+
print("woof")
288+
}
289+
"""
290+
let suggestion = CopilotCompletion(
291+
text: text,
292+
position: .init(line: 5, character: 34),
293+
uuid: "",
294+
range: .init(
295+
start: .init(line: 4, character: 7),
296+
end: .init(line: 5, character: 34)
297+
),
298+
displayText: ""
299+
)
300+
301+
var extraInfo = SuggestionInjector.ExtraInfo()
302+
var lines = content.breakLines()
303+
var cursor = CursorPosition(line: 0, character: 0)
304+
SuggestionInjector().acceptSuggestion(
305+
intoContentWithoutSuggestion: &lines,
306+
cursorPosition: &cursor,
307+
completion: suggestion,
308+
extraInfo: &extraInfo
309+
)
310+
311+
XCTAssertTrue(extraInfo.didChangeContent)
312+
XCTAssertTrue(extraInfo.didChangeCursorPosition)
313+
XCTAssertNil(extraInfo.suggestionRange)
314+
XCTAssertEqual(lines, content.breakLines().applying(extraInfo.modifications))
315+
XCTAssertEqual(cursor, .init(line: 7, character: 5))
316+
XCTAssertEqual(lines.joined(separator: ""), """
317+
protocol Animal {
318+
func speak()
319+
}
320+
321+
struct Dog {
322+
func speak() {
323+
print("woof")
324+
}
325+
}
326+
327+
func foo() {}
328+
""")
329+
}
228330
}

0 commit comments

Comments
 (0)