code for the article "Lenses and Prisms in Swift: a pragmatic approach"

lenses-and-prisms-in-swift-a-pragmatic-approach.swift

/// source: https://broomburgo.github.io/fun-ios/post/lenses-and-prisms-in-swift-a-pragmatic-approach/

protocol LensType {
	associatedtype WholeType
	associatedtype PartType

	var get: (WholeType) -> PartType { get }
	var set: (PartType,WholeType) -> WholeType { get }

	init(get: @escaping (WholeType) -> PartType, set: @escaping (PartType,WholeType) -> WholeType)
}

struct Lens<Whole,Part>: LensType {
	typealias WholeType = Whole
	typealias PartType = Part

	let get: (Whole) -> Part /// get the "focused" part
	let set: (Part,Whole) -> Whole /// set a new value for the "focused" part

	init(get: @escaping (Whole) -> Part, set: @escaping (Part,Whole) -> Whole) {
		self.get = get
		self.set = set
	}
}

extension LensType {
	func over(_ transform: @escaping (PartType) -> PartType) -> (WholeType) -> WholeType {
		return { whole in self.set(transform(self.get(whole)),whole) }
	}

	func join<Subpart, OtherLens>(_ other: OtherLens) -> Lens<WholeType,Subpart> where OtherLens: LensType, OtherLens.WholeType == PartType, OtherLens.PartType == Subpart {
		return Lens<WholeType,Subpart>(
			get: { other.get(self.get($0)) },
			set: { (subpart: Subpart, whole: WholeType) -> WholeType in
					self.set(other.set(subpart,self.get(whole)),whole)
		})
	}

	func zip<OtherPart, OtherLens>(_ other: OtherLens) -> Lens<WholeType,(PartType,OtherPart)> where OtherLens: LensType, OtherLens.WholeType == WholeType, OtherLens.PartType == OtherPart {
		return Lens<WholeType,(PartType,OtherPart)>(
			get: { (self.get($0),other.get($0)) },
			set: { other.set($0.1,self.set($0.0,$1)) })
	}
}

struct Company {
	var name: String
	var board: BoardOfDirectors

	struct lens {
		static let name = Lens<Company,String>(
			get: { $0.name },
			set: { (p,w) in
				var m_w = w
				m_w.name = p
				return m_w
		})

		static let board = Lens<Company,BoardOfDirectors>(
			get: { $0.board },
			set: { (p,w) in
				var m_w = w
				m_w.board = p
				return m_w
		})
	}
}

struct BoardOfDirectors {
	var ceo: Employee
	var cto: Employee
	var cfo: Employee
	var coo: Employee

	struct lens {
		static let ceo = Lens<BoardOfDirectors,Employee>(
			get: { $0.ceo },
			set: { BoardOfDirectors(ceo: $0, cto: $1.cto, cfo: $1.cfo, coo: $1.coo) })

		static let cto = Lens<BoardOfDirectors,Employee>(
			get: { $0.cto },
			set: { BoardOfDirectors(ceo: $1.ceo, cto: $0, cfo: $1.cfo, coo: $1.coo) })

		static let cfo = Lens<BoardOfDirectors,Employee>(
			get: { $0.cfo },
			set: { BoardOfDirectors(ceo: $1.ceo, cto: $1.cto, cfo: $0, coo: $1.coo) })

		static let coo = Lens<BoardOfDirectors,Employee>(
			get: { $0.coo },
			set: { BoardOfDirectors(ceo: $1.ceo, cto: $1.cto, cfo: $1.cfo, coo: $0) })
	}
}

struct Employee {
	var name: String
	var salary: Salary

	struct lens {
		static let name = Lens<Employee,String>(
			get: { $0.name },
			set: { (p,w) in
				var m_w = w
				m_w.name = p
				return m_w
		})

		static let salary = Lens<Employee,Salary>(
			get: { $0.salary },
			set: { (p,w) in
				var m_w = w
				m_w.salary = p
				return m_w
		})
	}
}

struct Salary {
	var amount: Double
	var bonus: Double

	struct lens {
		static let amount = Lens<Salary,Double>(
			get: { $0.amount },
			set: { Salary(amount: $0, bonus: $1.bonus) })

		static let bonus = Lens<Salary,Double>(
			get: { $0.bonus },
			set: { Salary(amount: $1.amount, bonus: $0) })
	}
}

