Intersperse.swift 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670
  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. // Intersperse
  13. //===----------------------------------------------------------------------===//
  14. /// A sequence that presents the elements of a base sequence of elements with a
  15. /// separator between each of those elements.
  16. public struct InterspersedSequence<Base: Sequence> {
  17. @usableFromInline
  18. internal let base: Base
  19. @usableFromInline
  20. internal let separator: Base.Element
  21. @inlinable
  22. internal init(base: Base, separator: Base.Element) {
  23. self.base = base
  24. self.separator = separator
  25. }
  26. }
  27. extension InterspersedSequence: Sequence {
  28. /// The iterator for an `InterspersedSequence` sequence.
  29. public struct Iterator: IteratorProtocol {
  30. @usableFromInline
  31. internal enum State {
  32. case start
  33. case element(Base.Element)
  34. case separator
  35. }
  36. @usableFromInline
  37. internal var iterator: Base.Iterator
  38. @usableFromInline
  39. internal let separator: Base.Element
  40. @usableFromInline
  41. internal var state = State.start
  42. @inlinable
  43. internal init(iterator: Base.Iterator, separator: Base.Element) {
  44. self.iterator = iterator
  45. self.separator = separator
  46. }
  47. @inlinable
  48. public mutating func next() -> Base.Element? {
  49. // After the start, the state flips between element and separator. Before
  50. // returning a separator, a check is made for the next element as a
  51. // separator is only returned between two elements. The next element is
  52. // stored to allow it to be returned in the next iteration.
  53. switch state {
  54. case .start:
  55. state = .separator
  56. return iterator.next()
  57. case .separator:
  58. guard let next = iterator.next() else { return nil }
  59. state = .element(next)
  60. return separator
  61. case .element(let element):
  62. state = .separator
  63. return element
  64. }
  65. }
  66. }
  67. @inlinable
  68. public func makeIterator() -> InterspersedSequence<Base>.Iterator {
  69. Iterator(iterator: base.makeIterator(), separator: separator)
  70. }
  71. }
  72. extension InterspersedSequence: Collection where Base: Collection {
  73. /// A position in an `InterspersedSequence` instance.
  74. public struct Index: Comparable {
  75. @usableFromInline
  76. internal enum Representation: Equatable {
  77. case element(Base.Index)
  78. case separator(next: Base.Index)
  79. }
  80. @usableFromInline
  81. internal let representation: Representation
  82. @inlinable
  83. internal init(representation: Representation) {
  84. self.representation = representation
  85. }
  86. @inlinable
  87. public static func < (lhs: Index, rhs: Index) -> Bool {
  88. switch (lhs.representation, rhs.representation) {
  89. case let (.element(li), .element(ri)),
  90. let (.separator(next: li), .separator(next: ri)),
  91. let (.element(li), .separator(next: ri)):
  92. return li < ri
  93. case let (.separator(next: li), .element(ri)):
  94. return li <= ri
  95. }
  96. }
  97. @inlinable
  98. internal static func element(_ index: Base.Index) -> Self {
  99. Self(representation: .element(index))
  100. }
  101. @inlinable
  102. internal static func separator(next: Base.Index) -> Self {
  103. Self(representation: .separator(next: next))
  104. }
  105. }
  106. @inlinable
  107. public var startIndex: Index {
  108. base.startIndex == base.endIndex ? endIndex : .element(base.startIndex)
  109. }
  110. @inlinable
  111. public var endIndex: Index {
  112. .separator(next: base.endIndex)
  113. }
  114. @inlinable
  115. public func index(after i: Index) -> Index {
  116. precondition(i != endIndex, "Can't advance past endIndex")
  117. switch i.representation {
  118. case let .element(index):
  119. return .separator(next: base.index(after: index))
  120. case let .separator(next):
  121. return .element(next)
  122. }
  123. }
  124. @inlinable
  125. public subscript(position: Index) -> Element {
  126. switch position.representation {
  127. case .element(let index): return base[index]
  128. case .separator: return separator
  129. }
  130. }
  131. @inlinable
  132. public func distance(from start: Index, to end: Index) -> Int {
  133. switch (start.representation, end.representation) {
  134. case let (.element(element), .separator(next: separator)):
  135. return 2 * base.distance(from: element, to: separator) - 1
  136. case let (.separator(next: separator), .element(element)):
  137. return 2 * base.distance(from: separator, to: element) + 1
  138. case let (.element(start), .element(end)),
  139. let (.separator(start), .separator(end)):
  140. return 2 * base.distance(from: start, to: end)
  141. }
  142. }
  143. @inlinable
  144. public func index(_ index: Index, offsetBy distance: Int) -> Index {
  145. distance >= 0
  146. ? offsetForward(index, by: distance)
  147. : offsetBackward(index, by: -distance)
  148. }
  149. @inlinable
  150. public func index(
  151. _ index: Index,
  152. offsetBy distance: Int,
  153. limitedBy limit: Index
  154. ) -> Index? {
  155. if distance >= 0 {
  156. return limit >= index
  157. ? offsetForward(index, by: distance, limitedBy: limit)
  158. : offsetForward(index, by: distance)
  159. } else {
  160. return limit <= index
  161. ? offsetBackward(index, by: -distance, limitedBy: limit)
  162. : offsetBackward(index, by: -distance)
  163. }
  164. }
  165. @inlinable
  166. internal func offsetForward(_ i: Index, by distance: Int) -> Index {
  167. guard let index = offsetForward(i, by: distance, limitedBy: endIndex)
  168. else { fatalError("Index is out of bounds") }
  169. return index
  170. }
  171. @inlinable
  172. internal func offsetBackward(_ i: Index, by distance: Int) -> Index {
  173. guard let index = offsetBackward(i, by: distance, limitedBy: startIndex)
  174. else { fatalError("Index is out of bounds") }
  175. return index
  176. }
  177. @inlinable
  178. internal func offsetForward(
  179. _ index: Index, by distance: Int, limitedBy limit: Index
  180. ) -> Index? {
  181. assert(distance >= 0)
  182. assert(limit >= index)
  183. switch (index.representation, limit.representation, distance.isMultiple(of: 2)) {
  184. case let (.element(index), .element(limit), true),
  185. let (.separator(next: index), .element(limit), false):
  186. return base.index(index, offsetBy: distance / 2, limitedBy: limit)
  187. .map { .element($0) }
  188. case let (.element(index), .element(limit), false),
  189. let (.element(index), .separator(next: limit), false),
  190. let (.separator(next: index), .element(limit), true),
  191. let (.separator(next: index), .separator(next: limit), true):
  192. return base.index(index, offsetBy: (distance + 1) / 2, limitedBy: limit)
  193. .map { .separator(next: $0) }
  194. case let (.element(index), .separator(next: limit), true),
  195. let (.separator(next: index), .separator(next: limit), false):
  196. return base.index(index, offsetBy: distance / 2, limitedBy: limit)
  197. .flatMap { $0 == limit ? nil : .element($0) }
  198. }
  199. }
  200. @inlinable
  201. internal func offsetBackward(
  202. _ index: Index, by distance: Int, limitedBy limit: Index
  203. ) -> Index? {
  204. assert(distance >= 0)
  205. assert(limit <= index)
  206. switch (index.representation, limit.representation, distance.isMultiple(of: 2)) {
  207. case let (.element(index), .element(limit), true),
  208. let (.element(index), .separator(next: limit), true),
  209. let (.separator(next: index), .element(limit), false),
  210. let (.separator(next: index), .separator(next: limit), false):
  211. return base.index(index, offsetBy: -((distance + 1) / 2), limitedBy: limit)
  212. .map { .element($0) }
  213. case let (.element(index), .separator(next: limit), false),
  214. let (.separator(next: index), .separator(next: limit), true):
  215. return base.index(index, offsetBy: -(distance / 2), limitedBy: limit)
  216. .map { .separator(next: $0) }
  217. case let (.element(index), .element(limit), false),
  218. let (.separator(next: index), .element(limit), true):
  219. return base.index(index, offsetBy: -(distance / 2), limitedBy: limit)
  220. .flatMap { $0 == limit ? nil : .separator(next: $0) }
  221. }
  222. }
  223. }
  224. extension InterspersedSequence: BidirectionalCollection
  225. where Base: BidirectionalCollection
  226. {
  227. @inlinable
  228. public func index(before i: Index) -> Index {
  229. precondition(i != startIndex, "Can't move before startIndex")
  230. switch i.representation {
  231. case let .element(index):
  232. return .separator(next: index)
  233. case let .separator(next):
  234. return .element(base.index(before: next))
  235. }
  236. }
  237. }
  238. extension InterspersedSequence: RandomAccessCollection
  239. where Base: RandomAccessCollection {}
  240. extension InterspersedSequence: LazySequenceProtocol
  241. where Base: LazySequenceProtocol {}
  242. extension InterspersedSequence: LazyCollectionProtocol
  243. where Base: LazySequenceProtocol & Collection {}
  244. //===----------------------------------------------------------------------===//
  245. // InterspersedMap
  246. //===----------------------------------------------------------------------===//
  247. /// A sequence over the results of applying a closure to the sequence's
  248. /// elements, with a separator that separates each pair of adjacent transformed
  249. /// values.
  250. @usableFromInline
  251. internal struct InterspersedMapSequence<Base: Sequence, Result> {
  252. @usableFromInline
  253. internal let base: Base
  254. @usableFromInline
  255. internal let transform: (Base.Element) -> Result
  256. @usableFromInline
  257. internal let separator: (Base.Element, Base.Element) -> Result
  258. }
  259. extension InterspersedMapSequence: Sequence {
  260. @usableFromInline
  261. internal struct Iterator: IteratorProtocol {
  262. @usableFromInline
  263. internal enum State {
  264. case start
  265. case element(Base.Element)
  266. case separator(previous: Base.Element)
  267. }
  268. @usableFromInline
  269. internal var base: Base.Iterator
  270. @usableFromInline
  271. internal let transform: (Base.Element) -> Result
  272. @usableFromInline
  273. internal let separator: (Base.Element, Base.Element) -> Result
  274. @usableFromInline
  275. internal var state = State.start
  276. @inlinable
  277. internal init(
  278. base: Base.Iterator,
  279. transform: @escaping (Base.Element) -> Result,
  280. separator: @escaping (Base.Element, Base.Element) -> Result
  281. ) {
  282. self.base = base
  283. self.transform = transform
  284. self.separator = separator
  285. }
  286. @inlinable
  287. internal mutating func next() -> Result? {
  288. switch state {
  289. case .start:
  290. guard let first = base.next() else { return nil }
  291. state = .separator(previous: first)
  292. return transform(first)
  293. case .separator(let previous):
  294. guard let next = base.next() else { return nil }
  295. state = .element(next)
  296. return separator(previous, next)
  297. case .element(let element):
  298. state = .separator(previous: element)
  299. return transform(element)
  300. }
  301. }
  302. }
  303. @inlinable
  304. internal func makeIterator() -> Iterator {
  305. Iterator(
  306. base: base.makeIterator(),
  307. transform: transform,
  308. separator: separator)
  309. }
  310. }
  311. extension InterspersedMapSequence: Collection where Base: Collection {
  312. @usableFromInline
  313. internal struct Index: Comparable {
  314. @usableFromInline
  315. internal enum Representation {
  316. case element(Base.Index)
  317. case separator(previous: Base.Index, next: Base.Index)
  318. }
  319. @usableFromInline
  320. internal let representation: Representation
  321. @inlinable
  322. internal init(representation: Representation) {
  323. self.representation = representation
  324. }
  325. @inlinable
  326. internal static func element(_ index: Base.Index) -> Self {
  327. Self(representation: .element(index))
  328. }
  329. @inlinable
  330. internal static func separator(previous: Base.Index, next: Base.Index) -> Self {
  331. Self(representation: .separator(previous: previous, next: next))
  332. }
  333. @inlinable
  334. internal static func == (lhs: Self, rhs: Self) -> Bool {
  335. switch (lhs.representation, rhs.representation) {
  336. case let (.element(lhs), .element(rhs)),
  337. let (.separator(_, next: lhs), .separator(_, next: rhs)):
  338. return lhs == rhs
  339. case (.element, .separator), (.separator, .element):
  340. return false
  341. }
  342. }
  343. @inlinable
  344. internal static func < (lhs: Self, rhs: Self) -> Bool {
  345. switch (lhs.representation, rhs.representation) {
  346. case let (.element(lhs), .element(rhs)),
  347. let (.separator(_, next: lhs), .separator(_, next: rhs)),
  348. let (.element(lhs), .separator(_, next: rhs)),
  349. let (.separator(previous: lhs, _), .element(rhs)):
  350. return lhs < rhs
  351. }
  352. }
  353. }
  354. @inlinable
  355. internal var startIndex: Index {
  356. base.isEmpty ? endIndex : .element(base.startIndex)
  357. }
  358. @inlinable
  359. internal var endIndex: Index {
  360. .separator(previous: base.endIndex, next: base.endIndex)
  361. }
  362. @inlinable
  363. internal func index(after index: Index) -> Index {
  364. switch index.representation {
  365. case .element(let index):
  366. let next = base.index(after: index)
  367. return .separator(previous: index, next: next)
  368. case .separator(_, let next):
  369. return .element(next)
  370. }
  371. }
  372. @inlinable
  373. internal subscript(position: Index) -> Result {
  374. switch position.representation {
  375. case .element(let index):
  376. return transform(base[index])
  377. case let .separator(previous, next):
  378. return separator(base[previous], base[next])
  379. }
  380. }
  381. @inlinable
  382. internal func distance(from start: Index, to end: Index) -> Int {
  383. switch (start.representation, end.representation) {
  384. case let (.element(lhs), .element(rhs)),
  385. let (.separator(_, next: lhs), .separator(_, next: rhs)):
  386. return 2 * base.distance(from: lhs, to: rhs)
  387. case let (.element(lhs), .separator(_, next: rhs)):
  388. return 2 * base.distance(from: lhs, to: rhs) - 1
  389. case let (.separator(_, next: lhs), .element(rhs)):
  390. return 2 * base.distance(from: lhs, to: rhs) + 1
  391. }
  392. }
  393. @inlinable
  394. internal func index(_ index: Index, offsetBy distance: Int) -> Index {
  395. guard distance != 0 else { return index }
  396. return distance > 0
  397. ? offsetForward(index, by: distance)
  398. : offsetBackward(index, by: -distance)
  399. }
  400. @inlinable
  401. internal func index(
  402. _ index: Index,
  403. offsetBy distance: Int,
  404. limitedBy limit: Index
  405. ) -> Index? {
  406. guard distance != 0 else { return index }
  407. if distance > 0 {
  408. return limit >= index
  409. ? offsetForward(index, by: distance, limitedBy: limit)
  410. : offsetForward(index, by: distance)
  411. } else {
  412. return limit <= index
  413. ? offsetBackward(index, by: -distance, limitedBy: limit)
  414. : offsetBackward(index, by: -distance)
  415. }
  416. }
  417. @inlinable
  418. internal func offsetForward(_ i: Index, by distance: Int) -> Index {
  419. guard let index = offsetForward(i, by: distance, limitedBy: endIndex)
  420. else { fatalError("Index is out of bounds") }
  421. return index
  422. }
  423. @inlinable
  424. internal func offsetBackward(_ i: Index, by distance: Int) -> Index {
  425. guard let index = offsetBackward(i, by: distance, limitedBy: startIndex)
  426. else { fatalError("Index is out of bounds") }
  427. return index
  428. }
  429. @inlinable
  430. internal func offsetForward(
  431. _ index: Index, by distance: Int, limitedBy limit: Index
  432. ) -> Index? {
  433. assert(distance > 0)
  434. assert(limit >= index)
  435. switch (index.representation, limit.representation, distance.isMultiple(of: 2)) {
  436. case let (.element(index), .element(limit), true),
  437. let (.separator(_, next: index), .element(limit), false):
  438. return base.index(index, offsetBy: distance / 2, limitedBy: limit)
  439. .map { .element($0) }
  440. case let (.element(index), .element(limit), false),
  441. let (.element(index), .separator(_, next: limit), false),
  442. let (.separator(_, next: index), .element(limit), true),
  443. let (.separator(_, next: index), .separator(_, next: limit), true):
  444. return base.index(index, offsetBy: (distance - 1) / 2, limitedBy: limit)
  445. .flatMap {
  446. guard $0 != limit else { return nil }
  447. let next = base.index(after: $0)
  448. return next == base.endIndex
  449. ? endIndex
  450. : .separator(previous: $0, next: next)
  451. }
  452. case let (.element(index), .separator(_, next: limit), true),
  453. let (.separator(_, next: index), .separator(_, next: limit), false):
  454. return base.index(index, offsetBy: distance / 2, limitedBy: limit)
  455. .flatMap { $0 == limit ? nil : .element($0) }
  456. }
  457. }
  458. @inlinable
  459. internal func offsetBackward(
  460. _ index: Index, by distance: Int, limitedBy limit: Index
  461. ) -> Index? {
  462. assert(distance > 0)
  463. assert(limit <= index)
  464. switch (index.representation, limit.representation, distance.isMultiple(of: 2)) {
  465. case let (.element(index), .element(limit), true),
  466. let (.element(index), .separator(_, next: limit), true),
  467. let (.separator(_, next: index), .element(limit), false),
  468. let (.separator(_, next: index), .separator(_, next: limit), false):
  469. return base.index(index, offsetBy: -((distance + 1) / 2), limitedBy: limit)
  470. .map { .element($0) }
  471. case let (.element(index), .separator(_, next: limit), false),
  472. let (.separator(_, next: index), .separator(_, next: limit), true):
  473. return base.index(index, offsetBy: -(distance / 2), limitedBy: limit)
  474. .map { .separator(previous: base.index($0, offsetBy: -1), next: $0) }
  475. case let (.element(index), .element(limit), false),
  476. let (.separator(_, next: index), .element(limit), true):
  477. return base.index(index, offsetBy: -(distance / 2), limitedBy: limit)
  478. .flatMap {
  479. $0 == limit
  480. ? nil
  481. : .separator(previous: base.index($0, offsetBy: -1), next: $0)
  482. }
  483. }
  484. }
  485. }
  486. extension InterspersedMapSequence: BidirectionalCollection
  487. where Base: BidirectionalCollection
  488. {
  489. @inlinable
  490. internal func index(before index: Index) -> Index {
  491. switch index.representation {
  492. case .element(let index):
  493. let previous = base.index(before: index)
  494. return .separator(previous: previous, next: index)
  495. case let .separator(previous, next):
  496. let index = next == base.endIndex ? base.index(before: next) : previous
  497. return .element(index)
  498. }
  499. }
  500. }
  501. extension InterspersedMapSequence.Index: Hashable
  502. where Base.Index: Hashable
  503. {
  504. @inlinable
  505. internal func hash(into hasher: inout Hasher) {
  506. switch representation {
  507. case .element(let base):
  508. hasher.combine(false)
  509. hasher.combine(base)
  510. case .separator(_, let next):
  511. hasher.combine(true)
  512. hasher.combine(next)
  513. }
  514. }
  515. }
  516. extension InterspersedMapSequence: LazySequenceProtocol {}
  517. extension InterspersedMapSequence: LazyCollectionProtocol
  518. where Base: Collection {}
  519. //===----------------------------------------------------------------------===//
  520. // interspersed(with:)
  521. //===----------------------------------------------------------------------===//
  522. extension Sequence {
  523. /// Returns a sequence containing elements of this sequence with the given
  524. /// separator inserted in between each element.
  525. ///
  526. /// Any value of the sequence's element type can be used as the separator.
  527. ///
  528. /// ```
  529. /// for value in [1,2,3].interspersed(with: 0) {
  530. /// print(value)
  531. /// }
  532. /// // 1
  533. /// // 0
  534. /// // 2
  535. /// // 0
  536. /// // 3
  537. /// ```
  538. ///
  539. /// The following shows a String being interspersed with a Character:
  540. /// ```
  541. /// let result = "ABCDE".interspersed(with: "-")
  542. /// print(String(result))
  543. /// // "A-B-C-D-E"
  544. /// ```
  545. ///
  546. /// - Parameter separator: Value to insert in between each of this sequence’s
  547. /// elements.
  548. /// - Returns: The interspersed sequence of elements.
  549. ///
  550. /// - Complexity: O(1)
  551. @inlinable
  552. public func interspersed(
  553. with separator: Element
  554. ) -> InterspersedSequence<Self> {
  555. InterspersedSequence(base: self, separator: separator)
  556. }
  557. }
  558. //===----------------------------------------------------------------------===//
  559. // lazy.interspersedMap(_:with:)
  560. //===----------------------------------------------------------------------===//
  561. extension LazySequenceProtocol {
  562. /// Returns a sequence over the results of applying a closure to the
  563. /// sequence's elements, with a separator that separates each pair of adjacent
  564. /// transformed values.
  565. ///
  566. /// The transformation closure lets you intersperse a sequence using a
  567. /// separator of a different type than the original's sequence's elements.
  568. /// Each separator is produced by a closure that is given access to the
  569. /// two elements in the original sequence right before and after it.
  570. ///
  571. /// let strings = [1, 2, 2].interspersedMap(String.init,
  572. /// with: { $0 == $1 ? " == " : " != " })
  573. /// print(strings.joined()) // "1 != 2 == 2"
  574. ///
  575. @usableFromInline
  576. internal func interspersedMap<Result>(
  577. _ transform: @escaping (Element) -> Result,
  578. with separator: @escaping (Element, Element) -> Result
  579. ) -> InterspersedMapSequence<Elements, Result> {
  580. InterspersedMapSequence(
  581. base: elements,
  582. transform: transform,
  583. separator: separator)
  584. }
  585. }