From 72532f6f032f286c594be125459657d7faf79c82 Mon Sep 17 00:00:00 2001
From: Ivan Gavrilovic <gavra@google.com>
Date: Sat, 19 Feb 2022 20:07:11 +0000
Subject: [PATCH] KT-41456: Incremental KAPT - optimize .class file counting

When counting .class files, stop once the required number is reached.
Also, implement simple counting logic directly, it is 25% faster than
using File.walk(). The way this information is used is a heuristic, so
reducing computation cost makes sense.

^KT-41456 Fixed
---
 .../kotlin/kapt3/base/incremental/cache.kt    | 33 +++++++++++++++++--
 1 file changed, 31 insertions(+), 2 deletions(-)

diff --git a/plugins/kapt3/kapt3-base/src/org/jetbrains/kotlin/kapt3/base/incremental/cache.kt b/plugins/kapt3/kapt3-base/src/org/jetbrains/kotlin/kapt3/base/incremental/cache.kt
index 94b658347a0..122d8487c3d 100644
--- a/plugins/kapt3/kapt3-base/src/org/jetbrains/kotlin/kapt3/base/incremental/cache.kt
+++ b/plugins/kapt3/kapt3-base/src/org/jetbrains/kotlin/kapt3/base/incremental/cache.kt
@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.kapt3.base.incremental
 
 import java.io.*
+import java.util.ArrayDeque
 
 // TODO(gavra): switch away from Java serialization
 class JavaClassCacheManager(val file: File) : Closeable {
@@ -59,9 +60,9 @@ class JavaClassCacheManager(val file: File) : Closeable {
          *
          * However, comparing total numbers is probably good enough, and it should eliminate most of the issues. See KT-41456 for details.
          */
-        val totalCompiledClasses = compiledSources.sumOf { it.walk().filter { it.extension == "class" }.count() }
         val totalDeclaredTypes = javaCache.getSourceFileDefinedTypesCount()
-        if (totalCompiledClasses < totalDeclaredTypes) {
+        val compileOutputHasEnoughClassFiles = checkMinNumberOfClassFiles(compiledSources, totalDeclaredTypes)
+        if (!compileOutputHasEnoughClassFiles) {
             return SourcesToReprocess.FullRebuild
         }
 
@@ -197,4 +198,32 @@ sealed class SourcesToReprocess {
     ) : SourcesToReprocess()
 
     object FullRebuild : SourcesToReprocess()
+}
+
+/** Returns if specified root dirs have at least [required] number of class files. */
+private fun checkMinNumberOfClassFiles(roots: List<File>, required: Int): Boolean {
+    var currentlyMissing = required
+    roots.filter { it.isDirectory }.forEach {
+        val inThisRoot = countClassFilesUpToLimit(it, currentlyMissing)
+        currentlyMissing -= inThisRoot
+    }
+    return currentlyMissing <= 0
+}
+
+private fun countClassFilesUpToLimit(root: File, limit: Int): Int {
+    var cnt = 0;
+    val toVisit = ArrayDeque<File>()
+    toVisit.add(root)
+    while (cnt < limit && toVisit.isNotEmpty()) {
+        val curr = toVisit.removeFirst()
+        if (curr.isFile && curr.toString().endsWith(".class")) {
+            cnt++
+        } else if (curr.isDirectory) {
+            curr.listFiles()?.forEach {
+                toVisit.addLast(it)
+            }
+        }
+    }
+
+    return cnt
 }
\ No newline at end of file