[clang] [clang][CodeGen][UBSan] Fixing shift-exponent generation for _BitInt (PR #80515)

[Adam Magier via cfe-commits]

https://github.com/AdamMagierFOSS created https://github.com/llvm/llvm-project/pull/80515

Testing the shift-exponent check with small width _BitInt values exposed a bug in ScalarExprEmitter::GetWidthMinusOneValue when using the result to determine valid exponent sizes. False positives were reported for some left shifts when width(LHS)-1 > range(RHS) and false negatives were reported for right shifts when value(RHS) > range(LHS). This patch caps the maximum value of GetWidthMinusOneValue to fit within range(RHS) to fix the issue with left shifts and fixes a code generation in EmitShr to fix the issue with right shifts.

>From 4e1c37ae83dec050fc9b7aa172db01fa0b2b6d68 Mon Sep 17 00:00:00 2001
From: Adam Magier <adam.magier at ericsson.com>
Date: Sat, 3 Feb 2024 00:38:54 +0100
Subject: [PATCH] [clang][CodeGen][UBSan] Fixing shift-exponent generation for
 _BitInt

Testing the shift-exponent check with small width _BitInt values exposed
a bug in ScalarExprEmitter::GetWidthMinusOneValue when using the result
to determine valid exponent sizes. False positives were reported for
some left shifts when width(LHS)-1 > range(RHS) and false negatives were
reported for right shifts when value(RHS) > range(LHS). This patch caps
the maximum value of GetWidthMinusOneValue to fit within range(RHS) to
fix the issue with left shifts and fixes a code generation in EmitShr to
fix the issue with right shifts.
---
 clang/lib/CodeGen/CGExprScalar.cpp      |  9 ++++++-
 clang/test/CodeGen/ubsan-shift-bitint.c | 36 +++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)
 create mode 100644 clang/test/CodeGen/ubsan-shift-bitint.c

diff --git a/clang/lib/CodeGen/CGExprScalar.cpp b/clang/lib/CodeGen/CGExprScalar.cpp
index 5502f685f6474..e2e3ed839714a 100644
--- a/clang/lib/CodeGen/CGExprScalar.cpp
+++ b/clang/lib/CodeGen/CGExprScalar.cpp
@@ -4121,6 +4121,13 @@ Value *ScalarExprEmitter::GetWidthMinusOneValue(Value* LHS,Value* RHS) {
     Ty = cast<llvm::IntegerType>(VT->getElementType());
   else
     Ty = cast<llvm::IntegerType>(LHS->getType());
+  // Testing with small _BitInt types has shown that Ty->getBitwidth() - 1
+  // can sometimes overflow the capacity of RHS->getType(), cap the value
+  // to be the largest RHS->getType() can hold
+  llvm::APInt RHSMax =
+      llvm::APInt::getMaxValue(RHS->getType()->getScalarSizeInBits());
+  if (RHSMax.ult(Ty->getBitWidth()))
+    return llvm::ConstantInt::get(RHS->getType(), RHSMax);
   return llvm::ConstantInt::get(RHS->getType(), Ty->getBitWidth() - 1);
 }
 
@@ -4235,7 +4242,7 @@ Value *ScalarExprEmitter::EmitShr(const BinOpInfo &Ops) {
            isa<llvm::IntegerType>(Ops.LHS->getType())) {
     CodeGenFunction::SanitizerScope SanScope(&CGF);
     llvm::Value *Valid =
-        Builder.CreateICmpULE(RHS, GetWidthMinusOneValue(Ops.LHS, RHS));
+        Builder.CreateICmpULE(Ops.RHS, GetWidthMinusOneValue(Ops.LHS, Ops.RHS));
     EmitBinOpCheck(std::make_pair(Valid, SanitizerKind::ShiftExponent), Ops);
   }
 
diff --git a/clang/test/CodeGen/ubsan-shift-bitint.c b/clang/test/CodeGen/ubsan-shift-bitint.c
new file mode 100644
index 0000000000000..8ca94b7de5a42
--- /dev/null
+++ b/clang/test/CodeGen/ubsan-shift-bitint.c
@@ -0,0 +1,36 @@
+// RUN: %clang_cc1 %s -O0 -fsanitize=shift-exponent -emit-llvm -o - | FileCheck %s
+
+// Checking that the code generation is using the unextended/untruncated
+// exponent values and capping the values accordingly
+
+// CHECK-LABEL: define{{.*}} i32 @test_left_variable
+int test_left_variable(unsigned _BitInt(5) b, unsigned _BitInt(2) e) {
+  // CHECK: [[E_REG:%.+]] = load [[E_SIZE:i2]]
+  // CHECK: icmp ule [[E_SIZE]] [[E_REG]], -1
+  return b << e;
+}
+
+// CHECK-LABEL: define{{.*}} i32 @test_right_variable
+int test_right_variable(unsigned _BitInt(2) b, unsigned _BitInt(3) e) {
+  // CHECK: [[E_REG:%.+]] = load [[E_SIZE:i3]]
+  // CHECK: icmp ule [[E_SIZE]] [[E_REG]], 1
+  return b >> e;
+}
+
+// Old code generation would give false positives on left shifts when:
+//   value(e) > (width(b) - 1 % 2 ** width(e))
+// CHECK-LABEL: define{{.*}} i32 @test_left_literal
+int test_left_literal(unsigned _BitInt(5) b) {
+  // CHECK-NOT: br i1 false, label %cont, label %handler.shift_out_of_bounds
+  // CHECK: br i1 true, label %cont, label %handler.shift_out_of_bounds
+  return b << 3uwb;
+}
+
+// Old code generation would give false positives on right shifts when:
+//   (value(e) % 2 ** width(b)) < width(b)
+// CHECK-LABEL: define{{.*}} i32 @test_right_literal
+int test_right_literal(unsigned _BitInt(2) b) {
+  // CHECK-NOT: br i1 true, label %cont, label %handler.shift_out_of_bounds
+  // CHECK: br i1 false, label %cont, label %handler.shift_out_of_bounds
+  return b >> 4uwb;
+}