[llvm] [ProfileData] Sink the length checks (PR #95604)

[Kazu Hirata via llvm-commits]

https://github.com/kazutakahirata created https://github.com/llvm/llvm-project/pull/95604

The new API getValueArrayForSite returns ArrayRef<InstrProfValueData>,
packaging the array length and contents together.

This patch sinks the array length checks just before we check the
contents.  This way, we check both the array length and contents
immediately after calling getValueArrayForSite.

>From c7a3ea1364168f015c61ffa02e1dcee98b8f2bcb Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu at google.com>
Date: Fri, 14 Jun 2024 14:03:21 -0700
Subject: [PATCH] [ProfileData] Sink the length checks

The new API getValueArrayForSite returns ArrayRef<InstrProfValueData>,
packaging the array length and contents together.

This patch sinks the array length checks just before we check the
contents.  This way, we check both the array length and contents
immediately after calling getValueArrayForSite.
---
 llvm/unittests/ProfileData/InstrProfTest.cpp | 30 +++++++++++---------
 1 file changed, 16 insertions(+), 14 deletions(-)

diff --git a/llvm/unittests/ProfileData/InstrProfTest.cpp b/llvm/unittests/ProfileData/InstrProfTest.cpp
index 749d35e02c286..9e3dc6a2ad13e 100644
--- a/llvm/unittests/ProfileData/InstrProfTest.cpp
+++ b/llvm/unittests/ProfileData/InstrProfTest.cpp
@@ -854,20 +854,15 @@ TEST_P(InstrProfReaderWriterTest, icall_and_vtable_data_read_write) {
   // Test the number of instrumented indirect call sites and the number of
   // profiled values at each site.
   ASSERT_EQ(4U, R->getNumValueSites(IPVK_IndirectCallTarget));
-  EXPECT_EQ(3U, R->getNumValueDataForSite(IPVK_IndirectCallTarget, 0));
-  EXPECT_EQ(0U, R->getNumValueDataForSite(IPVK_IndirectCallTarget, 1));
-  EXPECT_EQ(2U, R->getNumValueDataForSite(IPVK_IndirectCallTarget, 2));
-  EXPECT_EQ(2U, R->getNumValueDataForSite(IPVK_IndirectCallTarget, 3));
 
   // Test the number of instrumented vtable sites and the number of profiled
   // values at each site.
   ASSERT_EQ(R->getNumValueSites(IPVK_VTableTarget), 2U);
-  EXPECT_EQ(R->getNumValueDataForSite(IPVK_VTableTarget, 0), 3U);
-  EXPECT_EQ(R->getNumValueDataForSite(IPVK_VTableTarget, 1), 2U);
 
   // First indirect site.
   {
     auto VD = R->getValueArrayForSite(IPVK_IndirectCallTarget, 0);
+    ASSERT_THAT(VD, SizeIs(3));
 
     EXPECT_EQ(VD[0].Count, 3U * getProfWeight());
     EXPECT_EQ(VD[1].Count, 2U * getProfWeight());
@@ -878,9 +873,14 @@ TEST_P(InstrProfReaderWriterTest, icall_and_vtable_data_read_write) {
     EXPECT_STREQ((const char *)VD[2].Value, "callee1");
   }
 
+  EXPECT_THAT(R->getValueArrayForSite(IPVK_IndirectCallTarget, 1), SizeIs(0));
+  EXPECT_THAT(R->getValueArrayForSite(IPVK_IndirectCallTarget, 2), SizeIs(2));
+  EXPECT_THAT(R->getValueArrayForSite(IPVK_IndirectCallTarget, 3), SizeIs(2));
+
   // First vtable site.
   {
     auto VD = R->getValueArrayForSite(IPVK_VTableTarget, 0);
+    ASSERT_THAT(VD, SizeIs(3));
 
     EXPECT_EQ(VD[0].Count, 3U * getProfWeight());
     EXPECT_EQ(VD[1].Count, 2U * getProfWeight());
@@ -894,6 +894,7 @@ TEST_P(InstrProfReaderWriterTest, icall_and_vtable_data_read_write) {
   // Second vtable site.
   {
     auto VD = R->getValueArrayForSite(IPVK_VTableTarget, 1);
+    ASSERT_THAT(VD, SizeIs(2));
 
     EXPECT_EQ(VD[0].Count, 2U * getProfWeight());
     EXPECT_EQ(VD[1].Count, 1U * getProfWeight());
@@ -1112,20 +1113,13 @@ TEST_P(MaybeSparseInstrProfTest, icall_and_vtable_data_merge) {
   EXPECT_THAT_ERROR(R.takeError(), Succeeded());
   // For indirect calls.
   ASSERT_EQ(5U, R->getNumValueSites(IPVK_IndirectCallTarget));
-  ASSERT_EQ(4U, R->getNumValueDataForSite(IPVK_IndirectCallTarget, 0));
-  ASSERT_EQ(0U, R->getNumValueDataForSite(IPVK_IndirectCallTarget, 1));
-  ASSERT_EQ(4U, R->getNumValueDataForSite(IPVK_IndirectCallTarget, 2));
-  ASSERT_EQ(2U, R->getNumValueDataForSite(IPVK_IndirectCallTarget, 3));
-  ASSERT_EQ(3U, R->getNumValueDataForSite(IPVK_IndirectCallTarget, 4));
   // For vtables.
   ASSERT_EQ(R->getNumValueSites(IPVK_VTableTarget), 3U);
-  ASSERT_EQ(R->getNumValueDataForSite(IPVK_VTableTarget, 0), 4U);
-  ASSERT_EQ(R->getNumValueDataForSite(IPVK_VTableTarget, 1), 4U);
-  ASSERT_EQ(R->getNumValueDataForSite(IPVK_VTableTarget, 2), 3U);
 
   // Test the merged values for indirect calls.
   {
     auto VD = R->getValueArrayForSite(IPVK_IndirectCallTarget, 0);
+    ASSERT_THAT(VD, SizeIs(4));
     EXPECT_STREQ((const char *)VD[0].Value, "callee2");
     EXPECT_EQ(VD[0].Count, 7U);
     EXPECT_STREQ((const char *)VD[1].Value, "callee3");
@@ -1135,7 +1129,10 @@ TEST_P(MaybeSparseInstrProfTest, icall_and_vtable_data_merge) {
     EXPECT_STREQ((const char *)VD[3].Value, "callee1");
     EXPECT_EQ(VD[3].Count, 1U);
 
+    ASSERT_THAT(R->getValueArrayForSite(IPVK_IndirectCallTarget, 1), SizeIs(0));
+
     auto VD_2 = R->getValueArrayForSite(IPVK_IndirectCallTarget, 2);
+    ASSERT_THAT(VD_2, SizeIs(4));
     EXPECT_STREQ((const char *)VD_2[0].Value, "callee3");
     EXPECT_EQ(VD_2[0].Count, 6U);
     EXPECT_STREQ((const char *)VD_2[1].Value, "callee4");
@@ -1146,12 +1143,14 @@ TEST_P(MaybeSparseInstrProfTest, icall_and_vtable_data_merge) {
     EXPECT_EQ(VD_2[3].Count, 1U);
 
     auto VD_3 = R->getValueArrayForSite(IPVK_IndirectCallTarget, 3);
+    ASSERT_THAT(VD_3, SizeIs(2));
     EXPECT_STREQ((const char *)VD_3[0].Value, "callee8");
     EXPECT_EQ(VD_3[0].Count, 2U);
     EXPECT_STREQ((const char *)VD_3[1].Value, "callee7");
     EXPECT_EQ(VD_3[1].Count, 1U);
 
     auto VD_4 = R->getValueArrayForSite(IPVK_IndirectCallTarget, 4);
+    ASSERT_THAT(VD_4, SizeIs(3));
     EXPECT_STREQ((const char *)VD_4[0].Value, "callee3");
     EXPECT_EQ(VD_4[0].Count, 6U);
     EXPECT_STREQ((const char *)VD_4[1].Value, "callee2");
@@ -1163,6 +1162,7 @@ TEST_P(MaybeSparseInstrProfTest, icall_and_vtable_data_merge) {
   // Test the merged values for vtables
   {
     auto VD0 = R->getValueArrayForSite(IPVK_VTableTarget, 0);
+    ASSERT_THAT(VD0, SizeIs(4));
     EXPECT_EQ(VD0[0].Value, getCalleeAddress(vtable2));
     EXPECT_EQ(VD0[0].Count, 7U);
     EXPECT_EQ(VD0[1].Value, getCalleeAddress(vtable3));
@@ -1173,6 +1173,7 @@ TEST_P(MaybeSparseInstrProfTest, icall_and_vtable_data_merge) {
     EXPECT_EQ(VD0[3].Count, 1U);
 
     auto VD1 = R->getValueArrayForSite(IPVK_VTableTarget, 1);
+    ASSERT_THAT(VD1, SizeIs(4));
     EXPECT_EQ(VD1[0].Value, getCalleeAddress(vtable3));
     EXPECT_EQ(VD1[0].Count, 6U);
     EXPECT_EQ(VD1[1].Value, getCalleeAddress(vtable4));
@@ -1183,6 +1184,7 @@ TEST_P(MaybeSparseInstrProfTest, icall_and_vtable_data_merge) {
     EXPECT_EQ(VD1[3].Count, 1U);
 
     auto VD2 = R->getValueArrayForSite(IPVK_VTableTarget, 2);
+    ASSERT_THAT(VD2, SizeIs(3));
     EXPECT_EQ(VD2[0].Value, getCalleeAddress(vtable3));
     EXPECT_EQ(VD2[0].Count, 6U);
     EXPECT_EQ(VD2[1].Value, getCalleeAddress(vtable2));