8337795: Type annotation attached to incorrect type during class reading

Reviewed-by: vromero

src/jdk.compiler/share/classes/com/sun/tools/javac/jvm/ClassReader.java

@@ -36,9 +36,9 @@
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
-import java.util.function.BiFunction;
 import java.util.function.IntFunction;
 import java.util.function.Predicate;
+import java.util.stream.IntStream;
 
 import javax.lang.model.element.Modifier;
 import javax.lang.model.element.NestingKind;
@@ -2461,13 +2461,8 @@ private Type addTypeAnnotations(Type type, Predicate<TypeAnnotationPosition> fil
                         .add(attribute);
             }
 
-            // Search the structure of the type to find the contained types at each type path
-            Map<Type, List<Attribute.TypeCompound>> attributesByType = new HashMap<>();
-            new TypeAnnotationLocator(attributesByPath, attributesByType).visit(type, List.nil());
-
             // Rewrite the type and add the annotations
-            type = new TypeAnnotationTypeMapping(attributesByType).visit(type, null);
-            Assert.check(attributesByType.isEmpty(), "Failed to apply annotations to types");
+            type = new TypeAnnotationStructuralTypeMapping(attributesByPath).visit(type, List.nil());
 
             return type;
         }
@@ -2495,120 +2490,102 @@ private static Predicate<TypeAnnotationPosition> classExtends(int index) {
     }
 
     /**
-     * Visit all contained types, assembling a type path to represent the current location, and
-     * record the types at each type path that need to be annotated.
+     * A type mapping that rewrites the type to include type annotations.
+     *
+     * <p>This logic is similar to {@link Type.StructuralTypeMapping}, but also tracks the path to
+     * the contained types being rewritten, and so cannot easily share the existing logic.
      */
