From c037bf0a614e87d7b0aa428c13ba4210386c9447 Mon Sep 17 00:00:00 2001
From: Lorenzo Gabriele <lorenzolespaul@gmail.com>
Date: Mon, 30 Sep 2024 16:44:11 +0200
Subject: [PATCH] Cache Scalafix instances globally

---
 .../goyeau/mill/scalafix/ScalafixCache.scala  | 24 +++++++++++++++++++
 .../goyeau/mill/scalafix/ScalafixModule.scala |  4 ++--
 2 files changed, 26 insertions(+), 2 deletions(-)
 create mode 100644 mill-scalafix/src/com/goyeau/mill/scalafix/ScalafixCache.scala

diff --git a/mill-scalafix/src/com/goyeau/mill/scalafix/ScalafixCache.scala b/mill-scalafix/src/com/goyeau/mill/scalafix/ScalafixCache.scala
new file mode 100644
index 0000000..2fa3797
--- /dev/null
+++ b/mill-scalafix/src/com/goyeau/mill/scalafix/ScalafixCache.scala
@@ -0,0 +1,24 @@
+package com.goyeau.mill.scalafix
+
+import com.goyeau.mill.scalafix.CoursierUtils
+import coursier.core.Repository
+import scalafix.interfaces.Scalafix
+
+import scala.collection.mutable
+import scala.ref.SoftReference
+import scala.jdk.CollectionConverters._
+
+private[scalafix] object ScalafixCache {
+
+  private val cache = mutable.Map.empty[(String, Seq[Repository]), SoftReference[Scalafix]]
+
+  def getOrElseCreate(scalaVersion: String, repositories: Seq[Repository]) =
+    cache.get((scalaVersion, repositories)) match {
+      case Some(SoftReference(value)) => value
+      case _ =>
+        val newResult =
+          Scalafix.fetchAndClassloadInstance(scalaVersion, repositories.map(CoursierUtils.toApiRepository).asJava)
+        cache.update((scalaVersion, repositories), SoftReference(newResult))
+        newResult
+    }
+}
diff --git a/mill-scalafix/src/com/goyeau/mill/scalafix/ScalafixModule.scala b/mill-scalafix/src/com/goyeau/mill/scalafix/ScalafixModule.scala
index c55eefe..234b19d 100644
--- a/mill-scalafix/src/com/goyeau/mill/scalafix/ScalafixModule.scala
+++ b/mill-scalafix/src/com/goyeau/mill/scalafix/ScalafixModule.scala
@@ -77,8 +77,8 @@ object ScalafixModule {
       wd: os.Path
   ): Result[Unit] =
     if (sources.nonEmpty) {
-      val scalafix = Scalafix
-        .fetchAndClassloadInstance(scalaVersion, repositories.map(CoursierUtils.toApiRepository).asJava)
+      val scalafix = ScalafixCache
+        .getOrElseCreate(scalaVersion, repositories)
         .newArguments()
         .withParsedArguments(args.asJava)
         .withWorkingDirectory(wd.toNIO)