package org.checkerframework.checker.nullness;

import com.sun.source.tree.ExpressionTree;
import com.sun.source.tree.NewClassTree;
import com.sun.source.tree.VariableTree;
import javax.lang.model.element.ExecutableElement;
import javax.lang.model.type.TypeKind;
import org.checkerframework.checker.nullness.KeyForPropagator.PropagationDirection;
import org.checkerframework.framework.type.AnnotatedTypeFactory;
import org.checkerframework.framework.type.AnnotatedTypeMirror;
import org.checkerframework.framework.type.AnnotatedTypeMirror.AnnotatedDeclaredType;
import org.checkerframework.framework.type.treeannotator.TreeAnnotator;
import org.checkerframework.javacutil.TreeUtils;

/**
 * For the following initializations we wish to propagate the annotations from the left-hand side to
 * the right-hand side or vice versa:
 *
 * <p>1. If a keySet is being saved to a newly declared set, we transfer the annotations from the
 * keySet to the lhs. e.g.,
 *
 * <pre>{@code
 * // The user is not required to explicitly annotate the LHS's type argument with @KeyFor("m")
 * Set<String> keySet = m.keySet();
 * }</pre>
 *
 * 2. If a variable declaration contains type arguments with an @KeyFor annotation and its
 * initializer is a new class tree with corresponding type arguments that have an @UknownKeyFor
 * primary annotation, we transfer from the LHS to RHS. e.g.,
 *
 * <pre>{@code
 * // The user does not have to write @KeyFor("m") on both sides
 * List<@KeyFor("m") String> keys = new ArrayList<String>();
 * }</pre>
 *
 * 3. IMPORTANT NOTE: The following case must be (and is) handled in KeyForAnnotatedTypeFactory. In
 * BaseTypeVisitor we check to make sure that the constructor called in a NewClassTree is actually
 * compatible with the annotations placed on the NewClassTree. This requires that, prior to this
 * check we also propagate the annotations to this constructor in constructorFromUse so that the
 * constructor call matches the type given to the NewClassTree.
 *
 * @see
 *     org.checkerframework.checker.nullness.KeyForAnnotatedTypeFactory#constructorFromUse(com.sun.source.tree.NewClassTree)
 *     <p>Note propagation only occurs between two AnnotatedDeclaredTypes. If either side is not an
 *     AnnotatedDeclaredType then this class does nothing.
 */
public class KeyForPropagationTreeAnnotator extends TreeAnnotator {
    private final KeyForPropagator keyForPropagator;
    private final ExecutableElement keySetMethod;

    public KeyForPropagationTreeAnnotator(
            AnnotatedTypeFactory atypeFactory, KeyForPropagator propagationTreeAnnotator) {
        super(atypeFactory);
        this.keyForPropagator = propagationTreeAnnotator;
        keySetMethod =
                TreeUtils.getMethod(
                        java.util.Map.class.getName(),
                        "keySet",
                        0,
                        atypeFactory.getProcessingEnv());
    }

    /** @return true iff expression is a call to java.util.Map.KeySet */
    public boolean isCallToKeyset(ExpressionTree expression) {
        return TreeUtils.isMethodInvocation(
                expression, keySetMethod, atypeFactory.getProcessingEnv());
    }

    /**
     * Transfers annotations on type arguments from the initializer to the variableTree, if the
     * initializer is a call to java.util.Map.keySet.
     */
    @Override
    public Void visitVariable(VariableTree variableTree, AnnotatedTypeMirror type) {
        super.visitVariable(variableTree, type);

        // This should only happen on Map.keySet();
        if (type.getKind() == TypeKind.DECLARED) {
            final ExpressionTree initializer = variableTree.getInitializer();

            if (isCallToKeyset(initializer)) {
                final AnnotatedDeclaredType variableType = (AnnotatedDeclaredType) type;
                final AnnotatedTypeMirror initializerType =
                        atypeFactory.getAnnotatedType(initializer);

                // Propagate just for declared (class) types, not for array types, boxed primitives,
                // etc.
                if (variableType.getKind() == TypeKind.DECLARED) {
                    keyForPropagator.propagate(
                            (AnnotatedDeclaredType) initializerType,
                            variableType,
                            PropagationDirection.TO_SUPERTYPE,
                            atypeFactory);
                }
            }
        }

        return null;
    }

    /** Transfers annotations to type if the left hand side is a variable declaration. */
    @Override
    public Void visitNewClass(NewClassTree node, AnnotatedTypeMirror type) {
        keyForPropagator.propagateNewClassTree(
                node, type, (KeyForAnnotatedTypeFactory) atypeFactory);
        return super.visitNewClass(node, type);
    }
}
