Do not remove annotations that are members of annotations referenced …

…from code

Reviewed By: thezhangwei, ssj933

Differential Revision: D66399299

fbshipit-source-id: 020a3c88a73fe48dbe9d32101799c57bf986fda2

opt/annokill/AnnoKill.cpp

@@ -128,6 +128,43 @@ AnnoKill::AnnoKill(
   }
 }
 
+namespace {
+void gather_complete_referenced_annos(
+    const AnnoKill::AnnoSet& initial_referenced_annos,
+    DexType* type,
+    AnnoKill::AnnoSet* result) {
+  auto cls = type_class(type);
+  if (cls == nullptr || !is_annotation(cls) || cls->is_external()) {
+    return;
+  }
+  if (result->count(type) > 0) {
+    return;
+  }
+  result->emplace(type);
+  if (initial_referenced_annos.count(type) == 0) {
+    TRACE(ANNO, 3,
+          "Annotation type %s referenced indirectly from a referenced "
+          "annotation; keeping.",
+          SHOW(type));
+  }
+  auto process = [&](DexType* t) {
+    auto effective_type =
+        const_cast<DexType*>(type::get_element_type_if_array(t));
+    gather_complete_referenced_annos(initial_referenced_annos, effective_type,
+                                     result);
+  };
+  // gather_types too broad here (that will keep too much). Just check return
+  // type of members of the annotation (arguments should be disallowed) and
+  // static field types.
+  for (auto m : cls->get_all_methods()) {
+    process(m->get_proto()->get_rtype());
+  }
+  for (auto f : cls->get_sfields()) {
+    process(f->get_type());
+  }
+}
+} // namespace
+
 AnnoKill::AnnoSet AnnoKill::get_referenced_annos() {
   Timer timer{"get_referenced_annos"};
 
@@ -330,6 +367,13 @@ AnnoKill::AnnoSet AnnoKill::get_referenced_annos() {
       });
   referenced_annos.insert(concurrent_referenced_annos.begin(),
                           concurrent_referenced_annos.end());
+  // For each referenced annotation, make sure any annotations it references are
+  // also tracked as referenced, so we don't end up with a dangling ref.
+  AnnoKill::AnnoSet gathered;
+  for (auto referenced : referenced_annos) {
+    gather_complete_referenced_annos(referenced_annos, referenced, &gathered);
+  }
+  referenced_annos.insert(gathered.begin(), gathered.end());
   return referenced_annos;
 }
 

test/instr/AnnoKillTest.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+package com.redex;
+
+@interface Funny {
+
+}
+
+@interface VeryFunny {
+
+}
+
+@interface Zero {
+  // This is a field on an annotation. Why someone would do this, I don't know.
+  Funny f = null;
+  VeryFunny[] a = null;
+}
+
+@interface One {
+  int x() default 10;
+  Zero zero();
+}
+
+@interface Two {
+  int y();
+  One[] one() default {};
+}
+
+@interface Unused {
+  int z();
+}
+
+@Two(y = 1)
+class Use {}
+
+class Main {
+    public static void go() {
+        System.out.println(Two.class);
+    }
+}

test/instr/AnnoKillTestVerify.cpp

@@ -0,0 +1,29 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <gmock/gmock.h>
+
+#include "Show.h"
+#include "verify/VerifyUtil.h"
+
+using namespace testing;
+
+TEST_F(PostVerify, VerifyKeptAndRemoved) {
+  EXPECT_EQ(find_class_named(classes, "Lcom/redex/Unused;"), nullptr)
+      << "Should remove type Unused";
+  EXPECT_NE(find_class_named(classes, "Lcom/redex/Two;"), nullptr)
+      << "Should not remove Two! It has a code reference.";
+  EXPECT_NE(find_class_named(classes, "Lcom/redex/One;"), nullptr)
+      << "Should not remove One! Otherwise we'll have a torn enum, and that's "
+         "bad :p";
+  EXPECT_NE(find_class_named(classes, "Lcom/redex/Zero;"), nullptr)
+      << "Should not remove Zero!";
+  EXPECT_NE(find_class_named(classes, "Lcom/redex/Funny;"), nullptr)
+      << "Should not remove Funny! It is referenced from a static field.";
+  EXPECT_NE(find_class_named(classes, "Lcom/redex/VeryFunny;"), nullptr)
+      << "Should not remove VeryFunny! It is referenced from a static field.";
+}

test/instr/anno.config

@@ -0,0 +1,14 @@
+{
+    "AnnoKillPass": {
+        "keep_annos": [
+        ],
+        "kill_annos": [
+        ],
+        "kill_bad_signatures": true
+    },
+    "redex" : {
+        "passes" : [
+            "AnnoKillPass"
+        ]
+    }
+}