package org.checkerframework.common.aliasing;

import com.sun.source.tree.Tree;
import com.sun.source.tree.Tree.Kind;
import java.util.List;
import javax.lang.model.element.ExecutableElement;
import javax.lang.model.element.VariableElement;
import org.checkerframework.common.aliasing.qual.LeakedToResult;
import org.checkerframework.common.aliasing.qual.NonLeaked;
import org.checkerframework.common.aliasing.qual.Unique;
import org.checkerframework.dataflow.analysis.FlowExpressions;
import org.checkerframework.dataflow.analysis.FlowExpressions.Receiver;
import org.checkerframework.dataflow.analysis.RegularTransferResult;
import org.checkerframework.dataflow.analysis.TransferInput;
import org.checkerframework.dataflow.analysis.TransferResult;
import org.checkerframework.dataflow.cfg.node.AssignmentNode;
import org.checkerframework.dataflow.cfg.node.MethodInvocationNode;
import org.checkerframework.dataflow.cfg.node.Node;
import org.checkerframework.dataflow.cfg.node.ObjectCreationNode;
import org.checkerframework.framework.flow.CFAbstractAnalysis;
import org.checkerframework.framework.flow.CFStore;
import org.checkerframework.framework.flow.CFTransfer;
import org.checkerframework.framework.flow.CFValue;
import org.checkerframework.framework.type.AnnotatedTypeFactory;
import org.checkerframework.framework.type.AnnotatedTypeMirror;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedDeclaredType;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedExecutableType;
import org.checkerframework.javacutil.TreeUtils;

/**
 * Type refinement is treated in the usual way, except that at (pseudo-)assignments the RHS may lose
 * its type refinement, before the LHS is type-refined.
 *
 * <p>The RHS always loses its type refinement (it is widened to {@literal @}MaybeAliased, and its
 * declared type must have been {@literal @}MaybeAliased) except in the following cases:
 *
 * <ol>
 *   <li>The RHS is a fresh expression.
 *   <li>The LHS is a {@literal @}NonLeaked formal parameter and the RHS is an argument in a method
 *       call or constructor invocation.
 *   <li>The LHS is a {@literal @}LeakedToResult formal parameter, the RHS is an argument in a
 *       method call or constructor invocation, and the method's return value is discarded.
 * </ol>
 */
public class AliasingTransfer extends CFTransfer {

    private AnnotatedTypeFactory factory;

    public AliasingTransfer(CFAbstractAnalysis<CFValue, CFStore, CFTransfer> analysis) {
        super(analysis);
        factory = analysis.getTypeFactory();
    }

    /**
     * Case 1: For every assignment, the LHS is refined if the RHS has type {@literal @}Unique and
     * is a method invocation or a new class instance.
     */
    @Override
    public TransferResult<CFValue, CFStore> visitAssignment(
            AssignmentNode n, TransferInput<CFValue, CFStore> in) {
        Node rhs = n.getExpression();
        Tree treeRhs = rhs.getTree();
        AnnotatedTypeMirror rhsType = factory.getAnnotatedType(treeRhs);

        if (rhsType.hasAnnotation(Unique.class)
                && (rhs instanceof MethodInvocationNode || rhs instanceof ObjectCreationNode)) {
            return super.visitAssignment(n, in); // Do normal refinement.
        }
        // Widen the type of the rhs if the RHS's declared type wasn't @Unique.
        Receiver r = FlowExpressions.internalReprOf(factory, rhs);
        in.getRegularStore().clearValue(r);
        return new RegularTransferResult<>(null, in.getRegularStore());
    }

