RLE.swift

//
//  RLE.swift
//  RLE
//
//  Created by Mark Onyschuk on 2017-08-31.
//  Copyright © 2017 Mark Onyschuk. All rights reserved.
//

import Foundation

struct RLE<Element: Equatable>: BidirectionalCollection, Equatable, ExpressibleByArrayLiteral, CustomStringConvertible {
    fileprivate struct Run<Element: Equatable>: Equatable {
        var value: Element
        var count: Int

        init(_ value: Element, _ count: Int) {
            self.value = value
            self.count = count
        }
        
        static func ==(lhs: RLE.Run<Element>, rhs: RLE.Run<Element>) -> Bool {
            return lhs.value == rhs.value && lhs.count == rhs.count
        }
    }
    
    fileprivate var runs: [Run<Element>]
    
    var count: Int {
        return runs.reduce(0, { $0 + $1.count })
    }
    func contains(_ value: Element) -> Bool {
        return runs.contains(where: { $0.value == value })
    }
    
    mutating func append(_ value: Element) {
        if runs.isEmpty {
            runs.append(Run(value, 1))
        } else {
            if value == runs[runs.count - 1].value {
                runs[runs.count - 1].count += 1
            } else {
                runs.append(Run(value, 1))
            }
        }
    }
    
    mutating func insert(_ value: Element, at offset: Int) {
        guard let index = index(at: offset) else {
            fatalError()
        }
        
        if index == endIndex {
            self.append(value)
        } else {
            if runs[index.run].value == value {
                runs[index.run].count += 1
            } else {
                if index.offset == 0 {
                    runs.replaceSubrange(index.run..<index.run + 1, with: [
                        Run(value, 1),
                        runs[index.run]
                        ])
                } else {
                    runs.replaceSubrange(index.run..<index.run + 1, with: [
                        Run(runs[index.run].value, index.offset),
                        Run(value, 1),
                        Run(runs[index.run].value, runs[index.run].count - index.offset)
                        ])
                }
            }
        }
    }
    
    mutating func remove(at offset: Int) {
        guard let index = index(at: offset), index != endIndex else {
            fatalError()
        }
        
        if runs[index.run].count == 1 {
            runs.remove(at: index.run)
        } else {
            runs[index.run].count -= 1
        }
    }
    
    init() {
        self.runs = []
    }
    init(value: Element) {
        self.runs = [Run(value, 1)]
    }
    init<S: Sequence>(values: S) where S.Element == Element {
        self.runs = values.reduce(into: []) { runs, value in
            if runs.isEmpty {
                runs.append(Run(value, 1))
            } else {
                if runs[runs.count - 1].value == value {
                    runs[runs.count - 1].count += 1
                } else {
                    runs.append(Run(value, 1))
                }
            }
        }
    }
}

extension RLE {
    struct RLEIndex: Comparable {
        fileprivate let run, offset: Int;
        
        static func <(lhs: RLEIndex, rhs: RLEIndex) -> Bool {
            return (lhs.run, lhs.offset) < (rhs.run, rhs.offset)
        }
        static func ==(lhs: RLEIndex, rhs: RLEIndex) -> Bool {
            return lhs.run == rhs.run && lhs.offset == rhs.offset
        }
    }

    var startIndex: RLEIndex {
        return RLEIndex(run: 0, offset: 0)
    }
    var endIndex: RLEIndex {
        return RLEIndex(run: runs.count, offset: 0)
    }
    
    func index(after index: RLEIndex) -> RLEIndex {
        if index.offset < runs[index.run].count - 1 {
            return RLEIndex(run: index.run, offset: index.offset + 1)
        } else {
            return RLEIndex(run: index.run + 1, offset: 0)
        }
    }
    func index(before index: RLEIndex) -> RLEIndex {
        if index.offset > 0 {
            return RLEIndex(run: index.run, offset: index.offset - 1)
        } else {
            return RLEIndex(run: index.run - 1, offset: runs[index.run - 1].count - 1)
        }
    }
    
    subscript(index: RLEIndex) -> Element {
        return runs[index.run].value
    }
    
    fileprivate func index(at offset: Int) -> RLEIndex? {
        var idx = offset
        for ridx in 0..<runs.count {
            if idx < runs[ridx].count {
                return RLEIndex(run: ridx, offset: idx)
            }
            idx -= runs[ridx].count
        }
        return idx == 0 ? endIndex : nil
    }
    
    subscript(offset: Int) -> Element {
        return runs[index(at: offset)!.run].value
    }
}

extension RLE {
    static func ==(lhs: RLE<Element>, rhs: RLE<Element>) -> Bool {
        return lhs.runs == rhs.runs
    }
}

extension RLE {
    init(arrayLiteral values: Element...) {
        self.init(values: values)
    }
}

extension RLE {
    var description: String {
        return self.map({ $0 }).description
    }
}