From e6ec0c7715eb400ff92f891888968f87c8bb445f Mon Sep 17 00:00:00 2001
From: Alex Richardson <Alexander.Richardson@cl.cam.ac.uk>
Date: Fri, 12 May 2023 13:35:55 -0700
Subject: [PATCH] [InstSimplify] Add basic fold for @llvm.cheri.cap.flags.set

If we are setting the same value we just read we can simplify the
sequence to just the original argument.
---
 llvm/lib/Analysis/InstructionSimplify.cpp                   | 6 ++++++
 .../Transforms/InstSimplify/cheri-intrinsics-get-set.ll     | 6 +++---
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Analysis/InstructionSimplify.cpp b/llvm/lib/Analysis/InstructionSimplify.cpp
index 583e75a74066..8d7ca7c84b8b 100644
--- a/llvm/lib/Analysis/InstructionSimplify.cpp
+++ b/llvm/lib/Analysis/InstructionSimplify.cpp
@@ -6061,6 +6061,12 @@ static Value *simplifyBinaryIntrinsic(Function *F, Value *Op0, Value *Op1,
         return Op0;
     }
     break;
+  case Intrinsic::cheri_cap_flags_set:
+    // flags_set(x, flags_get(x)) -> x
+    if (match(Op1,
+              m_Intrinsic<Intrinsic::cheri_cap_flags_get>(m_Specific(Op0))))
+      return Op0;
+    break;
   case Intrinsic::cheri_cap_offset_set:
   case Intrinsic::cheri_cap_address_set: {
     Value *Base = getBasePtrIgnoringCapabilityAddressManipulation(Op0, Q.DL);
diff --git a/llvm/test/Transforms/InstSimplify/cheri-intrinsics-get-set.ll b/llvm/test/Transforms/InstSimplify/cheri-intrinsics-get-set.ll
index aa2a9be4129b..43ac026667eb 100644
--- a/llvm/test/Transforms/InstSimplify/cheri-intrinsics-get-set.ll
+++ b/llvm/test/Transforms/InstSimplify/cheri-intrinsics-get-set.ll
@@ -12,6 +12,8 @@ declare i8 addrspace(200)* @llvm.cheri.cap.INTRINSIC.set.i64(i8 addrspace(200)*,
 
 
 ;; This is a no-op and should be folded to ret %arg
+;; NB: The exception here is @llvm.cheri.cap.high.set since it is tag-clearing
+;; and therefore could only be used on known-untagged values.
 define i8 addrspace(200)* @fold_set_of_get(i8 addrspace(200)* %arg) nounwind {
 ; ADDRESS-LABEL: define {{[^@]+}}@fold_set_of_get
 ; ADDRESS-SAME: (i8 addrspace(200)* [[ARG:%.*]]) addrspace(200) #[[ATTR0:[0-9]+]] {
@@ -19,9 +21,7 @@ define i8 addrspace(200)* @fold_set_of_get(i8 addrspace(200)* %arg) nounwind {
 ;
 ; FLAGS-LABEL: define {{[^@]+}}@fold_set_of_get
 ; FLAGS-SAME: (i8 addrspace(200)* [[ARG:%.*]]) addrspace(200) #[[ATTR0:[0-9]+]] {
-; FLAGS-NEXT:    [[VALUE:%.*]] = tail call i64 @llvm.cheri.cap.flags.get.i64(i8 addrspace(200)* [[ARG]])
-; FLAGS-NEXT:    [[RET:%.*]] = tail call i8 addrspace(200)* @llvm.cheri.cap.flags.set.i64(i8 addrspace(200)* [[ARG]], i64 [[VALUE]])
-; FLAGS-NEXT:    ret i8 addrspace(200)* [[RET]]
+; FLAGS-NEXT:    ret i8 addrspace(200)* [[ARG]]
 ;
 ; HIGH-LABEL: define {{[^@]+}}@fold_set_of_get
 ; HIGH-SAME: (i8 addrspace(200)* [[ARG:%.*]]) addrspace(200) #[[ATTR0:[0-9]+]] {