Create a gist now

Instantly share code, notes, and snippets.

What would you like to do?
import scalaz.{LensFamily, Lens}
import Lens.{lensg, lensFamilyg}
object LensFns {
/**
* A lens that requires a key be provided to extract B from A.
*
* @tparam A The record type.
* @tparam K The field key type.
* @tparam B The field type.
*/
def KeyedLens[A, K, B](set: A => K => B => A, get: A => K => B) = KeyedLens2(set, get)
/**
* A lens that requires a key be provided to extract B from A.
*
* @tparam A1 The initial record type.
* @tparam A2 The final record type.
* @tparam K The field key type.
* @tparam B1 The initial field type.
* @tparam B2 The final field type.
*/
def KeyedLens2[A1, A2, K, B1, B2](set: A1 => K => B2 => A2, get: A1 => K => B1) =
lensFamilyg[(A1, K), A2, B1, B2](
ak => b => set(ak._1)(ak._2)(b),
ak => get(ak._1)(ak._2)
)
private def composeKK[A1, A2, K1, B1, B2, K2, C1, C2](g: LensFamily[(A1, K1), A2, B1, B2], f: LensFamily[(B1, K2), B2, C1, C2]) =
KeyedLens2[A1, A2, (K1, K2), C1, C2](
a => k12 => c => {
val (k1, k2) = k12
val b1 = g.get((a, k1))
val b2 = f.set((b1, k2), c)
g.set((a, k1), b2)
},
a => k12 => {
val (k1, k2) = k12
val b = g.get(a, k1)
f.get(b, k2)
}
)
private def composeLK[A1, A2, B1, B2, K, C1, C2](g: LensFamily[A1, A2, B1, B2], f: LensFamily[(B1, K), B2, C1, C2]) =
KeyedLens2[A1, A2, K, C1, C2](
a => k => c => {
val b1 = g.get(a)
val b2 = f.set((b1, k), c)
g.set(a, b2)
},
a => k => {
val b = g.get(a)
f.get(b, k)
}
)
private def composeKL[A1, A2, K, B1, B2, C1, C2](g: LensFamily[(A1, K), A2, B1, B2], f: LensFamily[B1, B2, C1, C2]) =
KeyedLens2[A1, A2, K, C1, C2](
a => k => c => {
val b1 = g.get(a, k)
val b2 = f.set(b1, c)
g.set((a, k), b2)
},
a => k => f.get(g.get(a, k))
)
implicit class KeyedLensExt[A1, A2, K, B1, B2](val l: LensFamily[(A1, K), A2, B1, B2]) extends AnyVal {
def >@=@>[K2, C1, C2](f: LensFamily[(B1, K2), B2, C1, C2]) = composeKK(l, f)
def >@=>[C1, C2](f: LensFamily[B1, B2, C1, C2]) = composeKL(l, f)
def <=@<[Z1, Z2](g: LensFamily[Z1, Z2, A1, A2]) = composeLK(g, l)
def <@=@<[Z1, Z2, K0](g: LensFamily[(Z1, K0), Z2, A1, A2]) = composeKK(g, l)
}
implicit class PlainLensExt[A1, A2, B1, B2](val l: LensFamily[A1, A2, B1, B2]) extends AnyVal {
def >=@>[K, C1, C2](f: LensFamily[(B1, K), B2, C1, C2]) = composeLK(l, f)
def <@=<[Z1, Z2, K](g: LensFamily[(Z1, K), Z2, A1, A2]) = composeKL(g, l)
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment