QualifierPolymorphism assignment context resolution

framework/src/main/java/org/checkerframework/framework/type/poly/AbstractQualifierPolymorphism.java

@@ -2,6 +2,7 @@
 
 import com.sun.source.tree.MethodInvocationTree;
 import com.sun.source.tree.NewClassTree;
+import com.sun.source.util.TreePath;
 import java.util.ArrayList;
 import java.util.Iterator;
 import java.util.List;
@@ -28,6 +29,7 @@
 import org.checkerframework.framework.util.AnnotatedTypes;
 import org.checkerframework.framework.util.AnnotationMirrorMap;
 import org.checkerframework.framework.util.AnnotationMirrorSet;
+import org.checkerframework.framework.util.QualifierPolymorphismUtil;
 import org.checkerframework.javacutil.AnnotationBuilder;
 import org.checkerframework.javacutil.TreeUtils;
 import org.checkerframework.javacutil.TypesUtils;
@@ -149,6 +151,21 @@ public void annotate(MethodInvocationTree tree, AnnotatedExecutableType type) {
                                     atypeFactory.getReceiverType(tree), type.getReceiverType()));
         }
 
+        TreePath path = atypeFactory.getPath(tree);
+        if (path != null) {
+            AnnotatedTypeMirror assignmentContext =
+                    QualifierPolymorphismUtil.assignedTo(atypeFactory, path);
+
+            if (assignmentContext != null) {
+                matchingMapping =
+                        collector.reduceWithUpperBounds(
+                                matchingMapping,
+                                collector.visit(
+                                        // Actual assignment lhs type
+                                        assignmentContext, type.getReturnType(), true));
+            }
+        }
+
         if (matchingMapping != null && !matchingMapping.isEmpty()) {
             replacer.visit(type, matchingMapping);
         } else {
@@ -353,6 +370,49 @@ public AnnotationMirrorMap<AnnotationMirrorSet> reduce(
             return res;
         }
 
+        /** Reduces lower bounds r1 with upper bounds r2 */
+        private AnnotationMirrorMap<AnnotationMirrorSet> reduceWithUpperBounds(
+                AnnotationMirrorMap<AnnotationMirrorSet> r1,
+                AnnotationMirrorMap<AnnotationMirrorSet> r2) {
+
+            if (r1 == null || r1.isEmpty()) {
+                return r2;
+            }
+            if (r2 == null || r2.isEmpty()) {
+                return r1;
+            }
+
+            AnnotationMirrorMap<AnnotationMirrorSet> res = new AnnotationMirrorMap<>();
+            // Ensure that all qualifiers from r1 and r2 are visited.
+            AnnotationMirrorSet r2remain = new AnnotationMirrorSet();
+            r2remain.addAll(r2.keySet());
+            for (Map.Entry<AnnotationMirror, AnnotationMirrorSet> kv1 : r1.entrySet()) {
+                AnnotationMirror key1 = kv1.getKey();
+                AnnotationMirrorSet a1Annos = kv1.getValue();
+                AnnotationMirrorSet a2Annos = r2.get(key1);
+                if (a2Annos != null && !a2Annos.isEmpty()) {
+                    r2remain.remove(key1);
+                    AnnotationMirrorSet subres = new AnnotationMirrorSet();
+                    for (AnnotationMirror top : topQuals) {
+                        AnnotationMirror a1 = qualHierarchy.findAnnotationInHierarchy(a1Annos, top);
+                        AnnotationMirror a2 = qualHierarchy.findAnnotationInHierarchy(a2Annos, top);
+                        if (a1 != null) {
+                            subres.add(a1);
+                        } else if (a2 != null) {
+                            subres.add(a2);
+                        }
+                    }
+                    res.put(key1, subres);
+                } else {
+                    res.put(key1, a1Annos);
+                }
+            }
+            for (AnnotationMirror key2 : r2remain) {
+                res.put(key2, r2.get(key2));
+            }
+            return res;
+        }
+
         /**
          * Calls {@link #visit(AnnotatedTypeMirror, AnnotatedTypeMirror)} for each type in types.
          */
@@ -378,10 +438,11 @@ private AnnotationMirrorMap<AnnotationMirrorSet> visit(
          *
          * @param type AnnotateTypeMirror used to find instantiations
          * @param polyType AnnotatedTypeMirror that may have polymorphich qualifiers
+         * @param polyIsSub boolean indicates whether polyType is subtype of type
          * @return a mapping of polymorphic qualifiers to their instantiations
          */
         private AnnotationMirrorMap<AnnotationMirrorSet> visit(
-                AnnotatedTypeMirror type, AnnotatedTypeMirror polyType) {
+                AnnotatedTypeMirror type, AnnotatedTypeMirror polyType, boolean polyIsSub) {
             if (type.getKind() == TypeKind.NULL) {
                 return mapQualifierToPoly(type, polyType);
             }
@@ -394,19 +455,35 @@ private AnnotationMirrorMap<AnnotationMirrorSet> visit(
 
                 switch (polyType.getKind()) {
                     case WILDCARD:
-                        AnnotatedTypeMirror asSuper =
-                                AnnotatedTypes.asSuper(atypeFactory, wildcardType, polyType);
-                        return visit(asSuper, polyType, null);
+                        if (polyIsSub) {
+                            return visit(
+                                    wildcardType,
+                                    AnnotatedTypes.asSuper(atypeFactory, polyType, wildcardType),
+                                    null);
+                        } else {
+                            return visit(
+                                    AnnotatedTypes.asSuper(atypeFactory, wildcardType, polyType),
+                                    polyType,
+                                    null);
+                        }
                     case TYPEVAR:
                         return mapQualifierToPoly(wildcardType.getExtendsBound(), polyType);
                     default:
                         return mapQualifierToPoly(wildcardType.getExtendsBound(), polyType);
                 }
             }
 
-            AnnotatedTypeMirror asSuper = AnnotatedTypes.asSuper(atypeFactory, type, polyType);
+            if (polyIsSub) {
+                return visit(type, AnnotatedTypes.asSuper(atypeFactory, polyType, type), null);
+            } else {
+                return visit(AnnotatedTypes.asSuper(atypeFactory, type, polyType), polyType, null);
+            }
+        }
 
-            return visit(asSuper, polyType, null);
+        /** Same visit but defaults polyIsSub to false */
+        private AnnotationMirrorMap<AnnotationMirrorSet> visit(
+                AnnotatedTypeMirror type, AnnotatedTypeMirror polyType) {
+            return visit(type, polyType, false);
         }
 
         /**

framework/src/main/java/org/checkerframework/framework/util/QualifierPolymorphismUtil.java

@@ -0,0 +1,113 @@
+package org.checkerframework.framework.util;
+
+import com.sun.source.tree.AssignmentTree;
+import com.sun.source.tree.CompoundAssignmentTree;
+import com.sun.source.tree.ExpressionTree;
+import com.sun.source.tree.LambdaExpressionTree;
+import com.sun.source.tree.MethodInvocationTree;
+import com.sun.source.tree.MethodTree;
+import com.sun.source.tree.NewArrayTree;
+import com.sun.source.tree.NewClassTree;
+import com.sun.source.tree.ReturnTree;
+import com.sun.source.tree.Tree;
+import com.sun.source.tree.Tree.Kind;
+import com.sun.source.tree.VariableTree;
+import com.sun.source.util.TreePath;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+import javax.lang.model.element.ExecutableElement;
+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.framework.util.typeinference.TypeArgInferenceUtil;
+import org.checkerframework.javacutil.BugInCF;
+import org.checkerframework.javacutil.Pair;
+import org.checkerframework.javacutil.TreeUtils;
+
+public class QualifierPolymorphismUtil {
+
+    private static final Set<TreePath> visitedPaths = new HashSet<>();
+
+    /**
+     * Returns the annotated type that the leaf of path is assigned to, if it is within an
+     * assignment context. Returns the annotated type that the method invocation at the leaf is
+     * assigned to.
+     *
+     * @return type that path leaf is assigned to
+     */
+    public static AnnotatedTypeMirror assignedTo(AnnotatedTypeFactory atypeFactory, TreePath path) {
+        if (visitedPaths.contains(path)) {
+            // inform the caller to skip assignment context resolution
+            return null;
+        }
+
+        visitedPaths.add(path);
+
+        Tree assignmentContext = TreeUtils.getAssignmentContext(path);
+        AnnotatedTypeMirror res;
+        if (assignmentContext == null) {
+            res = null;
+        } else if (assignmentContext instanceof AssignmentTree) {
+            ExpressionTree variable = ((AssignmentTree) assignmentContext).getVariable();
+            res = atypeFactory.getAnnotatedType(variable);
+        } else if (assignmentContext instanceof CompoundAssignmentTree) {
+            ExpressionTree variable = ((CompoundAssignmentTree) assignmentContext).getVariable();
+            res = atypeFactory.getAnnotatedType(variable);
+        } else if (assignmentContext instanceof MethodInvocationTree) {
+            MethodInvocationTree methodInvocation = (MethodInvocationTree) assignmentContext;
+            ExecutableElement methodElt = TreeUtils.elementFromUse(methodInvocation);
+            // TODO move to getAssignmentContext
+            AnnotatedTypeMirror receiver = atypeFactory.getReceiverType(methodInvocation);
+            res =
+                    TypeArgInferenceUtil.assignedToExecutable(
+                            atypeFactory,
+                            path,
+                            methodElt,
+                            receiver,
+                            methodInvocation.getArguments());
+        } else if (assignmentContext instanceof NewArrayTree) {
+            // TODO: I left the previous implementation below, it definitely caused infinite loops
+            // TODO: if you called it from places like the TreeAnnotator.
+            res = null;
+
+            // TODO: This may cause infinite loop
+            //            AnnotatedTypeMirror type =
+            //                    atypeFactory.getAnnotatedType((NewArrayTree)assignmentContext);
+            //            type = AnnotatedTypes.innerMostType(type);
+            //            return type;
+
+        } else if (assignmentContext instanceof NewClassTree) {
+            // This need to be basically like MethodTree
+            NewClassTree newClassTree = (NewClassTree) assignmentContext;
+            ExecutableElement constructorElt = TreeUtils.constructor(newClassTree);
+            AnnotatedTypeMirror receiver = atypeFactory.fromNewClass(newClassTree);
+            res =
+                    TypeArgInferenceUtil.assignedToExecutable(
+                            atypeFactory,
+                            path,
+                            constructorElt,
+                            receiver,
+                            newClassTree.getArguments());
+        } else if (assignmentContext instanceof ReturnTree) {
+            HashSet<Kind> kinds = new HashSet<>(Arrays.asList(Kind.LAMBDA_EXPRESSION, Kind.METHOD));
+            Tree enclosing = TreeUtils.enclosingOfKind(path, kinds);
+
+            if (enclosing.getKind() == Kind.METHOD) {
+                res = atypeFactory.getAnnotatedType((MethodTree) enclosing).getReturnType();
+            } else {
+                Pair<AnnotatedDeclaredType, AnnotatedExecutableType> fninf =
+                        atypeFactory.getFnInterfaceFromTree((LambdaExpressionTree) enclosing);
+                res = fninf.second.getReturnType();
+            }
+
+        } else if (assignmentContext instanceof VariableTree) {
+            res = TypeArgInferenceUtil.assignedToVariable(atypeFactory, assignmentContext);
+        } else {
+            throw new BugInCF("AnnotatedTypes.assignedTo: shouldn't be here");
+        }
+        visitedPaths.remove(path);
+        return res;
+    }
+}

framework/src/main/java/org/checkerframework/framework/util/typeinference/TypeArgInferenceUtil.java

@@ -203,7 +203,7 @@ public static AnnotatedTypeMirror assignedTo(AnnotatedTypeFactory atypeFactory,
         }
     }
 
