Swift challenge: Sequence.headAndTail

headAndTail.swift

extension Sequence {
  var headAndTail: (head: Element, tail: SubSequence)? {
    var head: Element?

    let tail = drop {
      if head == nil {
        head = $0

        return true
      } else {
        return false
      }
    }

    guard head != nil else {
      return nil
    }

    return (head: head!, tail: tail)
  }
}

nonConsumingTestCase: do {
  let (head, tail) = [1, 2, 3, 4].headAndTail!

  assert(head == 1)
  assert(tail == [2, 3, 4])
  assert(tail == [2, 3, 4])
}

consumingTestCase: do {
  var i = 0

  let sequence = AnySequence {
    AnyIterator { () -> Int? in
      i += 1

      return i
    }
  }

  var (head, tail) = sequence.headAndTail!

  assert(head == 1)
  assert(tail.prefix(3).elementsEqual([2, 3, 4]))

  (head, tail) = sequence.headAndTail!

  assert(head == 5)
  assert(tail.prefix(3).elementsEqual([6, 7, 8]))
}