Stride.swift 7.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284
  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. // striding(by:)
  13. //===----------------------------------------------------------------------===//
  14. extension Sequence {
  15. /// Returns a sequence stepping through the elements every `step` starting at
  16. /// the first value. Any remainders of the stride will be trimmed.
  17. ///
  18. /// (0...10).striding(by: 2) // == [0, 2, 4, 6, 8, 10]
  19. /// (0...10).striding(by: 3) // == [0, 3, 6, 9]
  20. ///
  21. /// - Complexity: O(1). Access to successive values is O(k) where _k_ is the
  22. /// striding `step`.
  23. ///
  24. /// - Parameter step: The amount to step with each iteration.
  25. /// - Returns: Returns a sequence for stepping through the elements by the
  26. /// specified amount.
  27. @inlinable
  28. public func striding(by step: Int) -> StridingSequence<Self> {
  29. StridingSequence(base: self, stride: step)
  30. }
  31. }
  32. extension Collection {
  33. /// Returns a sequence stepping through the elements every `step` starting at
  34. /// the first value. Any remainders of the stride will be trimmed.
  35. ///
  36. /// (0...10).striding(by: 2) // == [0, 2, 4, 6, 8, 10]
  37. /// (0...10).striding(by: 3) // == [0, 3, 6, 9]
  38. ///
  39. /// - Complexity: O(1). Access to successive values is O(1) if the collection
  40. /// conforms to `RandomAccessCollection`; otherwise, O(_k_), where _k_ is
  41. /// the striding `step`.
  42. ///
  43. /// - Parameter step: The amount to step with each iteration.
  44. /// - Returns: Returns a collection for stepping through the elements by the
  45. /// specified amount.
  46. @inlinable
  47. public func striding(by step: Int) -> StridingCollection<Self> {
  48. StridingCollection(base: self, stride: step)
  49. }
  50. }
  51. /// A wrapper that strides over a base sequence.
  52. public struct StridingSequence<Base: Sequence> {
  53. @usableFromInline
  54. internal let base: Base
  55. @usableFromInline
  56. internal let stride: Int
  57. @inlinable
  58. internal init(base: Base, stride: Int) {
  59. precondition(stride > 0, "Stride must be greater than zero")
  60. self.base = base
  61. self.stride = stride
  62. }
  63. }
  64. extension StridingSequence {
  65. @inlinable
  66. public func striding(by step: Int) -> Self {
  67. Self(base: base, stride: stride * step)
  68. }
  69. }
  70. extension StridingSequence: Sequence {
  71. /// An iterator over a `StridingSequence` instance.
  72. public struct Iterator: IteratorProtocol {
  73. @usableFromInline
  74. internal var iterator: Base.Iterator
  75. @usableFromInline
  76. internal let stride: Int
  77. @usableFromInline
  78. internal var striding: Bool = false
  79. @inlinable
  80. internal init(iterator: Base.Iterator, stride: Int) {
  81. self.iterator = iterator
  82. self.stride = stride
  83. }
  84. @inlinable
  85. public mutating func next() -> Base.Element? {
  86. guard striding else {
  87. striding = true
  88. return iterator.next()
  89. }
  90. for _ in 0..<stride - 1 {
  91. guard iterator.next() != nil else { break }
  92. }
  93. return iterator.next()
  94. }
  95. }
  96. @inlinable
  97. public func makeIterator() -> Iterator {
  98. Iterator(iterator: base.makeIterator(), stride: stride)
  99. }
  100. }
  101. extension StridingSequence: LazySequenceProtocol
  102. where Base: LazySequenceProtocol {}
  103. /// A wrapper that strides over a base collection.
  104. public struct StridingCollection<Base: Collection> {
  105. @usableFromInline
  106. internal let base: Base
  107. @usableFromInline
  108. internal let stride: Int
  109. @inlinable
  110. internal init(base: Base, stride: Int) {
  111. precondition(stride > 0, "striding must be greater than zero")
  112. self.base = base
  113. self.stride = stride
  114. }
  115. }
  116. extension StridingCollection {
  117. @inlinable
  118. public func striding(by step: Int) -> Self {
  119. Self(base: base, stride: stride * step)
  120. }
  121. }
  122. extension StridingCollection: Collection {
  123. /// A position in a `StridingCollection` instance.
  124. public struct Index: Comparable {
  125. @usableFromInline
  126. internal let base: Base.Index
  127. @usableFromInline
  128. internal init(_ base: Base.Index) {
  129. self.base = base
  130. }
  131. @inlinable
  132. public static func < (lhs: Index, rhs: Index) -> Bool {
  133. lhs.base < rhs.base
  134. }
  135. }
  136. @inlinable
  137. public var startIndex: Index {
  138. Index(base.startIndex)
  139. }
  140. @inlinable
  141. public var endIndex: Index {
  142. Index(base.endIndex)
  143. }
  144. @inlinable
  145. public subscript(i: Index) -> Base.Element {
  146. base[i.base]
  147. }
  148. @inlinable
  149. public func index(after i: Index) -> Index {
  150. precondition(i.base != base.endIndex, "Advancing past end index")
  151. return index(i, offsetBy: 1)
  152. }
  153. @inlinable
  154. public func index(
  155. _ i: Index,
  156. offsetBy n: Int,
  157. limitedBy limit: Index
  158. ) -> Index? {
  159. guard n != 0 else { return i }
  160. guard limit != i else { return nil }
  161. return n > 0
  162. ? offsetForward(i, offsetBy: n, limitedBy: limit)
  163. : offsetBackward(i, offsetBy: -n, limitedBy: limit)
  164. }
  165. @inlinable
  166. internal func offsetForward(
  167. _ i: Index,
  168. offsetBy n: Int,
  169. limitedBy limit: Index
  170. ) -> Index? {
  171. if limit < i {
  172. if let idx = base.index(
  173. i.base,
  174. offsetBy: n * stride,
  175. limitedBy: base.endIndex
  176. ) {
  177. return Index(idx)
  178. } else {
  179. assert(distance(from: i, to: endIndex) == n, "Advancing past end index")
  180. return endIndex
  181. }
  182. } else if let idx = base.index(
  183. i.base,
  184. offsetBy: n * stride,
  185. limitedBy: limit.base
  186. ) {
  187. return Index(idx)
  188. } else {
  189. return distance(from: i, to: limit) == n
  190. ? endIndex
  191. : nil
  192. }
  193. }
  194. @inlinable
  195. internal func offsetBackward(
  196. _ i: Index,
  197. offsetBy n: Int,
  198. limitedBy limit: Index
  199. ) -> Index? {
  200. // We typically use the ternary operator but this significantly increases
  201. // compile times when using Swift 5.3.2
  202. // https://github.com/apple/swift-algorithms/issues/146
  203. let distance: Int
  204. if i == endIndex {
  205. distance = -((base.count - 1) % stride + 1) + (n - 1) * -stride
  206. } else {
  207. distance = n * -stride
  208. }
  209. return base.index(
  210. i.base,
  211. offsetBy: distance,
  212. limitedBy: limit.base
  213. ).map(Index.init)
  214. }
  215. @inlinable
  216. public var count: Int {
  217. base.isEmpty ? 0 : (base.count - 1) / stride + 1
  218. }
  219. @inlinable
  220. public func distance(from start: Index, to end: Index) -> Int {
  221. let distance = base.distance(from: start.base, to: end.base)
  222. return distance / stride + (distance % stride).signum()
  223. }
  224. @inlinable
  225. public func index(_ i: Index, offsetBy distance: Int) -> Index {
  226. precondition(distance <= 0 || i.base != base.endIndex, "Advancing past end index")
  227. precondition(distance >= 0 || i.base != base.startIndex, "Incrementing past start index")
  228. let limit = distance > 0 ? endIndex : startIndex
  229. let idx = index(i, offsetBy: distance, limitedBy: limit)
  230. precondition(idx != nil, "The distance \(distance) is not valid for this collection")
  231. return idx!
  232. }
  233. }
  234. extension StridingCollection: BidirectionalCollection
  235. where Base: RandomAccessCollection {
  236. @inlinable
  237. public func index(before i: Index) -> Index {
  238. precondition(i.base != base.startIndex, "Incrementing past start index")
  239. return index(i, offsetBy: -1)
  240. }
  241. }
  242. extension StridingCollection: RandomAccessCollection
  243. where Base: RandomAccessCollection {}
  244. extension StridingCollection: LazySequenceProtocol, LazyCollectionProtocol
  245. where Base: LazySequenceProtocol {}
  246. extension StridingCollection.Index: Hashable where Base.Index: Hashable {}