Swift Array Masking Operations

Array Masking.swift

// Example

var x = [1, 2, 3, 4, 5]   // [1, 2, 3, 4, 5]
var y = [1, 1, 2, 3, 5]   // [1, 1, 2, 3, 5]
var z = x - y             // [0, 1, 1, 1, 0]

x[ x % 2 != 0 ] = x + 1   // [2, 2, 4, 4, 6]
y[ y > 2 * z ] = y * y    // [1, 1, 2, 9, 25]
z[ x > 2 && x < 6 ] = 0   // [0, 1, 0, 0, 0]

x[{ abs(sin(Float($0))) > 0.3 }] = 1 // [1, 1, 1, 1, 6]

// Helper functions

func map2<S : SequenceType, T : SequenceType, U>(lhs: S, rhs: T, transform: (S.Generator.Element, T.Generator.Element) -> U) -> [U] {
    return Array(Zip2(lhs, rhs)).map(transform)
}

func maskFilter<S : SequenceType, T : SequenceType where T.Generator.Element == Bool>(source: S, mask: T, defaultValue: S.Generator.Element) -> [S.Generator.Element] {
    return Array(Zip2(source, mask)).map({ (value, include) in include ? value : defaultValue })
}

func choose<S : SequenceType, T : SequenceType, U : SequenceType where U.Generator.Element == Bool, S.Generator.Element == T.Generator.Element>(onTrue: S, onFalse: T, picker: U) -> [S.Generator.Element]{
    return map2(Zip2(onTrue, onFalse), picker, { zipped, choice in choice ? zipped.0 : zipped.1 })
}

// Masking operations
typealias ArrayMask = [Bool]

extension Array {
    subscript (mask: ArrayMask) -> Array<T> {
        set {
            self = choose(newValue, self, mask)
        }
        get {
            fatalError("Array mask getters are disallowed")
        }
    }
    
    subscript (mask: ArrayMask) -> T {
        set {
            self = maskFilter(self, !mask, newValue)
        }
        get {
            fatalError("Array mask getters are disallowed")
        }
    }
    
    subscript (maskMap: T -> Bool) -> T {
        set {
            self[self.map(maskMap)] = newValue
        }
        get {
            fatalError("Array mask getters are disallowed")
        }
    }
}

// Boolean operations
prefix func !(arr: Array<Bool>) -> Array<Bool> {
    return arr.map({ b in !b })
}

func ||(lhs: Array<Bool>, rhs: Array<Bool>) -> Array<Bool> {
    return map2(lhs, rhs, { (lhs, rhs) in lhs || rhs })
}

func &&(lhs: Array<Bool>, rhs: Array<Bool>) -> Array<Bool> {
    return map2(lhs, rhs, { (lhs, rhs) in lhs && rhs })
}

// ==

func ==<T : Equatable>(lhs: [T], rhs: [T]) -> ArrayMask {
    return map2(lhs, rhs, { lhs, rhs in lhs == rhs })
}

func ==<T : Equatable>(lhs: T, rhs: [T]) -> ArrayMask {
    return rhs == lhs
}

func ==<T : Equatable>(lhs: [T], rhs: T) -> ArrayMask {
    return map(lhs, { lhs in lhs == rhs })
}

// !=

func !=<T : Equatable>(lhs: [T], rhs: [T]) -> ArrayMask {
    return map2(lhs, rhs, { lhs, rhs in lhs != rhs })
}

func !=<T : Equatable>(lhs: T, rhs: [T]) -> ArrayMask {
    return rhs != lhs
}

func !=<T : Equatable>(lhs: [T], rhs: T) -> ArrayMask {
    return map(lhs, { lhs in lhs != rhs })
}

// <

func <<T : Comparable>(lhs: [T], rhs: [T]) -> ArrayMask {
    return map2(lhs, rhs, { lhs, rhs in lhs < rhs })
}

func <<T : Comparable>(lhs: [T], rhs: T) -> ArrayMask {
    return map(lhs, { lhs in lhs < rhs })
}

