Creating an extension to filter nils from an Array in Swift

I'm trying to write an extension to Array which will allow an array of optional T's to be transformed into an array of non-optional T's.

e.g. this could be written as a free function like this:

func removeAllNils(array: [T?]) -> [T] {
    return array
        .filter({ $0 != nil })   // remove nils, still a [T?]
        .map({ $0! })            // convert each element from a T? to a T
}

But, I can't get this to work as an extension. I'm trying to tell the compiler that the extension only applies to Arrays of optional values. This is what I have so far:

extension Array {
    func filterNils<U, T: Optional<U>>() -> [U] {
        return filter({ $0 != nil }).map({ $0! })
    }
}

(it doesn't compile!)

Solution 1:

As of Swift 2.0, you don't need to write your own extension to filter nil values from an Array, you can use flatMap, which flattens the Array and filters nils:

let optionals : [String?] = ["a", "b", nil, "d"]
let nonOptionals = optionals.flatMap{$0}
print(nonOptionals)

Prints:

[a, b, d]

Note:

There are 2 flatMap functions:

One flatMap is used to remove non-nil values which is shown above. Refer - https://developer.apple.com/documentation/swift/sequence/2907182-flatmap
The other flatMap is used to concatenate results. Refer - https://developer.apple.com/documentation/swift/sequence/2905332-flatmap

Solution 2:

TL;DR

Swift 4

Use array.compactMap { $0 }. Apple updated the framework so that it'll no longer cause bugs/confusion.

Swift 3

To avoid potential bugs/confusion, don't use array.flatMap { $0 } to remove nils; use an extension method such as array.removeNils() instead (implementation below, updated for Swift 3.0).

Although array.flatMap { $0 } works most of the time, there are several reasons to favor an array.removeNils() extension:

removeNils describes exactly what you want to do: remove nil values. Someone not familiar with flatMap would need to look it up, and, when they do look it up, if they pay close attention, they'll come to the same conclusion as my next point;
flatMap has two different implementations which do two entirely different things. Based off of type-checking, the compiler is going to decide which one is invoked. This can be very problematic in Swift, since type-inference is used heavily. (E.g. to determine the actual type of a variable, you may need to inspect multiple files.) A refactor could cause your app to invoke the wrong version of flatMap which could lead to difficult-to-find bugs.
Since there are two completely different functions, it makes understanding flatMap much more difficult since you can easily conflate the two.
flatMap can be called on non-optional arrays (e.g. [Int]), so if you refactor an array from [Int?] to [Int] you may accidentally leave behind flatMap { $0 } calls which the compiler won't warn you about. At best it'll simply return itself, at worst it'll cause the other implementation to be executed, potentially leading to bugs.
In Swift 3, if you don't explicitly cast the return type, the compiler will choose the wrong version, which causes unintended consequences. (See Swift 3 section below)
Finally, it slows down the compiler because the type-checking system needs to determine which of the overloaded functions to call.

To recap, there are two versions of the function in question both, unfortunately, named flatMap.

Flatten sequences by removing a nesting level (e.g. [[1, 2], [3]] -> [1, 2, 3])

public struct Array<Element> : RandomAccessCollection, MutableCollection {
    /// Returns an array containing the concatenated results of calling the
    /// given transformation with each element of this sequence.
    ///
    /// Use this method to receive a single-level collection when your
    /// transformation produces a sequence or collection for each element.
    ///
    /// In this example, note the difference in the result of using `map` and
    /// `flatMap` with a transformation that returns an array.
    ///
    ///     let numbers = [1, 2, 3, 4]
    ///
    ///     let mapped = numbers.map { Array(count: $0, repeatedValue: $0) }
    ///     // [[1], [2, 2], [3, 3, 3], [4, 4, 4, 4]]
    ///
    ///     let flatMapped = numbers.flatMap { Array(count: $0, repeatedValue: $0) }
    ///     // [1, 2, 2, 3, 3, 3, 4, 4, 4, 4]
    ///
    /// In fact, `s.flatMap(transform)`  is equivalent to
    /// `Array(s.map(transform).joined())`.
    ///
    /// - Parameter transform: A closure that accepts an element of this
    ///   sequence as its argument and returns a sequence or collection.
    /// - Returns: The resulting flattened array.
    ///
    /// - Complexity: O(*m* + *n*), where *m* is the length of this sequence
    ///   and *n* is the length of the result.
    /// - SeeAlso: `joined()`, `map(_:)`
    public func flatMap<SegmentOfResult : Sequence>(_ transform: (Element) throws -> SegmentOfResult) rethrows -> [SegmentOfResult.Iterator.Element]
}

Remove elements from a sequence (e.g. [1, nil, 3] -> [1, 3])

public struct Array<Element> : RandomAccessCollection, MutableCollection {
    /// Returns an array containing the non-`nil` results of calling the given
    /// transformation with each element of this sequence.
    ///
    /// Use this method to receive an array of nonoptional values when your
    /// transformation produces an optional value.
    ///
    /// In this example, note the difference in the result of using `map` and
    /// `flatMap` with a transformation that returns an optional `Int` value.
    ///
    ///     let possibleNumbers = ["1", "2", "three", "///4///", "5"]
    ///
    ///     let mapped: [Int?] = numbers.map { str in Int(str) }
    ///     // [1, 2, nil, nil, 5]
    ///
    ///     let flatMapped: [Int] = numbers.flatMap { str in Int(str) }
    ///     // [1, 2, 5]
    ///
    /// - Parameter transform: A closure that accepts an element of this
    ///   sequence as its argument and returns an optional value.
    /// - Returns: An array of the non-`nil` results of calling `transform`
    ///   with each element of the sequence.
    ///
    /// - Complexity: O(*m* + *n*), where *m* is the length of this sequence
    ///   and *n* is the length of the result.
    public func flatMap<ElementOfResult>(_ transform: (Element) throws -> ElementOfResult?) rethrows -> [ElementOfResult]
}

#2 is the one that people use to remove nils by passing { $0 } as transform. This works since the method performs a map, then filters out all nil elements.

You may be wondering "Why did Apple not rename #2 to removeNils()"? One thing to keep in mind is that using flatMap to remove nils is not the only usage of #2. In fact, since both versions take in a transform function, they can be much more powerful than those examples above.

For example, #1 could easily split an array of strings into individual characters (flatten) and capitalize each letter (map):

["abc", "d"].flatMap { $0.uppercaseString.characters } == ["A", "B", "C", "D"]

While number #2 could easily remove all even numbers (flatten) and multiply each number by -1 (map):

[1, 2, 3, 4, 5, 6].flatMap { ($0 % 2 == 0) ? nil : -$0 } == [-1, -3, -5]

(Note that this last example can cause Xcode 7.3 to spin for a very long time because there are no explicit types stated. Further proof as to why the methods should have different names.)

The real danger of blindly using flatMap { $0 } to remove nils comes not when you call it on [1, 2], but rather when you call it on something like [[1], [2]]. In the former case, it'll call invocation #2 harmlessly and return [1, 2]. In the latter case, you may think it would do the same (harmlessly return [[1], [2]] since there are no nil values), but it will actually return [1, 2] since it's using invocation #1.

The fact that flatMap { $0 } is used to remove nils seems to be more of Swift community recommendation rather than one coming from Apple. Perhaps if Apple notices this trend, they'll eventually provide a removeNils() function or something similar.

Until then, we're left with coming up with our own solution.

Solution

// Updated for Swift 3.0
protocol OptionalType {
    associatedtype Wrapped
    func map<U>(_ f: (Wrapped) throws -> U) rethrows -> U?
}

extension Optional: OptionalType {}

extension Sequence where Iterator.Element: OptionalType {
    func removeNils() -> [Iterator.Element.Wrapped] {
        var result: [Iterator.Element.Wrapped] = []
        for element in self {
            if let element = element.map({ $0 }) {
                result.append(element)
            }
        }
        return result
    }
}

(Note: Don't get confused with element.map... it has nothing to do with the flatMap discussed in this post. It's using Optional's map function to get an optional type that can be unwrapped. If you omit this part, you'll get this syntax error: "error: initializer for conditional binding must have Optional type, not 'Self.Generator.Element'." For more information about how map() helps us, see this answer I wrote about adding an extension method on SequenceType to count non-nils.)

Usage

let a: [Int?] = [1, nil, 3]
a.removeNils() == [1, 3]

Example

var myArray: [Int?] = [1, nil, 2]
assert(myArray.flatMap { $0 } == [1, 2], "Flat map works great when it's acting on an array of optionals.")
assert(myArray.removeNils() == [1, 2])

var myOtherArray: [Int] = [1, 2]
assert(myOtherArray.flatMap { $0 } == [1, 2], "However, it can still be invoked on non-optional arrays.")
assert(myOtherArray.removeNils() == [1, 2]) // syntax error: type 'Int' does not conform to protocol 'OptionalType'

var myBenignArray: [[Int]?] = [[1], [2, 3], [4]]
assert(myBenignArray.flatMap { $0 } == [[1], [2, 3], [4]], "Which can be dangerous when used on nested SequenceTypes such as arrays.")
assert(myBenignArray.removeNils() == [[1], [2, 3], [4]])

var myDangerousArray: [[Int]] = [[1], [2, 3], [4]]
assert(myDangerousArray.flatMap { $0 } == [1, 2, 3, 4], "If you forget a single '?' from the type, you'll get a completely different function invocation.")
assert(myDangerousArray.removeNils() == [[1], [2, 3], [4]]) // syntax error: type '[Int]' does not conform to protocol 'OptionalType'

(Notice how on the last one, flatMap returns [1, 2, 3, 4] while removeNils() would have been expected to return [[1], [2, 3], [4]].)

The solution is similar to the answer @fabb linked to.

However, I made a few modifications:

I didn't name the method flatten, since there is already a flatten method for sequence types, and giving the same name to entirely different methods is what got us in this mess in the first place. Not to mention that it's much easier to misinterpret what flatten does than it is removeNils.
Rather than creating a new type T on OptionalType, it uses the same name that Optional uses (Wrapped).
Instead of performing map{}.filter{}.map{}, which leads to O(M + N) time, I loop through the array once.
Instead of using flatMap to go from Generator.Element to Generator.Element.Wrapped?, I use map. There's no need to return nil values inside the map function, so map will suffice. By avoiding the flatMap function, it's harder to conflate yet another (i.e. 3rd) method with the same name which has an entirely different function.

The one drawback to using removeNils vs. flatMap is that the type-checker may need a bit more hinting:

[1, nil, 3].flatMap { $0 } // works
[1, nil, 3].removeNils() // syntax error: type of expression is ambiguous without more context

// but it's not all bad, since flatMap can have similar problems when a variable is used:
let a = [1, nil, 3] // syntax error: type of expression is ambiguous without more context
a.flatMap { $0 }
a.removeNils()

I haven't looked into it much, but it seems you can add:

extension SequenceType {
  func removeNils() -> Self {
    return self
  }
}

if you want to be able to call the method on arrays that contain non-optional elements. This could make a massive rename (e.g. flatMap { $0 } -> removeNils()) easier.

Assigning to self is different than assigning to a new variable?!

Take a look at the following code:

var a: [String?] = [nil, nil]

var b = a.flatMap{$0}
b // == []

a = a.flatMap{$0}
a // == [nil, nil]

Surprisingly, a = a.flatMap { $0 } does not remove nils when you assign it to a, but it does remove nils when you assign it to b! My guess is that this has something to do with the overloaded flatMap and Swift choosing the one we didn't mean to use.

You could temporarily resolve the problem by casting it to the expected type:

a = a.flatMap { $0 } as [String]
a // == []

But this can be easy to forget. Instead, I would recommend using the removeNils() method above.

Update

Seems like there's a proposal to deprecate at least one of the (3) overloads of flatMap: https://github.com/apple/swift-evolution/blob/master/proposals/0187-introduce-filtermap.md