From 617f307d177829a80df5f08566a8314ff8b9c80c Mon Sep 17 00:00:00 2001
From: Lorenzo Gabriele <lorenzolespaul@gmail.com>
Date: Mon, 30 Sep 2024 23:11:33 +0200
Subject: [PATCH] Cache also `withToolClasspath` using two-layered cache

---
 .../goyeau/mill/scalafix/ScalafixCache.scala  | 47 ++++++++++++++-----
 .../goyeau/mill/scalafix/ScalafixModule.scala |  8 +---
 2 files changed, 35 insertions(+), 20 deletions(-)

diff --git a/mill-scalafix/src/com/goyeau/mill/scalafix/ScalafixCache.scala b/mill-scalafix/src/com/goyeau/mill/scalafix/ScalafixCache.scala
index 72681a9..58aeebc 100644
--- a/mill-scalafix/src/com/goyeau/mill/scalafix/ScalafixCache.scala
+++ b/mill-scalafix/src/com/goyeau/mill/scalafix/ScalafixCache.scala
@@ -2,7 +2,10 @@ package com.goyeau.mill.scalafix
 
 import com.goyeau.mill.scalafix.CoursierUtils
 import coursier.core.Repository
+import mill.Agg
+import mill.scalalib.Dep
 import scalafix.interfaces.Scalafix
+import scalafix.interfaces.ScalafixArguments
 
 import java.util.concurrent.ConcurrentHashMap
 import scala.jdk.CollectionConverters._
@@ -10,17 +13,35 @@ import scala.ref.SoftReference
 
 private[scalafix] object ScalafixCache {
 
-  private val cache = new ConcurrentHashMap[(String, Seq[Repository]), SoftReference[Scalafix]]
-
-  def getOrElseCreate(scalaVersion: String, repositories: Seq[Repository]) =
-    cache.compute(
-      (scalaVersion, repositories),
-      {
-        case (_, v @ SoftReference(_)) => v
-        case _ =>
-          SoftReference(
-            Scalafix.fetchAndClassloadInstance(scalaVersion, repositories.map(CoursierUtils.toApiRepository).asJava)
-          )
-      }
-    )()
+  private val scalafixCache = new Cache[(String, java.util.List[coursierapi.Repository]), Scalafix](createFunction = {
+    case (scalaVersion, repositories) =>
+      Scalafix.fetchAndClassloadInstance(scalaVersion, repositories)
+  })
+
+  private val scalafixArgumentsCache =
+    new Cache[(String, Seq[Repository], Agg[Dep]), ScalafixArguments](createFunction = {
+      case (scalaVersion, repositories, scalafixIvyDeps) =>
+        val repos = repositories.map(CoursierUtils.toApiRepository).asJava
+        val deps  = scalafixIvyDeps.map(CoursierUtils.toCoordinates).iterator.toSeq.asJava
+        scalafixCache
+          .getOrElseCreate((scalaVersion, repos))
+          .newArguments()
+          .withToolClasspath(Seq.empty.asJava, deps, repos)
+    })
+
+  def getOrElseCreate(scalaVersion: String, repositories: Seq[Repository], scalafixIvyDeps: Agg[Dep]) =
+    scalafixArgumentsCache.getOrElseCreate((scalaVersion, repositories, scalafixIvyDeps))
+
+  private class Cache[A, B <: AnyRef](createFunction: A => B) {
+    private val cache = new ConcurrentHashMap[A, SoftReference[B]]
+
+    def getOrElseCreate(a: A) =
+      cache.compute(
+        a,
+        {
+          case (_, v @ SoftReference(_)) => v
+          case _                         => SoftReference(createFunction(a))
+        }
+      )()
+  }
 }
diff --git a/mill-scalafix/src/com/goyeau/mill/scalafix/ScalafixModule.scala b/mill-scalafix/src/com/goyeau/mill/scalafix/ScalafixModule.scala
index e63c61c..db2a6c2 100644
--- a/mill-scalafix/src/com/goyeau/mill/scalafix/ScalafixModule.scala
+++ b/mill-scalafix/src/com/goyeau/mill/scalafix/ScalafixModule.scala
@@ -77,8 +77,7 @@ object ScalafixModule {
   ): Result[Unit] =
     if (sources.nonEmpty) {
       val scalafix = ScalafixCache
-        .getOrElseCreate(scalaVersion, repositories)
-        .newArguments()
+        .getOrElseCreate(scalaVersion, repositories, scalafixIvyDeps)
         .withParsedArguments(args.asJava)
         .withWorkingDirectory(wd.toNIO)
         .withConfig(scalafixConfig.map(_.toNIO).asJava)
@@ -86,11 +85,6 @@ object ScalafixModule {
         .withScalaVersion(scalaVersion)
         .withScalacOptions(scalacOptions.asJava)
         .withPaths(sources.map(_.toNIO).asJava)
-        .withToolClasspath(
-          Seq.empty.asJava,
-          scalafixIvyDeps.map(CoursierUtils.toCoordinates).iterator.toSeq.asJava,
-          repositories.map(CoursierUtils.toApiRepository).asJava
-        )
 
       log.info(s"Rewriting and linting ${sources.size} Scala sources against ${scalafix.rulesThatWillRun.size} rules")
       val errors = scalafix.run()