Cache word indices

Author: Azoy

Sources/_StringProcessing/ByteCodeGen.swift

@@ -118,12 +118,13 @@ fileprivate extension Compiler.ByteCodeGen {
     // need to supply both a slice bounds and a per-search bounds.
     switch kind {
     case .startOfSubject:
-      builder.buildAssert { (input, pos, subjectBounds) in
+      builder.buildAssert { (_, _, input, pos, subjectBounds) in
         pos == subjectBounds.lowerBound
       }
 
     case .endOfSubjectBeforeNewline:
-      builder.buildAssert { [semanticLevel = options.semanticLevel] (input, pos, subjectBounds) in
+      builder.buildAssert { [semanticLevel = options.semanticLevel]
+          (_, _, input, pos, subjectBounds) in
         if pos == subjectBounds.upperBound { return true }
         switch semanticLevel {
         case .graphemeCluster:
@@ -136,7 +137,7 @@ fileprivate extension Compiler.ByteCodeGen {
       }
 
     case .endOfSubject:
-      builder.buildAssert { (input, pos, subjectBounds) in
+      builder.buildAssert { (_, _, input, pos, subjectBounds) in
         pos == subjectBounds.upperBound
       }
 
@@ -149,16 +150,16 @@ fileprivate extension Compiler.ByteCodeGen {
 
       // FIXME: This needs to be based on `searchBounds`,
       // not the `subjectBounds` given as an argument here
-      builder.buildAssert { (input, pos, subjectBounds) in false }
+      builder.buildAssert { (_, _, input, pos, subjectBounds) in false }
 
     case .textSegment:
-      builder.buildAssert { (input, pos, _) in
+      builder.buildAssert { (_, _, input, pos, _) in
         // FIXME: Grapheme or word based on options
         input.isOnGraphemeClusterBoundary(pos)
       }
 
     case .notTextSegment:
-      builder.buildAssert { (input, pos, _) in
+      builder.buildAssert { (_, _, input, pos, _) in
         // FIXME: Grapheme or word based on options
         !input.isOnGraphemeClusterBoundary(pos)
       }
@@ -169,7 +170,8 @@ fileprivate extension Compiler.ByteCodeGen {
       // the DSL-based `.startOfLine` anchor should always match the start
       // of a line. Right now we don't distinguish between those anchors.
       if options.anchorsMatchNewlines {
-        builder.buildAssert { [semanticLevel = options.semanticLevel] (input, pos, subjectBounds) in
+        builder.buildAssert { [semanticLevel = options.semanticLevel]
+            (_, _, input, pos, subjectBounds) in
           if pos == subjectBounds.lowerBound { return true }
           switch semanticLevel {
           case .graphemeCluster:
@@ -179,7 +181,7 @@ fileprivate extension Compiler.ByteCodeGen {
           }
         }
       } else {
-        builder.buildAssert { (input, pos, subjectBounds) in
+        builder.buildAssert { (_, _, input, pos, subjectBounds) in
           pos == subjectBounds.lowerBound
         }
       }
@@ -190,7 +192,8 @@ fileprivate extension Compiler.ByteCodeGen {
       // the DSL-based `.endOfLine` anchor should always match the end
       // of a line. Right now we don't distinguish between those anchors.
       if options.anchorsMatchNewlines {
-        builder.buildAssert { [semanticLevel = options.semanticLevel] (input, pos, subjectBounds) in
+        builder.buildAssert { [semanticLevel = options.semanticLevel]
+            (_, _, input, pos, subjectBounds) in
           if pos == subjectBounds.upperBound { return true }
           switch semanticLevel {
           case .graphemeCluster:
@@ -200,13 +203,14 @@ fileprivate extension Compiler.ByteCodeGen {
           }
         }
       } else {
-        builder.buildAssert { (input, pos, subjectBounds) in
+        builder.buildAssert { (_, _, input, pos, subjectBounds) in
           pos == subjectBounds.upperBound
         }
       }
 
     case .wordBoundary:
-      builder.buildAssert { [options] (input, pos, subjectBounds) in
+      builder.buildAssert { [options]
+          (cache, maxIndex, input, pos, subjectBounds) in
         if options.usesSimpleUnicodeBoundaries {
           // TODO: How should we handle bounds?
           return _CharacterClassModel.word.isBoundary(
@@ -216,12 +220,13 @@ fileprivate extension Compiler.ByteCodeGen {
             with: options
           )
         } else {
-          return input.isOnWordBoundary(at: pos)
+          return input.isOnWordBoundary(at: pos, using: &cache, &maxIndex)
         }
       }
 
     case .notWordBoundary:
-      builder.buildAssert { [options] (input, pos, subjectBounds) in
+      builder.buildAssert { [options]
+          (cache, maxIndex, input, pos, subjectBounds) in
         if options.usesSimpleUnicodeBoundaries {
           // TODO: How should we handle bounds?
           return !_CharacterClassModel.word.isBoundary(
@@ -231,7 +236,7 @@ fileprivate extension Compiler.ByteCodeGen {
             with: options
           )
         } else {
-          return !input.isOnWordBoundary(at: pos)
+          return !input.isOnWordBoundary(at: pos, using: &cache, &maxIndex)
         }
       }
     }

Sources/_StringProcessing/Engine/MEProgram.swift

@@ -16,7 +16,13 @@ struct MEProgram {
 
   typealias ConsumeFunction = (Input, Range<Input.Index>) -> Input.Index?
   typealias AssertionFunction =
-    (Input, Input.Index, Range<Input.Index>) throws -> Bool
+    (
+      inout Set<String.Index>?,
+      inout String.Index?,
+      Input,
+      Input.Index,
+      Range<Input.Index>
+    ) throws -> Bool
   typealias TransformFunction =
     (Input, Processor._StoredCapture) throws -> Any?
   typealias MatcherFunction =

Sources/_StringProcessing/Engine/Processor.swift

@@ -78,6 +78,9 @@ struct Processor {
 
   var storedCaptures: Array<_StoredCapture>
 
+  var wordIndexCache: Set<String.Index>? = nil
+  var wordIndexMaxIndex: String.Index? = nil
+  
   var state: State = .inProgress
 
   var failureReason: Error? = nil
@@ -401,7 +404,13 @@ extension Processor {
       let reg = payload.assertion
       let assertion = registers[reg]
       do {
-        guard try assertion(input, currentPosition, subjectBounds) else {
+        guard try assertion(
+          &wordIndexCache,
+          &wordIndexMaxIndex,
+          input,
+          currentPosition,
+          subjectBounds
+        ) else {
           signalFailure()
           return
         }

Sources/_StringProcessing/Unicode/WordBreaking.swift

@@ -13,20 +13,42 @@
 import Swift
 
 extension String {
-  func isOnWordBoundary(at i: String.Index) -> Bool {
+  func isOnWordBoundary(
+    at i: String.Index,
+    using cache: inout Set<String.Index>?,
+    _ maxIndex: inout String.Index?
+  ) -> Bool {
     guard i != startIndex, i != endIndex else {
       return true
     }
 
+    // If our index is already in our cache, then this is obviously on a
+    // boundary.
+    if let cache = cache, cache.contains(i) {
+      return true
+    }
+    
+    // If its not in the cache AND our max index is larger than our index, it
+    // means this index is never on a word boundary in our string. If our index
+    // is larger than max index, we may need to still do work to determine if
+    // i is on a boundary. If it's equal to max index, then it should've been
+    // taken the cache path.
+    if let maxIndex = maxIndex, i < maxIndex {
+      return false
+    }
+    
     if #available(SwiftStdlib 5.7, *) {
       var indices: Set<String.Index> = []
-      var j = startIndex
+      var j = maxIndex ?? startIndex
 
       while j < endIndex, j <= i {
         indices.insert(j)
         j = _wordIndex(after: j)
       }
 
+      cache = indices
+      maxIndex = j
+      
       return indices.contains(i)
     } else {
       return false