[C2y] Fix _Countof handling of VLAs (#141621)

It turns out that getVLASize() does not get you the size of a single
dimension of the VLA, it gets you the full count of all elements. This
caused _Countof to return invalid values on VLA ranks. Now switched to
using getVLAElements1D() instead, which only gets a single dimension.

Fixes #141409

Author: AaronBallman

clang/lib/CodeGen/CGExprScalar.cpp

@@ -3574,20 +3574,20 @@ ScalarExprEmitter::VisitUnaryExprOrTypeTraitExpr(
           CGF.EmitIgnoredExpr(E->getArgumentExpr());
         }
 
-        auto VlaSize = CGF.getVLASize(VAT);
-        llvm::Value *size = VlaSize.NumElts;
+        // For _Countof, we just want to return the size of a single dimension.
+        if (Kind == UETT_CountOf)
+          return CGF.getVLAElements1D(VAT).NumElts;
 
         // For sizeof and __datasizeof, we need to scale the number of elements
-        // by the size of the array element type. For _Countof, we just want to
-        // return the size directly.
-        if (Kind != UETT_CountOf) {
-          // Scale the number of non-VLA elements by the non-VLA element size.
-          CharUnits eltSize = CGF.getContext().getTypeSizeInChars(VlaSize.Type);
-          if (!eltSize.isOne())
-            size = CGF.Builder.CreateNUWMul(CGF.CGM.getSize(eltSize), size);
-        }
+        // by the size of the array element type.
+        auto VlaSize = CGF.getVLASize(VAT);
 
-        return size;
+        // Scale the number of non-VLA elements by the non-VLA element size.
+        CharUnits eltSize = CGF.getContext().getTypeSizeInChars(VlaSize.Type);
+        if (!eltSize.isOne())
+          return CGF.Builder.CreateNUWMul(CGF.CGM.getSize(eltSize),
+                                          VlaSize.NumElts);
+        return VlaSize.NumElts;
       }
     }
   } else if (E->getKind() == UETT_OpenMPRequiredSimdAlign) {

clang/test/C/C2y/n3369_3.c

@@ -164,3 +164,57 @@ void test4(int n) {
   int y = _Countof(int[n++][7]);
 }
 
+// CHECK-64-LABEL: define dso_local void @test5(
+// CHECK-64-SAME: ) #[[ATTR0]] {
+// CHECK-64-NEXT:  [[ENTRY:.*:]]
+// CHECK-64-NEXT:    [[J:%.*]] = alloca i32, align
+// CHECK-64-NEXT:    [[SAVED_STACK:%.*]] = alloca ptr, align
+// CHECK-64-NEXT:    [[A:%.*]] = alloca i32, align
+// CHECK-64-NEXT:    [[B:%.*]] = alloca i32, align
+// CHECK-64-NEXT:    [[C:%.*]] = alloca i32, align
+// CHECK-64-NEXT:    [[D:%.*]] = alloca i32, align
+// CHECK-64-NEXT:    store i32 2, ptr [[J]], align
+// CHECK-64-NEXT:    [[TMP0:%.*]] = call ptr @llvm.stacksave.p0()
+// CHECK-64-NEXT:    store ptr [[TMP0]], ptr [[SAVED_STACK]], align
+// CHECK-64-NEXT:    [[VLA:%.*]] = alloca i32, i64 120, align
+// CHECK-64-NEXT:    store i32 1, ptr [[A]], align
+// CHECK-64-NEXT:    store i32 4, ptr [[B]], align
+// CHECK-64-NEXT:    store i32 5, ptr [[C]], align
+// CHECK-64-NEXT:    store i32 6, ptr [[D]], align
+// CHECK-64-NEXT:    [[TMP1:%.*]] = load ptr, ptr [[SAVED_STACK]], align
+// CHECK-64-NEXT:    call void @llvm.stackrestore.p0(ptr [[TMP1]])
+// CHECK-64-NEXT:    ret void
+//
+// CHECK-32-LABEL: define dso_local void @test5(
+// CHECK-32-SAME: ) #[[ATTR0]] {
+// CHECK-32-NEXT:  [[ENTRY:.*:]]
+// CHECK-32-NEXT:    [[J:%.*]] = alloca i32, align
+// CHECK-32-NEXT:    [[SAVED_STACK:%.*]] = alloca ptr, align
+// CHECK-32-NEXT:    [[A:%.*]] = alloca i32, align
+// CHECK-32-NEXT:    [[B:%.*]] = alloca i32, align
+// CHECK-32-NEXT:    [[C:%.*]] = alloca i32, align
+// CHECK-32-NEXT:    [[D:%.*]] = alloca i32, align
+// CHECK-32-NEXT:    store i32 2, ptr [[J]], align
+// CHECK-32-NEXT:    [[TMP0:%.*]] = call ptr @llvm.stacksave.p0()
+// CHECK-32-NEXT:    store ptr [[TMP0]], ptr [[SAVED_STACK]], align
+// CHECK-32-NEXT:    [[VLA:%.*]] = alloca i32, i32 120, align
+// CHECK-32-NEXT:    store i32 1, ptr [[A]], align
+// CHECK-32-NEXT:    store i32 4, ptr [[B]], align
+// CHECK-32-NEXT:    store i32 5, ptr [[C]], align
+// CHECK-32-NEXT:    store i32 6, ptr [[D]], align
+// CHECK-32-NEXT:    [[TMP1:%.*]] = load ptr, ptr [[SAVED_STACK]], align
+// CHECK-32-NEXT:    call void @llvm.stackrestore.p0(ptr [[TMP1]])
+// CHECK-32-NEXT:    ret void
+//
+void test5() {
+  // Ensure we're getting the variable-length dimensions correctly. We were
+  // previously returning the size of all VLA dimensions multiplied together
+  // to get the total element count rather than the element count for a single
+  // array rank. See GH141409
+  const int j = 2;
+  int arr[1][j + 2][j + 3][j + 4];
+  int a = _Countof arr;
+  int b = _Countof *arr;
+  int c = _Countof **arr;
+  int d = _Countof ***arr;
+}