prefixed postfixed sequence

prepostseq.swift

// Released under the terms of the MIT license.

struct PrefixedPostfixedSequence<T : SequenceType, U where U == T.Generator.Element> : SequenceType {
  private let sequence : T
  private let initial : U?
  private let final : U?

  init(_ sequence: T, initial: U? = nil, final: U? = nil) {
    self.sequence = sequence
    self.initial = initial
    self.final = final
  }

  func generate() -> PrefixedPostfixedSeqGenerator<T.Generator, U> {
    return PrefixedPostfixedSeqGenerator(generator: sequence.generate(), initial: initial, final: final)
  }
}

struct PrefixedPostfixedSeqGenerator<T : GeneratorType, U where U == T.Element> : GeneratorType {
  private var generator: T?
  private var initial : U?
  private var final : U?

  private init(generator: T, initial: U?, final: U?) {
    self.generator = generator
    self.initial = initial
    self.final = final
  }

  mutating func next() -> U? {
    // 1. Try the initial element
    if let actualInitial = initial {
      initial = nil
      return actualInitial
    }
    // 2. Try the actual sequence
    if let g = generator?.next() {
      return g
    }
    else {
      generator = nil
    }
    // 3. Try the final element
    if let actualFinal = final {
      final = nil
      return actualFinal
    }
    return nil
  }
}

// Try it out
var buffer : [Character] = []
for item in PrefixedPostfixedSequence("hello world", initial: Character("?"), final: Character("!")) {
  buffer.append(item)
}
println(buffer)

var buffer2 : [Int] = []
for item in PrefixedPostfixedSequence([10, 11, 12, 13], initial: 1) {
  buffer2.append(item)
}
println(buffer2)