Add group tracking to UniqueMethodTracker

Summary:
Extends `UniqueMethodTracker` to track groups of methods with identical code.
When `insert()` is called, the method is now also recorded in an internal
`ConcurrentMap` under its representative method. A representative method is the first method with this code inserted to this group.

A new `groups()` accessor returns a const reference to this map.

Reviewed By: thezhangwei

Differential Revision: D91266062

fbshipit-source-id: eaea68d40f2741681fddf8f154adbcb609d70ad4

Author: xuhdev

libredex/UniqueMethodTracker.cpp

@@ -19,6 +19,19 @@ std::pair<const DexMethod*, bool> UniqueMethodTracker::insert(
 
 std::pair<const DexMethod*, bool> UniqueMethodTracker::insert(
     const DexMethod* method, size_t code_hash) {
-  const auto [ptr, inserted] = m_unique_methods.insert(Key{code_hash, method});
-  return {ptr->method, inserted};
+  const DexMethod* representative = nullptr;
+  bool was_new = false;
+
+  // The third parameter to the updater is true if the entry already existed,
+  // false if it was newly created.
+  m_groups.update(
+      Key{code_hash, method},
+      [method, &representative, &was_new](
+          const Key& key, UnorderedSet<const DexMethod*>& group, bool existed) {
+        was_new = !existed;
+        representative = existed ? key.method : method;
+        group.insert(method);
+      });
+
+  return {representative, was_new};
 }

libredex/UniqueMethodTracker.h

@@ -11,26 +11,12 @@
 
 #include "ConcurrentContainers.h"
 #include "ControlFlow.h"
+#include "DeterministicContainers.h"
 #include "DexClass.h"
 
-// Tracks unique method code patterns using hash + string equality.
+// Tracks unique method code patterns using hash + CFG equality.
 // Thread-safe for concurrent inserts.
 class UniqueMethodTracker {
- public:
-  // Inserts a method and returns (representative, was_inserted).
-  // - If the code is new, returns (method, true).
-  // - If the code was seen before, returns (first_method_with_same_code,
-  //   false).
-  // - If the method has no code or CFG, returns (nullptr, false).
-  std::pair<const DexMethod*, bool> insert(const DexMethod* method);
-
-  size_t size() const { return m_unique_methods.size(); }
-
- protected:
-  // Insert with a specific hash value. Exposed for testing collision handling.
-  std::pair<const DexMethod*, bool> insert(const DexMethod* method,
-                                           size_t code_hash);
-
  private:
   // Compare two CFGs instruction by instruction.
   static bool cfg_code_equals(const cfg::ControlFlowGraph& a,
@@ -43,6 +29,7 @@ class UniqueMethodTracker {
                       });
   }
 
+ public:
   struct Key {
     size_t code_hash;
     const DexMethod* method;
@@ -64,9 +51,37 @@ class UniqueMethodTracker {
     }
   };
 
+ private:
   struct KeyHash {
     size_t operator()(const Key& key) const { return key.code_hash; }
   };
 
-  InsertOnlyConcurrentSet<Key, KeyHash> m_unique_methods;
+ public:
+  using GroupMap = ConcurrentMap<Key, UnorderedSet<const DexMethod*>, KeyHash>;
+
+  // Inserts a method and returns (representative, was_inserted).
+  // - If the code is new, returns (method, true).
+  // - If the code was seen before, returns (first_method_with_same_code,
+  //   false).
+  // - If the method has no code or CFG, returns (nullptr, false).
+  std::pair<const DexMethod*, bool> insert(const DexMethod* method);
+
+  size_t size() const { return m_groups.size(); }
+
+  // Returns the groups of methods with identical code.
+  // Use UnorderedIterable to iterate over the map.
+  // Each group's key.method is the representative, value is all methods with
+  // that code (including the representative).
+  const GroupMap& groups() const { return m_groups; }
+
+ protected:
+  // Insert with a specific hash value. Exposed for testing collision handling.
+  std::pair<const DexMethod*, bool> insert(const DexMethod* method,
+                                           size_t code_hash);
+
+ private:
+  // Maps code pattern (by Key) to all methods with that code.
+  // The key.method is the representative (first method inserted with this
+  // code).
+  GroupMap m_groups;
 };

test/unit/Makefile.am

@@ -531,6 +531,7 @@ type_system_test_LDADD = $(COMMON_MOCK_TEST_LIBS)
 type_util_test_SOURCES = TypeUtilTest.cpp
 
 unique_method_tracker_test_SOURCES = UniqueMethodTrackerTest.cpp
+unique_method_tracker_test_LDADD = $(COMMON_MOCK_TEST_LIBS)
 
 up_code_motion_test_SOURCES = UpCodeMotionTest.cpp
 up_code_motion_test_LDADD = $(COMMON_MOCK_TEST_LIBS)

test/unit/UniqueMethodTrackerTest.cpp

@@ -9,10 +9,12 @@
 
 #include <gmock/gmock.h>
 
+#include "DexHasher.h"
 #include "IRAssembler.h"
 #include "RedexTest.h"
 
 using ::testing::IsNull;
