Implementing Foldable, Map and FlatMap

Thanks to Lukas for the illustration

I have been spending a good portion of my time studying SICP. SICP avoids focusing on the terms such as Functor (Map) and Monad (FlatMap) and instead focuses solely on their implementations and the power they provide us.

You can find all the code below in the following playground:

TYRONEMICHAEL/foldableMapFlatMap

The implementations are surprisingly simple. Look at the implementation of Optional’s Map in Swift’s very own source code.

public func map<U>(
    _ transform: (Wrapped) throws -> U) rethrows -> U? {    
    switch self {    
    case .some(let y):      
        return .some(try transform(y))    
    case .none:      
        return .none    
    }  
}

Most developers I know make heavy use of Map and FlatMap. I do think it is of utmost importance to understand their implementations in order to understand how they actually work across different Types. As an exercise try code Array’s implementation of Map and FlatMap before reading further.

In order to implement FlatMap, we will use two functions as building blocks, reduce (sometimes referred to as Accumulate or FoldRight) and Map. So let’s implement reduce first.

func reduce<A, B>(
    _ fn: (A, B) -> B, _ initial: B, _ list: [A]) -> B {
    guard let value = list.first else {
       return initial
    }
    
    return fn(value, reduce(fn, initial, Array(list.dropFirst(1))))
}

The above function is recursive. It executes in the following way:

let sum: (Int, Int) -> Int = { a, b in a + b }
reduce(sum, 0, [1, 2, 3])
// return fn(value, reduce(fn, initial, Array(list.dropFirst(1))))
// fn(1, fn(2, fn(3, 0)))
// fn(1, fn(2, 3))
// fn(1, 5)
// 6

Now you can see why it is sometimes referred to as FoldRight. It executes from right to left. If we want to write an implementation for FoldLeft, we can write the same function iteratively.

func foldLeft<A, B>(
    _ fn: (A, B) -> B, _ initial: B, _ list: [A]) -> B {
    
    func iter(_ fn: (A, B) -> B, _ result: B, _ list: [A]) -> B {
        guard let value = list.first else {
           return result
        }
        return iter(fn, fn(value, result), Array(list.dropFirst(1)))
    }
    
    return iter(fn, initial, list)
}

Now let’s implement Map.

func map<A, B>(_ fn: (A) -> B, _ list: [A]) -> [B] {
    guard let value = list.first else {
        return []
    }
    
    return [fn(value)] + map(fn, dropFirst(list))
}

Pretty simple. Do you see the correlation between Optional’s implementation and the Array’s implementation? Both unwrap the value and pass it to the function, take the result produced by the function, and finally wrap the altered value back up in its corresponding type: .some(fn(value)) & [fn(value)]. If you look at any other Type’s implementation of Map, you will find the exact same thing SomeType(fn(value)).

Now let’s implement the final function, FlatMap. The simplest of the three, but the most dense.

func flatMap<A, B>(_ fn: (A) -> [B], _ list: [A]) -> [B] {
    return foldRight(+, [], map(fn, list))
}
let square: (Int) -> [Int] = { x in [x * x] } 
map(square, [1, 2, 3]) // [[1], [4], [9]]
flatMap(square, [1, 2, 3]) // [1, 4, 9]

Spend time thinking about how the above using both FoldRight (the reduce method we implemented in the beginning) and Map produces the correct result.

The above is exactly why functional programming, to me, is so powerful. The combination of FoldRight and Map to produce FlatMap is pure elegance.

FlatMap just takes the value and unlike Map, does not re-wrap the value in its corresponding type. That’s it. So if you have a nested type you solve it with the following use of FlatMap:

let value = Optional.some(Optional.some(5))
print(value.flatMap { $0 }) // Optional.some(5)

The main inspiration for the article came from the following problem introduced to me in SICP, and is a fitting conclusion:

Suppose we wish to generate all the permutations of a set S; that is, all the ways of ordering the items in the set. For instance, the permutations [1, 2, 3] are [1, 2, 3], [1, 3, 2], [2, 1, 3], [2, 3 1], [3, 1, 2], [3, 2, 1].

We can solve the problem using only FlatMap and Map (no loops or assignment). Something that took me quite some time to wrap my head around and I hope that it inspires you the same way it did me.

func remove<A: Equatable>(_ a: A, _ arr: [A]) -> [A] {
    return arr.filter { x in a != x }
}
func permutations(_ list: [Int]) -> [[Int]] {
    if list.isEmpty {
        return [[]]
    }
    return flatMap({ x in
        return map({ y in
            [x] + y
        }, permutations(remove(x, list)))
    }, list)
}
permutations([1, 2, 3])
// [
//  [1, 2, 3], [1, 3, 2], [2, 1, 3], [2, 3, 1], [3, 1, 2], [3, 2, 1]
// ]

If you enjoyed this article, please get in touch with me below:

Tyrone Avnit

All the code above can be found here:

TYRONEMICHAEL/foldableMapFlatMap