func <<T : Comparable>(lhs: T, rhs: [T]) -> ArrayMask {
    return map(rhs, { rhs in lhs < rhs })
}

// <=

func <=<T : Comparable>(lhs: [T], rhs: [T]) -> ArrayMask {
    return map2(lhs, rhs, { lhs, rhs in lhs <= rhs })
}

func <=<T : Comparable>(lhs: [T], rhs: T) -> ArrayMask {
    return map(lhs, { lhs in lhs <= rhs })
}

func <=<T : Comparable>(lhs: T, rhs: [T]) -> ArrayMask {
    return map(rhs, { rhs in lhs <= rhs })
}

// >

func ><T : Comparable>(lhs: [T], rhs: [T]) -> ArrayMask {
    return rhs < lhs
}

func ><T : Comparable>(lhs: [T], rhs: T) -> ArrayMask {
    return rhs < lhs
}

func ><T : Comparable>(lhs: T, rhs: [T]) -> ArrayMask {
    return rhs < lhs
}

// >=

func >=<T : Comparable>(lhs: [T], rhs: [T]) -> ArrayMask {
    return rhs <= lhs
}

func >=<T : Comparable>(lhs: [T], rhs: T) -> ArrayMask {
    return rhs <= lhs
}

func >=<T : Comparable>(lhs: T, rhs: [T]) -> ArrayMask {
    return rhs <= lhs
}

// Integer arithmetic operations

// +

func +<T : IntegerArithmeticType>(lhs: [T], rhs: [T]) -> [T] {
    return map2(lhs, rhs, { lhs, rhs in lhs + rhs })
}

func +<T : IntegerArithmeticType>(lhs: [T], rhs: T) -> [T] {
    return map(lhs, { lhs in lhs + rhs })
}

func +<T : IntegerArithmeticType>(lhs: T, rhs: [T]) -> [T] {
    return rhs + lhs
}

// -

func -<T : IntegerArithmeticType>(lhs: [T], rhs: [T]) -> [T] {
    return map2(lhs, rhs, { lhs, rhs in lhs - rhs })
}

func -<T : IntegerArithmeticType>(lhs: [T], rhs: T) -> [T] {
    return map(lhs, { lhs in lhs - rhs })
}

func -<T : IntegerArithmeticType>(lhs: T, rhs: [T]) -> [T] {
    return map(rhs, { rhs in lhs - rhs })
}

prefix func -<T : SignedNumberType>(arr: [T]) -> [T] {
    return arr.map({ v in -v })
}

// *

func *<T : IntegerArithmeticType>(lhs: [T], rhs: [T]) -> [T] {
    return map2(lhs, rhs, { lhs, rhs in lhs * rhs })
}

func *<T : IntegerArithmeticType>(lhs: [T], rhs: T) -> [T] {
    return map(lhs, { lhs in lhs * rhs })
}

func *<T : IntegerArithmeticType>(lhs: T, rhs: [T]) -> [T] {
    return rhs * lhs
}

// /

func /<T : IntegerArithmeticType>(lhs: [T], rhs: [T]) -> [T] {
    return map2(lhs, rhs, { lhs, rhs in lhs / rhs })
}

func /<T : IntegerArithmeticType>(lhs: [T], rhs: T) -> [T] {
    return map(lhs, { lhs in lhs / rhs })
}

func /<T : IntegerArithmeticType>(lhs: T, rhs: [T]) -> [T] {
    return map(rhs, { rhs in lhs / rhs })
}

// %

func %<T : IntegerArithmeticType>(lhs: [T], rhs: [T]) -> [T] {
    return map2(lhs, rhs, { lhs, rhs in lhs % rhs })
}

func %<T : IntegerArithmeticType>(lhs: [T], rhs: T) -> [T] {
    return map(lhs, { lhs in lhs % rhs })
}

func %<T : IntegerArithmeticType>(lhs: T, rhs: [T]) -> [T] {
    return map(rhs, { rhs in lhs % rhs })
}