var currentCompany = Company(
	name: "Something inc.",
	board: BoardOfDirectors(
		ceo: Employee(
			name: "Jane Doe",
			salary: Salary(
				amount: 100,
				bonus: 10)),
		cto: Employee(
			name: "John Doe",
			salary: Salary(
				amount: 80,
				bonus: 8)),
		cfo: Employee(
			name: "Jane Doe Jr.",
			salary: Salary(
				amount: 80,
				bonus: 4)),
		coo: Employee(
			name: "John Doe Jr.",
			salary: Salary(
				amount: 80,
				bonus: 4))))

struct Repository {
	static var getCompany: Company {
		print("Retrieving company: \(currentCompany)\n")
		return currentCompany
	}

	static func setCompany(_ company: Company) {
		print("Saving company: \(company)\n")
		currentCompany = company
	}
}

func example1() {
	var com = Repository.getCompany
	com.board.ceo.salary.bonus *= 0.5
	Repository.setCompany(com)
}

extension Double {
	static func average(_ values: Double...) -> Double {
		return values.reduce(0, +)/Double(values.count)
	}
}

func example2() {
	var com = Repository.getCompany
	let ceoSalary = com.board.ceo.salary
	let cfoSalary = com.board.cfo.salary
	com.board.cto.salary.amount = Double.average(ceoSalary.amount,cfoSalary.amount)
	Repository.setCompany(com)
}

func example3() {
	var com = Repository.getCompany
	let cfoAmount = com.board.cfo.salary.amount
	var cfoBonus = com.board.cfo.salary.bonus
	if cfoBonus/cfoAmount < 0.06 {
		cfoBonus = cfoAmount*0.06
	}
	com.board.cfo.salary.bonus = cfoBonus
	Repository.setCompany(com)
}

let onCEO = Company.lens.board.join(BoardOfDirectors.lens.ceo)
let onCTO = Company.lens.board.join(BoardOfDirectors.lens.cto)
let onCFO = Company.lens.board.join(BoardOfDirectors.lens.cfo)
let onSalaryAmount = Employee.lens.salary.join(Salary.lens.amount)
let onSalaryBonus = Employee.lens.salary.join(Salary.lens.bonus)

func updateCompany(with action: (Company) -> Company) {
	Repository.setCompany(action(Repository.getCompany))
}

func example1_lenses() {
	updateCompany(with: onCEO.join(onSalaryAmount).over { $0*0.5 })
}

func example2_lenses() {
	updateCompany {
		onCTO.join(onSalaryAmount)
			.set(Double.average(
				onCEO.join(onSalaryAmount).get($0),
				onCFO.join(onSalaryAmount).get($0)),
			     $0)
	}
}

func example3_lenses() {
	updateCompany(with: onCFO.join(onSalaryAmount)
		.zip(onCFO.join(onSalaryBonus))
		.over { (amount,bonus) in
			if bonus/amount < 0.06 {
				return (amount,amount*0.06)
			} else {
				return (amount,bonus)
			}
	})
}

typealias AnyDict = [String:Any]

/* this is wrong
func anyDictLens<Part>(at key: String, as: Part) -> Lens<AnyDict,Part> {
	return Lens<AnyDict,Part>(
		get: { (whole: AnyDict) -> Part in whole[key] as! Part },
		set: { (part: Part, whole: AnyDict) -> AnyDict in
			var m_dict = whole
			m_dict[key] = part
			return m_dict
	})
}
*/

func anyDictLens<Part>(at key: String, as: Part) -> Lens<AnyDict,Part?> {
	return Lens<AnyDict,Part?>(
		get: { (whole: AnyDict) -> Part? in whole[key] as? Part },
		set: { (part: Part?, whole: AnyDict) -> AnyDict in
			var m_dict = whole
			m_dict[key] = part
			return m_dict
	})
}

let dict: AnyDict = ["user" : ["name" : "Mr. Creosote"]]

let lens1 = anyDictLens(at: "user", as: AnyDict.self)
let lens2 = anyDictLens(at: "name", as: String.self)

/* this won't compile
let nameLens = lens1.join(lens2)
*/

protocol PrismType {
	associatedtype WholeType
	associatedtype PartType

	var tryGet: (WholeType) -> PartType? { get }
	var inject: (PartType) -> WholeType { get }

	init(tryGet: @escaping (WholeType) -> PartType?, inject: @escaping (PartType) -> WholeType)
}

struct Prism<Whole,Part>: PrismType {
	typealias WholeType = Whole
	typealias PartType = Part

	let tryGet: (Whole) -> Part? /// get the part, if possible
	let inject: (Part) -> Whole /// changes the value to reflect the part that's injected in

	init(tryGet: @escaping (Whole) -> Part?, inject: @escaping (Part) -> Whole) {
		self.tryGet = tryGet
		self.inject = inject
	}
}

