centralize ScalafixInterface memoization within a setting key

src/main/scala/scalafix/internal/sbt/BlockingCache.scala

@@ -2,21 +2,32 @@ package scalafix.internal.sbt
 
 import java.{util => ju}
 
+import scala.util.Try
+
 /** A basic thread-safe cache without any eviction. */
 class BlockingCache[K, V] {
 
-  // Number of keys is expected to be very small so the global lock should not be a bottleneck
-  private val underlying = ju.Collections.synchronizedMap(new ju.HashMap[K, V])
+  private val underlying = new ju.concurrent.ConcurrentHashMap[K, V]
+
+  private case class SkipUpdate(prev: V) extends Exception
 
-  /** @param value
-    *   By-name parameter evaluated when the key if missing. Value computation
-    *   is guaranteed to be called only once per key across all invocations.
-    */
-  def getOrElseUpdate(key: K, value: => V): V =
-    underlying.computeIfAbsent(
-      key,
-      new ju.function.Function[K, V] {
-        override def apply(k: K): V = value
-      }
+  // Evaluations of `transform` are guaranteed to be sequential for the same key
+  def compute(key: K, transform: Option[V] => Option[V]): V = {
+    Try(
+      underlying.compute(
+        key,
+        new ju.function.BiFunction[K, V, V] {
+          override def apply(key: K, prev: V): V = {
+            transform(Option(prev)).getOrElse(throw SkipUpdate(prev))
+          }
+        }
+      )
+    ).fold(
+      {
+        case SkipUpdate(prev) => prev
+        case ex => throw ex
+      },
+      identity
     )
+  }
 }

src/main/scala/scalafix/internal/sbt/ScalafixCompletions.scala

@@ -10,8 +10,6 @@ import sbt.complete.DefaultParsers._
 
 class ScalafixCompletions(
     workingDirectory: Path,
-    // Unfortunately, local rules will not show up as completions in the parser, as that parser can only
-    // depend on settings, while local rules classpath must be looked up via tasks
     loadedRules: () => Seq[ScalafixRule],
     terminalWidth: Option[Int],
     allowedTargetFilesPrefixes: Seq[Path], // Nil means all values are valid

src/main/scala/scalafix/internal/sbt/ScalafixInterface.scala

@@ -4,6 +4,7 @@ import java.nio.file.Path
 import java.{util => jutil}
 import coursierapi.Repository
 import sbt._
+import sbt.io.RegularFileFilter
 import scalafix.interfaces.{Scalafix => ScalafixAPI, _}
 import scalafix.sbt.InvalidArgument
 
@@ -31,6 +32,16 @@ object Arg {
         customDependencies.map(_.asCoursierCoordinates).asJava,
         repositories.asJava
       )
+
+    def mutableFiles: Seq[File] =
+      customURIs
+        .map(uri => java.nio.file.Paths.get(uri).toFile)
+        .flatMap {
+          case classDirectory if classDirectory.isDirectory =>
+            classDirectory.**(RegularFileFilter).get
+          case jar =>
+            Seq(jar)
+        }
   }
 
   case class Rules(rules: Seq[String]) extends Arg with CacheKey {
@@ -81,7 +92,7 @@ object Arg {
 
 class ScalafixInterface private (
     scalafixArguments: ScalafixArguments, // hide it to force usage of withArgs so we can intercept arguments
-    val args: Seq[Arg]
+    val args: Seq[Arg] = Nil
 ) {
 
   def withArgs(args: Arg*): ScalafixInterface = {
@@ -112,31 +123,42 @@ class ScalafixInterface private (
 }
 
 object ScalafixInterface {
-  private class LazyValue[T](thunk: () => T) extends (() => T) {
-    private lazy val _value = scala.util.Try(thunk())
-    override def apply(): T = _value.get
-  }
-  private val fromToolClasspathMemo: BlockingCache[
-    (String, Seq[ModuleID], Seq[Repository]),
-    ScalafixInterface
-  ] = new BlockingCache
+
+  type Cache = BlockingCache[
+    (
+        String, // scalafixScalaBinaryVersion
+        Option[Arg.ToolClasspath]
+    ),
+    (
+        ScalafixInterface,
+        Seq[Long] // mutable files stamping
+    )
+  ]
 
   private[scalafix] val defaultLogger: Logger = ConsoleLogger(System.out)
 
-  def fromToolClasspath(
+  def apply(
+      cache: Cache,
       scalafixScalaBinaryVersion: String,
-      scalafixDependencies: Seq[ModuleID],
-      scalafixCustomResolvers: Seq[Repository],
+      toolClasspath: Arg.ToolClasspath,
       logger: Logger,
       callback: ScalafixMainCallback
-  ): () => ScalafixInterface =
-    new LazyValue({ () =>
-      fromToolClasspathMemo.getOrElseUpdate(
-        (
-          scalafixScalaBinaryVersion,
-          scalafixDependencies,
-          scalafixCustomResolvers
-        ), {
+  ): ScalafixInterface = {
+
+    // Build or retrieve from the cache the scalafix interface that can run
+    // built-in rules. This is the most costly instantiation, which should be
+    // shared as much as possible.
+    val (buildinRulesInterface, _) = cache.compute(
+      (
+        scalafixScalaBinaryVersion,
+        None
+      ),
+      {
+        case Some(_) =>
+          // cache hit, don't update
+          None
+        case None =>
+          // cache miss, resolve scalafix artifacts and classload them
           if (scalafixScalaBinaryVersion.startsWith("3"))
             logger.error(
               "To use Scalafix on Scala 3 projects, you must unset `scalafixScalaBinaryVersion`. " +
@@ -151,27 +173,54 @@ object ScalafixInterface {
           val scalafixArguments = ScalafixAPI
             .fetchAndClassloadInstance(
               scalafixScalaBinaryVersion,
-              scalafixCustomResolvers.asJava
+              toolClasspath.repositories.asJava
             )
             .newArguments()
             .withMainCallback(callback)
-          val printStream =
+
+          val printStream = Arg.PrintStream(
             new PrintStream(
               LoggingOutputStream(
                 logger,
                 Level.Info
               )
             )
-          new ScalafixInterface(scalafixArguments, Nil)
-            .withArgs(
-              Arg.PrintStream(printStream),
-              Arg.ToolClasspath(
-                Nil,
-                scalafixDependencies,
-                scalafixCustomResolvers
-              )
+          )
+
+          Some(
+            (
+              new ScalafixInterface(scalafixArguments).withArgs(printStream),
+              Nil
             )
-        }
-      )
-    })
+          )
+      }
+    )
+
+    // stamp the files that might change, to potentialy force invalidation
+    // of the corresponding cache entry
+    val currentStamps =
+      toolClasspath.mutableFiles.map(IO.getModifiedTimeOrZero)
+
+    val (toolClasspathInterface, _) = cache.compute(
+      (
+        scalafixScalaBinaryVersion,
+        Some(toolClasspath)
+      ),
+      {
+        case Some((_, oldStamps)) if (currentStamps == oldStamps) =>
+          // cache hit, don't update
+          None
+        case _ =>
+          // cache miss or stale stamps, resolve custom rules artifacts and classload them
+          Some(
+            (
+              buildinRulesInterface.withArgs(toolClasspath),
+              currentStamps
+            )
+          )
+      }
+    )
+
+    toolClasspathInterface
+  }
 }