-    private static class TypeAnnotationLocator
-            extends Types.DefaultTypeVisitor<Void, List<TypeAnnotationPosition.TypePathEntry>> {
+    private static final class TypeAnnotationStructuralTypeMapping
+            extends Types.TypeMapping<List<TypeAnnotationPosition.TypePathEntry>> {
+
         private final Map<List<TypeAnnotationPosition.TypePathEntry>,
-                          ListBuffer<Attribute.TypeCompound>> attributesByPath;
-        private final Map<Type, List<Attribute.TypeCompound>> attributesByType;
+                ListBuffer<Attribute.TypeCompound>> attributesByPath;
 
-        private TypeAnnotationLocator(
+        private TypeAnnotationStructuralTypeMapping(
                 Map<List<TypeAnnotationPosition.TypePathEntry>, ListBuffer<Attribute.TypeCompound>>
-                        attributesByPath,
-                Map<Type, List<Attribute.TypeCompound>> attributesByType) {
+                    attributesByPath) {
             this.attributesByPath = attributesByPath;
-            this.attributesByType = attributesByType;
         }
 
+
         @Override
-        public Void visitClassType(ClassType t, List<TypeAnnotationPosition.TypePathEntry> path) {
+        public Type visitClassType(ClassType t, List<TypeAnnotationPosition.TypePathEntry> path) {
             // As described in JVMS 4.7.20.2, type annotations on nested types are located with
             // 'left-to-right' steps starting on 'the outermost part of the type for which a type
             // annotation is admissible'. So the current path represents the outermost containing
             // type of the type being visited, and we add type path steps for every contained nested
             // type.
-            List<ClassType> enclosing = List.nil();
-            for (Type curr = t;
-                    curr != null && curr != Type.noType;
+            Type outer = t.getEnclosingType();
+            Type outer1 = outer != Type.noType ? visit(outer, path) : outer;
+            for (Type curr = t.getEnclosingType();
+                    curr != Type.noType;
                     curr = curr.getEnclosingType()) {
-                enclosing = enclosing.prepend((ClassType) curr);
-            }
-            for (ClassType te : enclosing) {
-                if (te.typarams_field != null) {
-                    int i = 0;
-                    for (Type typaram : te.typarams_field) {
-                        visit(typaram, path.append(new TypeAnnotationPosition.TypePathEntry(
-                                TypeAnnotationPosition.TypePathEntryKind.TYPE_ARGUMENT, i++)));
-                    }
-                }
-                visitType(te, path);
                 path = path.append(TypeAnnotationPosition.TypePathEntry.INNER_TYPE);
             }
-            return null;
+            List<Type> typarams = t.getTypeArguments();
+            List<Type> typarams1 = rewriteTypeParams(path, typarams);
+            if (outer1 != outer || typarams != typarams1) {
+                t = new ClassType(outer1, typarams1, t.tsym, t.getMetadata());
+            }
+            return reannotate(t, path);
         }
 
-        @Override
-        public Void visitWildcardType(
-                WildcardType t, List<TypeAnnotationPosition.TypePathEntry> path) {
-            visit(t.type, path.append(TypeAnnotationPosition.TypePathEntry.WILDCARD));
-            return super.visitWildcardType(t, path);
+        private List<Type> rewriteTypeParams(
+                List<TypeAnnotationPosition.TypePathEntry> path, List<Type> typarams) {
+            var i = IntStream.iterate(0, x -> x + 1).iterator();
+            return typarams.map(typaram -> visit(typaram,
+                    path.append(new TypeAnnotationPosition.TypePathEntry(
+                            TypeAnnotationPosition.TypePathEntryKind.TYPE_ARGUMENT, i.nextInt()))));
         }
 
         @Override
-        public Void visitArrayType(ArrayType t, List<TypeAnnotationPosition.TypePathEntry> path) {
-            visit(t.elemtype, path.append(TypeAnnotationPosition.TypePathEntry.ARRAY));
-            return super.visitArrayType(t, path);
+        public Type visitWildcardType(
+                WildcardType wt, List<TypeAnnotationPosition.TypePathEntry> path) {
+            Type t = wt.type;
+            if (t != null) {
+                t = visit(t, path.append(TypeAnnotationPosition.TypePathEntry.WILDCARD));
+            }
+            if (t != wt.type) {
+                wt = new WildcardType(t, wt.kind, wt.tsym, wt.bound, wt.getMetadata());
+            }
+            return reannotate(wt, path);
         }
 
         @Override
-        public Void visitType(Type t, List<TypeAnnotationPosition.TypePathEntry> path) {
-            ListBuffer<Attribute.TypeCompound> attributes = attributesByPath.remove(path);
-            if (attributes != null) {
-                attributesByType.put(t, attributes.toList());
+        public Type visitArrayType(ArrayType t, List<TypeAnnotationPosition.TypePathEntry> path) {
+            Type elemtype = t.elemtype;
+            Type elemtype1 =
+                    visit(elemtype, path.append(TypeAnnotationPosition.TypePathEntry.ARRAY));
+            if (elemtype1 != elemtype)  {
+                t = new ArrayType(elemtype1, t.tsym, t.getMetadata());
             }
-            return null;
+            return reannotate(t, path);
         }
-    }
-
-    /** A type mapping that rewrites the type to include type annotations. */
-    private static class TypeAnnotationTypeMapping extends Type.StructuralTypeMapping<Void> {
 
-        private final Map<Type, List<Attribute.TypeCompound>> attributesByType;
-
-        private TypeAnnotationTypeMapping(
-                Map<Type, List<Attribute.TypeCompound>> attributesByType) {
-            this.attributesByType = attributesByType;
+        @Override
+        public Type visitType(Type t, List<TypeAnnotationPosition.TypePathEntry> path) {
+            return reannotate(t, path);
         }
 
-        private <T extends Type> Type reannotate(T t, BiFunction<T, Void, Type> f) {
-            // We're relying on object identify of Type instances to record where the annotations
-            // need to be added, so we have to retrieve the annotations for each type before
-            // rewriting it, and then add them after its contained types have been rewritten.
-            List<Attribute.TypeCompound> attributes = attributesByType.remove(t);
-            Type mapped = f.apply(t, null);
-            if (attributes == null) {
-                return mapped;
+        Type reannotate(Type type, List<TypeAnnotationPosition.TypePathEntry> path) {
+            List<Attribute.TypeCompound> attributes = attributesForPath(path);
+            if (attributes.isEmpty()) {
+                return type;
             }
             // Runtime-visible and -invisible annotations are completed separately, so if the same
             // type has annotations from both it will get annotated twice.
-            TypeMetadata.Annotations existing = mapped.getMetadata(TypeMetadata.Annotations.class);
+            TypeMetadata.Annotations existing = type.getMetadata(TypeMetadata.Annotations.class);
             if (existing != null) {
                 existing.annotationBuffer().addAll(attributes);
-                return mapped;
+                return type;
             }
-            return mapped.annotatedType(attributes);
+            return type.annotatedType(attributes);
         }
 
-        @Override
-        public Type visitClassType(ClassType t, Void unused) {
-            return reannotate(t, super::visitClassType);
-        }
-
-        @Override
-        public Type visitWildcardType(WildcardType t, Void unused) {
-            return reannotate(t, super::visitWildcardType);
-        }
-
-        @Override
-        public Type visitArrayType(ArrayType t, Void unused) {
-            return reannotate(t, super::visitArrayType);
-        }
-
-        @Override
-        public Type visitType(Type t, Void unused) {
-            return reannotate(t, (x, u) -> x);
+        List<Attribute.TypeCompound> attributesForPath(
+                List<TypeAnnotationPosition.TypePathEntry> path) {
+            ListBuffer<Attribute.TypeCompound> attributes = attributesByPath.remove(path);
+            return attributes != null ? attributes.toList() : List.nil();
         }
     }
 

test/langtools/tools/javac/processing/model/type/BasicAnnoTests.java

@@ -717,4 +717,8 @@ private class GenericNested<Y> {
             GenericNested(@TA(120) GenericInner120<X> GenericInner120.this) {}
         }
     }
+
+    @Test(posn=1, annoType=TA.class, expect="130")
+    @Test(posn=23, annoType=TA.class, expect="131")
+    public Map<@TA(130) String, @TA(131) String> f130;
 }