Windows.swift 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361
  1. //===----------------------------------------------------------------------===//
  2. //
  3. // This source file is part of the Swift Algorithms open source project
  4. //
  5. // Copyright (c) 2020 Apple Inc. and the Swift project authors
  6. // Licensed under Apache License v2.0 with Runtime Library Exception
  7. //
  8. // See https://swift.org/LICENSE.txt for license information
  9. //
  10. //===----------------------------------------------------------------------===//
  11. //===----------------------------------------------------------------------===//
  12. // windows(ofCount:)
  13. //===----------------------------------------------------------------------===//
  14. extension Collection {
  15. /// Returns a collection of all the overlapping slices of a given size.
  16. ///
  17. /// Use this method to iterate over overlapping subsequences of this
  18. /// collection. This example prints every five character substring of `str`:
  19. ///
  20. /// let str = "Hello, world!"
  21. /// for substring in str.windows(ofCount: 5) {
  22. /// print(substring)
  23. /// }
  24. /// // "Hello"
  25. /// // "ello,"
  26. /// // "llo, "
  27. /// // "lo, W"
  28. /// // ...
  29. /// // "orld!"
  30. ///
  31. /// - Parameter count: The number of elements in each window subsequence.
  32. /// - Returns: A collection of subsequences of this collection, each with
  33. /// length `count`. If this collection is shorter than `count`, the
  34. /// resulting collection is empty.
  35. ///
  36. /// - Complexity: O(1) if the collection conforms to
  37. /// `RandomAccessCollection`, otherwise O(*k*) where `k` is `count`.
  38. /// Access to successive windows is O(1).
  39. @inlinable
  40. public func windows(ofCount count: Int) -> WindowsOfCountCollection<Self> {
  41. WindowsOfCountCollection(base: self, windowSize: count)
  42. }
  43. }
  44. /// A collection wrapper that presents a sliding window over the elements of
  45. /// a collection.
  46. public struct WindowsOfCountCollection<Base: Collection> {
  47. @usableFromInline
  48. internal let base: Base
  49. @usableFromInline
  50. internal let windowSize: Int
  51. @usableFromInline
  52. internal var endOfFirstWindow: Base.Index?
  53. @inlinable
  54. internal init(base: Base, windowSize: Int) {
  55. precondition(windowSize > 0, "Windows size must be greater than zero")
  56. self.base = base
  57. self.windowSize = windowSize
  58. self.endOfFirstWindow =
  59. base.index(base.startIndex, offsetBy: windowSize, limitedBy: base.endIndex)
  60. }
  61. }
  62. extension WindowsOfCountCollection: Collection {
  63. /// A position in a `WindowsOfCountCollection` instance.
  64. public struct Index: Comparable {
  65. @usableFromInline
  66. internal var lowerBound: Base.Index
  67. @usableFromInline
  68. internal var upperBound: Base.Index
  69. @inlinable
  70. internal init(lowerBound: Base.Index, upperBound: Base.Index) {
  71. self.lowerBound = lowerBound
  72. self.upperBound = upperBound
  73. }
  74. @inlinable
  75. public static func == (lhs: Index, rhs: Index) -> Bool {
  76. lhs.lowerBound == rhs.lowerBound
  77. }
  78. @inlinable
  79. public static func < (lhs: Index, rhs: Index) -> Bool {
  80. lhs.lowerBound < rhs.lowerBound
  81. }
  82. }
  83. @inlinable
  84. public var startIndex: Index {
  85. if let upperBound = endOfFirstWindow {
  86. return Index(lowerBound: base.startIndex, upperBound: upperBound)
  87. } else {
  88. return endIndex
  89. }
  90. }
  91. @inlinable
  92. public var endIndex: Index {
  93. Index(lowerBound: base.endIndex, upperBound: base.endIndex)
  94. }
  95. @inlinable
  96. public subscript(index: Index) -> Base.SubSequence {
  97. precondition(
  98. index.lowerBound != index.upperBound,
  99. "Windows index is out of range")
  100. return base[index.lowerBound..<index.upperBound]
  101. }
  102. @inlinable
  103. public func index(after index: Index) -> Index {
  104. precondition(index != endIndex, "Advancing past end index")
  105. guard index.upperBound < base.endIndex else { return endIndex }
  106. let lowerBound = windowSize == 1
  107. ? index.upperBound
  108. : base.index(after: index.lowerBound)
  109. let upperBound = base.index(after: index.upperBound)
  110. return Index(lowerBound: lowerBound, upperBound: upperBound)
  111. }
  112. @inlinable
  113. public func index(_ i: Index, offsetBy distance: Int) -> Index {
  114. guard distance != 0 else { return i }
  115. return distance > 0
  116. ? offsetForward(i, by: distance)
  117. : offsetBackward(i, by: -distance)
  118. }
  119. @inlinable
  120. public func index(
  121. _ i: Index,
  122. offsetBy distance: Int,
  123. limitedBy limit: Index
  124. ) -> Index? {
  125. guard distance != 0 else { return i }
  126. guard limit != i else { return nil }
  127. if distance > 0 {
  128. return limit > i
  129. ? offsetForward(i, by: distance, limitedBy: limit)
  130. : offsetForward(i, by: distance)
  131. } else {
  132. return limit < i
  133. ? offsetBackward(i, by: -distance, limitedBy: limit)
  134. : offsetBackward(i, by: -distance)
  135. }
  136. }
  137. @inlinable
  138. internal func offsetForward(_ i: Index, by distance: Int) -> Index {
  139. guard let index = offsetForward(i, by: distance, limitedBy: endIndex)
  140. else { fatalError("Index is out of bounds") }
  141. return index
  142. }
  143. @inlinable
  144. internal func offsetBackward(_ i: Index, by distance: Int) -> Index {
  145. guard let index = offsetBackward(i, by: distance, limitedBy: startIndex)
  146. else { fatalError("Index is out of bounds") }
  147. return index
  148. }
  149. @inlinable
  150. internal func offsetForward(
  151. _ i: Index, by distance: Int, limitedBy limit: Index
  152. ) -> Index? {
  153. assert(distance > 0)
  154. assert(limit > i)
  155. // `endIndex` and the index before it both have `base.endIndex` as their
  156. // upper bound, so we first advance to the base index _before_ the upper
  157. // bound of the output, in order to avoid advancing past the end of `base`
  158. // when advancing to `endIndex`.
  159. //
  160. // Advancing by 4:
  161. //
  162. // input: [x|x x x x x|x x x x] [x x|x x x x x|x x x]
  163. // |> > >|>| or |> > >|
  164. // output: [x x x x x|x x x x x] [x x x x x x x x x x] (`endIndex`)
  165. if distance >= windowSize {
  166. // Avoid traversing `self[i.lowerBound..<i.upperBound]` when the lower
  167. // bound of the output is greater than or equal to the upper bound of the
  168. // input.
  169. // input: [x|x x x x|x x x x x x x]
  170. // |> >|> > >|>|
  171. // output: [x x x x x x x|x x x x|x]
  172. guard limit.lowerBound >= i.upperBound,
  173. let lowerBound = base.index(
  174. i.upperBound,
  175. offsetBy: distance - windowSize,
  176. limitedBy: limit.lowerBound),
  177. let indexBeforeUpperBound = base.index(
  178. lowerBound,
  179. offsetBy: windowSize - 1,
  180. limitedBy: limit.upperBound)
  181. else { return nil }
  182. // If `indexBeforeUpperBound` equals `base.endIndex`, we're advancing to
  183. // `endIndex`.
  184. guard indexBeforeUpperBound != base.endIndex else { return endIndex }
  185. return Index(
  186. lowerBound: lowerBound,
  187. upperBound: base.index(after: indexBeforeUpperBound))
  188. } else {
  189. // input: [x|x x x x x x|x x x x x]
  190. // |> > > >| |> > >|>|
  191. // output: [x x x x x|x x x x x x|x]
  192. guard let indexBeforeUpperBound = base.index(
  193. i.upperBound,
  194. offsetBy: distance - 1,
  195. limitedBy: limit.upperBound)
  196. else { return nil }
  197. // If `indexBeforeUpperBound` equals the limit, the upper bound itself
  198. // exceeds it.
  199. guard indexBeforeUpperBound != limit.upperBound || limit == endIndex
  200. else { return nil }
  201. // If `indexBeforeUpperBound` equals `base.endIndex`, we're advancing to
  202. // `endIndex`.
  203. guard indexBeforeUpperBound != base.endIndex else { return endIndex }
  204. return Index(
  205. lowerBound: base.index(i.lowerBound, offsetBy: distance),
  206. upperBound: base.index(after: indexBeforeUpperBound))
  207. }
  208. }
  209. @inlinable
  210. internal func offsetBackward(
  211. _ i: Index, by distance: Int, limitedBy limit: Index
  212. ) -> Index? {
  213. assert(distance > 0)
  214. assert(limit < i)
  215. if i == endIndex {
  216. // Advance `base.endIndex` by `distance - 1`, because the index before
  217. // `endIndex` also has `base.endIndex` as its upper bound.
  218. //
  219. // Advancing by 4:
  220. //
  221. // input: [x x x x x x x x x x] (`endIndex`)
  222. // |< < < < <|< < <|
  223. // output: [x x|x x x x x|x x x]
  224. guard let upperBound = base.index(
  225. base.endIndex,
  226. offsetBy: -(distance - 1),
  227. limitedBy: limit.upperBound)
  228. else { return nil }
  229. return Index(
  230. lowerBound: base.index(upperBound, offsetBy: -windowSize),
  231. upperBound: upperBound)
  232. } else if distance >= windowSize {
  233. // Avoid traversing `self[i.lowerBound..<i.upperBound]` when the upper
  234. // bound of the output is less than or equal to the lower bound of the
  235. // input.
  236. //
  237. // input: [x x x x x x x|x x x x|x]
  238. // |< < < <|< <|
  239. // output: [x|x x x x|x x x x x x x]
  240. guard limit.upperBound <= i.lowerBound,
  241. let upperBound = base.index(
  242. i.lowerBound,
  243. offsetBy: -(distance - windowSize),
  244. limitedBy: limit.upperBound)
  245. else { return nil }
  246. return Index(
  247. lowerBound: base.index(upperBound, offsetBy: -windowSize),
  248. upperBound: upperBound)
  249. } else {
  250. // input: [x x x x x|x x x x x x|x]
  251. // |< < < <| |< < < <|
  252. // output: [x|x x x x x x|x x x x x]
  253. guard let lowerBound = base.index(
  254. i.lowerBound,
  255. offsetBy: -distance,
  256. limitedBy: limit.lowerBound)
  257. else { return nil }
  258. return Index(
  259. lowerBound: lowerBound,
  260. upperBound: base.index(i.lowerBound, offsetBy: -distance))
  261. }
  262. }
  263. @inlinable
  264. public func distance(from start: Index, to end: Index) -> Int {
  265. guard start <= end else { return -distance(from: end, to: start) }
  266. guard start != end else { return 0 }
  267. guard end != endIndex else {
  268. // We add 1 here because the index before `endIndex` also has
  269. // `base.endIndex` as its upper bound.
  270. return base[start.upperBound...].count + 1
  271. }
  272. if start.upperBound <= end.lowerBound {
  273. // The distance between `start.lowerBound` and `start.upperBound` is
  274. // already known.
  275. //
  276. // start: [x|x x x x|x x x x x x x]
  277. // |- - - -|> >|
  278. // end: [x x x x x x x|x x x x|x]
  279. return windowSize + base[start.upperBound..<end.lowerBound].count
  280. } else {
  281. // start: [x|x x x x x x|x x x x x]
  282. // |> > > >|
  283. // end: [x x x x x|x x x x x x|x]
  284. return base[start.lowerBound..<end.lowerBound].count
  285. }
  286. }
  287. }
  288. extension WindowsOfCountCollection: BidirectionalCollection
  289. where Base: BidirectionalCollection
  290. {
  291. @inlinable
  292. public func index(before index: Index) -> Index {
  293. precondition(index != startIndex, "Incrementing past start index")
  294. if index == endIndex {
  295. return Index(
  296. lowerBound: base.index(index.lowerBound, offsetBy: -windowSize),
  297. upperBound: index.upperBound
  298. )
  299. } else {
  300. return Index(
  301. lowerBound: base.index(before: index.lowerBound),
  302. upperBound: base.index(before: index.upperBound)
  303. )
  304. }
  305. }
  306. }
  307. extension WindowsOfCountCollection: RandomAccessCollection
  308. where Base: RandomAccessCollection {}
  309. extension WindowsOfCountCollection: LazySequenceProtocol, LazyCollectionProtocol
  310. where Base: LazySequenceProtocol {}
  311. extension WindowsOfCountCollection.Index: Hashable where Base.Index: Hashable {}