Improve performance of `ClassStamper`.

- Use `Builder`s to avoid building intermediate collections
- Use plain `Set` instead of `SortedSet`
  - Sorting only needs to happen at the end of the computation in `transitiveStamp`
  - This is now done with a call to `.toSeq.sorted`
This commit is contained in:
Matt Dziuban 2026-05-20 10:08:42 -04:00
parent d9b912288a
commit 57e1b84997
1 changed files with 37 additions and 27 deletions

View File

@ -23,7 +23,6 @@ import sbt.util.CacheImplicits
import sbt.util.CacheImplicits.given
import scala.collection.concurrent
import scala.collection.mutable
import scala.collection.SortedSet
import xsbti.{ FileConverter, HashedVirtualFileRef, VirtualFileRef }
object IncrementalTest:
@ -175,8 +174,7 @@ class ClassStamper private[sbt] (
converter,
)
private val stamps = mutable.Map.empty[String, SortedSet[Digest]]
private val internalStamps = mutable.Map.empty[String, SortedSet[Digest]]
private val stamps = mutable.Map.empty[String, Set[Digest]]
// Cached so by-name `analyses0` is only evaluated once
private lazy val analyses = analyses0
private val stampVf: VirtualFileRef => Digest =
@ -192,53 +190,65 @@ class ClassStamper private[sbt] (
extraHashes: Seq[Digest],
log: Logger,
): Option[Digest] =
val digests = SortedSet(analyses.flatMap(internalStamp(javaClassName, _, Set.empty, log))*)
log.debug(s"test: transitiveStamp($javaClassName, $extraHashes) = $digests")
if digests.nonEmpty then Some(Digest.sha256Hash(digests.toSeq ++ extraHashes*))
val builder = Set.newBuilder[Digest]
analyses.foreach: analysis =>
internalStamp(builder, javaClassName, analysis, Set.empty, log)
val digests = builder.result().toSeq.sorted
// log.debug(s"test: transitiveStamp($javaClassName, $extraHashes) = $digests")
if digests.nonEmpty then Some(Digest.sha256Hash(digests ++ extraHashes*))
else None
private def internalStamp(
builder: mutable.Builder[Digest, Set[Digest]],
javaClassName: String,
analysis: Analysis,
alreadySeen: Set[String],
log: Logger,
): SortedSet[Digest] =
): Unit =
import analysis.relations
// log.debug(s"test: internalStamp($javaClassName)")
def internalStamp0(className: String): SortedSet[Digest] =
def internalStamp0(className: String): Unit =
// Use a new builder so we can cache the result in `stamps`
val newBuilder = Set.newBuilder[Digest]
// Zinc doesn't fully track the transitive dependencies
val internalDeps = relations
relations
.internalClassDeps(className)
.flatMap: otherCN =>
internalStamp(otherCN, analysis, alreadySeen + javaClassName, log)
.foreach: otherCN =>
internalStamp(newBuilder, otherCN, analysis, alreadySeen + javaClassName, log)
// log.debug(s" internalStamp: internalDeps: $className = $internalDeps")
val internalJarDeps = relations
relations
.externalDeps(className)
.flatMap: libClassName =>
transitiveStamp(libClassName, Nil, log)
val externalDeps = relations
.foreach: libClassName =>
newBuilder ++= transitiveStamp(libClassName, Nil, log)
relations
.externalDeps(className)
.flatMap: libClassName =>
.foreach: libClassName =>
relations.libraryClassName
.reverse(libClassName)
.map(stampVf)
val classDigests = analysis.apis.internal
.foreach: vf =>
newBuilder += stampVf(vf)
analysis.apis.internal
.get(className)
.toSet
.map: analyzed =>
Digest.dummy(37 * (17 + analyzed.transitiveBytecodeHash) + analyzed.bytecodeHash)
val xs =
(internalDeps union internalJarDeps union externalDeps union classDigests)
.to(SortedSet)
.foreach: analyzed =>
newBuilder += Digest.dummy(
37 * (17 + analyzed.transitiveBytecodeHash) + analyzed.bytecodeHash
)
val xs = newBuilder.result()
if xs.nonEmpty then stamps(className) = xs
else ()
xs
if alreadySeen.contains(javaClassName) then SortedSet.empty
builder ++= xs
if alreadySeen.contains(javaClassName) then ()
else
stamps.get(javaClassName) match
case Some(xs) => xs
case Some(xs) => builder ++= xs
case _ =>
// Note: internalClassDeps uses Scala-encoded class name for companion objects
val classNames = relations.productClassName.reverse(javaClassName)
SortedSet(classNames.toSeq*).flatMap(internalStamp0)
classNames.foreach(internalStamp0)
end ClassStamper