Skip to content

Commit 143d766

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

File tree

1 file changed

+229
-10
lines changed

1 file changed

+229
-10
lines changed

‎Sources/Algorithms/Product.swift

Lines changed: 229 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,219 @@ 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+
if distance >= 0 {
211+
return limit >= i
212+
? offsetForward(i, by: distance, limitedBy: limit)
213+
: offsetForward(i, by: distance)
214+
} else {
215+
return limit <= i
216+
? offsetBackward(i, by: -distance, limitedBy: limit)
217+
: offsetBackward(i, by: -distance)
218+
}
219+
}
186220

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

0 commit comments

Comments
 (0)