-    private static AnnotatedTypeMirror assignedToExecutable(
+    public static AnnotatedTypeMirror assignedToExecutable(
             AnnotatedTypeFactory atypeFactory,
             TreePath path,
             ExecutableElement methodElt,

javacutil/src/main/java/org/checkerframework/javacutil/TreeUtils.java

@@ -344,7 +344,7 @@ public static ExpressionTree skipParens(final ExpressionTree tree) {
     }
 
     /**
-     * Returns the tree with the assignment context for the treePath leaf node. (Does not handle
+     * Returns the tree with the assignment context for the treePath leaf node. (Handles
      * pseudo-assignment of an argument to a parameter or a receiver expression to a receiver.)
      *
      * <p>The assignment context for the {@code treePath} is the leaf of its parent, if the parent
@@ -398,6 +398,18 @@ public static Tree getAssignmentContext(final TreePath treePath) {
             case RETURN:
             case VARIABLE:
                 return parent;
+            case MEMBER_SELECT:
+                // Also check case when treepath's leaf tree is used as methohd
+                // invocation's actual receiver
+                // If so, return that method invocation tree too as the assignment
+                // context tree rather than null as we did before
+                TreePath grandParentPath = parentPath.getParentPath();
+                if (grandParentPath != null
+                        && grandParentPath.getLeaf() instanceof MethodInvocationTree) {
+                    return grandParentPath.getLeaf();
+                } else {
+                    return null;
+                }
             default:
                 // 11 Tree.Kinds are CompoundAssignmentTrees,
                 // so use instanceof rather than listing all 11.