/*
 * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

package org.jetbrains.kotlin.powerassert.diagram

import org.jetbrains.kotlin.constant.EvaluatedConstTracker
import org.jetbrains.kotlin.ir.builders.IrBlockBuilder
import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
import org.jetbrains.kotlin.ir.declarations.IrVariable
import org.jetbrains.kotlin.ir.expressions.IrExpression

/**
 * Represents a variable generated by the power-assert compiler plugin when it expands an
 * [IrExpression] to explain said expression.
 */
@Suppress("unused") // Public API
class ExplainVariable(
    /** The temporary variable, which was generated for the expression. */
    val variable: IrVariable,
    /** The original expression, which was extracted to a variable. */
    val original: IrExpression,
    /** The *file* offset suggested for diagramming the result of the expression. */
    val displayOffset: Int,
)

/**
 * Expand and explain each sub-expression of the specified [IrExpression].
 * [ExplainVariable]s are always provided to [builder] in evaluated order.
 *
 * When provided, a [EvaluatedConstTracker] will be used to determine if a
 * [org.jetbrains.kotlin.ir.expressions.IrConst] was evaluated by the compiler, and include it
 * as a variable if it was. Otherwise, all constants are ignored.
 *
 * For example, the following expression,
 *
 * ```
 * a + b
 * ```
 *
 * Will be transformed into the following:
 *
 * ```
 * run {
 *     val tmp0 = a
 *     val tmp1 = b
 *     val tmp2 = tmp0 + tmp1
 *     // call to `builder` with `tmp0`, `tmp1`, and `tmp2` as `ExplainVariable`s.
 *     tmp2
 * }
 * ```
 *
 * Complex expressions may result in [builder] being called multiple times with different variables.
 *
 * ```
 * // Result of transforming: `a && b`
 * run {
 *     val tmp0 = a
 *     if (tmp0) {
 *         val tmp1 = b
 *         // call to `builder` with `tmp0` and `tmp1`
 *         tmp1
 *     } else {
 *         // call to `builder` with just `tmp0`
 *         tmp0
 *     }
 * }
 * ```
 *
 * > Note that results of `&&` and `||` are not included in the list of variables. See KT-73070 for
 * details.
 *
 * If the entire expression is a constant, the expression will not be transformed and [builder]
 * will not be called. In this case, the original expression is simply returned.
 */
@Suppress("unused") // Public API
fun IrBuilderWithScope.irExplain(
    expression: IrExpression,
    sourceFile: SourceFile,
    constTracker: EvaluatedConstTracker? = null,
    builder: IrBlockBuilder.(List<ExplainVariable>) -> Unit,
): IrExpression {
    val root = buildTree(constTracker, sourceFile, expression)
    if (root == null || !root.isVisible()) return expression
    return buildDiagramNesting(sourceFile, root) { value, variables ->
        val explainedVariables = variables
            .filterIsInstance<IrTemporaryVariable.Displayable>()
            .map {
                val offset = it.sourceRangeInfo.startOffset + findDisplayOffset(it.original, it.sourceRangeInfo, it.text)
                ExplainVariable(
                    variable = it.temporary,
                    original = it.original,
                    displayOffset = offset,
                )
            }
        builder(explainedVariables)
        value
    }
}
