/* 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.util.boundary
import scala.util.boundary.break

import org.opalj.br.ArrayType
import org.opalj.br.ClassType
import org.opalj.br.DeclaredMethod
import org.opalj.br.DefinedMethod
import org.opalj.br.Field
import org.opalj.br.PCAndInstruction
import org.opalj.br.ReferenceType
import org.opalj.br.Type
import org.opalj.br.analyses.DeclaredFieldsKey
import org.opalj.br.analyses.ProjectInformationKeys
import org.opalj.br.analyses.SomeProject
import org.opalj.br.analyses.cg.ClosedPackagesKey
import org.opalj.br.analyses.cg.InitialEntryPointsKey
import org.opalj.br.analyses.cg.InitialInstantiatedTypesKey
import org.opalj.br.fpcf.BasicFPCFTriggeredAnalysisScheduler
import org.opalj.br.fpcf.ContextProviderKey
import org.opalj.br.fpcf.FPCFAnalysis
import org.opalj.br.fpcf.analyses.ContextProvider
import org.opalj.br.fpcf.properties.Context
import org.opalj.br.fpcf.properties.cg.Callers
import org.opalj.br.fpcf.properties.cg.InstantiatedTypes
import org.opalj.br.fpcf.properties.cg.NoCallers
import org.opalj.br.instructions.INVOKESPECIAL
import org.opalj.br.instructions.NEW
import org.opalj.collection.immutable.UIDSet
import org.opalj.fpcf.Entity
import org.opalj.fpcf.EOptionP
import org.opalj.fpcf.EPK
import org.opalj.fpcf.EPS
import org.opalj.fpcf.FinalP
import org.opalj.fpcf.InterimEUBP
import org.opalj.fpcf.InterimPartialResult
import org.opalj.fpcf.NoResult
import org.opalj.fpcf.PartialResult
import org.opalj.fpcf.PropertyBounds
import org.opalj.fpcf.PropertyComputationResult
import org.opalj.fpcf.PropertyKey
import org.opalj.fpcf.PropertyStore
import org.opalj.fpcf.Results
import org.opalj.fpcf.SomeEPS
import org.opalj.fpcf.UBP
import org.opalj.util.elidedAssert

/**
 * Marks types as instantiated if their constructor is invoked. Constructors invoked by subclass
 * constructors do not result in additional instantiated types.
 * The analysis does not just looks for "new" instructions, in order to support reflection.
 *
 * This analysis is adapted from the RTA version. Instead of adding the instantiations to the type
 * set of the Project, they are added to the type set of the calling method. Which entity the type
 * is attached to depends on the call graph variant used.
 *
 * TODO: Refactor this and the rta version in order to provide a common base-class.
 *
 * @author Florian Kuebler
 * @author Andreas Bauer
 */
