/* BSD 2-Clause License - see OPAL/LICENSE for details. */
package org.opalj
package tac
package fpcf
package analyses
package cg
package xta

import scala.collection.mutable.ArrayBuffer
import scala.jdk.CollectionConverters.*

import org.opalj.br.ClassType
import org.opalj.br.Code
import org.opalj.br.DeclaredMethod
import org.opalj.br.DefinedField
import org.opalj.br.Method
import org.opalj.br.ReferenceType
import org.opalj.br.analyses.DeclaredFields
import org.opalj.br.analyses.DeclaredFieldsKey
import org.opalj.br.analyses.ProjectInformationKeys
import org.opalj.br.analyses.SomeProject
import org.opalj.br.fpcf.BasicFPCFTriggeredAnalysisScheduler
import org.opalj.br.fpcf.FPCFAnalysis
import org.opalj.br.fpcf.properties.cg.Callees
import org.opalj.br.fpcf.properties.cg.Callers
import org.opalj.br.fpcf.properties.cg.InstantiatedTypes
import org.opalj.br.fpcf.properties.fieldaccess.MethodFieldReadAccessInformation
import org.opalj.br.fpcf.properties.fieldaccess.MethodFieldWriteAccessInformation
import org.opalj.br.instructions.CHECKCAST
import org.opalj.collection.immutable.UIDSet
import org.opalj.fpcf.Entity
import org.opalj.fpcf.EPS
import org.opalj.fpcf.EUBP
import org.opalj.fpcf.FinalEP
import org.opalj.fpcf.InterimPartialResult
import org.opalj.fpcf.PartialResult
import org.opalj.fpcf.ProperPropertyComputationResult
import org.opalj.fpcf.PropertyBounds
import org.opalj.fpcf.PropertyKind
import org.opalj.fpcf.PropertyStore
import org.opalj.fpcf.Results
import org.opalj.fpcf.SomeEPS
import org.opalj.fpcf.SomePartialResult
import org.opalj.tac.cg.TypeIteratorKey
import org.opalj.tac.fpcf.properties.NoTACAI
import org.opalj.tac.fpcf.properties.TACAI
import org.opalj.util.elidedAssert

/**
 * This analysis handles the type propagation of XTA, MTA, FTA and CTA call graph
 * algorithms.
 *
 * @param project         Project under analysis
 * @param selectTypeSetEntity Function which, for each entity, selects which entity its type set is attached to.
 *
 * @author Andreas Bauer
 */