    /**
     * Handling pseudo-assignments. Called by {@code CFAbstractTransfer.visitMethodInvocation()}.
     *
     * <p>Case 2: Given a method call, traverses all formal parameters of the method declaration,
     * and if it doesn't have the {@literal @}NonLeaked or {@literal @}LeakedToResult annotations,
     * we remove the node of the respective argument in the method call from the store. If parameter
     * has {@literal @}LeakedToResult, {@code visitMethodInvocation()} handles it.
     */
    @Override
    protected void processPostconditions(
            MethodInvocationNode n, CFStore store, ExecutableElement methodElement, Tree tree) {
        super.processPostconditions(n, store, methodElement, tree);
        if (TreeUtils.isEnumSuper(n.getTree())) {
            // Skipping the init() method for enums.
            return;
        }
        List<Node> args = n.getArguments();
        List<? extends VariableElement> params = methodElement.getParameters();
        assert (args.size() == params.size())
                : "Number of arguments in "
                        + "the method call "
                        + n
                        + " is different from the"
                        + " number of parameters for the method declaration: "
                        + methodElement.getSimpleName();

        AnnotatedExecutableType annotatedType = factory.getAnnotatedType(methodElement);
        List<AnnotatedTypeMirror> paramTypes = annotatedType.getParameterTypes();
        for (int i = 0; i < args.size(); i++) {
            Node arg = args.get(i);
            AnnotatedTypeMirror paramType = paramTypes.get(i);
            if (!paramType.hasAnnotation(NonLeaked.class)
                    && !paramType.hasAnnotation(LeakedToResult.class)) {
                store.clearValue(FlowExpressions.internalReprOf(factory, arg));
            }
        }

        // Now, doing the same as above for the receiver parameter
        Node receiver = n.getTarget().getReceiver();
        AnnotatedDeclaredType receiverType = annotatedType.getReceiverType();
        if (receiverType != null
                && !receiverType.hasAnnotation(LeakedToResult.class)
                && !receiverType.hasAnnotation(NonLeaked.class)) {
            store.clearValue(FlowExpressions.internalReprOf(factory, receiver));
        }
    }

    /**
     * Case 3: Given a method invocation expression, if the parent of the expression is not a
     * statement, check if there are any arguments of the method call annotated as
     * {@literal @}LeakedToResult and remove it from the store, since it might be leaked.
     */
    @Override
    public TransferResult<CFValue, CFStore> visitMethodInvocation(
            MethodInvocationNode n, TransferInput<CFValue, CFStore> in) {
        Tree parent = n.getTreePath().getParentPath().getLeaf();
        boolean parentIsStatement = parent.getKind() == Kind.EXPRESSION_STATEMENT;

        if (!parentIsStatement) {

            ExecutableElement methodElement = TreeUtils.elementFromUse(n.getTree());
            List<Node> args = n.getArguments();
            List<? extends VariableElement> params = methodElement.getParameters();
            assert (args.size() == params.size())
                    : "Number of arguments in "
                            + "the method call "
                            + n
                            + " is different from the"
                            + " number of parameters for the method declaration: "
                            + methodElement.getSimpleName();
            CFStore store = in.getRegularStore();

            for (int i = 0; i < args.size(); i++) {
                Node arg = args.get(i);
                VariableElement param = params.get(i);
                if (factory.getAnnotatedType(param).hasAnnotation(LeakedToResult.class)) {
                    // If argument can leak to result, and parent is not a
                    // single statement, remove that node from store.
                    store.clearValue(FlowExpressions.internalReprOf(factory, arg));
                }
            }

            // Now, doing the same as above for the receiver parameter
            Node receiver = n.getTarget().getReceiver();
            AnnotatedExecutableType annotatedType = factory.getAnnotatedType(methodElement);
            AnnotatedDeclaredType receiverType = annotatedType.getReceiverType();
            if (receiverType != null && receiverType.hasAnnotation(LeakedToResult.class)) {
                store.clearValue(FlowExpressions.internalReprOf(factory, receiver));
            }
        }
        // If parent is a statement, processPostconditions will handle the
        // pseudo-assignments.
        return super.visitMethodInvocation(n, in);
    }
}