class InstantiatedTypesAnalysis private[analyses] (
    final val project:     SomeProject,
    val setEntitySelector: TypeSetEntitySelector
) extends FPCFAnalysis {

    private implicit val contextProvider: ContextProvider = project.get(ContextProviderKey)

    def analyze(declaredMethod: DeclaredMethod): PropertyComputationResult = {

        val callersEOptP = propertyStore(declaredMethod, Callers.key)

        val callersUB: Callers = (callersEOptP: @unchecked) match {
            case FinalP(NoCallers) =>
                // nothing to do, since there is no caller
                return NoResult;

            case eps: EPS[_, _] =>
                if (eps.ub eq NoCallers) {
                    // we can not create a dependency here, so the analysis is not allowed to create
                    // such a result
                    throw new IllegalStateException("illegal immediate result for callers")
                } else {
                    eps.ub
                }
            // the method is reachable, so we analyze it!
        }

        val declaredType = declaredMethod.declaringClassType.asClassType
        val loadConstantTypes = getLoadConstantTypes(declaredMethod)

        val instantiatedTypes = PartialResult[TypeSetEntity, InstantiatedTypes](
            declaredMethod,
            InstantiatedTypes.key,
            InstantiatedTypes.update(declaredMethod, loadConstantTypes)
        )

        val cfOpt = project.classFile(declaredType)

        // only constructors may initialize a class; abstract classes can never be instantiated
        if (declaredMethod.name != "<init>" || cfOpt.isDefined && cfOpt.get.isAbstract) {
            if (loadConstantTypes.isEmpty)
                return NoResult;
            else
                return Results(instantiatedTypes)
        }

        processCallers(declaredMethod, declaredType, ArrayBuffer(instantiatedTypes), callersEOptP, callersUB, null)
    }

    private def processCallers(
        declaredMethod: DeclaredMethod,
        declaredType:   ClassType,
        partialResults: ArrayBuffer[PartialResult[TypeSetEntity, InstantiatedTypes]],
        callersEOptP:   EOptionP[DeclaredMethod, Callers],
        callersUB:      Callers,
        seenCallers:    Callers
    ): PropertyComputationResult = {
        callersUB.forNewCallerContexts(seenCallers, callersEOptP.e) {
            (_, callerContext, _, isDirect) =>
                processCaller(declaredMethod, declaredType, callerContext, isDirect, partialResults)
        }

        if (callersEOptP.isFinal) {
            Results(partialResults.iterator)
        } else {
            val reRegistration =
                InterimPartialResult(
                    Set(callersEOptP),
                    continuation(declaredMethod, declaredType, callersUB)
                )

            Results(reRegistration, partialResults.iterator)
        }
    }

    private def processCaller(
        declaredMethod: DeclaredMethod,
        declaredType:   ClassType,
        callContext:    Context,
        isDirect:       Boolean,
        partialResults: ArrayBuffer[PartialResult[TypeSetEntity, InstantiatedTypes]]
    ): Unit = boundary {
        // A constructor is called from an unknown context, there could be an initialization.
        if (!callContext.hasContext) {
            partialResults += partialResult(declaredType, ExternalWorld)
            return;
        }

        val caller = callContext.method

        // Indirect calls, e.g. via reflection, are to be treated as instantiations as well
        if (!isDirect) {
            partialResults += partialResult(declaredType, caller)
            return;
        }

        // A constructor is called by a non-constructor method, there will be an initialization.
        if (caller.name != "<init>") {
            partialResults += partialResult(declaredType, caller)
            return;
        }

        // The constructor is called from another constructor. It is only a newly instantiated
        // type if it was no super call. Thus, the caller must be a direct subtype.
        if (!classHierarchy.isSubtypeOf(caller.declaringClassType, declaredType)) {
            partialResults += partialResult(declaredType, caller)
            return;
        }

        // Actually it must be the direct subtype! -- we did the first check to return early
        project.classFile(caller.declaringClassType.asClassType).foreach { cf =>
            cf.superclassType.foreach { supertype =>
                if (supertype != declaredType) {
                    partialResults += partialResult(declaredType, caller)
                    break();
                }
            }
        }

        // If the caller is not available, we have to assume that it was no super call
        if (!caller.hasSingleDefinedMethod) {
            partialResults += partialResult(declaredType, caller)
            return;
        }

        val callerMethod = caller.definedMethod

        // If the caller has no body, we have to assume that it was no super call
        if (callerMethod.body.isEmpty) {
            partialResults += partialResult(declaredType, caller)
            return;
        }

        val supercall = INVOKESPECIAL(
            declaredType,
            isInterface = false,
            "<init>",
            declaredMethod.descriptor
        )

        val pcsOfSuperCalls = callerMethod.body.get.collectInstructionsWithPC {
            case pcAndInstr @ PCAndInstruction(_, `supercall`) => pcAndInstr
        }

        elidedAssert(pcsOfSuperCalls.nonEmpty)

        // There can be only one super call, so there must be an explicit call
        if (pcsOfSuperCalls.size > 1) {
            partialResults += partialResult(declaredType, caller)
            return;
        }

        // There is exactly the current call as potential super call, it still might be no super
        // call if the class has another constructor that calls the super. In that case,
        // there must either be a NEW of the `declaredType` or it is a super call.
        val newInstr = NEW(declaredType)
        val hasNew = callerMethod.body.get.exists(pcInst => pcInst.instruction == newInstr)
        if (hasNew)
            partialResults += partialResult(declaredType, caller)
    }

    private def continuation(
        declaredMethod: DeclaredMethod,
        declaredType:   ClassType,
        seenCallers:    Callers
    )(someEPS: SomeEPS): PropertyComputationResult = {
        val eps = someEPS.asInstanceOf[EPS[DeclaredMethod, Callers]]
        processCallers(declaredMethod, declaredType, ArrayBuffer.empty, eps, eps.ub, seenCallers)
    }

    private def partialResult(
        declaredType: ClassType,
        entity:       Entity
    ): PartialResult[TypeSetEntity, InstantiatedTypes] = {

        // Subtypes of Throwable are tracked globally.
        val setEntity =
            if (classHierarchy.isSubtypeOf(declaredType, ClassType.Throwable))
                project
            else
                setEntitySelector(entity)

        PartialResult[TypeSetEntity, InstantiatedTypes](
            setEntity,
            InstantiatedTypes.key,
            InstantiatedTypes.update(setEntity, UIDSet(declaredType))
        )
    }

    def getInstantiatedTypesUB(
        instantiatedTypesEOptP: EOptionP[SomeProject, InstantiatedTypes]
    ): UIDSet[ReferenceType] = {
        instantiatedTypesEOptP match {
            case eps: EPS[_, _] => eps.ub.types
            case _              => UIDSet.empty
        }
    }
}

