Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -143,43 +143,71 @@ extension VocabularyRescorer {
return Array(repeating: nil, count: baseWords.count)
}

var alignments: [[Range<String.Index>]] = []

func search(
wordIndex: Int,
cursor: String.Index,
ranges: [Range<String.Index>]
) {
guard alignments.count < 2 else { return }
guard wordIndex < baseWords.count else {
guard isDelimiterOnly(baseText[cursor..<baseText.endIndex]) else { return }
alignments.append(ranges)
return
}
// Iterative backtracking: recursion depth and per-level path copies scaled with word count (#961).
// A word may only start inside the delimiter run after the cursor, so each level has few candidates.
struct State: Hashable {
let wordIndex: Int
let cursor: String.Index
}
struct Frame {
let state: State
let limit: String.Index
var next: String.Index?
let alignmentsAtEntry: Int
}

let word = baseWords[wordIndex]
var searchStart = cursor
while searchStart < baseText.endIndex,
let match = baseText.range(
of: word,
options: .literal,
range: searchStart..<baseText.endIndex
)
{
if isDelimiterOnly(baseText[cursor..<match.lowerBound]) {
search(
wordIndex: wordIndex + 1,
cursor: match.upperBound,
ranges: ranges + [match]
)
var alignmentCount = 0
var alignment: [Range<String.Index>] = []
var path: [Range<String.Index>] = []
var deadStates = Set<State>()

func frame(wordIndex: Int, cursor: String.Index) -> Frame {
let limit = baseText[cursor...].firstIndex(where: { !isInterWordDelimiter($0) }) ?? baseText.endIndex
return Frame(
state: State(wordIndex: wordIndex, cursor: cursor),
limit: limit,
next: cursor,
alignmentsAtEntry: alignmentCount
)
}

var stack = [frame(wordIndex: 0, cursor: baseText.startIndex)]
path.reserveCapacity(baseWords.count)

while alignmentCount < 2, let top = stack.last {
let wordIndex = top.state.wordIndex
if wordIndex == baseWords.count {
if top.limit == baseText.endIndex {
alignmentCount += 1
if alignmentCount == 1 { alignment = path }
}
guard alignments.count < 2, match.lowerBound < baseText.endIndex else { return }
searchStart = baseText.index(after: match.lowerBound)
stack.removeLast()
if !path.isEmpty { path.removeLast() }
continue
}

guard let start = top.next, start < baseText.endIndex else {
if alignmentCount == top.alignmentsAtEntry { deadStates.insert(top.state) }
stack.removeLast()
if !path.isEmpty { path.removeLast() }
continue
}
stack[stack.count - 1].next = start < top.limit ? baseText.index(after: start) : nil

guard
let match = baseText.range(
of: baseWords[wordIndex],
options: [.literal, .anchored],
range: start..<baseText.endIndex
),
!deadStates.contains(State(wordIndex: wordIndex + 1, cursor: match.upperBound))
else { continue }

path.append(match)
stack.append(frame(wordIndex: wordIndex + 1, cursor: match.upperBound))
}

search(wordIndex: 0, cursor: baseText.startIndex, ranges: [])
guard alignments.count == 1, let alignment = alignments.first else {
guard alignmentCount == 1 else {
return Array(repeating: nil, count: baseWords.count)
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -301,6 +301,23 @@ final class VocabularyCandidateEvidenceTests: XCTestCase {
XCTAssertNil(ambiguous[0])
}

func testLongTranscriptAlignsOffTheMainThreadWithoutStackOverflow() async {
let vocabulary = ["so", "we", "should", "um", "go", "go", "C++,", "re-enter", "don't", "“remote”"]
let baseWords = (0..<20_000).map { vocabulary[$0 * 7 % vocabulary.count] }
let baseText = baseWords.joined(separator: " ")

let ranges = await Task.detached {
VocabularyRescorer.alignBaseWordsToUTF8Ranges(baseText: baseText, baseWords: baseWords)
}.value
XCTAssertEqual(ranges.count, baseWords.count)
XCTAssertTrue(ranges.allSatisfy { $0 != nil })

let truncated = await Task.detached {
VocabularyRescorer.alignBaseWordsToUTF8Ranges(baseText: baseText, baseWords: Array(baseWords.dropLast()))
}.value
XCTAssertTrue(truncated.allSatisfy { $0 == nil }, "A trailing unmatched word must fail closed")
}

func testMissingReorderedAndNormalizationOnlyMatchesFailClosed() {
let cases: [(String, [String])] = [
("one two", ["two", "one"]),
Expand Down
Loading