Skip to content

Commit e4989f2

Browse files
timvermeulentimvermeulen
committed
Index offsetting for Product2
1 parent d872b64 commit e4989f2

File tree

1 file changed

+232
-10
lines changed

1 file changed

+232
-10
lines changed

‎Sources/Algorithms/Product.swift

Lines changed: 232 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
//
1010
//===----------------------------------------------------------------------===//
1111

12-
/// A sequence that represents the product of two sequence's elements.
12+
/// A sequence that represents the product of two sequences' elements.
1313
public struct Product2<Base1: Sequence, Base2: Collection> {
1414
/// The outer sequence in the product.
1515
public let base1: Base1
@@ -119,17 +119,27 @@ extension Product2: Collection where Base1: Collection {
119119
Index(i1: base1.endIndex, i2: base2.startIndex)
120120
}
121121

122+
@inlinable
123+
public subscript(position: Index) -> (Base1.Element, Base2.Element) {
124+
(base1[position.i1], base2[position.i2])
125+
}
126+
127+
/// Forms an index from a pair of base indices, converting
128+
/// `(i, base2.endIndex)` to `(base1.index(after: i), base2.startIndex)` if
129+
/// necessary.
130+
@usableFromInline
131+
internal func convertIndex(_ i1: Base1.Index, _ i2: Base2.Index) -> Index {
132+
i2 == base2.endIndex
133+
? Index(i1: base1.index(after: i1), i2: base2.startIndex)
134+
: Index(i1: i1, i2: i2)
135+
}
136+
122137
@inlinable
123138
public func index(after i: Index) -> Index {
124139
precondition(i.i1 != base1.endIndex, "Can't advance past endIndex")
125-
let newIndex2 = base2.index(after: i.i2)
126-
return newIndex2 == base2.endIndex
127-
? Index(i1: base1.index(after: i.i1), i2: base2.startIndex)
128-
: Index(i1: i.i1, i2: newIndex2)
140+
return convertIndex(i.i1, base2.index(after: i.i2))
129141
}
130142

131-
// TODO: Implement index(_:offsetBy:) and index(_:offsetBy:limitedBy:)
132-
133143
@inlinable
134144
public func distance(from start: Index, to end: Index) -> Int {
135145
guard start.i1 <= end.i1
@@ -183,10 +193,222 @@ extension Product2: Collection where Base1: Collection {
183193
+ left
184194
}
185195
}
196+
197+
public func index(_ i: Index, offsetBy distance: Int) -> Index {
198+
guard distance != 0 else { return i }
199+
200+
return distance > 0
201+
? offsetForward(i, by: distance)
202+
: offsetBackward(i, by: -distance)
203+
}
204+
205+
public func index(
206+
_ i: Index,
207+
offsetBy distance: Int,
208+
limitedBy limit: Index
209+
) -> Index? {
210+
guard distance != 0 else { return i }
211+
guard limit != i else { return nil }
212+
213+
if distance > 0 {
214+
return limit > i
215+
? offsetForward(i, by: distance, limitedBy: limit)
216+
: offsetForward(i, by: distance)
217+
} else {
218+
return limit < i
219+
? offsetBackward(i, by: -distance, limitedBy: limit)
220+
: offsetBackward(i, by: -distance)
221+
}
222+
}
186223

187-
@inlinable
188-
public subscript(position: Index) -> (Base1.Element, Base2.Element) {
189-
return (base1[position.i1], base2[position.i2])
224+
@usableFromInline
225+
internal func offsetForward(_ i: Index, by distance: Int) -> Index {
226+
guard let index = offsetForward(i, by: distance, limitedBy: endIndex)
227+
else { fatalError("Index is out of bounds") }
228+
return index
229+
}
230+
231+
@usableFromInline
232+
internal func offsetBackward(_ i: Index, by distance: Int) -> Index {
233+
guard let index = offsetBackward(i, by: distance, limitedBy: startIndex)
234+
else { fatalError("Index is out of bounds") }
235+
return index
236+
}
237+
238+
@usableFromInline
239+
internal func offsetForward(
240+
_ i: Index, by distance: Int, limitedBy limit: Index
241+
) -> Index? {
242+
assert(distance > 0)
243+
assert(limit > i)
244+
245+
if limit.i1 == i.i1 {
246+
// Delegate to `base2` if the offset is limited to `i.i1`.
247+
//
248+
// i.i2 limit.i2
249+
// v v
250+
// i.i1 > [x x x|x x x x x x|x x x]
251+
252+
return base2.index(i.i2, offsetBy: distance, limitedBy: limit.i2)
253+
.map { i2 in Index(i1: i.i1, i2: i2) }
254+
}
255+
256+
257+
if let i2 = base2.index(i.i2, offsetBy: distance, limitedBy: base2.endIndex) {
258+
// `distance` does not overflow `base2[i.i2...]`.
259+
//
260+
// i.i2 i2
261+
// v v
262+
// i.i1 > [x x x|x x x x x x|x x x]
263+
// [ |> > > > > >| ] (`distance`)
264+
265+
return convertIndex(i.i1, i2)
266+
}
267+
268+
let suffixCount = base2[i.i2...].count
269+
let remaining = distance - suffixCount
270+
let nextI1 = base1.index(after: i.i1)
271+
272+
if limit.i1 == nextI1 {
273+
// Delegate to `base2` if the offset is limited to `nextI1`.
274+
//
275+
// i.i2
276+
// v
277+
// i.i1 > [x x x|x x x x x x x x x]
278+
// nextI1 > [x x x x x x x x x|x x x]
279+
// ^
280+
// limit.i2
281+
282+
return base2.index(base2.startIndex, offsetBy: remaining, limitedBy: limit.i2)
283+
.map { i2 in Index(i1: nextI1, i2: i2) }
284+
}
285+
286+
if let i2 = base2.index(base2.startIndex, offsetBy: remaining, limitedBy: i.i2) {
287+
// `remaining` does not overflow `base2[..<i.i2]`.
288+
//
289+
// i.i2
290+
// v
291+
// i.i1 > [x x x x x x x x x|x x x]
292+
// [ |> > >] (`suffixCount`)
293+
// [> > >| ] (`remaining`)
294+
// nextI1 > [x x x|x x x x x x x x x]
295+
// ^
296+
// i2
297+
298+
return Index(i1: nextI1, i2: i2)
299+
}
300+
301+
let prefixCount = base2[..<i.i2].count
302+
let base2Count = prefixCount + suffixCount
303+
let base1Distance = remaining / base2Count
304+
305+
guard let i1 = base1.index(nextI1, offsetBy: base1Distance, limitedBy: limit.i1)
306+
else { return nil }
307+
308+
// The distance from `base2.startIndex` to the target.
309+
let base2Distance = remaining % base2Count
310+
311+
let base2Limit = limit.i1 == i1 ? limit.i2 : base2.endIndex
312+
return base2.index(base2.startIndex, offsetBy: base2Distance, limitedBy: base2Limit)
313+
.map { i2 in Index(i1: i1, i2: i2) }
314+
}
315+
316+
@usableFromInline
317+
internal func offsetBackward(
318+
_ i: Index, by distance: Int, limitedBy limit: Index
319+
) -> Index? {
320+
assert(distance > 0)
321+
assert(limit < i)
322+
323+
if limit.i1 == i.i1 {
324+
// Delegate to `base2` if the offset is limited to `i.i1`.
325+
//
326+
// limit.i2 i.i2
327+
// v v
328+
// i.i1 > [x x x|x x x x x x|x x x]
329+
330+
return base2.index(i.i2, offsetBy: -distance, limitedBy: limit.i2)
331+
.map { i2 in Index(i1: i.i1, i2: i2) }
332+
}
333+
334+
if let i2 = base2.index(i.i2, offsetBy: -distance, limitedBy: base2.startIndex) {
335+
// `distance` does not underflow `base2[..<i.i2]`.
336+
//
337+
// i2 i.i2
338+
// v v
339+
// i.i1 > [x x x|x x x x x x|x x x]
340+
// [ |< < < < < <| ] (`distance`)
341+
342+
return Index(i1: i.i1, i2: i2)
343+
}
344+
345+
let prefixCount = base2[..<i.i2].count
346+
let remaining = distance - prefixCount
347+
let previousI1 = base1.index(i.i1, offsetBy: -1)
348+
349+
if limit.i1 == previousI1 {
350+
// Delegate to `base2` if the offset is limited to `previousI1`.
351+
//
352+
// limit.i2
353+
// v
354+
// previousI1 > [x x x|x x x x x x x x x]
355+
// i.i1 > [x x x x x x x x x|x x x]
356+
// ^
357+
// i.i2
358+
359+
return base2.index(base2.endIndex, offsetBy: -remaining, limitedBy: limit.i2)
360+
.map { i2 in Index(i1: previousI1, i2: i2) }
361+
}
362+
363+
if let i2 = base2.index(base2.endIndex, offsetBy: -remaining, limitedBy: i.i2) {
364+
// `remaining` does not underflow `base2[i.i2...]`.
365+
//
366+
// i2
367+
// v
368+
// previousI1 > [x x x x x x x x x|x x x]
369+
// [ |< < <] (`remaining`)
370+
// [< < <| ] (`prefixCount`)
371+
// i.i1 > [x x x|x x x x x x x x x]
372+
// ^
373+
// i.i2
374+
375+
return Index(i1: previousI1, i2: i2)
376+
}
377+
378+
let suffixCount = base2[i.i2...].count
379+
let base2Count = prefixCount + suffixCount
380+
let base1Distance = remaining / base2Count
381+
382+
// The distance from `base2.endIndex` to the target.
383+
let base2Distance = remaining % base2Count
384+
385+
if base2Distance == 0 {
386+
// We end up exactly between two cycles, so `base1Distance` would
387+
// overshoot the target by 1.
388+
//
389+
// base2.startIndex
390+
// v
391+
// i1 > |x x x x x x x x x x x x] >
392+
// ... > `base1Distance` times
393+
// previousI1 > [x x x x x x x x x x x x] >
394+
// i.i1 > [x x x|x x x x x x x x x]
395+
// ^
396+
// i.i2
397+
398+
if let i1 = base1.index(previousI1, offsetBy: -(base1Distance - 1), limitedBy: limit.i1) {
399+
let index = Index(i1: i1, i2: base2.startIndex)
400+
return index < limit ? nil : index
401+
} else {
402+
return nil
403+
}
404+
}
405+
406+
guard let i1 = base1.index(previousI1, offsetBy: -base1Distance, limitedBy: limit.i1)
407+
else { return nil }
408+
409+
let base2Limit = limit.i1 == i1 ? limit.i2 : base2.startIndex
410+
return base2.index(base2.endIndex, offsetBy: -base2Distance, limitedBy: base2Limit)
411+
.map { i2 in Index(i1: i1, i2: i2) }
190412
}
191413
}
192414

0 commit comments

Comments
 (0)