Lens.swift

import Prelude

public struct Lens<A, B> {
    private let get: A -> B
    private let set: (A, B) -> A

    public init(get: A -> B, set: (A, B) -> A) {
        self.get = get
        self.set = set
    }

    public init(get: A -> B, set: (inout A, B) -> ()) {
        self.get = get
        self.set = {
            var copy = $0; set(&copy, $1); return copy
        }
    }
}

// MARK: - Basics

public func get<A, B>(lens: Lens<A, B>, a: A) -> B {
    return lens.get(a)
}

public func get<A, B>(lens: Lens<A, B>)(a: A) -> B {
    return lens.get(a)
}

public func get<A, B>(lens: Lens<A, B>, a: A?) -> B? {
    return map(a, lens.get)
}

public func get<A, B>(lens: Lens<A, B>)(a: A?) -> B? {
    return a.map(lens.get)
}

public func set<A, B>(lens: Lens<A, B>, a: A, b: B) -> A {
    return lens.set(a, b)
}

public func set<A, B>(lens: Lens<A, B>, a: A)(b: B) -> A {
    return lens.set(a, b)
}

public func mod<A, B>(lens: Lens<A, B>, a: A, f: B -> B) -> A {
    return get(lens, a) |> f |> set(lens, a)
}

// MARK: - Compose

public func compose<A, B, C>(left: Lens<A, B>, right: Lens<B, C>) -> Lens<A, C> {
    let get: A -> C = { a in
        right.get(left.get(a))
    }

    let set: (A, C) -> A = { a, c in
        left.set(a, right.set(left.get(a), c))
    }

    return Lens(get: get, set: set)
}

public func >>><A, B, C>(lhs: Lens<A, B>, rhs: Lens<B, C>) -> Lens<A, C> {
    return compose(lhs, rhs)
}

// MARK: - Lift

public func lift<A, B>(lens: Lens<A, B>) -> Lens<[A], [B]> {
    let get: [A] -> [B] = { input in
        return map(input, lens.get)
    }

    let set: ([A], [B]) -> [A] = { xs, ys in
        return zip(xs, ys, lens.set)
    }

    return Lens(get: get, set: set)
}

// MARK: - Split

public func split<A, B, C, D>(left: Lens<A, B>, right: Lens<C, D>) -> Lens<(A, C), (B, D)> {
    let get: (A, C) -> (B, D) = { (a, c) in
        (left.get(a), right.get(c))
    }

    let set: ((A, C), (B, D)) -> (A, C) = {
        (left.set($0.0, $1.0), right.set($0.1, $1.1))
    }

    return Lens(get: get, set: set)
}

infix operator *** {
    associativity left
    precedence 150
}

public func ***<A, B, C, D>(lhs: Lens<A, B>, rhs: Lens<C, D>) -> Lens<(A, C), (B, D)> {
    return split(lhs, rhs)
}

// MARK: - Fanout

public func fanout<A, B, C>(left: Lens<A, B>, right: Lens<A, C>) -> Lens<A, (B, C)> {
    let get: A -> (B, C) = { a in
        (left.get(a), right.get(a))
    }

    let set: (A, (B, C)) -> A = { (a, input) in
        right.set(left.set(a, input.0), input.1)
    }

    return Lens(get: get, set: set)
}

infix operator &&& {
    associativity left
    precedence 120
}

public func &&&<A, B, C>(lhs: Lens<A, B>, rhs: Lens<A, C>) -> Lens<A, (B, C)> {
    return fanout(lhs, rhs)
}

LensSpec.swift

import Prelude

import Nimble
import Quick

// MARK: Example struct

struct Outer: Equatable {
    var count: Int

    var inner: Inner

    init(count: Int = 0, inner: Inner = Inner(count: 0)) {
        self.count = count
        self.inner = inner
    }
}

func ==(lhs: Outer, rhs: Outer) -> Bool {
    return lhs.count == rhs.count
        && lhs.inner == rhs.inner
}

struct Inner: Equatable {
    var count: Int
}

func ==(lhs: Inner, rhs: Inner) -> Bool {
    return lhs.count == rhs.count
}

struct OuterLenses {
    static let count = Lens<Outer, Int>(get: { $0.count }, set: { (inout outer: Outer, count) in
        outer.count = count
    })

    static let inner = Lens<Outer, Inner>(get: { $0.inner }, set: { (inout outer: Outer, inner) in
        outer.inner = inner
    })
}

struct InnerLenses {
    static let count = Lens<Inner, Int>(get: { $0.count }, set: { (inout inner: Inner, count) in
        inner.count = count
    })
}

class LensSpec: QuickSpec {
    override func spec() {
        describe("A Lens") {
            let example: Outer = Outer(count: 2)

            let count = OuterLenses.count

            it("should get values") {
                expect(get(count, example)!).to(equal(2))
            }

            it("should set values") {
                expect(set(count, example, 4)).to(equal(Outer(count: 4)))
            }

            it("should modify values") {
                expect(mod(count, example, { $0 + 2 })).to(equal(Outer(count: 4)))
            }
        }

        describe("A composed Lens") {
            let example = Outer(count: 0, inner: Inner(count: 2))

            let innerCount = OuterLenses.inner >>> InnerLenses.count

            it("should get values") {
                expect(get(innerCount, example)!).to(equal(2))
            }

            it("should set values") {
                expect(set(innerCount, example, 4)).to(equal(Outer(count: 0, inner: Inner(count: 4))))
            }

            it("should modify values") {
                expect(mod(innerCount, example, { $0 + 2 })).to(equal(Outer(count: 0, inner: Inner(count: 4))))
            }
        }

        describe("Lifted lenses") {
            let inner = [
                Inner(count: 1),
                Inner(count: 2),
                Inner(count: 3),
                Inner(count: 4)
            ]

            let lifted = lift(InnerLenses.count)

            it("should get values") {
                let result = get(lifted, inner)

                expect(result).to(equal([ 1, 2, 3, 4 ]))
            }

            it("should set values") {
                let result = set(lifted, inner, [ 2, 4, 6, 8 ])

                expect(result).to(equal([
                    Inner(count: 2),
                    Inner(count: 4),
                    Inner(count: 6),
                    Inner(count: 8)
                ]))
            }

            it("should reduce the resulting array size accordingly") {
                // Does this make sense?
                let result = set(lifted, inner, [ 42 ])

                expect(result).to(equal([
                    Inner(count: 42)
                ]))
            }
        }

        describe("Split lenses") {
            let outer = Outer(count: 2, inner: Inner(count: 4))
            let inner = Inner(count: 9)

            let both = OuterLenses.count *** InnerLenses.count

            it("should get values") {
                let result = get(both, (outer, inner))

                expect(result.0).to(equal(2))
                expect(result.1).to(equal(9))
            }

            it("should set values") {
                let result = set(both, (outer, inner)) <| (12, 34)

                expect(result.0.count).to(equal(12))
                expect(result.0.inner.count).to(equal(4))
                expect(result.1.count).to(equal(34))
            }
        }

        describe("Fanned out lenses") {
            let example = Outer(count: 0, inner: Inner(count: 2))

            let both = OuterLenses.count &&& (OuterLenses.inner >>> InnerLenses.count)

            it("should get values") {
                let result = get(both, example)

                expect(result.0).to(equal(0))
                expect(result.1).to(equal(2))
            }

            it("should set values") {
                let result = set(both, example) <| (12, 34)

                expect(result.count).to(equal(12))
                expect(result.inner.count).to(equal(34))
            }
        }
    }
}