Sequence of Range<Int>s

RangeStrider.swift

// known-good: Swift 5.0.1

/*

# how to use:

instead of:

  for i in 0 ..< (n-1)/bs {
        let startIdx = i * bs
        let endIdx = startIdx + bs
        let xb = xTrain[startIdx..<endIdx]
        let yb = yTrain[startIdx..<endIdx]
  }
  
Do:

  for indices in batchedIndices(0,n, bs) {
          let xb = xTrain[indices]
          let yb = yTrain[indices]
          // do stuff with xb and yb here...
  }

Or to be safer, if c is an Int-indexed Collection:

  for indices in c.rangeBatches(size:bs) {
          let xb = xTrain[indices]
          let yb = yTrain[indices]
          // do stuff with xb and yb here...
  }
  
*/

//
// shorter version
// 

/// Returns a sequence of ranges, of a specified batch size, over (startIndex ..< endIndex)
private func batchedRanges(startIndex:Int, endIndex:Int, batchSize:Int) -> UnfoldSequence<Range<Int>,Int>
{
  return sequence(state: startIndex) { (batchStartIndex) -> Range<Int>? in
    let remaining = endIndex - batchStartIndex
    guard remaining > 0 else { return nil }
    let currentBatchSize = min(batchSize,remaining)
    let batchEndIndex = batchStartIndex.advanced(by: currentBatchSize)
    defer {  batchStartIndex = batchEndIndex  }
    return batchStartIndex ..< batchEndIndex
  }
}

extension Collection
  where Self.Index == Int {
  /// Returns a sequence of ranges over the collection, of or below the specified batch size
  func rangeBatches(size:Int) -> AnySequence<Range<Self.Index>> {
    return batchedRanges(startIndex: self.startIndex, endIndex: self.endIndex, batchSize: size)
  }
}


//
// long-winded version
// 

/// An iterator for a RangeStrideThrough
struct RangeIterator : IteratorProtocol
{
  var batchStartIndex:Int
  let collectionEndIndex:Int
  let batchSize:Int

  typealias Element = Range<Int>

  mutating func next() -> Range<Int>? {
    let distanceRemaining = collectionEndIndex - batchStartIndex
    if distanceRemaining <= 0 {
      return nil
    }
    else {
      let thisBatchSize = min(batchSize,distanceRemaining)
      let batchEnd = batchStartIndex + thisBatchSize
      let batchRange = batchStartIndex..<batchEnd
      batchStartIndex = batchEnd
      return batchRange
    }
  }
}

/// A Sequence of Range<Int>, starting at 0.
struct RangeStrideThrough : Sequence
{
  typealias Element = Range<Int>
  typealias Iterator = RangeIterator

  let startIndex:Int
  let endIndex:Int
  let batchSize:Int

  __consuming func makeIterator() -> RangeIterator {
    return RangeIterator(batchStartIndex: startIndex, collectionEndIndex: endIndex, batchSize: batchSize)
  }
}

/// Returns a sequence of Int index ranges, starting at 0, with each of batchSize or less.
func batchedIndices(_ collectionSize:Int, _ batchSize:Int) -> RangeStrideThrough {
  return RangeStrideThrough(startIndex:0,
                            endIndex: collectionSize,
                            batchSize: batchSize)
}

func batchedIndices<C:Collection>(_ coll:C, _ batchSize:Int) -> RangeStrideThrough
  where C.Index == Int
{
  return RangeStrideThrough(startIndex:coll.startIndex,
                            endIndex: coll.endIndex,
                            batchSize: batchSize)
}