Merge pull request #81 from KacperFKorban/optimize-modifyAll

Optimize `modifyAll`

quicklens/src/main/scala-3/com/softwaremill/quicklens/QuicklensMacros.scala

@@ -15,39 +15,46 @@ object QuicklensMacros {
 
   def from[T: Type, R: Type](f: Expr[T => R])(using Quotes): Expr[T] => Expr[R] = (x: Expr[T]) => '{ $f($x) }
 
-  def modifyLensApplyImpl[T, U](path: Expr[T => U])(using Quotes, Type[T], Type[U]): Expr[PathLazyModify[T, U]] =
-    '{ PathLazyModify((t, mod) => ${ modifyImpl('t, path) }.using(mod)) }
-
-  def modifyAllLensApplyImpl[T, U](path1: Expr[T => U], paths: Expr[Seq[T => U]])(using
-      Quotes,
-      Type[T],
-      Type[U]
-  ): Expr[PathLazyModify[T, U]] =
+  def modifyLensApplyImpl[T, U](path: Expr[T => U])(using Quotes, Type[T], Type[U]): Expr[PathLazyModify[T, U]] = '{
+    PathLazyModify { (t, mod) =>
+      ${
+        toPathModify('t, modifyImpl('t, path))
+      }.using(mod)
+    }
+  }
+
+  def modifyAllLensApplyImpl[T: Type, U: Type](
+      path1: Expr[T => U],
+      paths: Expr[Seq[T => U]]
+  )(using Quotes): Expr[PathLazyModify[T, U]] =
     '{ PathLazyModify((t, mod) => ${ modifyAllImpl('t, path1, paths) }.using(mod)) }
 
-  def modifyAllImpl[S, A](obj: Expr[S], focus: Expr[S => A], focusesExpr: Expr[Seq[S => A]])(using
-      Quotes,
-      Type[S],
-      Type[A]
-  ): Expr[PathModify[S, A]] = {
+  def modifyAllImpl[S: Type, A: Type](
+      obj: Expr[S],
+      focus: Expr[S => A],
+      focusesExpr: Expr[Seq[S => A]]
+  )(using Quotes): Expr[PathModify[S, A]] = {
     import quotes.reflect.*
 
     val focuses = focusesExpr match {
       case Varargs(args) => args
     }
 
-    val modF1 = fromPathModify(modifyImpl(obj, focus))
-    val modF = to[(A => A), S] { (mod: Expr[A => A]) =>
+    val modF1 = modifyImpl(obj, focus)
+    val modF = { (mod: Expr[A => A]) =>
       focuses.foldLeft(from[(A => A), S](modF1).apply(mod)) { case (objAcc, focus) =>
-        val modCur = fromPathModify(modifyImpl(objAcc, focus))
+        val modCur = modifyImpl(objAcc, focus)
         from[(A => A), S](modCur).apply(mod)
       }
     }
 
-    toPathModify(obj, modF)
+    toPathModify(obj, to(modF))
   }
 
-  def modifyImpl[S, A](obj: Expr[S], focus: Expr[S => A])(using Quotes, Type[S], Type[A]): Expr[PathModify[S, A]] = {
+  def toPathModifyFromFocus[S: Type, A: Type](obj: Expr[S], focus: Expr[S => A])(using Quotes): Expr[PathModify[S, A]] =
+    toPathModify(obj, modifyImpl(obj, focus))
+
+  private def modifyImpl[S: Type, A: Type](obj: Expr[S], focus: Expr[S => A])(using Quotes): Expr[(A => A) => S] = {
     import quotes.reflect.*
 
     def unsupportedShapeInfo(tree: Tree) =
@@ -197,6 +204,6 @@ object QuicklensMacros {
 
     val res: (Expr[A => A] => Expr[S]) = (mod: Expr[A => A]) =>
       mapToCopy(Symbol.spliceOwner, mod, obj.asTerm, path).asExpr.asInstanceOf[Expr[S]]
-    toPathModify(obj, to(res))
+    to(res)
   }
 }

quicklens/src/main/scala-3/com/softwaremill/quicklens/package.scala

@@ -14,7 +14,7 @@ package object quicklens {
       *
       * You can use `.each` to traverse options, lists, etc.
       */
-    inline def modify(inline path: S => A): PathModify[S, A] = ${ modifyImpl('obj, 'path) }
+    inline def modify(inline path: S => A): PathModify[S, A] = ${ toPathModifyFromFocus('obj, 'path) }
 
     /** Create an object allowing modifying the given (deeply nested) fields accessible in a `case class` hierarchy via
       * `paths` on the given `obj`.

quicklens/src/test/scala/com/softwaremill/quicklens/HugeModifyTest.scala

@@ -0,0 +1,64 @@
+package com.softwaremill.quicklens
+
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+
+class HugeModifyTest extends AnyFlatSpec with Matchers {
+  import HugeModifyTestData._
+
+  it should "expand a huge function" in {
+    val c5 = C5(1)
+    val c4 = C4(c5, c5, c5, c5)
+    val c3 = C3(c4, c4, c4, c4)
+    val c2 = C2(c3, c3, c3, c3)
+    val c1 = C1(c2, c2, c2, c2)
+
+    val c5e = C5(2)
+    val c4e = C4(c5e, c5, c5, c5)
+    val c3e = C3(c4e, c4, c4, c4)
+    val c2e = C2(c3e, c3, c3, c3)
+    val c1e = C1(c2e, c2e, c2e, c2)
+
+    val res = c1
+      .modifyAll(
+        _.a.a.a.a.a,
+        _.b.a.a.a.a,
+        _.c.a.a.a.a
+      ).using(_ + 1)
+    res should be(c1e)
+  }
+}
+
+object HugeModifyTestData {
+  case class C1(
+    a: C2,
+    b: C2,
+    c: C2,
+    d: C2
+  )
+
+  case class C2(
+    a: C3,
+    b: C3,
+    c: C3,
+    d: C3
+  )
+
+  case class C3(
+    a: C4,
+    b: C4,
+    c: C4,
+    d: C4
+  )
+
+  case class C4(
+    a: C5,
+    b: C5,
+    c: C5,
+    d: C5
+  )
+
+  case class C5(
+    a: Int
+  )
+}