+using ::testing::NotNull;
 using ::testing::SizeIs;
 
 class UniqueMethodTrackerTest : public RedexTest, public UniqueMethodTracker {
@@ -51,6 +53,22 @@ class UniqueMethodTrackerTest : public RedexTest, public UniqueMethodTracker {
 
     ASSERT_THAT(*this, SizeIs(2u));
   }
+
+  // Helper to find a group by its representative method.
+  const UnorderedSet<const DexMethod*>* find_group(
+      const DexMethod* representative) const {
+    const auto* code = representative->get_code();
+    if (code == nullptr || !code->cfg_built()) {
+      return nullptr;
+    }
+    size_t hash = hashing::DexMethodHasher(representative).run().code_hash;
+    Key key{hash, representative};
+    auto it = this->groups().find(key);
+    if (it == this->groups().end()) {
+      return nullptr;
+    }
+    return &it->second;
+  }
 };
 
 TEST_F(UniqueMethodTrackerTest, UniqueMethodTrackerIdenticalCode) {
@@ -196,3 +214,117 @@ TEST_F(UniqueMethodTrackerTest, UniqueMethodTrackerHashCollision) {
 
   EXPECT_THAT(*this, SizeIs(4u));
 }
+
+TEST_F(UniqueMethodTrackerTest, GroupsTracksDuplicates) {
+  // Verify that groups() correctly tracks methods with identical code.
+  auto* method1 = assembler::method_from_string(R"(
+    (method (public static) "LFoo;.dup1:()I"
+      (
+        (const v0 100)
+        (return v0)
+      )
+    )
+  )");
+
+  auto* method2 = assembler::method_from_string(R"(
+    (method (public static) "LFoo;.dup2:()I"
+      (
+        (const v0 100)
+        (return v0)
+      )
+    )
+  )");
+
+  auto* method3 = assembler::method_from_string(R"(
+    (method (public static) "LFoo;.dup3:()I"
+      (
+        (const v0 100)
+        (return v0)
+      )
+    )
+  )");
+
+  method1->get_code()->build_cfg();
+  method2->get_code()->build_cfg();
+  method3->get_code()->build_cfg();
+
+  this->insert(method1);
+  this->insert(method2);
+  this->insert(method3);
+
+  // Find the group for method1 (the representative).
+  const auto* group = find_group(method1);
+  ASSERT_THAT(group, NotNull()) << "Expected group for method1 to exist";
+
+  EXPECT_THAT(*group, SizeIs(3u)) << "Expected all 3 methods in the same group";
+  EXPECT_NE(group->find(method1), group->end());
+  EXPECT_NE(group->find(method2), group->end());
+  EXPECT_NE(group->find(method3), group->end());
+}
+
+TEST_F(UniqueMethodTrackerTest, GroupsDistinctForDifferentCode) {
+  // Verify that methods with different code are in separate groups.
+  auto* method1 = assembler::method_from_string(R"(
+    (method (public static) "LFoo;.unique1:()I"
+      (
+        (const v0 200)
+        (return v0)
+      )
+    )
+  )");
+
+  auto* method2 = assembler::method_from_string(R"(
+    (method (public static) "LFoo;.unique2:()I"
+      (
+        (const v0 201)
+        (return v0)
+      )
+    )
+  )");
+
+  method1->get_code()->build_cfg();
+  method2->get_code()->build_cfg();
+
+  this->insert(method1);
+  this->insert(method2);
+
+  // Each method should be in its own group.
+  const auto* group1 = find_group(method1);
+  ASSERT_THAT(group1, NotNull());
+  EXPECT_THAT(*group1, SizeIs(1u));
+  EXPECT_NE(group1->find(method1), group1->end());
+
+  const auto* group2 = find_group(method2);
+  ASSERT_THAT(group2, NotNull());
+  EXPECT_THAT(*group2, SizeIs(1u));
+  EXPECT_NE(group2->find(method2), group2->end());
+}
+
+TEST_F(UniqueMethodTrackerTest, DuplicateInsertionIgnored) {
+  // Verify that inserting the same method twice does not add duplicates.
+  auto* method = assembler::method_from_string(R"(
+    (method (public static) "LFoo;.duplicate:()I"
+      (
+        (const v0 999)
+        (return v0)
+      )
+    )
+  )");
+
+  method->get_code()->build_cfg();
+
+  const auto [rep1, inserted1] = this->insert(method);
+  ASSERT_TRUE(inserted1);
+  ASSERT_EQ(rep1, method);
+
+  // Insert same method again.
+  const auto [rep2, inserted2] = this->insert(method);
+  EXPECT_FALSE(inserted2) << "Re-inserting same method should return false";
+  EXPECT_EQ(rep2, method) << "Representative should still be the same method";
+
+  // The group should contain exactly one entry.
+  const auto* group = find_group(method);
+  ASSERT_THAT(group, NotNull());
+  EXPECT_THAT(*group, SizeIs(1u))
+      << "Group should have exactly 1 method, not duplicates";
+}