lens + prism + affine

optics.swift

// MARK: - Original material

// Brandon Williams - Lenses in Swift: https://youtu.be/ofjehH9f-CU
// Lenses and Prisms in Swift: a pragmatic approach: https://broomburgo.github.io/fun-ios/post/lenses-and-prisms-in-swift-a-pragmatic-approach/
// Lenses and Prisms in Swift - Elviro Rocca: https://youtu.be/8VhYFEAQ0FY
// Elviro Rocca - Advanced Swift Optics: https://youtu.be/ki2WSw2WXV4

// MARK: - Either

enum Either<Left, Right> {
    case left(Left)
    case right(Right)
}

// MARK: - Lens

struct Lens<Whole, Part> {
    let get: (Whole) -> Part
    let set: (Part, Whole) -> Whole
}

extension Lens {
    init(_ keyPath: WritableKeyPath<Whole, Part>) {
        get = { whole in whole[keyPath: keyPath] }
        set = { newPart, whole in
            var whole = whole
            whole[keyPath: keyPath] = newPart
            return whole
        }
    }
}

extension Lens {
    func modify(_ transform: @escaping (Part) -> Part) -> (Whole) -> Whole {
        return { whole in
            self.set(transform(self.get(whole)), whole)
        }
    }
}

extension Lens {
    static func zip<Part1, Part2>(
        _ lens1: Lens<Whole, Part1>,
        _ lens2: Lens<Whole, Part2>
    ) -> Lens<Whole, (Part1, Part2)> where Part == (Part1, Part2) {
        return Lens<Whole, (Part1, Part2)>(
            get: { whole in (lens1.get(whole), lens2.get(whole)) },
            set: { parts, whole in lens2.set(parts.1, lens1.set(parts.0, whole)) }
        )
    }
}

extension Lens {
    func compose<Subpart>(_ other: Lens<Part, Subpart>) -> Lens<Whole, Subpart> {
        return Lens<Whole,Subpart>(
            get: { whole in other.get(self.get(whole)) },
            set: { subpart, whole in self.set(other.set(subpart, self.get(whole)), whole) }
        )
    }
}

// MARK: - Prism

struct Prism<Whole, Part> {
    let tryGet: (Whole) -> Part?
    let set: (Part) -> Whole
}

extension Prism {
    func tryModify(_ transform: @escaping (Part) -> Part) -> (Whole) -> Whole {
        return { whole in
            self.tryGet(whole).map { self.set(transform($0)) } ?? whole
        }
    }
}

extension Prism {
    static func zip<Part1, Part2>(
        _ prism1: Prism<Whole, Part1>,
        _ prism2: Prism<Whole, Part2>
    ) -> Prism<Whole, Either<Part1, Part2>> where Part == Either<Part1, Part2> {
        return Prism<Whole, Either<Part1, Part2>>(
            tryGet: { whole in prism1.tryGet(whole).map(Either.left) ?? prism2.tryGet(whole).map(Either.right) },
            set: { part in
                switch part {
                case let .left(value):  return prism1.set(value)
                case let .right(value): return prism2.set(value)
                }
            }
        )
    }
}

extension Prism {
    func compose<Subpart>(_ other: Prism<Part, Subpart>) -> Prism<Whole, Subpart> {
        return Prism<Whole, Subpart>(
            tryGet: { whole in self.tryGet(whole).flatMap(other.tryGet) },
            set: { subpart in self.set(other.set(subpart)) }
        )
    }
}

extension Prism {
    func isCase(_ whole: Whole) -> Bool {
        return tryGet(whole) != nil
    }
}

// MARK: - Affine

struct Affine<Whole, Part> {
    let tryGet: (Whole) -> Part?
    let trySet: (Part, Whole) -> Whole?
}

extension Affine {
    func compose<Subpart>(_ other: Affine<Part, Subpart>) -> Affine<Whole, Subpart> {
        return Affine<Whole, Subpart>(
            tryGet: { whole in self.tryGet(whole).flatMap(other.tryGet) },
            trySet: { subpart, whole in self.tryGet(whole).flatMap { other.trySet(subpart, $0) }.flatMap { self.trySet($0, whole) } }
        )
    }
}

extension Lens {
    var affine: Affine<Whole, Part> {
        return Affine<Whole, Part>(
            tryGet: self.get,
            trySet: self.set
        )
    }
}

extension Prism {
    var affine: Affine<Whole, Part> {
        return Affine<Whole, Part>(
            tryGet: self.tryGet,
            trySet: { part, _ in self.set(part) }
        )
    }
}

// MARK: - Operators

precedencegroup LeftCompositionPrecedence {
    associativity: left
}

infix operator .. : LeftCompositionPrecedence

extension Lens {
    static func .. <Subpart>(lhs: Lens<Whole, Part>, rhs: Lens<Part, Subpart>) -> Lens<Whole, Subpart> {
        return lhs.compose(rhs)
    }
}

extension Prism {
    static func .. <Subpart>(lhs: Prism<Whole, Part>, rhs: Prism<Part, Subpart>) -> Prism<Whole, Subpart> {
        return lhs.compose(rhs)
    }
}

extension Affine {
    static func .. <Subpart>(lhs: Affine<Whole, Part>, rhs: Affine<Part, Subpart>) -> Affine<Whole, Subpart> {
        return lhs.compose(rhs)
    }
}