extension PrismType {
	func tryOver(_ transform: @escaping (PartType) -> PartType) -> (WholeType) -> WholeType? {
		return { whole in self.tryGet(whole).map { self.inject(transform($0)) } }
	}

	func join<OtherPart, OtherPrism>(_ other: OtherPrism) -> Prism<WholeType,OtherPart> where OtherPrism: PrismType, OtherPrism.WholeType == PartType, OtherPrism.PartType == OtherPart  {
		return Prism<WholeType,OtherPart>(
			tryGet: { self.tryGet($0).flatMap(other.tryGet) },
			inject: { self.inject(other.inject($0)) })
	}

	func zip<OtherWhole, OtherPrism>(_ other: OtherPrism) -> Prism<(WholeType,OtherWhole),PartType> where OtherPrism: PrismType, OtherPrism.PartType == PartType, OtherPrism.WholeType == OtherWhole {
		return Prism<(WholeType,OtherWhole),PartType>(
			tryGet: { (whole,otherWhole) in
				self.tryGet(whole) ?? other.tryGet(otherWhole)
		},
			inject: { part in
				(self.inject(part),other.inject(part))
		})
	}
}

func anyDictPrism<Part>(at key: String, as: Part.Type) -> Prism<AnyDict,Part> {
	return Prism<AnyDict,Part>(
		tryGet: { $0[key] as? Part },
		inject: { [key:$0] })
}

let prism1 = anyDictPrism(at: "user", as: AnyDict.self)
let prism2 = anyDictPrism(at: "name", as: String.self)

let namePrism = prism1.join(prism2)
let name = namePrism.tryGet(dict) /// it's Mr. Creosote!

let prism3 = anyDictPrism(at: "weight", as: String.self)

let weightPrism = prism1.join(prism3)
let weight = weightPrism.tryGet(dict) /// this is nil

let otherDict = prism1.join(prism3).inject("200 kg") /// this is ["user": ["weight": "200 kg"]]

func dictOverride(_ first: AnyDict, _ second: AnyDict) -> AnyDict {
	var m_dict = first
	for (key,value) in second {
		if let current = m_dict[key] {
			if let firstAnyDict = current as? AnyDict, let secondAnyDict = value as? AnyDict {
				m_dict[key] = dictOverride(firstAnyDict, secondAnyDict)
			} else {
				m_dict[key] = value
			}
		} else {
			m_dict[key] = value
		}
	}
	return m_dict
}

let fullDict = dictOverride(dict, otherDict) /// this is ["user": ["name": "Mr. Creosote", "weight": "200 kg"]]

let jsonExample1 = [
	"info" : [
		"image" : [
			"icon_urls" : [
				"https://image.org/1.png",
				"https://image.org/2.png",
				"https://image.org/3.png"]]]]

let jsonExample2 = [
	"info" : [
		"image" : [
			"icon_urls" : [String]()]]]

let jsonExample3 = [
	"info" : [
		"image" : [
			"WRONG_KEY" : [
				"https://image.org/1.png",
				"https://image.org/2.png",
				"https://image.org/3.png"]]]]

let infoDictPrism = anyDictPrism(at: "info", as: AnyDict.self)
let imageDictPrism = anyDictPrism(at: "image", as: AnyDict.self)
let iconURLsPrism = anyDictPrism(at: "icon_urls", as: [String].self)
let firstURLPrism = Prism<[String],String>(
	tryGet: { $0.first },
	inject: { [$0] })
let imageURLPrism1 = infoDictPrism
	.join(imageDictPrism)
	.join(iconURLsPrism)
	.join(firstURLPrism)

let firstURL = imageURLPrism1.tryGet(jsonExample1) /// this is "https://image.org/1.png"
let noFirstURL = imageURLPrism1.tryGet(jsonExample2) /// this is nil
let stillNoFirstURL = imageURLPrism1.tryGet(jsonExample3) /// this is nil

let jsonExample4 = [
	"info" : [
		"main_image_icon_url" : "https://image.org/3.png",
		"image" : [
			"icon_urls" : [
				"https://image.org/1.png",
				"https://image.org/2.png",
				"https://image.org/3.png"]]]]

let imageURLPrism2 = infoDictPrism.join(anyDictPrism(at: "main_image_icon_url", as: String.self))

let finalURL1 = imageURLPrism2.zip(imageURLPrism1).tryGet(jsonExample4,jsonExample4) /// this is "https://image.org/3.png"
let finalURL2 = imageURLPrism2.zip(imageURLPrism1).tryGet(jsonExample1,jsonExample1) /// this is "https://image.org/1.png"