Refactor generated IR classes caching

Author: lupuuss

mokkery-plugin/src/main/kotlin/dev/mokkery/plugin/Utils.kt

@@ -10,9 +10,4 @@ object Caches {
     val classReferences by cacheKey<IrClassReferencer, IrClass>()
     val functionReferences by cacheKey<IrFunctionReferencer, IrSimpleFunction>()
     val propertyReferences by cacheKey<IrPropertyReferencer, IrProperty>()
-
-
-    val mockClasses by cacheKey<IrClass, IrClass>()
-    val spyClasses by cacheKey<IrClass, IrClass>()
-    val mockManyClasses by cacheKey<Set<IrClass>, IrClass>()
 }

mokkery-plugin/src/main/kotlin/dev/mokkery/plugin/ir/Caches.kt

@@ -1,8 +1,6 @@
 package dev.mokkery.plugin.ir.transformer
 
 import dev.mokkery.plugin.Mokkery
-import dev.mokkery.plugin.caches
-import dev.mokkery.plugin.ir.Caches
 import dev.mokkery.plugin.ir.IrMokkeryPluginScope
 import dev.mokkery.plugin.ir.MokkeryIr
 import dev.mokkery.plugin.ir.applyTransformChildrenVoid
@@ -20,7 +18,6 @@ import dev.mokkery.plugin.ir.transformer.templating.replaceVerify
 import dev.mokkery.plugin.ir.transformer.templating.replaceVerifySuspend
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.declarations.IrClass
-import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
@@ -66,13 +63,6 @@ class MokkeryRootTransformer(pluginScope: IrMokkeryPluginScope) : CoreTransforme
         }
     }
 
