swift-indices-where-wrapper.swift

struct SparseSlice<T> {
    var data: Dictionary<Int, T>
    
    mutating func update(_ f: (inout T) -> Void) {
        for (k, v) in data {
            var v = v
            f(&v)
            data[k] = v
        }
    }
}

struct ArrayWhere<Element> {
    subscript (wh: @escaping (Element) -> Bool) -> SparseSlice<Element> {
        get {
            var data = [Int: Element]()
            for i in self.data.startIndex..<self.data.endIndex {
                if wh(self.data[i]) {
                    data[i] = self.data[i]
                }
            }
            return SparseSlice(data: data)
        }
        set {
            for (i, v) in newValue.data {
                self.data[i] = v
            }
        }
    }
    
    var data: Array<Element>
}

extension Array {
    var `where`: ArrayWhere<Element> {
        get {
            return ArrayWhere(data: self)
        }
        set {
            self = newValue.data
        }
    }
    
    subscript(`where` wh: @escaping (Element) -> Bool) -> SparseSlice<Element> {
        get {
            var data = [Int: Element]()
            for i in startIndex..<endIndex {
                if wh(self[i]) {
                    data[i] = self[i]
                }
            }
            return SparseSlice(data: data)
        }
        set {
            for (i, v) in newValue.data {
                self[i] = v
            }
        }
    }
}

var s = [0, 1, 2, 3, 4]
print(s)
s.where[]{ $0 % 2 == 0 }.update { (x: inout Int) in x += 100  }
print(s)