Cache also withToolClasspath using two-layered cache

joan38 · Sep 30, 2024 · 617f307 · 617f307
1 parent c8798c6
commit 617f307
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 20 deletions.
diff --git a/mill-scalafix/src/com/goyeau/mill/scalafix/ScalafixCache.scala b/mill-scalafix/src/com/goyeau/mill/scalafix/ScalafixCache.scala
@@ -2,25 +2,46 @@ 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._
 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
@@ -77,20 +77,14 @@ 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)
         .withClasspath(classpath.map(_.toNIO).asJava)
         .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()