final class TypePropagationAnalysis private[analyses] (
    val project:         SomeProject,
    selectTypeSetEntity: TypeSetEntitySelector
) extends ReachableMethodAnalysis {

    private val debug = false
    private val _trace: TypePropagationTrace = new TypePropagationTrace()

    private type State = TypePropagationState[ContextType]

    private implicit val declaredFields: DeclaredFields = project.get(DeclaredFieldsKey)

    // We need to also propagate types if the method has no body, e.g. for native methods with configured data.
    // Those methods set the TAC EPS to null. Access to state.tac will always be guarded by if(state.methodHasBody).
    override def processMethodWithoutBody(callContext: ContextType): ProperPropertyComputationResult =
        processMethod(callContext, FinalEP(null, NoTACAI))

    override def processMethod(
        callContext: ContextType,
        tacEP:       EPS[Method, TACAI]
    ): ProperPropertyComputationResult = {

        val declaredMethod = callContext.method

        val methodOpt = if (declaredMethod.hasSingleDefinedMethod) {
            Some(declaredMethod.definedMethod)
        } else None

        val typeSetEntity = selectTypeSetEntity(declaredMethod)
        val instantiatedTypesEOptP = propertyStore(typeSetEntity, InstantiatedTypes.key)
        val calleesEOptP = propertyStore(declaredMethod, Callees.key)
        val readAccessEOptP = methodOpt.map(m => propertyStore(m, MethodFieldReadAccessInformation.key))
        val writeAccessEOptP = methodOpt.map(m => propertyStore(m, MethodFieldWriteAccessInformation.key))

        if (debug) _trace.traceInit(declaredMethod)

        implicit val state: TypePropagationState[ContextType] = new TypePropagationState(
            callContext,
            typeSetEntity,
            tacEP,
            instantiatedTypesEOptP,
            calleesEOptP,
            readAccessEOptP,
            writeAccessEOptP
        )
        implicit val partialResults: ArrayBuffer[SomePartialResult] = ArrayBuffer.empty[SomePartialResult]

        if (calleesEOptP.hasUBP)
            processCallees(calleesEOptP.ub)
        if (readAccessEOptP.isDefined && readAccessEOptP.get.hasUBP)
            processReadAccesses(readAccessEOptP.get.ub)
        if (writeAccessEOptP.isDefined && writeAccessEOptP.get.hasUBP)
            processWriteAccesses(writeAccessEOptP.get.ub)

        if (state.methodHasBody) {
            processTACStatements
            processArrayTypes(state.ownInstantiatedTypes)
        }

        returnResults(partialResults.iterator)
    }

    /**
     * Processes the method upon initialization. Finds array accesses for further processing in [[processArrayTypes]].
     */
    private def processTACStatements(implicit state: State): Unit = {
        val tac = state.tac
        tac.stmts.foreach {
            case Assignment(_, _, expr) if expr.astID == ArrayLoad.ASTID =>
                state.methodReadsArrays = true

            case stmt: Stmt[?] if stmt.astID == ArrayStore.ASTID =>
                state.methodWritesArrays = true

            case _ =>
        }
    }

    private def c(state: State)(eps: SomeEPS): ProperPropertyComputationResult = eps match {

        case EUBP(e: DeclaredMethod, _: Callees) =>
            if (debug) {
                elidedAssert(e == state.callContext.method)
                _trace.traceCalleesUpdate(e)
            }
            handleUpdateOfCallees(eps.asInstanceOf[EPS[DeclaredMethod, Callees]])(using state)

        case EUBP(e: Method, _: MethodFieldReadAccessInformation) =>
            if (debug) {
                elidedAssert(e == state.callContext.method.definedMethod)
                _trace.traceReadAccessUpdate(e)
            }
            handleUpdateOfReadAccesses(eps.asInstanceOf[EPS[Method, MethodFieldReadAccessInformation]])(using state)

        case EUBP(e: Method, _: MethodFieldWriteAccessInformation) =>
            if (debug) {
                elidedAssert(e == state.callContext.method.definedMethod)
                _trace.traceWriteAccessUpdate(e)
            }
            handleUpdateOfWriteAccesses(eps.asInstanceOf[EPS[Method, MethodFieldWriteAccessInformation]])(using state)

        case EUBP(e: TypeSetEntity, t: InstantiatedTypes) if e == state.typeSetEntity =>
            if (debug) _trace.traceTypeUpdate(state.callContext.method, e, t.types)
            handleUpdateOfOwnTypeSet(eps.asInstanceOf[EPS[TypeSetEntity, InstantiatedTypes]])(using state)

        case EUBP(e: TypeSetEntity, t: InstantiatedTypes) =>
            if (debug) _trace.traceTypeUpdate(state.callContext.method, e, t.types)
            handleUpdateOfBackwardPropagationTypeSet(eps.asInstanceOf[EPS[TypeSetEntity, InstantiatedTypes]])(
                using state
            )

        case _ =>
            sys.error("received unexpected update")
    }

    private def handleUpdateOfCallees(
        eps: EPS[DeclaredMethod, Callees]
    )(
        implicit state: State
    ): ProperPropertyComputationResult = {
        state.updateCalleeDependee(eps)
        implicit val partialResults: ArrayBuffer[SomePartialResult] = ArrayBuffer.empty[SomePartialResult]
        processCallees(eps.ub)
        returnResults(partialResults.iterator)
    }

    private def handleUpdateOfReadAccesses(
        eps: EPS[Method, MethodFieldReadAccessInformation]
    )(implicit state: State): ProperPropertyComputationResult = {
        state.updateReadAccessDependee(eps)
        implicit val partialResults: ArrayBuffer[SomePartialResult] = ArrayBuffer.empty[SomePartialResult]
        processReadAccesses(eps.ub)
        returnResults(partialResults.iterator)
    }

    private def handleUpdateOfWriteAccesses(
        eps: EPS[Method, MethodFieldWriteAccessInformation]
    )(implicit state: State): ProperPropertyComputationResult = {
        state.updateWriteAccessDependee(eps)
        implicit val partialResults: ArrayBuffer[SomePartialResult] = ArrayBuffer.empty[SomePartialResult]
        processWriteAccesses(eps.ub)
        returnResults(partialResults.iterator)
    }

    private def handleUpdateOfOwnTypeSet(
        eps: EPS[TypeSetEntity, InstantiatedTypes]
    )(
        implicit state: State
    ): ProperPropertyComputationResult = {
        val previouslySeenTypes = state.ownInstantiatedTypes.size
        state.updateOwnInstantiatedTypesDependee(eps)
        val unseenTypes = UIDSet(eps.ub.dropOldest(previouslySeenTypes).toSeq*)

        implicit val partialResults: ArrayBuffer[SomePartialResult] = ArrayBuffer.empty[SomePartialResult]
        for (fpe <- state.forwardPropagationEntities.iterator().asScala) {
            val filters = state.forwardPropagationFilters(fpe)
            val propagation = propagateTypes(fpe, unseenTypes, filters)
            if (propagation.isDefined)
                partialResults += propagation.get
        }

        processArrayTypes(unseenTypes)

        returnResults(partialResults.iterator)
    }

    private def handleUpdateOfBackwardPropagationTypeSet(
        eps: EPS[TypeSetEntity, InstantiatedTypes]
    )(
        implicit state: State
    ): ProperPropertyComputationResult = {
        val typeSetEntity = eps.e
        val previouslySeenTypes = state.seenTypes(typeSetEntity)
        state.updateBackwardPropagationDependee(eps)
        val unseenTypes = UIDSet(eps.ub.dropOldest(previouslySeenTypes).toSeq*)

        val filters = state.backwardPropagationFilters(typeSetEntity)
        val propagationResult = propagateTypes(state.typeSetEntity, unseenTypes, filters)

        returnResults(propagationResult)
    }

    private def processArrayTypes(
        unseenTypes: UIDSet[ReferenceType]
    )(
        implicit
        state:          State,
        partialResults: ArrayBuffer[SomePartialResult]
    ): Unit = {
        for (t <- unseenTypes if t.isArrayType; at = t.asArrayType if at.elementType.isReferenceType) {
            if (state.methodWritesArrays) {
                registerEntityForForwardPropagation(at, UIDSet(at.componentType.asReferenceType))
            }
            if (state.methodReadsArrays) {
                registerEntityForBackwardPropagation(at, at.componentType.asReferenceType)
            }
        }
    }

    private def isIgnoredCallee(callee: DeclaredMethod): Boolean = {
        // Special case: Object.<init> is implicitly called as a super call by any method X.<init> when X does
        // not have a supertype.
        // The "this" type X will flow to the type set of Object.<init>. Since Object.<init> is usually
        // part of the external world, the external world type set is then polluted with any type which
        // was constructed anywhere in the program.
        callee.declaringClassType == ClassType.Object && callee.name == "<init>"
    }

    private def processCallees(
        callees: Callees
    )(
        implicit
        state:          State,
        partialResults: ArrayBuffer[SomePartialResult]
    ): Unit = {
        val bytecodeOpt = if (state.methodHasBody) Some(state.callContext.method.definedMethod.body.get) else None
        for {
            pc <- callees.callSitePCs(state.callContext)
            calleeContext <- callees.callees(state.callContext, pc)
            callee = calleeContext.method
            if !state.isSeenCallee(pc, callee) && !isIgnoredCallee(callee)
        } {
            // Some sanity checks ...
            // Methods with multiple defined methods should never appear as callees.
            elidedAssert(!callee.hasMultipleDefinedMethods)
            // Instances of DefinedMethod we see should only be those where the method is defined in the class file of
            // the declaring class type (i.e., it is not a DefinedMethod instance of some inherited method).  However,
            // in inconsistent bytecode scenarios, the call graph may resolve to a non-implemented abstract method.
            elidedAssert(!callee.hasSingleDefinedMethod ||
                (callee.declaringClassType == callee.asDefinedMethod.definedMethod.classFile.thisType) ||
                callee.asDefinedMethod.definedMethod.isAbstract)

            // Remember callee (with PC) so we don't have to process it again later.
            state.addSeenCallee(pc, callee)

            maybeRegisterMethodForForwardPropagation(callee, pc, bytecodeOpt)

            maybeRegisterMethodForBackwardPropagation(callee, pc, bytecodeOpt)
        }
    }

    private def processReadAccesses(
        readAccesses: MethodFieldReadAccessInformation
    )(implicit state: State, partialResults: ArrayBuffer[SomePartialResult]): Unit = {
        val bytecodeOpt = if (state.methodHasBody) Some(state.callContext.method.definedMethod.body.get) else None
        for {
            pc <- readAccesses.getAccessSites(state.callContext)
            numDirectAccess = readAccesses.numDirectAccesses(state.callContext, pc)
            numIndirectAccess = readAccesses.numIndirectAccesses(state.callContext, pc)
            declaredField <- readAccesses.getNewestAccessedFields(
                state.callContext,
                pc,
                numDirectAccess - state.seenDirectReadAccesses.getOrElse(pc, 0),
                numIndirectAccess - state.seenIndirectReadAccesses.getOrElse(pc, 0)
            )
        } {
            if (declaredField.fieldType.isReferenceType) {
                // Internally, generic fields have type "Object" due to type erasure. In many cases
                // (but not all!), the Java compiler will place the "actual" return type within a checkcast
                // instruction right after the field read instruction.
                val nextInstructionOpt =
                    bytecodeOpt.map(bytecode => bytecode.instructions(bytecode.pcOfNextInstruction(pc)))
                val mostPreciseFieldType =
                    if (nextInstructionOpt.isDefined && nextInstructionOpt.get.isCheckcast)
                        nextInstructionOpt.get.asInstanceOf[CHECKCAST].referenceType
                    else
                        declaredField.fieldType.asReferenceType

                declaredField match {
                    case DefinedField(f) if project.isProjectType(f.classFile.thisType) =>
                        registerEntityForBackwardPropagation(declaredField.asDefinedField, mostPreciseFieldType)

                    case _ =>
                        registerEntityForBackwardPropagation(declaredField, mostPreciseFieldType)
                }
            }

            state.seenDirectReadAccesses += pc -> numDirectAccess
            state.seenIndirectReadAccesses += pc -> numIndirectAccess
        }
    }

    private def processWriteAccesses(
        writeAccesses: MethodFieldWriteAccessInformation
    )(implicit state: State, partialResults: ArrayBuffer[SomePartialResult]): Unit = {
        for {
            pc <- writeAccesses.getAccessSites(state.callContext)
            numDirectAccess = writeAccesses.numDirectAccesses(state.callContext, pc)
            numIndirectAccess = writeAccesses.numIndirectAccesses(state.callContext, pc)
            declaredField <- writeAccesses.getNewestAccessedFields(
                state.callContext,
                pc,
                numDirectAccess - state.seenDirectWriteAccesses.getOrElse(pc, 0),
                numIndirectAccess - state.seenIndirectWriteAccesses.getOrElse(pc, 0)
            )
        } {
            val fieldType = declaredField.fieldType
            if (fieldType.isReferenceType) {
                declaredField match {
                    case DefinedField(f) if project.isProjectType(f.classFile.thisType) =>
                        registerEntityForForwardPropagation(
                            declaredField.asDefinedField,
                            UIDSet(fieldType.asReferenceType)
                        )

                    case _ =>
                        registerEntityForForwardPropagation(declaredField, UIDSet(fieldType.asReferenceType))
                }
            }

            state.seenDirectWriteAccesses += pc -> numDirectAccess
            state.seenIndirectWriteAccesses += pc -> numIndirectAccess
        }
    }

    private def maybeRegisterMethodForForwardPropagation(
        callee:      DeclaredMethod,
        pc:          Int,
        bytecodeOpt: Option[Code]
    )(
        implicit
        state:          State,
        partialResults: ArrayBuffer[SomePartialResult]
    ): Unit = {
        val params = UIDSet.newBuilder[ReferenceType]

        for (param <- callee.descriptor.parameterTypes) {
            if (param.isReferenceType) {
                params += param.asReferenceType
            }
        }

        // If the call is not static, we need to take the implicit "this" parameter into account.
        if (callee.hasSingleDefinedMethod && !callee.definedMethod.isStatic ||
            // If we can't ask the target if it is static, we look in the bytecode. If we do not have any bytecode,
            // we soundly assume that the call might be virtual and need to add "this".
            !callee.hasSingleDefinedMethod && !bytecodeOpt.exists(_.instructions(pc).isInvokeStatic)
        ) {
            params += callee.declaringClassType
        }

        // If we do not have any params at this point, there is no forward propagation!
        val typeFilters = params.result()
        if (typeFilters.isEmpty) {
            return;
        }

        registerEntityForForwardPropagation(callee, typeFilters)
    }

    private def maybeRegisterMethodForBackwardPropagation(
        callee:      DeclaredMethod,
        pc:          Int,
        bytecodeOpt: Option[Code]
    )(
        implicit
        state:          State,
        partialResults: ArrayBuffer[SomePartialResult]
    ): Unit = {
        // If the method body is not available, we have to assume the return value might be used
        val returnValueIsUsed = if (!state.methodHasBody) true
        else {
            val tacIndex = state.tac.properStmtIndexForPC(pc)
            val tacInstr = state.tac.instructions(tacIndex)
            tacInstr.isAssignment
        }

        if (returnValueIsUsed) {
            // Internally, generic methods have return type "Object" due to type erasure. In many cases
            // (but not all!), the Java compiler will place the "actual" return type within a checkcast
            // instruction right after the call.
            val mostPreciseReturnType = {
                val nextInstructionOpt = bytecodeOpt.map { bytecode =>
                    val nextPc = bytecode.pcOfNextInstruction(pc)
                    bytecode.instructions(nextPc)
                }
                // If method body is not available, we will assume the callee descriptor return type as well
                if (nextInstructionOpt.exists(_.isCheckcast)) {
                    nextInstructionOpt.get.asInstanceOf[CHECKCAST].referenceType
                } else {
                    callee.descriptor.returnType
                }
            }

            // Return type could also be a basic type (i.e., int). We don't care about those.
            if (mostPreciseReturnType.isReferenceType) {
                registerEntityForBackwardPropagation(callee, mostPreciseReturnType.asReferenceType)
            }
        }
    }

    private def registerEntityForForwardPropagation(
        e:       Entity,
        filters: UIDSet[ReferenceType]
    )(
        implicit
        state:          State,
        partialResults: ArrayBuffer[SomePartialResult]
    ): Unit = {
        // Propagation from and to the same entity can be ignored.
        val typeSetEntity = selectTypeSetEntity(e)
        if (typeSetEntity == state.typeSetEntity) {
            return;
        }

        val filterSetHasChanged = state.registerForwardPropagationEntity(typeSetEntity, filters)
        if (filterSetHasChanged) {
            val propagationResult =
                propagateTypes(typeSetEntity, state.ownInstantiatedTypes, state.forwardPropagationFilters(typeSetEntity))
            if (propagationResult.isDefined)
                partialResults += propagationResult.get
        }
    }

    private def registerEntityForBackwardPropagation(
        e:                     Entity,
        mostPreciseUpperBound: ReferenceType
    )(
        implicit
        state:          State,
        partialResults: ArrayBuffer[SomePartialResult]
    ): Unit = {
        val typeSetEntity = selectTypeSetEntity(e)
        if (typeSetEntity == state.typeSetEntity) {
            return;
        }

        val filter = UIDSet(mostPreciseUpperBound)

        if (!state.backwardPropagationDependeeIsRegistered(typeSetEntity)) {
            val dependee = propertyStore(typeSetEntity, InstantiatedTypes.key)

            state.updateBackwardPropagationDependee(dependee)
            state.updateBackwardPropagationFilters(typeSetEntity, filter)

            if (dependee.hasNoUBP) {
                return;
            }

            val propagation = propagateTypes(state.typeSetEntity, dependee.ub.types, filter)
            if (propagation.isDefined) {
                partialResults += propagation.get
            }
        } else {
            val filterSetHasChanged = state.updateBackwardPropagationFilters(typeSetEntity, filter)
            if (filterSetHasChanged) {
                // Since the filters were updated, it is possible that types which were previously seen but not
                // propagated are now relevant for back propagation. Therefore, we need to propagate from the
                // entire dependee type set.
                val allDependeeTypes = state.backwardPropagationDependeeInstantiatedTypes(typeSetEntity)
                val propagation = propagateTypes(state.typeSetEntity, allDependeeTypes, filter)
                if (propagation.isDefined) {
                    partialResults += propagation.get
                }
            }
        }
    }

    private def propagateTypes[E >: Null <: TypeSetEntity](
        targetSetEntity: E,
        newTypes:        UIDSet[ReferenceType],
        filters:         Set[ReferenceType]
    ): Option[PartialResult[E, InstantiatedTypes]] = {

        if (newTypes.isEmpty) {
            return None;
        }

        val filteredTypes = newTypes.foldLeft(UIDSet.newBuilder[ReferenceType]) { (builder, nt) =>
            val fitr = filters.iterator
            var candidateMatches = false
            while (!candidateMatches && fitr.hasNext) {
                val tf = fitr.next()
                if (candidateMatchesTypeFilter(nt, tf)) {
                    candidateMatches = true
                    builder += nt
                }
            }
            builder
        }.result()

        if (filteredTypes.nonEmpty) {
            if (debug) _trace.traceTypePropagation(targetSetEntity, filteredTypes)
            val partialResult = PartialResult[E, InstantiatedTypes](
                targetSetEntity,
                InstantiatedTypes.key,
                InstantiatedTypes.update(targetSetEntity, filteredTypes)
            )

            Some(partialResult)
        } else {
            None
        }
    }

    private def returnResults(
        partialResults: IterableOnce[SomePartialResult]
    )(implicit state: State): ProperPropertyComputationResult = {
        // Always re-register the continuation. It is impossible for all dependees to be final in XTA/...
        Results(
            InterimPartialResult(state.dependees, c(state)),
            partialResults
        )
    }
}

