Move out compile time function checker logic to EvaluationMode.kt

compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/checker/EvaluationMode.kt

@@ -5,6 +5,81 @@
 
 package org.jetbrains.kotlin.ir.interpreter.checker
 
-enum class EvaluationMode {
+import org.jetbrains.kotlin.ir.declarations.*
-    FULL, WITH_ANNOTATIONS, ONLY_BUILTINS
+import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.interpreter.compileTimeAnnotation
+import org.jetbrains.kotlin.ir.interpreter.contractsDslAnnotation
+import org.jetbrains.kotlin.ir.interpreter.evaluateIntrinsicAnnotation
+import org.jetbrains.kotlin.ir.interpreter.hasAnnotation
+import org.jetbrains.kotlin.ir.interpreter.isUnsigned
+import org.jetbrains.kotlin.ir.types.isAny
+import org.jetbrains.kotlin.ir.types.isPrimitiveType
+import org.jetbrains.kotlin.ir.types.isString
+import org.jetbrains.kotlin.ir.util.*
+import org.jetbrains.kotlin.name.FqName
+
+enum class EvaluationMode(protected val mustCheckBody: Boolean) {
+    FULL(mustCheckBody = true) {
+        override fun canEvaluateFunction(function: IrFunction, expression: IrCall?): Boolean {
+            return true
+        }
+    },
+
+    WITH_ANNOTATIONS(mustCheckBody = false) {
+        override fun canEvaluateFunction(function: IrFunction, expression: IrCall?): Boolean {
+            if (function.isCompileTimeProperty()) return true
+            return function.isMarkedWith(compileTimeAnnotation) || function.origin == IrBuiltIns.BUILTIN_OPERATOR ||
+                    (function is IrSimpleFunction && function.isOperator && function.name.asString() == "invoke") ||
+                    (function is IrSimpleFunction && function.isFakeOverride && function.overriddenSymbols.any { canEvaluateFunction(it.owner) }) ||
+                    function.isCompileTimeTypeAlias()
+        }
+
+        private fun IrDeclaration?.isCompileTimeProperty(): Boolean {
+            val property = (this as? IrSimpleFunction)?.correspondingPropertySymbol?.owner ?: return false
+            if (property.isConst) return true
+            if (property.isMarkedWith(compileTimeAnnotation) || property.isCompileTimeTypeAlias()) return true
+
+            val backingField = property.backingField
+            val backingFieldExpression = backingField?.initializer?.expression as? IrGetValue
+            return backingFieldExpression?.origin == IrStatementOrigin.INITIALIZE_PROPERTY_FROM_PARAMETER
+        }
+    },
+
+    ONLY_BUILTINS(mustCheckBody = false) {
+        override fun canEvaluateFunction(function: IrFunction, expression: IrCall?): Boolean {
+            if ((function as? IrSimpleFunction)?.correspondingPropertySymbol?.owner?.isConst == true) return true
+
+            val parent = function.parentClassOrNull ?: return false
+            val parentType = parent.defaultType
+            return when {
+                parentType.isPrimitiveType() -> function.name.asString() !in setOf("inc", "dec", "rangeTo", "hashCode")
+                parentType.isString() -> function.name.asString() !in setOf("subSequence", "hashCode")
+                parentType.isAny() -> function.name.asString() == "toString" && expression?.dispatchReceiver !is IrGetObjectValue
+                parent.isObject -> parent.parentClassOrNull?.defaultType?.let { it.isPrimitiveType() || it.isUnsigned() } == true
+                parentType?.isUnsignedType() == true && function is IrConstructor -> true
+                else -> false
+            }
+        }
+    };
+
+    abstract fun canEvaluateFunction(function: IrFunction, expression: IrCall? = null): Boolean
+
+    fun canEvaluateBody(function: IrFunction): Boolean {
+        return (mustCheckBody || function.isLocal) && !function.isContract() && !function.isMarkedAsEvaluateIntrinsic()
+    }
+
+    protected val compileTimeTypeAliases = setOf(
+        "java.lang.StringBuilder", "java.lang.IllegalArgumentException", "java.util.NoSuchElementException"
+    )
+
+    private fun IrDeclaration.isContract() = isMarkedWith(contractsDslAnnotation)
+    private fun IrDeclaration.isMarkedAsEvaluateIntrinsic() = isMarkedWith(evaluateIntrinsicAnnotation)
+    protected fun IrDeclaration.isCompileTimeTypeAlias() = this.parentClassOrNull?.fqNameWhenAvailable?.asString() in compileTimeTypeAliases
+
+    protected fun IrDeclaration.isMarkedWith(annotation: FqName): Boolean {
+        if (this is IrClass && this.isCompanion) return false
+        if (this.hasAnnotation(annotation)) return true
+        return (this.parent as? IrClass)?.isMarkedWith(annotation) ?: false
+    }
 }

compiler/ir/ir.interpreter/src/org/jetbrains/kotlin/ir/interpreter/checker/IrCompileTimeChecker.kt

@@ -8,13 +8,7 @@ package org.jetbrains.kotlin.ir.interpreter.checker
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns
 import org.jetbrains.kotlin.ir.expressions.*
-import org.jetbrains.kotlin.ir.interpreter.builtins.compileTimeAnnotation
-import org.jetbrains.kotlin.ir.interpreter.builtins.contractsDslAnnotation
-import org.jetbrains.kotlin.ir.interpreter.builtins.evaluateIntrinsicAnnotation
-import org.jetbrains.kotlin.ir.interpreter.hasAnnotation
-import org.jetbrains.kotlin.ir.interpreter.isUnsigned
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
@@ -25,48 +19,6 @@ class IrCompileTimeChecker(
 ) : IrElementVisitor<Boolean, Nothing?> {
     private val visitedStack = mutableListOf<IrElement>().apply { if (containingDeclaration != null) add(containingDeclaration) }
 
-    private val compileTimeTypeAliases = setOf(
-        "java.lang.StringBuilder", "java.lang.IllegalArgumentException", "java.util.NoSuchElementException"
-    )
-
-    private fun IrDeclaration.isContract() = isMarkedWith(contractsDslAnnotation)
-    private fun IrDeclaration.isMarkedAsEvaluateIntrinsic() = isMarkedWith(evaluateIntrinsicAnnotation)
-    private fun IrDeclaration.isMarkedAsCompileTime(expression: IrCall? = null): Boolean {
-        if (mode == EvaluationMode.WITH_ANNOTATIONS) {
-            return isMarkedWith(compileTimeAnnotation) || this.origin == IrBuiltIns.BUILTIN_OPERATOR ||
-                    (this is IrSimpleFunction && this.isOperator && this.name.asString() == "invoke") ||
-                    (this is IrSimpleFunction && this.isFakeOverride && this.overriddenSymbols.any { it.owner.isMarkedAsCompileTime() }) ||
-                    this.parentClassOrNull?.fqNameWhenAvailable?.asString() in compileTimeTypeAliases
-        }
-
-        val parent = this.parentClassOrNull
-        val parentType = parent?.defaultType
-        return when {
-            parentType?.isPrimitiveType() == true -> (this as IrFunction).name.asString() !in setOf("inc", "dec", "rangeTo", "hashCode")
-            parentType?.isString() == true -> (this as IrDeclarationWithName).name.asString() !in setOf("subSequence", "hashCode")
-            parentType?.isAny() == true -> (this as IrFunction).name.asString() == "toString" && expression?.dispatchReceiver !is IrGetObjectValue
-            parent?.isObject == true -> parent.parentClassOrNull?.defaultType?.let { it.isPrimitiveType() || it.isUnsignedType() } == true
-            parentType?.isUnsignedType() == true && (this is IrConstructor) -> true
-            else -> false
-        }
-    }
-
-    private fun IrDeclaration.isMarkedWith(annotation: FqName): Boolean {
-        if (this is IrClass && this.isCompanion) return false
-        if (this.hasAnnotation(annotation)) return true
-        return (this.parent as? IrClass)?.isMarkedWith(annotation) ?: false
-    }
-
-    private fun IrProperty?.isCompileTime(): Boolean {
-        if (this == null) return false
-        if (this.isConst) return true
-        if (this.isMarkedAsCompileTime()) return true
-
-        val backingField = this.backingField
-        val backingFieldExpression = backingField?.initializer?.expression as? IrGetValue
-        return backingFieldExpression?.origin == IrStatementOrigin.INITIALIZE_PROPERTY_FROM_PARAMETER
-    }
-
     private fun IrElement.asVisited(block: () -> Boolean): Boolean {
         visitedStack += this
         val result = block()
@@ -86,26 +38,21 @@ class IrCompileTimeChecker(
 
     private fun visitConstructor(expression: IrFunctionAccessExpression): Boolean {
         return when {
-            expression.symbol.owner.isMarkedAsEvaluateIntrinsic() -> true
             !visitValueParameters(expression, null) -> false
-            else -> expression.symbol.owner.isMarkedAsCompileTime()
+            else -> mode.canEvaluateFunction(expression.symbol.owner)
         }
     }
 
     override fun visitCall(expression: IrCall, data: Nothing?): Boolean {
-        if (expression.symbol.owner.isContract()) return false
+        val owner = expression.symbol.owner
+        if (!mode.canEvaluateFunction(owner, expression)) return false
 
-        val property = (expression.symbol.owner as? IrSimpleFunction)?.correspondingPropertySymbol?.owner
+        val dispatchReceiverComputable = expression.dispatchReceiver?.accept(this, null) ?: true
-        if (expression.symbol.owner.isMarkedAsCompileTime(expression) || property.isCompileTime()) {
+        val extensionReceiverComputable = expression.extensionReceiver?.accept(this, null) ?: true
-            val dispatchReceiverComputable = expression.dispatchReceiver?.accept(this, null) ?: true
+        if (!visitValueParameters(expression, null)) return false
-            val extensionReceiverComputable = expression.extensionReceiver?.accept(this, null) ?: true
+        val bodyComputable = if (mode.canEvaluateBody(owner)) owner.body?.accept(this, null) ?: true else true
-            if (!visitValueParameters(expression, null)) return false
-            val bodyComputable = if (expression.symbol.owner.isLocal) expression.symbol.owner.body?.accept(this, null) ?: true else true
 
-            return dispatchReceiverComputable && extensionReceiverComputable && bodyComputable
+        return dispatchReceiverComputable && extensionReceiverComputable && bodyComputable
-        }
-
-        return false
     }
 
     override fun visitVariable(declaration: IrVariable, data: Nothing?): Boolean {
@@ -206,13 +153,13 @@ class IrCompileTimeChecker(
 
     override fun visitFunctionReference(expression: IrFunctionReference, data: Nothing?): Boolean {
         return expression.asVisited {
-            expression.symbol.owner.isMarkedAsCompileTime() && expression.symbol.owner.body?.accept(this, data) == true
+            mode.canEvaluateFunction(expression.symbol.owner) && expression.symbol.owner.body?.accept(this, data) == true
         }
     }
 
     override fun visitFunctionExpression(expression: IrFunctionExpression, data: Nothing?): Boolean {
         val isLambda = expression.origin == IrStatementOrigin.LAMBDA || expression.origin == IrStatementOrigin.ANONYMOUS_FUNCTION
-        val isCompileTime = expression.function.isMarkedAsCompileTime()
+        val isCompileTime = mode.canEvaluateFunction(expression.function)
         return expression.function.asVisited {
             if (isLambda || isCompileTime) expression.function.body?.accept(this, data) == true else false
         }