Swift Offset Generator

Swift Offset Generator.swift

/*
 * Allows you to offset offset a generator
 * and loop back around to the initial values
 */
struct OffsetGenerator<G : GeneratorType> : GeneratorType {
    var generator: G
    var skipped = Array<G.Element>()
    
    init(generator: G, var offset: Int){
        self.generator = generator
        while offset > 0, let next = self.generator.next() {
            skipped.append(next)
            offset--
        }
    }
    
    mutating func next() -> G.Element? {
        if let next = generator.next() {
            return next
        }
        else if let next = skipped.first {
            skipped.removeAtIndex(0)
            return next
        }
        else {
            return nil
        }
    }
}

struct OffsetSequence<S : SequenceType> : SequenceType {
    
    let sequence: S
    let offset: Int
    init(sequence: S, offset: Int){
        self.sequence = sequence
        self.offset = offset
    }
    
    func generate() -> OffsetGenerator<S.Generator> {
        return OffsetGenerator(generator: sequence.generate(), offset: offset)
    }
}

func offset<S : SequenceType>(sequence: S, offset: Int) -> OffsetSequence<S> {
    return OffsetSequence(sequence: sequence, offset: offset)
}

// Example

var str = "Hey World"

for c in offset(str, 2){
    println(c) // -> "y"
}              //    " "
               //    "W"
               //    "o"
               //    "r"
               //    "l"
               //    "d"
               //    "H"
               //    "e"