RunLengthEncoding.swift

RunLengthEncoding.swift

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

import Foundation

struct RunLengthEncoding<Element: Equatable>: Equatable {
    fileprivate struct Run<Element: Equatable>: Equatable {
        var item:  Element
        var count: Int

        init(_ item: Element) {
            self.item  = item
            self.count = 1
        }
        
        init(_ count: Int, _ item: Element) {
            self.item  = item
            self.count = count
        }
    }
    
    fileprivate var runs: [Run<Element>]
    
    var count: Int {
        return runs.reduce(0) { $0 + $1.count }
    }
    func contains(_ item: Element) -> Bool {
        return runs.contains { $0.item == item }
    }
    
    mutating func append(_ item: Element) {
        if runs.isEmpty {
            runs.append(Run(item))
        } else {
            if item == runs[runs.count - 1].item {
                runs[runs.count - 1].count += 1
            } else {
                runs.append(Run(item))
            }
        }
    }
    
    mutating func insert(_ item: Element, at offset: Int) {
        guard let index = index(at: offset) else {
            fatalError()
        }
        
        if index == endIndex {
            self.append(item)
            
        } else {
            if runs[index.run].item == item {
                runs[index.run].count += 1
                
            } else {
                if index.offset == 0 {
                    runs.replaceSubrange(index.run..<index.run + 1, with: [
                        Run(item),
                        runs[index.run]
                        ])
                } else {
                    runs.replaceSubrange(index.run..<index.run + 1, with: [
                        Run(index.offset, runs[index.run].item),
                        Run(item),
                        Run(runs[index.run].count - index.offset, runs[index.run].item)
                        ])
                }
            }
        }
    }
    
    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(1, value)]
    }
    init<C: Collection>(items: C) where C.Element == Element {
        self.runs = items.reduce(into: []) { runs, item in
            if runs.isEmpty || runs[runs.count - 1].item != item {
                runs.append(Run(item))
            } else {
                runs[runs.count - 1].count += 1
            }
        }
    }
}

extension RunLengthEncoding: BidirectionalCollection {
    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].item
    }
    
    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].item
    }
}

extension RunLengthEncoding: ExpressibleByArrayLiteral {
    init(arrayLiteral values: Element...) {
        self.init(items: values)
    }
}

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

extension RunLengthEncoding: Codable where Element: Codable {}
extension RunLengthEncoding.Run: Codable where Element: Codable {}

extension RunLengthEncoding: Hashable where Element: Hashable {}
extension RunLengthEncoding.Run: Hashable where Element: Hashable {}

extension Collection where Element: Equatable {
    func runLengthEncoded() -> RunLengthEncoding<Element> {
        return RunLengthEncoding<Element>(items: self)
    }
}