final class TypePropagationAnalysisScheduler(
    val selectSetEntity: TypeSetEntitySelector
) extends BasicFPCFTriggeredAnalysisScheduler {

    override def requiredProjectInformation: ProjectInformationKeys = Seq(TypeIteratorKey, DeclaredFieldsKey)

    override type InitializationData = Null

    override def triggeredBy: PropertyKind = Callers.key

    override def init(p: SomeProject, ps: PropertyStore): Null = null

    override def register(project: SomeProject, propertyStore: PropertyStore, i: Null): FPCFAnalysis = {
        val analysis = new TypePropagationAnalysis(project, selectSetEntity)
        propertyStore.registerTriggeredComputation(Callers.key, analysis.analyze)
        analysis
    }

    override def uses: Set[PropertyBounds] = PropertyBounds.ubs(
        InstantiatedTypes,
        Callees,
        TACAI,
        MethodFieldReadAccessInformation,
        MethodFieldWriteAccessInformation
    )

    override def uses(p: SomeProject, ps: PropertyStore): Set[PropertyBounds] =
        p.get(TypeIteratorKey).usedPropertyKinds

    override def derivesEagerly: Set[PropertyBounds] = Set.empty

    override def derivesCollaboratively: Set[PropertyBounds] = PropertyBounds.ubs(InstantiatedTypes)
}