class InstantiatedTypesAnalysisScheduler(
    val selectSetEntity: TypeSetEntitySelector
) extends BasicFPCFTriggeredAnalysisScheduler {

    override def requiredProjectInformation: ProjectInformationKeys = Seq(
        ContextProviderKey,
        ClosedPackagesKey,
        InitialEntryPointsKey,
        InitialInstantiatedTypesKey
    )

    override type InitializationData = Null

    override def uses: Set[PropertyBounds] = PropertyBounds.ubs(InstantiatedTypes, Callers)

    override def triggeredBy: PropertyKey[Callers] = Callers.key

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

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

    override def register(p: SomeProject, ps: PropertyStore, unused: Null): FPCFAnalysis = {
        val analysis = new InstantiatedTypesAnalysis(p, selectSetEntity)
        ps.registerTriggeredComputation(triggeredBy, analysis.analyze)
        analysis
    }

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

    def assignInitialTypeSets(p: SomeProject, ps: PropertyStore): Unit = {
        val packageIsClosed = p.get(ClosedPackagesKey)
        val declaredFields = p.get(DeclaredFieldsKey)
        val entryPoints = p.get(InitialEntryPointsKey)
        val initialInstantiatedTypes = UIDSet[ReferenceType](p.get(InitialInstantiatedTypesKey).toSeq*)

        // While processing entry points and fields, we keep track of all array types we see, as
        // well as subtypes and lower-dimensional types. These types also need to be
        // pre-initialized. Note: This set only contains ArrayTypes whose element-type is an
        // ClassType. Arrays of primitive types can be ignored.
        val seenArrayTypes = UIDSet.newBuilder[ArrayType]

        import p.classHierarchy

        def initialize(setEntity: TypeSetEntity, types: UIDSet[ReferenceType]): Unit = {
            ps.preInitialize(setEntity, InstantiatedTypes.key) { pc =>
                (pc: @unchecked) match
                    case UBP(typeSet: InstantiatedTypes) =>
                        InterimEUBP(setEntity, typeSet.updated(types))
                    case _: EPK[_, _] =>
                        InterimEUBP(setEntity, InstantiatedTypes(types))
            }
        }

        // The external world will need some initial types to be instantiated. This is required especially when the
        // Base JAR is not loaded, as the TypeSetEntity of fields like "System.out" will then be "ExternalWorld"
        initialize(ExternalWorld, initialInstantiatedTypes)

        def isRelevantArrayType(rt: Type): Boolean =
            rt.isArrayType && rt.asArrayType.elementType.isClassType

        // For each method which is also an entry point, we assume that the caller has passed all subtypes of the
        // method's parameter types to the method.
        for {
            ep <- entryPoints
        } {
            val typeFilters = UIDSet.newBuilder[ReferenceType]
            val arrayTypeAssignments = UIDSet.newBuilder[ArrayType]

            // If the entry point is not static (or if we do not know whether it may be static), we add the type
            // to the type filters
            val expandTypeFilter = ep match {
                case defM: DefinedMethod => !defM.definedMethod.isStatic
                case _                   => true
            }
            if (expandTypeFilter) {
                typeFilters += ep.declaringClassType
            }

            for (pt <- ep.descriptor.parameterTypes) {
                if (pt.isClassType) {
                    typeFilters += pt.asClassType
                } else if (isRelevantArrayType(pt)) {
                    seenArrayTypes += pt.asArrayType

                    val dim = pt.asArrayType.dimensions
                    val et = pt.asArrayType.elementType.asClassType
                    if (initialInstantiatedTypes.contains(et)) {
                        arrayTypeAssignments += ArrayType(dim, et)
                    }
                }
            }

            val typeFilterSet = typeFilters.result()

            // Initial assignments of ClassTypes
            val classTypeAssignments = initialInstantiatedTypes.filter(iit =>
                typeFilterSet.exists(tf => classHierarchy.isSubtypeOf(iit, tf))
            )

            val initialAssignment = classTypeAssignments ++ arrayTypeAssignments.result()

            val dmSetEntity = selectSetEntity(ep)

            initialize(dmSetEntity, initialAssignment)
        }

        // Returns true if the field's type indicates that the field should be pre-initialized.
        @inline def fieldIsRelevant(f: Field): Boolean = {
            // Only fields which are ArrayType or ClassType are relevant.
            f.fieldType.isReferenceType &&
            // If the field is an ArrayType, then the array's element-type must be a ClassType.
            // In other words: We don't care about arrays of primitive types (e.g. int[]) which
            // do not have to be pre-initialized.
            (!f.fieldType.isArrayType || f.fieldType.asArrayType.elementType.isClassType)
        }

        // Returns true if a field can be written by the user of a library containing that field.
        def fieldIsAccessible(f: Field): Boolean = {
            // Public fields can always be accessed.
            f.isPublic ||
            // Protected fields can only be accessed by subclasses. In that case, the library
            // user can create a subclass of the type containing the field and add a setter method.
            // This only applies if the field's type can be extended in the first place.
            (f.isProtected && !f.classFile.isEffectivelyFinal) ||
            // If the field is package private, it can only be written if the package is
            // open for modification. In that case, the library user can put a method
            // writing that field into the field's type's namespace.
            (f.isPackagePrivate && !packageIsClosed(f.classFile.thisType.packageName))
        }

        for {
            iit <- initialInstantiatedTypes;
            // Only class types should be initially instantiated.
            ct = iit.asClassType
        } {
            // Assign initial types to all accessible fields.
            p.classFile(ct) match {
                case Some(cf) =>
                    for (f <- cf.fields if f.isNotFinal && fieldIsRelevant(f) && fieldIsAccessible(f)) {
                        val fieldType = f.fieldType.asReferenceType

                        val initialAssignments = if (fieldType.isClassType) {
                            val ct = fieldType.asClassType
                            initialInstantiatedTypes.foldLeft(UIDSet.newBuilder[ReferenceType]) {
                                (assignments, iit) =>
                                    if (classHierarchy.isSubtypeOf(iit, ct)) {
                                        assignments += iit
                                    }
                                    assignments
                            }.result()
                        } else {
                            val at = fieldType.asArrayType
                            seenArrayTypes += at
                            val dim = at.dimensions
                            val et = at.elementType.asClassType
                            val allSubtypes = classHierarchy.allSubtypes(et, reflexive = true)
                            initialInstantiatedTypes.foldLeft(UIDSet.newBuilder[ReferenceType]) {
                                (assignments, iit) =>
                                    if (allSubtypes.contains(iit.asClassType)) {
                                        assignments += ArrayType(dim, iit)
                                    }
                                    assignments
                            }.result()
                        }

                        val fieldSetEntity = selectSetEntity(declaredFields(f))
                        initialize(fieldSetEntity, initialAssignments)
                    }
                case None =>
                // Nothing to do here, no classfile => no fields
            }
        }

        // Next, process all ArrayTypes that have been seen while processing entry points and fields,
        // and initialize their type sets.

        // Remember which ArrayTypes were processed, so we don't do it twice.
        val initializedArrayTypes = new java.util.HashSet[ArrayType]()

        def initializeArrayType(at: ArrayType): Unit = {
            // If this type has already been initialized, we skip it.
            if (initializedArrayTypes.contains(at)) {
                return;
            }

            initializedArrayTypes.add(at)

            val et = at.elementType.asClassType
            val allSubtypes = p.classHierarchy.allSubtypes(et, reflexive = true)
            val subtypes =
                initialInstantiatedTypes.foldLeft(UIDSet.newBuilder[ReferenceType]) { (builder, iit) =>
                    if (allSubtypes.contains(iit.asClassType)) {
                        builder += iit
                    }
                    builder
                }.result()

            val dim = at.dimensions
            if (dim > 1) {
                // Initialize multidimensional ArrayType. E.g., if at == A[][] and A is a supertype of A1,
                // we need to assign A[] and A1[] to the type set of A[][].
                val assignedArrayTypes: UIDSet[ArrayType] = UIDSet.fromSpecific(subtypes.map(ArrayType(dim - 1, _)))
                initialize(at, assignedArrayTypes.asInstanceOf[UIDSet[ReferenceType]])

                // After that, we also need to initialize the ArrayTypes which were just assigned. It is possible
                // that these were types which were not initially seen when processing entry points and fields.
                assignedArrayTypes foreach initializeArrayType
            } else {
                // If dim == 1, we just need to assign the "pure" ClassTypes to the ArrayType.
                initialize(at, subtypes)
            }
        }

        seenArrayTypes.result() foreach initializeArrayType
    }
}