extension Lens {
    static func .. <Subpart>(lhs: Lens<Whole, Part>, rhs: Affine<Part, Subpart>) -> Affine<Whole, Subpart> {
        return lhs.affine .. rhs
    }

    static func .. <Subpart>(lhs: Affine<Whole, Part>, rhs: Lens<Part, Subpart>) -> Affine<Whole, Subpart> {
        return lhs .. rhs.affine
    }

    static func .. <Subpart>(lhs: Lens<Whole, Part>, rhs: Prism<Part, Subpart>) -> Affine<Whole, Subpart> {
        return lhs .. rhs.affine
    }
}

extension Prism {
    static func .. <Subpart>(lhs: Prism<Whole, Part>, rhs: Affine<Part, Subpart>) -> Affine<Whole, Subpart> {
        return lhs.affine .. rhs
    }

    static func .. <Subpart>(lhs: Affine<Whole, Part>, rhs: Prism<Part, Subpart>) -> Affine<Whole, Subpart> {
        return lhs .. rhs.affine
    }

    static func .. <Subpart>(lhs: Prism<Whole, Part>, rhs: Lens<Part, Subpart>) -> Affine<Whole, Subpart> {
        return lhs .. rhs.affine
    }
}

Example

enum Event {
    case application(Application, String)
    case login(Login)

    enum Application {
        case didBecomeActive
        case openURL
    }

    enum Login {
        case tryLogin(outcome: LoginOutcome)
        case logout(motivation: LogoutMotivation)

        enum LoginOutcome {
            case success
            case failure(message: String)
        }

        enum LogoutMotivation {
            case manual
            case sessionExpired
        }
    }
}

// MARK: - Lenses

extension Event {
    static func lens<Value>(_ keyPath: WritableKeyPath<Event, Value>) -> Lens<Event, Value> {
        return Lens(keyPath)
    }
}

extension Event.Application {
    static func lens<Value>(_ keyPath: WritableKeyPath<Event.Application, Value>) -> Lens<Event.Application, Value> {
        return Lens(keyPath)
    }
}

extension Event.Login {
    static func lens<Value>(_ keyPath: WritableKeyPath<Event.Login, Value>) -> Lens<Event.Login, Value> {
        return Lens(keyPath)
    }
}

extension Event.Login.LoginOutcome {
    static func lens<Value>(_ keyPath: WritableKeyPath<Event.Login.LoginOutcome, Value>) -> Lens<Event.Login.LoginOutcome, Value> {
        return Lens(keyPath)
    }
}

extension Event.Login.LogoutMotivation {
    static func lens<Value>(_ keyPath: WritableKeyPath<Event.Login.LogoutMotivation, Value>) -> Lens<Event.Login.LogoutMotivation, Value> {
        return Lens(keyPath)
    }
}

// MARK: - Prisms

extension Event {
    static var loginPrism: Prism<Event, Event.Login> {
        return Prism<Event, Event.Login>(
            tryGet: { whole in
                guard case let .login(part) = whole else { return nil }
                return part
            },
            set: { part in .login(part) }
        )
    }

    static var applicationPrism: Prism<Event, (Event.Application, String)> {
        return Prism<Event, (Event.Application, String)>(
            tryGet: { whole in
                guard case let .application(part) = whole else { return nil }
                return part
            },
            set: { part in .application(part.0, part.1) }
        )
    }
}

extension Event.Application {
    static var didBecomeActivePrism: Prism<Event.Application, Void> {
        return Prism<Event.Application, Void>(
            tryGet: { whole in
                guard case .didBecomeActive = whole else { return nil }
                return ()
            },
            set: { part in .didBecomeActive }
        )
    }
}

extension Event.Login {
    static var logoutPrism: Prism<Event.Login, Event.Login.LogoutMotivation> {
        return Prism<Event.Login, Event.Login.LogoutMotivation>(
            tryGet: { whole in
                guard case let .logout(part) = whole else { return nil }
                return part
            },
            set: { part in .logout(motivation: part) }
        )
    }
}

extension Event.Login.LogoutMotivation {
    static var sessionExpiredPrism: Prism<Event.Login.LogoutMotivation, Void> {
        return Prism<Event.Login.LogoutMotivation, Void>(
            tryGet: { whole in
                guard case .sessionExpired = whole else { return nil }
                return ()
            },
            set: { part in .sessionExpired }
        )
    }
}

// MARK: - Helpers

func firstLens<A, B>() -> Lens<(A, B), A> {
    return Lens<(A, B), A>(
        get: { whole in whole.0 },
        set: { part, whole in (part, whole.1) }
    )
}

func secondLens<A, B>() -> Lens<(A, B), B> {
    return Lens<(A, B), B>(
        get: { whole in whole.1 },
        set: { part, whole in (whole.0, part) }
    )
}

// MARK: - Usage

let event = Event.login(.logout(motivation: .sessionExpired))

(Event.loginPrism .. Event.Login.logoutPrism .. Event.Login.LogoutMotivation.sessionExpiredPrism).isCase(event) // true
(Event.applicationPrism .. firstLens() .. Event.Application.didBecomeActivePrism).tryGet(event) // nil