[CIR] CountOf VLA with Array element type (llvm#169404)

Implement CountOf on VariableArrayType with IntegerConstant SizeExpr

Author: AmrDeveloper

clang/lib/CIR/CodeGen/CIRGenExprScalar.cpp

@@ -2325,14 +2325,45 @@ mlir::Value ScalarExprEmitter::VisitUnaryExprOrTypeTraitExpr(
   const QualType typeToSize = e->getTypeOfArgument();
   const mlir::Location loc = cgf.getLoc(e->getSourceRange());
   if (auto kind = e->getKind();
-      kind == UETT_SizeOf || kind == UETT_DataSizeOf) {
-    if (cgf.getContext().getAsVariableArrayType(typeToSize)) {
-      cgf.getCIRGenModule().errorNYI(e->getSourceRange(),
-                                     "sizeof operator for VariableArrayType",
-                                     e->getStmtClassName());
-      return builder.getConstant(
-          loc, cir::IntAttr::get(cgf.cgm.uInt64Ty,
-                                 llvm::APSInt(llvm::APInt(64, 1), true)));
+      kind == UETT_SizeOf || kind == UETT_DataSizeOf || kind == UETT_CountOf) {
+    if (const VariableArrayType *vat =
+            cgf.getContext().getAsVariableArrayType(typeToSize)) {
+      // For _Countof, we only want to evaluate if the extent is actually
+      // variable as opposed to a multi-dimensional array whose extent is
+      // constant but whose element type is variable.
+      bool evaluateExtent = true;
+      if (kind == UETT_CountOf && vat->getElementType()->isArrayType()) {
+        evaluateExtent =
+            !vat->getSizeExpr()->isIntegerConstantExpr(cgf.getContext());
+      }
+
+      if (evaluateExtent) {
+        if (e->isArgumentType()) {
+          // sizeof(type) - make sure to emit the VLA size.
+          cgf.emitVariablyModifiedType(typeToSize);
+        } else {
+          // C99 6.5.3.4p2: If the argument is an expression of type
+          // VLA, it is evaluated.
+          cgf.getCIRGenModule().errorNYI(
+              e->getSourceRange(),
+              "sizeof operator for VariableArrayType & evaluateExtent "
+              "ignoredExpr",
+              e->getStmtClassName());
+          return {};
+        }
+
+        // For _Countof, we just want to return the size of a single dimension.
+        if (kind == UETT_CountOf)
+          return cgf.getVLAElements1D(vat).numElts;
+
+        cgf.getCIRGenModule().errorNYI(
+            e->getSourceRange(),
+            "sizeof operator for VariableArrayType & evaluateExtent",
+            e->getStmtClassName());
+        return builder.getConstant(
+            loc, cir::IntAttr::get(cgf.cgm.uInt64Ty,
+                                   -llvm::APSInt(llvm::APInt(64, 1), true)));
+      }
     }
   } else if (e->getKind() == UETT_OpenMPRequiredSimdAlign) {
     cgf.getCIRGenModule().errorNYI(

clang/lib/CIR/CodeGen/CIRGenFunction.cpp

@@ -1138,6 +1138,14 @@ CIRGenFunction::getVLASize(const VariableArrayType *type) {
   return {numElements, elementType};
 }
 
+CIRGenFunction::VlaSizePair
+CIRGenFunction::getVLAElements1D(const VariableArrayType *vla) {
+  mlir::Value vlaSize = vlaSizeMap[vla->getSizeExpr()];
+  assert(vlaSize && "no size for VLA!");
+  assert(vlaSize->getType() == sizeTy);
+  return {vlaSize, vla->getElementType()};
+}
+
 // TODO(cir): Most of this function can be shared between CIRGen
 // and traditional LLVM codegen
 void CIRGenFunction::emitVariablyModifiedType(QualType type) {

clang/lib/CIR/CodeGen/CIRGenFunction.h

@@ -498,6 +498,10 @@ class CIRGenFunction : public CIRGenTypeCache {
     VlaSizePair(mlir::Value num, QualType ty) : numElts(num), type(ty) {}
   };
 
+  /// Return the number of elements for a single dimension
+  /// for the given array type.
+  VlaSizePair getVLAElements1D(const VariableArrayType *vla);
+
   /// Returns an MLIR::Value+QualType pair that corresponds to the size,
   /// in non-variably-sized elements, of a variable length array type,
   /// plus that largest non-variably-sized element type.  Assumes that

clang/test/CIR/CodeGen/count-of.c

@@ -0,0 +1,52 @@
+// RUN: %clang_cc1 -std=c2y -triple x86_64-unknown-linux-gnu -Wno-unused-value -fclangir -emit-cir %s -o %t.cir
+// RUN: FileCheck --input-file=%t.cir %s -check-prefix=CIR
+// RUN: %clang_cc1 -std=c2y -triple x86_64-unknown-linux-gnu -Wno-unused-value -fclangir -emit-llvm %s -o %t-cir.ll
+// RUN: FileCheck --input-file=%t-cir.ll %s -check-prefix=LLVM
+// RUN: %clang_cc1 -std=c2y -triple x86_64-unknown-linux-gnu -Wno-unused-value -emit-llvm %s -o %t.ll
+// RUN: FileCheck --input-file=%t.ll %s -check-prefix=OGCG
+
+unsigned long vla_with_array_element_type_with_const_size() {
+  long size;
+  return _Countof(int[5][size]);
+}
+
+// CIR: %[[RET_ADDR:.*]] = cir.alloca !u64i, !cir.ptr<!u64i>, ["__retval"]
+// CIR: %[[SIZE_ADDR:.*]] = cir.alloca !s64i, !cir.ptr<!s64i>, ["size"]
+// CIR: %[[CONST_5:.*]] = cir.const #cir.int<5> : !u64i
+// CIR: cir.store %[[CONST_5]], %[[RET_ADDR]] : !u64i, !cir.ptr<!u64i>
+// CIR: %[[RET_VAL:.*]] = cir.load %[[RET_ADDR]] : !cir.ptr<!u64i>, !u64i
+// CIR: cir.return %[[RET_VAL]] : !u64i
+
+// LLVM: %[[RET_ADDR:.*]] = alloca i64, i64 1, align 8
+// LLVM: %[[SIZE_ADDR:.*]] = alloca i64, i64 1, align 8
+// LLVM: store i64 5, ptr %[[RET_ADDR]], align 8
+// LLVM: %[[RET_VAL:.*]] = load i64, ptr %[[RET_ADDR]], align 8
+// LLVM: ret i64 %[[RET_VAL]]
+
+// OGCG: %[[SIZE_ADDR:.*]] = alloca i64, align 8
+// OGCG: ret i64 5
+
+unsigned long vla_with_array_element_type_non_const_size() {
+  long size;
+  return _Countof(int[size][size]);
+}
+
+// CIR: %[[REET_ADDR:.*]] = cir.alloca !u64i, !cir.ptr<!u64i>, ["__retval"]
+// CIR: %[[SIZE_ADDR:.*]] = cir.alloca !s64i, !cir.ptr<!s64i>, ["size"]
+// CIR: %[[TMP_SIZE:.*]] = cir.load {{.*}} %[[SIZE_ADDR]] : !cir.ptr<!s64i>, !s64i
+// CIR: %[[TMP_SIZE_U64:.*]] = cir.cast integral %[[TMP_SIZE]] : !s64i -> !u64i
+// CIR: cir.store %[[TMP_SIZE_U64]], %[[RET_ADDR]] : !u64i, !cir.ptr<!u64i>
+// CIR: %[[TMP_RET:.*]] = cir.load %[[RET_ADDR]] : !cir.ptr<!u64i>, !u64i
+// CIR: cir.return %[[TMP_RET]] : !u64i
+
+// LLVM: %[[RET_ADDR:.*]] = alloca i64, i64 1, align 8
+// LLVM: %[[SIZE_ADDR:.*]] = alloca i64, i64 1, align 8
+// LLVM: %[[TMP_SIZE:.*]] = load i64, ptr %[[SIZE_ADDR]], align 8
+// LLVM: store i64 %[[TMP_SIZE]], ptr %[[RET_ADDR]], align 8
+// LLVM: %[[TMP_RET:.*]] = load i64, ptr %[[RET_ADDR]], align 8
+// LLVM: ret i64 %[[TMP_RET]]
+
+// OGCG: %[[SIZE_ADDR:.*]] = alloca i64, align 8
+// OGCG: %[[TMP_SIZE:.*]] = load i64, ptr %[[SIZE_ADDR]], align 8
+// OGCG: %[[TMP_SIZE_2:.*]] = load i64, ptr %[[SIZE_ADDR]], align 8
+// OGCG: ret i64 %[[TMP_SIZE]]