Ensure correct pruning of type annotations

Bug: b/271543766
Change-Id: Ibb200fa2456504cd5cbeab2f369143173f92448c
diff --git a/src/main/java/com/android/tools/r8/shaking/AnnotationRemover.java b/src/main/java/com/android/tools/r8/shaking/AnnotationRemover.java
index f6959c6..9321484 100644
--- a/src/main/java/com/android/tools/r8/shaking/AnnotationRemover.java
+++ b/src/main/java/com/android/tools/r8/shaking/AnnotationRemover.java
@@ -129,10 +129,12 @@
           if (!options.isKeepRuntimeVisibleParameterAnnotationsEnabled()) {
             return false;
           }
-        } else {
-          if (!options.isKeepRuntimeVisibleAnnotationsEnabled()) {
-            return false;
-          }
+        } else if (!annotation.isTypeAnnotation()
+            && !options.isKeepRuntimeVisibleAnnotationsEnabled()) {
+          return false;
+        } else if (annotation.isTypeAnnotation()
+            && !options.isKeepRuntimeVisibleTypeAnnotationsEnabled()) {
+          return false;
         }
         return isAnnotationTypeLive;
 
@@ -147,13 +149,14 @@
           if (!options.isKeepRuntimeInvisibleParameterAnnotationsEnabled()) {
             return false;
           }
-        } else {
-          if (!options.isKeepRuntimeInvisibleAnnotationsEnabled()) {
-            return false;
-          }
+        } else if (!annotation.isTypeAnnotation()
+            && !options.isKeepRuntimeInvisibleAnnotationsEnabled()) {
+          return false;
+        } else if (annotation.isTypeAnnotation()
+            && !options.isKeepRuntimeInvisibleTypeAnnotationsEnabled()) {
+          return false;
         }
         return isAnnotationTypeLive;
-
       default:
         throw new Unreachable("Unexpected annotation visibility.");
     }
diff --git a/src/main/java/com/android/tools/r8/shaking/GlobalKeepInfoConfiguration.java b/src/main/java/com/android/tools/r8/shaking/GlobalKeepInfoConfiguration.java
index d10b35c..7ff0803 100644
--- a/src/main/java/com/android/tools/r8/shaking/GlobalKeepInfoConfiguration.java
+++ b/src/main/java/com/android/tools/r8/shaking/GlobalKeepInfoConfiguration.java
@@ -35,4 +35,8 @@
   boolean isKeepRuntimeVisibleAnnotationsEnabled();
 
   boolean isKeepRuntimeVisibleParameterAnnotationsEnabled();
+
+  boolean isKeepRuntimeVisibleTypeAnnotationsEnabled();
+
+  boolean isKeepRuntimeInvisibleTypeAnnotationsEnabled();
 }
diff --git a/src/main/java/com/android/tools/r8/utils/InternalOptions.java b/src/main/java/com/android/tools/r8/utils/InternalOptions.java
index 0d7a71c..8275b55 100644
--- a/src/main/java/com/android/tools/r8/utils/InternalOptions.java
+++ b/src/main/java/com/android/tools/r8/utils/InternalOptions.java
@@ -785,6 +785,16 @@
     return proguardConfiguration.getKeepAttributes().runtimeVisibleParameterAnnotations;
   }
 
+  @Override
+  public boolean isKeepRuntimeVisibleTypeAnnotationsEnabled() {
+    return proguardConfiguration.getKeepAttributes().runtimeVisibleTypeAnnotations;
+  }
+
+  @Override
+  public boolean isKeepRuntimeInvisibleTypeAnnotationsEnabled() {
+    return proguardConfiguration.getKeepAttributes().runtimeInvisibleTypeAnnotations;
+  }
+
   /**
    * If any non-static class merging is enabled, information about types referred to by instanceOf
    * and check cast instructions needs to be collected.
diff --git a/src/test/java/com/android/tools/r8/annotations/TypeUseAnnotationPruneTest.java b/src/test/java/com/android/tools/r8/annotations/TypeUseAnnotationPruneTest.java
index e9fcba0..3a72818 100644
--- a/src/test/java/com/android/tools/r8/annotations/TypeUseAnnotationPruneTest.java
+++ b/src/test/java/com/android/tools/r8/annotations/TypeUseAnnotationPruneTest.java
@@ -44,18 +44,10 @@
     return StringUtils.joinLines(
         "printAnnotation - Class: " + notNullTestRuntimeTypeName,
         "printAnnotation - Class: NULL",
-        "printAnnotation - Extends(0): " + notNullTestRuntimeTypeName,
-        "printAnnotation - Implements(0): " + notNullTestRuntimeTypeName,
         "printAnnotation - Field: NULL",
         "printAnnotation - Field: NULL",
-        "printAnnotation - Field(0): " + notNullTestRuntimeTypeName,
         "printAnnotation - Method: NULL",
         "printAnnotation - Method: NULL",
-        "printAnnotation - MethodReturnType(0): " + notNullTestRuntimeTypeName,
-        "printAnnotation - MethodParameter at 0(0): " + notNullTestRuntimeTypeName,
-        "printAnnotation - MethodParameter at 1(0): " + notNullTestRuntimeTypeName,
-        "printAnnotation - MethodException at 0(0): " + notNullTestRuntimeTypeName,
-        "printAnnotation - MethodException at 1(0): " + notNullTestRuntimeTypeName,
         "Hello World!");
   }
 
@@ -88,30 +80,15 @@
               assertThat(notNullTestRuntime, isPresent());
               ClassSubject clazz = inspector.clazz(TestClassWithTypeAndGenericAnnotations.class);
               assertThat(clazz, isPresent());
-              if (parameters.isDexRuntime()) {
-                assertTrue(
-                    clazz.annotations().stream()
-                        .noneMatch(FoundAnnotationSubject::isTypeAnnotation));
-              } else {
-                // TODO(b/271543766): Assert that we have no type annotations
-              }
+              assertTrue(
+                  clazz.annotations().stream().noneMatch(FoundAnnotationSubject::isTypeAnnotation));
               FieldSubject field = clazz.uniqueFieldWithOriginalName("field");
               assertThat(field, isPresent());
-              if (parameters.isDexRuntime()) {
-                assertEquals(0, field.annotations().size());
-              } else {
-                // TODO(b/271543766): Assert that we have no annotations
-                assertEquals(4, field.annotations().size());
-              }
+              assertEquals(0, field.annotations().size());
               MethodSubject method = clazz.uniqueMethodWithOriginalName("method");
               assertThat(method, isPresent());
               // We create a dex annotation for the checked exception.
-              if (parameters.isDexRuntime()) {
-                assertEquals(1, method.annotations().size());
-              } else {
-                // TODO(b/271543766): Assert that we have no annotations
-                assertEquals(17, method.annotations().size());
-              }
+              assertEquals(1, method.annotations().size());
             },
             options -> options.programConsumer = ClassFileConsumer.emptyConsumer())
         .run(parameters.getRuntime(), MainWithTypeAndGeneric.class)