-    override fun visitFileNew(declaration: IrFile): IrFile {
-        caches[Caches.mockClasses].clear()
-        caches[Caches.spyClasses].clear()
-        caches[Caches.mockManyClasses].clear()
-        return declaration.applyTransformChildrenVoid()
-    }
-
     override fun visitModuleFragment(declaration: IrModuleFragment): IrModuleFragment {
         val time = TimeSource.Monotonic.markNow()
         declaration.transformChildrenVoid()

mokkery-plugin/src/main/kotlin/dev/mokkery/plugin/ir/transformer/MokkeryRootTransformer.kt

@@ -11,12 +11,18 @@ import org.jetbrains.kotlin.ir.builders.IrBuilder
 import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
 import org.jetbrains.kotlin.ir.builders.IrGeneratorContext
 import org.jetbrains.kotlin.ir.builders.irGetObject
+import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.name
 import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrVarargElement
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.types.typeWith
+import org.jetbrains.kotlin.ir.util.findDeclaration
+import org.jetbrains.kotlin.ir.util.kotlinFqName
+import org.jetbrains.kotlin.name.Name
 
 context(scope: TransformerScope)
 inline fun <T> declarationIrBuilder(
@@ -78,3 +84,36 @@ fun IrBuilderWithScope.irGetMokkeryScopeGlobal(): IrCall {
         .requirePropertyGetterOwner("global")
         .let { irCall(it) { arguments[0] = irGetObject(scopeCompanion.symbol) } }
 }
+
+context(scope: TransformerScope)
+inline fun findOrBuildClassInCurrentFile(
+    nameBase: String,
+    nameHashSource: List<IrClass> = emptyList(),
+    builder: (Name) -> IrClass
+): IrClass {
+    val name = generatedClassNameInCurrentFile(nameBase, nameHashSource)
+    return currentFileValue.findDeclaration<IrClass> { it.name == name } ?: builder(name)
+}
+
+context(scope: TransformerScope)
+fun generatedClassNameInCurrentFile(base: String, implements: List<IrClass>): Name {
+    val prefix = when {
+        implements.size == 1 -> "${implements[0].name}_${base}"
+        else -> base
+    }
+    val hash = implements
+        .map { it.kotlinFqName.asString() }
+        .plus(currentFileValue.kotlinFqNameStringWithName())
+        .hexHashString()
+    return Name.identifier("${prefix}_${hash}")
+}
+
+private fun List<String>.hexHashString(): String {
+    var hash = 0xcbf29ce484222325UL
+    forEach {
+        it.forEach { c -> hash = (hash xor c.code.toULong()) * 0x100000001b3UL }
+    }
+    return hash.toString(36)
+}
+
+private fun IrFile.kotlinFqNameStringWithName() = kotlinFqName.asString() + "." + this.name

mokkery-plugin/src/main/kotlin/dev/mokkery/plugin/ir/transformer/core/BuildersApi.kt

@@ -20,6 +20,7 @@ import dev.mokkery.plugin.ir.requireSimpleFunctionOwner
 import dev.mokkery.plugin.ir.transformer.core.TransformerScope
 import dev.mokkery.plugin.ir.transformer.core.addToCurrentFile
 import dev.mokkery.plugin.ir.transformer.core.declarationIrBuilder
+import dev.mokkery.plugin.ir.transformer.core.findOrBuildClassInCurrentFile
 import dev.mokkery.plugin.ir.transformer.core.referenced
 import dev.mokkery.plugin.ir.transformer.mock.stubs.irDelegatingConstructorWithStubs
 import dev.mokkery.plugin.ir.typeWith
@@ -46,31 +47,36 @@ import org.jetbrains.kotlin.ir.util.primaryConstructor
 import org.jetbrains.kotlin.name.Name
 
 context(scope: TransformerScope)
-fun buildDefaultsExtractorFactoryIfRequired(
-    className: Name,
+fun findOrBuildDefaultsExtractorFactoryIfRequired(
     classesToIntercept: List<IrClass>,
     bodyBuilder: IrBlockBodyBuilder,
 ): IrExpression? {
-    val defaultsExtractor = buildDefaultsExtractorOrNull(className, classesToIntercept) ?: return null
-    val defaultsFactoryClass = buildDefaultsFactoryClassWith(className, defaultsExtractor.primaryConstructor!!)
-    return bodyBuilder.run {
-        irCallConstructor(defaultsFactoryClass.primaryConstructor!!)
+    val anyDefaults = classesToIntercept.any { clazz ->
+        clazz.functions.any { it.parameters.any(IrValueParameter::hasDefaultValue) }
     }
+    if (!anyDefaults) return null
+    val defaultsExtractor = findOrBuildClassInCurrentFile(
+        nameBase = "DE",
+        nameHashSource = classesToIntercept,
+        builder = { buildDefaultsExtractor(it, classesToIntercept) }
+    )
+    val defaultsFactoryClass = findOrBuildClassInCurrentFile(
+        nameBase = "DEF",
+        nameHashSource = classesToIntercept,
+        builder = { buildDefaultsExtractorFactoryClass(it, defaultsExtractor.primaryConstructor!!) }
+    )
+    return bodyBuilder.run { irCallConstructor(defaultsFactoryClass.primaryConstructor!!) }
 }
 
 context(scope: TransformerScope)
-private fun buildDefaultsExtractorOrNull(
-    className: Name,
+private fun buildDefaultsExtractor(
+    name: Name,
     classesToIntercept: List<IrClass>
-): IrClass? {
-    val anyDefaults = classesToIntercept.any { clazz ->
-        clazz.functions.any { it.parameters.any(IrValueParameter::hasDefaultValue) }
-    }
-    if (!anyDefaults) return null
+): IrClass {
     val throwArgumentsFunction = referenced(MokkeryIr.Function.throwArguments)
     val methodWithoutDefaultFunction = referenced(MokkeryIr.Function.methodWithoutDefaultsError)
     val irBuiltIns = irBuiltIns
-    val defaultsExtractorClass = irFactory.buildClass { name = Name.identifier("DE${className.asString()}") }
+    val defaultsExtractorClass = irFactory.buildClass { this.name = name }
     classesToIntercept.forEach(defaultsExtractorClass::copyTypeParametersFrom)
     defaultsExtractorClass.addToCurrentFile()
     defaultsExtractorClass.createThisReceiverParameter()
@@ -121,11 +127,9 @@ private fun buildDefaultsExtractorOrNull(
 }
 
 context(scope: TransformerScope)
-private fun buildDefaultsFactoryClassWith(className: Name, constructor: IrConstructor): IrClass {
+private fun buildDefaultsExtractorFactoryClass(name: Name, constructor: IrConstructor): IrClass {
     val defaultsExtractorFactoryInterface = referenced(MokkeryIr.Class.DefaultsExtractorFactory)
-    val defaultsExtractorFactoryClassImpl = irFactory.buildClass {
-        name = Name.identifier("DEF${className.asString()}")
-    }
+    val defaultsExtractorFactoryClassImpl = irFactory.buildClass { this.name = name }
     defaultsExtractorFactoryClassImpl.superTypes = listOf(
         defaultsExtractorFactoryInterface.defaultType,
         irBuiltIns.anyType

mokkery-plugin/src/main/kotlin/dev/mokkery/plugin/ir/transformer/mock/BuildDefaultsExtractorFactory.kt

@@ -32,7 +32,6 @@ import dev.mokkery.plugin.ir.transformer.core.referenced
 import dev.mokkery.plugin.ir.transformer.core.referencedGetterSymbol
 import dev.mokkery.plugin.ir.transformer.mock.stubs.irDelegatingConstructorWithStubs
 import dev.mokkery.plugin.ir.typeWith
-import dev.mokkery.plugin.randomString
 import org.jetbrains.kotlin.ir.builders.IrBlockBodyBuilder
 import org.jetbrains.kotlin.ir.builders.declarations.addConstructor
 import org.jetbrains.kotlin.ir.builders.declarations.addValueParameter
@@ -58,17 +57,17 @@ import org.jetbrains.kotlin.ir.util.isClass
 import org.jetbrains.kotlin.ir.util.isInterface
 import org.jetbrains.kotlin.ir.util.kotlinFqName
 import org.jetbrains.kotlin.ir.util.parentAsClass
-import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.utils.memoryOptimizedMap
 
 context(scope: TransformerScope)
 fun buildMockClass(
+    name: Name,
     mokkeryKind: IrMokkeryKind,
     classToMock: IrClass,
 ): IrClass {
     val instanceScopeClass = referenced(MokkeryIr.Class.MokkeryInstanceScope)
-    val mockedClass = irFactory.buildClass { name = classToMock.name.createUniqueMockName(mokkeryKind.name) }
+    val mockedClass = irFactory.buildClass { this.name = name }
     mockedClass.addToCurrentFile()
     mockedClass.copyTypeParametersFrom(classToMock)
     val typedClassToMock = classToMock.symbol.typeWithParameters(mockedClass.typeParameters)
@@ -102,12 +101,10 @@ fun buildMockClass(
 }
 
 context(scope: TransformerScope)
-fun buildManyMockClass(classesToMock: List<IrClass>): IrClass {
+fun buildManyMockClass(name: Name, classesToMock: List<IrClass>): IrClass {
     val manyMocksMarkerClass = referenced(MokkeryIr.Class.mockMany(classesToMock.size))
     val mokkeryInstanceClass = referenced(MokkeryIr.Class.MokkeryInstanceScope)
-    val mockedClass = irFactory.buildClass {
-        name = manyMocksMarkerClass.kotlinFqName.createUniqueManyMockName()
-    }
+    val mockedClass = irFactory.buildClass { this.name = name }
     mockedClass.addToCurrentFile()
     classesToMock.forEach(mockedClass::copyTypeParametersFrom)
     mockedClass.createThisReceiverParameter()
@@ -223,8 +220,7 @@ private fun IrClass.addMockClassConstructor(
                     arguments[4] = irGet(thisReceiver!!)
                     arguments[5] = irGet(parameters[1])
                     arguments[6] = spyParam?.let(::irGet) ?: irNull()
-                    arguments[7] = buildDefaultsExtractorFactoryIfRequired(
-                        className = this@addMockClassConstructor.name,
+                    arguments[7] = findOrBuildDefaultsExtractorFactoryIfRequired(
                         classesToIntercept = classesToIntercept,
                         bodyBuilder = this@irBlockBody
                     )
@@ -252,12 +248,3 @@ private fun IrConstructor.addSpyParameter(classesToIntercept: List<IrClass>): Ir
     val classToSpy = classesToIntercept.singleOrNull() ?: error("Spy is not supported for intercepting multiple types!")
     return addValueParameter("obj", classToSpy.symbol.typeWithParameters(parentAsClass.typeParameters))
 }
-
-private fun Name.createUniqueMockName(type: String) = asString()
-    .plus("_$type}_${randomString()}")
-    .let(Name::identifier)
-
-private fun FqName.createUniqueManyMockName() = shortName()
-    .asString()
-    .plus("_${randomString()}")
-    .let(Name::identifier)

mokkery-plugin/src/main/kotlin/dev/mokkery/plugin/ir/transformer/mock/BuildMockClass.kt

@@ -1,10 +1,7 @@
 package dev.mokkery.plugin.ir.transformer.mock
 
-import dev.mokkery.plugin.Cache
-import dev.mokkery.plugin.caches
 import dev.mokkery.plugin.context.configuration
 import dev.mokkery.plugin.defaultMockMode
-import dev.mokkery.plugin.ir.Caches
 import dev.mokkery.plugin.ir.IrMokkeryKind.Mock
 import dev.mokkery.plugin.ir.IrMokkeryKind.Spy
 import dev.mokkery.plugin.ir.MokkeryIr
@@ -19,6 +16,7 @@ import dev.mokkery.plugin.ir.kClassReference
 import dev.mokkery.plugin.ir.platform
 import dev.mokkery.plugin.ir.transformer.core.TransformerScope
 import dev.mokkery.plugin.ir.transformer.core.declarationIrBuilder
+import dev.mokkery.plugin.ir.transformer.core.findOrBuildClassInCurrentFile
 import dev.mokkery.plugin.ir.transformer.core.irGetMokkeryScopeGlobal
 import dev.mokkery.plugin.ir.transformer.core.referenced
 import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
@@ -36,10 +34,10 @@ fun IrCall.replaceMockCall(): IrExpression {
     val typeToMock = typeArguments.firstOrNull() ?: return this
     val classToMock = typeToMock.getClass() ?: return this
     if (platform.isJs() && classToMock.defaultType.isAnyFunction()) return buildMockJsFunction(this, Mock)
-    val mockedClass = cachedMockImplementationClass(
-        typeToMock = classToMock,
-        cacheKey = Caches.mockClasses,
-        builder = { buildMockClass(Mock, it) }
+    val mockedClass = findOrBuildClassInCurrentFile(
+        nameBase = "Mock",
+        nameHashSource = listOf(classToMock),
+        builder = { buildMockClass(it, Mock, classToMock) }
     )
     return declarationIrBuilder {
         irMockConstructorCall(mockedClass, this@replaceMockCall)
@@ -48,12 +46,12 @@ fun IrCall.replaceMockCall(): IrExpression {
 
 context(scope: TransformerScope)
 fun IrCall.replaceMockManyCall(): IrExpression {
-    val classes = typeArguments.mapNotNullTo(mutableSetOf()) { it?.getClass() }
-    if (classes.isEmpty()) return this
-    val mockedClass = cachedMockImplementationClass(
-        typeToMock = classes,
-        cacheKey = Caches.mockManyClasses,
-        builder = { buildManyMockClass(it.toList()) }
+    val classesToMock = typeArguments.mapNotNull { it?.getClass() }
+    if (classesToMock.isEmpty()) return this
+    val mockedClass = findOrBuildClassInCurrentFile(
+        nameBase = "MockMany${classesToMock.size}",
+        nameHashSource = classesToMock,
+        builder = { buildManyMockClass(it, classesToMock) }
     )
     return declarationIrBuilder {
         irMockConstructorCall(mockedClass, this@replaceMockManyCall)
@@ -63,12 +61,12 @@ fun IrCall.replaceMockManyCall(): IrExpression {
 context(scope: TransformerScope)
 fun IrCall.replaceSpyCall(): IrExpression {
     val typeToMock = typeArguments.firstOrNull() ?: return this
-    val classToMock = typeToMock.getClass() ?: return this
-    if (platform.isJs() && classToMock.defaultType.isAnyFunction()) return buildMockJsFunction(this, Spy)
-    val spiedClass = cachedMockImplementationClass(
-        typeToMock = classToMock,
-        cacheKey = Caches.spyClasses,
-        builder = { buildMockClass(Spy, it) }
+    val classToSpy = typeToMock.getClass() ?: return this
+    if (platform.isJs() && classToSpy.defaultType.isAnyFunction()) return buildMockJsFunction(this, Spy)
+    val spiedClass = findOrBuildClassInCurrentFile(
+        nameBase = "Spy",
+        nameHashSource = listOf(classToSpy),
+        builder = { buildMockClass(it, Spy, classToSpy) }
     )
     return declarationIrBuilder {
         irSpyConstructorCall(spiedClass, this@replaceSpyCall)
@@ -118,17 +116,3 @@ private fun IrBuilderWithScope.irSpyConstructorCall(
             arguments[4 + index] = kClassReference(it ?: anyType)
         }
 }
-
-context(scope: TransformerScope)
-private fun <T> cachedMockImplementationClass(
-    typeToMock: T,
-    cacheKey: Cache.Key<T, IrClass>,
-    builder: (T) -> IrClass
-): IrClass {
-    val cache = caches[cacheKey]
-    val cached = cache[typeToMock]
-    if (cached != null) return cached
-    val newClass = builder(typeToMock)
-    cache[typeToMock] = newClass
-    return newClass
-}