[llvm] [llubi] Add support for captures (PR #201170)

Yingwei Zheng via llvm-commits llvm-commits at lists.llvm.org
Sun Jun 14 08:39:23 PDT 2026


https://github.com/dtcxzyw updated https://github.com/llvm/llvm-project/pull/201170

>From 35fc116a2e32895f19347baf41cea5c99acdd2b0 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Thu, 4 Jun 2026 01:31:10 +0800
Subject: [PATCH] [llubi] Add support for captures

---
 llvm/test/tools/llubi/captures.ll             | 120 ++++++++++++++++++
 ...dereferenceable_capture_read_provenance.ll |  24 ++++
 .../llubi/load_capture_read_provenance.ll     |  24 ++++
 .../tools/llubi/store_capture_provenance.ll   |  24 ++++
 .../tools/llubi/volatile_capture_address.ll   |  24 ++++
 llvm/test/tools/llubi/wildcard_provenance.ll  |  38 ++++++
 ...dcard_provenance_angelic_nondeterminism.ll |  31 +++++
 .../llubi/wildcard_provenance_gep_inbounds.ll |  29 +++++
 llvm/tools/llubi/lib/ExecutorBase.cpp         |  33 +++--
 llvm/tools/llubi/lib/ExecutorBase.h           |  20 +--
 llvm/tools/llubi/lib/Interpreter.cpp          |  82 ++++++++++--
 llvm/tools/llubi/lib/Library.cpp              |   6 +-
 llvm/tools/llubi/lib/Value.cpp                |   4 +-
 llvm/tools/llubi/lib/Value.h                  |  18 ++-
 14 files changed, 442 insertions(+), 35 deletions(-)
 create mode 100644 llvm/test/tools/llubi/captures.ll
 create mode 100644 llvm/test/tools/llubi/dereferenceable_capture_read_provenance.ll
 create mode 100644 llvm/test/tools/llubi/load_capture_read_provenance.ll
 create mode 100644 llvm/test/tools/llubi/store_capture_provenance.ll
 create mode 100644 llvm/test/tools/llubi/volatile_capture_address.ll
 create mode 100644 llvm/test/tools/llubi/wildcard_provenance.ll
 create mode 100644 llvm/test/tools/llubi/wildcard_provenance_angelic_nondeterminism.ll
 create mode 100644 llvm/test/tools/llubi/wildcard_provenance_gep_inbounds.ll

diff --git a/llvm/test/tools/llubi/captures.ll b/llvm/test/tools/llubi/captures.ll
new file mode 100644
index 0000000000000..4ac808c4248aa
--- /dev/null
+++ b/llvm/test/tools/llubi/captures.ll
@@ -0,0 +1,120 @@
+; NOTE: Assertions have been autogenerated by utils/update_llubi_test_checks.py UTC_ARGS: --version 6
+; RUN: llubi --verbose < %s 2>&1 | FileCheck %s
+
+define void @write(ptr %p) {
+  store i8 1, ptr %p
+  ret void
+}
+
+define ptr @address_only(ptr captures(address) %p) {
+  ret ptr %p
+}
+
+define ptr @address_is_null_only(ptr captures(address_is_null) %p) {
+  ret ptr %p
+}
+
+define ptr @provenance_only(ptr captures(provenance) %p) {
+  ret ptr %p
+}
+
+define ptr @read_provenance_only(ptr captures(read_provenance) %p) {
+  ret ptr %p
+}
+
+ at g = global ptr null
+
+define ptr @ret_capture_read_provenance_other_capture_address(ptr captures(address, ret: read_provenance) %p) {
+  store ptr %p, ptr @g
+  %p2 = load ptr, ptr @g
+  ret ptr %p2
+}
+
+define void @main() {
+  %p = alloca i8
+  call void @write(ptr captures(none) %p)
+  %p_poison = call ptr @address_only(ptr poison)
+  %p_address_only = call ptr @address_only(ptr %p)
+  %p_address_is_null_only = call ptr @address_is_null_only(ptr %p)
+  %cmp1 = icmp eq ptr %p_address_only, %p_address_only
+  %cmp2 = icmp eq ptr %p_address_is_null_only, %p_address_is_null_only
+  %cmp3 = icmp eq ptr %p_address_is_null_only, null
+  %cmp4 = icmp eq ptr null, %p_address_is_null_only
+  %cmp5 = icmp eq ptr null, null
+
+  %p_read_provenance = call ptr @read_provenance_only(ptr %p)
+  %p_provenance = call ptr @provenance_only(ptr %p)
+
+  %p_read_provenance_2 = call ptr @read_provenance_only(ptr %p_provenance)
+  %p_none = call ptr @provenance_only(ptr %p_address_only)
+
+  %p_read_provenance2 = call ptr @ret_capture_read_provenance_other_capture_address(ptr %p)
+  %p_address_only2 = load ptr, ptr @g
+
+  call void @llvm.assume(i1 true) ["dereferenceable"(ptr %p_read_provenance, i32 1)]
+  %val = load i8, ptr %p_read_provenance
+  store i8 0, ptr %p_provenance
+
+  ret void
+}
+; CHECK: Entering function: main
+; CHECK-NEXT:   %p = alloca i8, align 1 => ptr 0x11 [p]
+; CHECK-NEXT: Entering function: write
+; CHECK-NEXT:   ptr %p = ptr 0x11 [p]
+; CHECK-NEXT:   store i8 1, ptr %p, align 1
+; CHECK-NEXT:   ret void
+; CHECK-NEXT: Exiting function: write
+; CHECK-NEXT:   call void @write(ptr captures(none) %p)
+; CHECK-NEXT: Entering function: address_only
+; CHECK-NEXT:   ptr %p = poison
+; CHECK-NEXT:   ret ptr %p
+; CHECK-NEXT: Exiting function: address_only
+; CHECK-NEXT:   %p_poison = call ptr @address_only(ptr poison) => poison
+; CHECK-NEXT: Entering function: address_only
+; CHECK-NEXT:   ptr %p = ptr 0x11 [p]
+; CHECK-NEXT:   ret ptr %p
+; CHECK-NEXT: Exiting function: address_only
+; CHECK-NEXT:   %p_address_only = call ptr @address_only(ptr %p) => ptr 0x11 [p address]
+; CHECK-NEXT: Entering function: address_is_null_only
+; CHECK-NEXT:   ptr %p = ptr 0x11 [p]
+; CHECK-NEXT:   ret ptr %p
+; CHECK-NEXT: Exiting function: address_is_null_only
+; CHECK-NEXT:   %p_address_is_null_only = call ptr @address_is_null_only(ptr %p) => ptr 0x11 [p address_is_null]
+; CHECK-NEXT:   %cmp1 = icmp eq ptr %p_address_only, %p_address_only => T
+; CHECK-NEXT:   %cmp2 = icmp eq ptr %p_address_is_null_only, %p_address_is_null_only => poison
+; CHECK-NEXT:   %cmp3 = icmp eq ptr %p_address_is_null_only, null => F
+; CHECK-NEXT:   %cmp4 = icmp eq ptr null, %p_address_is_null_only => F
+; CHECK-NEXT:   %cmp5 = icmp eq ptr null, null => T
+; CHECK-NEXT: Entering function: read_provenance_only
+; CHECK-NEXT:   ptr %p = ptr 0x11 [p]
+; CHECK-NEXT:   ret ptr %p
+; CHECK-NEXT: Exiting function: read_provenance_only
+; CHECK-NEXT:   %p_read_provenance = call ptr @read_provenance_only(ptr %p) => ptr 0x11 [p read_provenance]
+; CHECK-NEXT: Entering function: provenance_only
+; CHECK-NEXT:   ptr %p = ptr 0x11 [p]
+; CHECK-NEXT:   ret ptr %p
+; CHECK-NEXT: Exiting function: provenance_only
+; CHECK-NEXT:   %p_provenance = call ptr @provenance_only(ptr %p) => ptr 0x11 [p provenance]
+; CHECK-NEXT: Entering function: read_provenance_only
+; CHECK-NEXT:   ptr %p = ptr 0x11 [p provenance]
+; CHECK-NEXT:   ret ptr %p
+; CHECK-NEXT: Exiting function: read_provenance_only
+; CHECK-NEXT:   %p_read_provenance_2 = call ptr @read_provenance_only(ptr %p_provenance) => ptr 0x11 [p read_provenance]
+; CHECK-NEXT: Entering function: provenance_only
+; CHECK-NEXT:   ptr %p = ptr 0x11 [p address]
+; CHECK-NEXT:   ret ptr %p
+; CHECK-NEXT: Exiting function: provenance_only
+; CHECK-NEXT:   %p_none = call ptr @provenance_only(ptr %p_address_only) => ptr 0x11 [p none]
+; CHECK-NEXT: Entering function: ret_capture_read_provenance_other_capture_address
+; CHECK-NEXT:   ptr %p = ptr 0x11 [p]
+; CHECK-NEXT:   store ptr %p, ptr @g, align 8
+; CHECK-NEXT:   %p2 = load ptr, ptr @g, align 8 => ptr 0x11 [p]
+; CHECK-NEXT:   ret ptr %p2
+; CHECK-NEXT: Exiting function: ret_capture_read_provenance_other_capture_address
+; CHECK-NEXT:   %p_read_provenance2 = call ptr @ret_capture_read_provenance_other_capture_address(ptr %p) => ptr 0x11 [p read_provenance]
+; CHECK-NEXT:   %p_address_only2 = load ptr, ptr @g, align 8 => ptr 0x11 [p address]
+; CHECK-NEXT:   call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %p_read_provenance, i32 1) ]
+; CHECK-NEXT:   %val = load i8, ptr %p_read_provenance, align 1 => i8 1
+; CHECK-NEXT:   store i8 0, ptr %p_provenance, align 1
+; CHECK-NEXT:   ret void
+; CHECK-NEXT: Exiting function: main
diff --git a/llvm/test/tools/llubi/dereferenceable_capture_read_provenance.ll b/llvm/test/tools/llubi/dereferenceable_capture_read_provenance.ll
new file mode 100644
index 0000000000000..809756499b8c8
--- /dev/null
+++ b/llvm/test/tools/llubi/dereferenceable_capture_read_provenance.ll
@@ -0,0 +1,24 @@
+; NOTE: Assertions have been autogenerated by utils/update_llubi_test_checks.py UTC_ARGS: --version 6
+; RUN: not llubi --verbose < %s 2>&1 | FileCheck %s
+
+define ptr @address_only(ptr captures(address) %p) {
+  ret ptr %p
+}
+
+ at g = global i8 0
+
+define void @main() {
+  %g_address_only = call ptr @address_only(ptr @g)
+  call void @llvm.assume(i1 true) ["dereferenceable"(ptr %g_address_only, i32 1)]
+  ret void
+}
+; CHECK: Entering function: main
+; CHECK-NEXT: Entering function: address_only
+; CHECK-NEXT:   ptr %p = ptr 0x8 [@g]
+; CHECK-NEXT:   ret ptr %p
+; CHECK-NEXT: Exiting function: address_only
+; CHECK-NEXT:   %g_address_only = call ptr @address_only(ptr @g) => ptr 0x8 [@g address]
+; CHECK-NEXT: Stacktrace:
+; CHECK-NEXT: #0   call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %g_address_only, i32 1) ] at @main <stdin>:12
+; CHECK-NEXT: Immediate UB detected: The pointer ptr 0x8 [@g address] violates dereferenceable(1) assumption.
+; CHECK-NEXT: error: Execution of function 'main' failed.
diff --git a/llvm/test/tools/llubi/load_capture_read_provenance.ll b/llvm/test/tools/llubi/load_capture_read_provenance.ll
new file mode 100644
index 0000000000000..ae4cf63a2dbb1
--- /dev/null
+++ b/llvm/test/tools/llubi/load_capture_read_provenance.ll
@@ -0,0 +1,24 @@
+; NOTE: Assertions have been autogenerated by utils/update_llubi_test_checks.py UTC_ARGS: --version 6
+; RUN: not llubi --verbose < %s 2>&1 | FileCheck %s
+
+define ptr @address_only(ptr captures(address) %p) {
+  ret ptr %p
+}
+
+ at g = global i8 0
+
+define void @main() {
+  %g_address_only = call ptr @address_only(ptr @g)
+  %res = load i8, ptr %g_address_only
+  ret void
+}
+; CHECK: Entering function: main
+; CHECK-NEXT: Entering function: address_only
+; CHECK-NEXT:   ptr %p = ptr 0x8 [@g]
+; CHECK-NEXT:   ret ptr %p
+; CHECK-NEXT: Exiting function: address_only
+; CHECK-NEXT:   %g_address_only = call ptr @address_only(ptr @g) => ptr 0x8 [@g address]
+; CHECK-NEXT: Stacktrace:
+; CHECK-NEXT: #0   %res = load i8, ptr %g_address_only, align 1 at @main <stdin>:12
+; CHECK-NEXT: Immediate UB detected: Invalid memory access via a pointer with nullary provenance.
+; CHECK-NEXT: error: Execution of function 'main' failed.
diff --git a/llvm/test/tools/llubi/store_capture_provenance.ll b/llvm/test/tools/llubi/store_capture_provenance.ll
new file mode 100644
index 0000000000000..a5fd778d7a200
--- /dev/null
+++ b/llvm/test/tools/llubi/store_capture_provenance.ll
@@ -0,0 +1,24 @@
+; NOTE: Assertions have been autogenerated by utils/update_llubi_test_checks.py UTC_ARGS: --version 6
+; RUN: not llubi --verbose < %s 2>&1 | FileCheck %s
+
+define ptr @address_read_provenance(ptr captures(address, read_provenance) %p) {
+  ret ptr %p
+}
+
+ at g = global i8 0
+
+define void @main() {
+  %g_address_read_provenance = call ptr @address_read_provenance(ptr @g)
+  store i8 1, ptr %g_address_read_provenance
+  ret void
+}
+; CHECK: Entering function: main
+; CHECK-NEXT: Entering function: address_read_provenance
+; CHECK-NEXT:   ptr %p = ptr 0x8 [@g]
+; CHECK-NEXT:   ret ptr %p
+; CHECK-NEXT: Exiting function: address_read_provenance
+; CHECK-NEXT:   %g_address_read_provenance = call ptr @address_read_provenance(ptr @g) => ptr 0x8 [@g address, read_provenance]
+; CHECK-NEXT: Stacktrace:
+; CHECK-NEXT: #0   store i8 1, ptr %g_address_read_provenance, align 1 at @main <stdin>:12
+; CHECK-NEXT: Immediate UB detected: Invalid memory access via a pointer with nullary provenance.
+; CHECK-NEXT: error: Execution of function 'main' failed.
diff --git a/llvm/test/tools/llubi/volatile_capture_address.ll b/llvm/test/tools/llubi/volatile_capture_address.ll
new file mode 100644
index 0000000000000..28a84e1b893b8
--- /dev/null
+++ b/llvm/test/tools/llubi/volatile_capture_address.ll
@@ -0,0 +1,24 @@
+; NOTE: Assertions have been autogenerated by utils/update_llubi_test_checks.py UTC_ARGS: --version 6
+; RUN: not llubi --verbose < %s 2>&1 | FileCheck %s
+
+define ptr @provenance_only(ptr captures(provenance) %p) {
+  ret ptr %p
+}
+
+ at g = global i8 0
+
+define void @main() {
+  %g_provenance_only = call ptr @provenance_only(ptr @g)
+  %res = load volatile i8, ptr %g_provenance_only
+  ret void
+}
+; CHECK: Entering function: main
+; CHECK-NEXT: Entering function: provenance_only
+; CHECK-NEXT:   ptr %p = ptr 0x8 [@g]
+; CHECK-NEXT:   ret ptr %p
+; CHECK-NEXT: Exiting function: provenance_only
+; CHECK-NEXT:   %g_provenance_only = call ptr @provenance_only(ptr @g) => ptr 0x8 [@g provenance]
+; CHECK-NEXT: Stacktrace:
+; CHECK-NEXT: #0   %res = load volatile i8, ptr %g_provenance_only, align 1 at @main <stdin>:12
+; CHECK-NEXT: Immediate UB detected: Invalid memory access via a pointer with nullary provenance.
+; CHECK-NEXT: error: Execution of function 'main' failed.
diff --git a/llvm/test/tools/llubi/wildcard_provenance.ll b/llvm/test/tools/llubi/wildcard_provenance.ll
new file mode 100644
index 0000000000000..806c67bb1edd6
--- /dev/null
+++ b/llvm/test/tools/llubi/wildcard_provenance.ll
@@ -0,0 +1,38 @@
+; NOTE: Assertions have been autogenerated by utils/update_llubi_test_checks.py UTC_ARGS: --version 6
+; RUN: llubi --verbose < %s 2>&1 | FileCheck %s
+
+target datalayout = "e-p:64:64:64"
+
+define ptr @identity(ptr captures(provenance) %p) {
+  ret ptr %p
+}
+
+define void @main() {
+  %p = alloca i32
+  %p_read_provenance_only = call ptr @identity(ptr captures(read_provenance) %p)
+  %p_provenance_only = call ptr @identity(ptr captures(provenance) %p)
+  %addr1 = ptrtoint ptr %p_read_provenance_only to i64
+  %addr2 = ptrtoint ptr %p_provenance_only to i64
+  %p_wildcard = inttoptr i64 %addr1 to ptr
+  store i32 0, ptr %p_wildcard
+
+  ret void
+}
+; CHECK: Entering function: main
+; CHECK-NEXT:   %p = alloca i32, align 4 => ptr 0x8 [p]
+; CHECK-NEXT: Entering function: identity
+; CHECK-NEXT:   ptr %p = ptr 0x8 [p]
+; CHECK-NEXT:   ret ptr %p
+; CHECK-NEXT: Exiting function: identity
+; CHECK-NEXT:   %p_read_provenance_only = call ptr @identity(ptr captures(read_provenance) %p) => ptr 0x8 [p read_provenance]
+; CHECK-NEXT: Entering function: identity
+; CHECK-NEXT:   ptr %p = ptr 0x8 [p]
+; CHECK-NEXT:   ret ptr %p
+; CHECK-NEXT: Exiting function: identity
+; CHECK-NEXT:   %p_provenance_only = call ptr @identity(ptr captures(provenance) %p) => ptr 0x8 [p provenance]
+; CHECK-NEXT:   %addr1 = ptrtoint ptr %p_read_provenance_only to i64 => i64 8
+; CHECK-NEXT:   %addr2 = ptrtoint ptr %p_provenance_only to i64 => i64 8
+; CHECK-NEXT:   %p_wildcard = inttoptr i64 %addr1 to ptr => ptr 0x8 [wildcard]
+; CHECK-NEXT:   store i32 0, ptr %p_wildcard, align 4
+; CHECK-NEXT:   ret void
+; CHECK-NEXT: Exiting function: main
diff --git a/llvm/test/tools/llubi/wildcard_provenance_angelic_nondeterminism.ll b/llvm/test/tools/llubi/wildcard_provenance_angelic_nondeterminism.ll
new file mode 100644
index 0000000000000..70c7747b2c2cd
--- /dev/null
+++ b/llvm/test/tools/llubi/wildcard_provenance_angelic_nondeterminism.ll
@@ -0,0 +1,31 @@
+; NOTE: Assertions have been autogenerated by utils/update_llubi_test_checks.py UTC_ARGS: --version 6
+; RUN: not llubi --verbose < %s 2>&1 | FileCheck %s
+
+target datalayout = "e-p:64:64:64"
+
+define void @main() {
+  %p1 = alloca i32
+  %p2 = alloca i32
+  %addr1 = ptrtoint ptr %p1 to i64
+  %addr2 = ptrtoint ptr %p2 to i64
+  %diff = sub i64 %addr2, %addr1
+  %p = inttoptr i64 %addr1 to ptr
+  %p_off = getelementptr i8, ptr %p, i64 %diff
+  store i32 0, ptr %p_off
+  store i32 0, ptr %p
+
+  ret void
+}
+; CHECK: Entering function: main
+; CHECK-NEXT:   %p1 = alloca i32, align 4 => ptr 0x8 [p1]
+; CHECK-NEXT:   %p2 = alloca i32, align 4 => ptr 0x10 [p2]
+; CHECK-NEXT:   %addr1 = ptrtoint ptr %p1 to i64 => i64 8
+; CHECK-NEXT:   %addr2 = ptrtoint ptr %p2 to i64 => i64 16
+; CHECK-NEXT:   %diff = sub i64 %addr2, %addr1 => i64 8
+; CHECK-NEXT:   %p = inttoptr i64 %addr1 to ptr => ptr 0x8 [wildcard]
+; CHECK-NEXT:   %p_off = getelementptr i8, ptr %p, i64 %diff => ptr 0x10 [wildcard]
+; CHECK-NEXT:   store i32 0, ptr %p_off, align 4
+; CHECK-NEXT: Stacktrace:
+; CHECK-NEXT: #0   store i32 0, ptr %p, align 4 at @main <stdin>:15
+; CHECK-NEXT: Immediate UB detected: Invalid memory access via a pointer with nullary provenance.
+; CHECK-NEXT: error: Execution of function 'main' failed.
diff --git a/llvm/test/tools/llubi/wildcard_provenance_gep_inbounds.ll b/llvm/test/tools/llubi/wildcard_provenance_gep_inbounds.ll
new file mode 100644
index 0000000000000..9718b4743e46c
--- /dev/null
+++ b/llvm/test/tools/llubi/wildcard_provenance_gep_inbounds.ll
@@ -0,0 +1,29 @@
+; NOTE: Assertions have been autogenerated by utils/update_llubi_test_checks.py UTC_ARGS: --version 6
+; RUN: not llubi --verbose < %s 2>&1 | FileCheck %s
+
+target datalayout = "e-p:64:64:64"
+
+define void @main() {
+  %p1 = alloca i32
+  %p2 = alloca i32
+  %addr1 = ptrtoint ptr %p1 to i64
+  %addr2 = ptrtoint ptr %p2 to i64
+  %diff = sub i64 %addr2, %addr1
+  %p = inttoptr i64 %addr1 to ptr
+  %p_off = getelementptr inbounds i8, ptr %p, i64 %diff
+  store i32 0, ptr %p_off
+
+  ret void
+}
+; CHECK: Entering function: main
+; CHECK-NEXT:   %p1 = alloca i32, align 4 => ptr 0x8 [p1]
+; CHECK-NEXT:   %p2 = alloca i32, align 4 => ptr 0x10 [p2]
+; CHECK-NEXT:   %addr1 = ptrtoint ptr %p1 to i64 => i64 8
+; CHECK-NEXT:   %addr2 = ptrtoint ptr %p2 to i64 => i64 16
+; CHECK-NEXT:   %diff = sub i64 %addr2, %addr1 => i64 8
+; CHECK-NEXT:   %p = inttoptr i64 %addr1 to ptr => ptr 0x8 [wildcard]
+; CHECK-NEXT:   %p_off = getelementptr inbounds i8, ptr %p, i64 %diff => poison
+; CHECK-NEXT: Stacktrace:
+; CHECK-NEXT: #0   store i32 0, ptr %p_off, align 4 at @main <stdin>:14
+; CHECK-NEXT: Immediate UB detected: Invalid memory access with a poison pointer.
+; CHECK-NEXT: error: Execution of function 'main' failed.
diff --git a/llvm/tools/llubi/lib/ExecutorBase.cpp b/llvm/tools/llubi/lib/ExecutorBase.cpp
index 482930af81b5b..ded9ac798cd9c 100644
--- a/llvm/tools/llubi/lib/ExecutorBase.cpp
+++ b/llvm/tools/llubi/lib/ExecutorBase.cpp
@@ -15,9 +15,11 @@
 namespace llvm::ubi {
 Frame::Frame(Function &F, CallBase *CallSite, Frame *LastFrame,
              ArrayRef<AnyValue> Args, AnyValue &RetVal,
-             const TargetLibraryInfoImpl &TLIImpl)
+             const TargetLibraryInfoImpl &TLIImpl,
+             CapturedProvenanceList CapturedProvenances)
     : Func(F), LastFrame(LastFrame), CallSite(CallSite), Args(Args),
-      RetVal(RetVal), TLI(TLIImpl, &F) {
+      RetVal(RetVal), TLI(TLIImpl, &F),
+      CapturedProvenances(std::move(CapturedProvenances)) {
   assert((Args.size() == F.arg_size() ||
           (F.isVarArg() && Args.size() >= F.arg_size())) &&
          "Expected enough arguments to call the function.");
@@ -55,12 +57,19 @@ void ExecutorBase::reportErrorString(StringRef Msg) {
 
 std::pair<MemoryObject *, uint64_t>
 ExecutorBase::verifyMemAccess(const Pointer &Ptr, uint64_t AccessSize,
-                              Align Alignment, bool IsStore) {
-  auto *MO = Ctx.checkProvenance(Ptr, [](const Provenance &) {
-    // TODO: check provenance
-    // TODO: check inrange(S, E)
-    return true;
-  });
+                              Align Alignment, bool IsStore, bool IsVolatile) {
+  auto *MO =
+      Ctx.checkProvenance(Ptr, [IsStore, IsVolatile](const Provenance &Prov) {
+        CaptureComponents CC = Prov.capability();
+        if (IsVolatile &&
+            (CC & CaptureComponents::Address) != CaptureComponents::Address)
+          return false;
+        if (IsStore ? !capturesFullProvenance(CC) : !capturesAnyProvenance(CC))
+          return false;
+
+        // TODO: check inrange(S, E)
+        return true;
+      });
   if (!MO) {
     reportImmediateUB()
         << "Invalid memory access via a pointer with nullary provenance.";
@@ -110,7 +119,7 @@ ExecutorBase::verifyMemAccess(const Pointer &Ptr, uint64_t AccessSize,
 }
 
 AnyValue ExecutorBase::load(const AnyValue &Ptr, Align Alignment, Type *ValTy,
-                            bool NoUndef) {
+                            bool NoUndef, bool IsVolatile) {
   if (Ptr.isPoison()) {
     reportImmediateUB() << "Invalid memory access with a poison pointer.";
     return AnyValue::getPoisonValue(Ctx, ValTy);
@@ -118,7 +127,7 @@ AnyValue ExecutorBase::load(const AnyValue &Ptr, Align Alignment, Type *ValTy,
   auto &PtrVal = Ptr.asPointer();
   if (auto [MO, Offset] = verifyMemAccess(
           PtrVal, Ctx.getEffectiveTypeStoreSize(ValTy), Alignment,
-          /*IsStore=*/false);
+          /*IsStore=*/false, IsVolatile);
       MO) {
     // Load from a dead stack object yields poison value.
     if (MO->getState() == MemoryObjectState::Dead)
@@ -135,7 +144,7 @@ AnyValue ExecutorBase::load(const AnyValue &Ptr, Align Alignment, Type *ValTy,
 }
 
 void ExecutorBase::store(const AnyValue &Ptr, Align Alignment,
-                         const AnyValue &Val, Type *ValTy) {
+                         const AnyValue &Val, Type *ValTy, bool IsVolatile) {
   if (Ptr.isPoison()) {
     reportImmediateUB() << "Invalid memory access with a poison pointer.";
     return;
@@ -143,7 +152,7 @@ void ExecutorBase::store(const AnyValue &Ptr, Align Alignment,
   auto &PtrVal = Ptr.asPointer();
   if (auto [MO, Offset] = verifyMemAccess(
           PtrVal, Ctx.getEffectiveTypeStoreSize(ValTy), Alignment,
-          /*IsStore=*/true);
+          /*IsStore=*/true, IsVolatile);
       MO) {
     if (MO->isConstant()) {
       reportImmediateUB() << "Try to write to a constant memory object: "
diff --git a/llvm/tools/llubi/lib/ExecutorBase.h b/llvm/tools/llubi/lib/ExecutorBase.h
index d165192b55d63..91a5a18f076fe 100644
--- a/llvm/tools/llubi/lib/ExecutorBase.h
+++ b/llvm/tools/llubi/lib/ExecutorBase.h
@@ -60,6 +60,10 @@ struct Frame {
   SmallVector<IntrusiveRefCntPtr<MemoryObject>> Allocas;
   // Values of arguments and executed instructions in this function.
   DenseMap<Value *, AnyValue> ValueMap;
+  // The provenances to be updated after the function returns or unwinds.
+  using CapturedProvenanceList =
+      SmallVector<std::pair<IntrusiveRefCntPtr<Provenance>, CaptureInfo>>;
+  CapturedProvenanceList CapturedProvenances;
 
   // Reserved for in-flight subroutines.
   Function *ResolvedCallee = nullptr;
@@ -68,7 +72,8 @@ struct Frame {
 
   Frame(Function &F, CallBase *CallSite, Frame *LastFrame,
         ArrayRef<AnyValue> Args, AnyValue &RetVal,
-        const TargetLibraryInfoImpl &TLIImpl);
+        const TargetLibraryInfoImpl &TLIImpl,
+        CapturedProvenanceList CapturedProvenance);
 };
 
 enum class DiagnosticKind {
@@ -101,15 +106,14 @@ class ExecutorBase {
 
   /// Check if the upcoming memory access is valid. Returns the resolved memory
   /// object and offset if it is valid.
-  std::pair<MemoryObject *, uint64_t> verifyMemAccess(const Pointer &Ptr,
-                                                      uint64_t AccessSize,
-                                                      Align Alignment,
-                                                      bool IsStore);
+  std::pair<MemoryObject *, uint64_t>
+  verifyMemAccess(const Pointer &Ptr, uint64_t AccessSize, Align Alignment,
+                  bool IsStore, bool IsVolatile);
 
-  AnyValue load(const AnyValue &Ptr, Align Alignment, Type *ValTy,
-                bool NoUndef);
+  AnyValue load(const AnyValue &Ptr, Align Alignment, Type *ValTy, bool NoUndef,
+                bool IsVolatile);
   void store(const AnyValue &Ptr, Align Alignment, const AnyValue &Val,
-             Type *ValTy);
+             Type *ValTy, bool IsVolatile);
 
   void requestProgramExit(ProgramExitInfo::ProgramExitKind Kind,
                           uint64_t ExitCode = 0);
diff --git a/llvm/tools/llubi/lib/Interpreter.cpp b/llvm/tools/llubi/lib/Interpreter.cpp
index 5f7c48794245d..1b079ca64f0fa 100644
--- a/llvm/tools/llubi/lib/Interpreter.cpp
+++ b/llvm/tools/llubi/lib/Interpreter.cpp
@@ -94,8 +94,10 @@ static bool violatesDereferenceableBytesAttr(const AnyValue &V, uint64_t Bytes,
     return true;
   }
 
-  auto *MO = Ctx.checkProvenance(Ptr, [&](const Provenance &) {
-    // TODO: check read_provenance
+  auto *MO = Ctx.checkProvenance(Ptr, [&](const Provenance &Prov) {
+    CaptureComponents CC = Prov.capability();
+    if (!capturesAnyProvenance(CC))
+      return false;
     // TODO: check nofree for attributes/metadata.
     return true;
   });
@@ -703,7 +705,8 @@ class InstExecutor : public InstVisitor<InstExecutor, void>,
       : ExecutorBase(C, H), DL(Ctx.getDataLayout()),
         Lib(Ctx, Handler, DL, static_cast<ExecutorBase &>(*this)) {
     CallStack.emplace_back(F, /*CallSite=*/nullptr, /*LastFrame=*/nullptr, Args,
-                           RetVal, Ctx.getTLIImpl());
+                           RetVal, Ctx.getTLIImpl(),
+                           Frame::CapturedProvenanceList{});
   }
 
   void visitReturnInst(ReturnInst &RI) {
@@ -1709,6 +1712,7 @@ class InstExecutor : public InstVisitor<InstExecutor, void>,
 
     // Handle parameter attributes (Attributes from resolved callee should be
     // applied if available).
+    Frame::CapturedProvenanceList Captures;
     for (auto [I, Arg] : enumerate(CB.args())) {
       Type *ArgTy = Arg->getType();
       AnyValue &ArgVal = CalleeArgs[I];
@@ -1717,6 +1721,17 @@ class InstExecutor : public InstVisitor<InstExecutor, void>,
       AttributeSet AttrsAtCallSite = CB.getParamAttributes(I);
       AttributeSet AttrsAtCallee = Callee->getAttributes().getParamAttrs(I);
       handleAttributes(ArgTy, ArgVal, AttrsAtCallSite, AttrsAtCallee);
+
+      // Handle captures
+      if (ArgVal.isPointer() && !CB.isByValArgument(I)) {
+        CaptureInfo CI =
+            AttrsAtCallSite.getCaptureInfo() & AttrsAtCallee.getCaptureInfo();
+        if (CI != CaptureInfo::all()) {
+          auto NewProv = ArgVal.asPointer().provenance().clone();
+          Captures.emplace_back(NewProv, CI);
+          ArgVal = ArgVal.asPointer().getWithNewProvenance(std::move(NewProv));
+        }
+      }
     }
 
     CurrentFrame->ResolvedCallee = Callee;
@@ -1740,7 +1755,7 @@ class InstExecutor : public InstVisitor<InstExecutor, void>,
       AnyValue &RetVal = CurrentFrame->CalleeRetVal;
       CurrentFrame->State = FrameState::Pending;
       CallStack.emplace_back(*Callee, &CB, CurrentFrame, Args, RetVal,
-                             Ctx.getTLIImpl());
+                             Ctx.getTLIImpl(), std::move(Captures));
     }
   }
 
@@ -2107,9 +2122,37 @@ class InstExecutor : public InstVisitor<InstExecutor, void>,
     visitBinOp(I, [&](const AnyValue &LHS, const AnyValue &RHS) -> AnyValue {
       if (LHS.isPoison() || RHS.isPoison())
         return AnyValue::poison();
-      // TODO: handle pointer comparison.
-      const APInt &LHSVal = LHS.asInteger();
-      const APInt &RHSVal = RHS.asInteger();
+      if (LHS.isPointer()) {
+        auto &LHSPtr = LHS.asPointer();
+        auto &RHSPtr = RHS.asPointer();
+        unsigned AS = I.getOperand(0)->getType()->getPointerAddressSpace();
+        // Check provenance
+        bool LHSIsNull = LHSPtr.isNullPtr(AS, DL);
+        bool RHSIsNull = RHSPtr.isNullPtr(AS, DL);
+
+        auto IsValidCompare = [&](const Pointer &Ptr,
+                                  bool OtherIsNull) -> bool {
+          return Ctx.checkProvenance(
+              Ptr,
+              [&](const Provenance &Prov) {
+                CaptureComponents CC = Prov.capability();
+                CaptureComponents Mask = OtherIsNull
+                                             ? CaptureComponents::AddressIsNull
+                                             : CaptureComponents::Address;
+                return (CC & Mask) == Mask;
+              },
+              /*HasSideEffect=*/false);
+        };
+
+        if (!LHSIsNull && !IsValidCompare(LHSPtr, RHSIsNull))
+          return AnyValue::poison();
+        if (!RHSIsNull && !IsValidCompare(RHSPtr, LHSIsNull))
+          return AnyValue::poison();
+      }
+      const APInt &LHSVal =
+          LHS.isPointer() ? LHS.asPointer().address() : LHS.asInteger();
+      const APInt &RHSVal =
+          RHS.isPointer() ? RHS.asPointer().address() : RHS.asInteger();
       if (I.hasSameSign() && LHSVal.isNonNegative() != RHSVal.isNonNegative())
         return AnyValue::poison();
       return AnyValue::boolean(
@@ -2268,8 +2311,9 @@ class InstExecutor : public InstVisitor<InstExecutor, void>,
   }
 
   void visitLoadInst(LoadInst &LI) {
-    auto RetVal = load(getValue(LI.getPointerOperand()), LI.getAlign(),
-                       LI.getType(), LI.hasMetadata(LLVMContext::MD_noundef));
+    auto RetVal =
+        load(getValue(LI.getPointerOperand()), LI.getAlign(), LI.getType(),
+             LI.hasMetadata(LLVMContext::MD_noundef), LI.isVolatile());
     // TODO: track volatile loads
     handleMetadata(LI.getType(), RetVal, LI);
     setResult(LI, std::move(RetVal));
@@ -2280,7 +2324,8 @@ class InstExecutor : public InstVisitor<InstExecutor, void>,
     auto &Val = getValue(SI.getValueOperand());
     // TODO: track volatile stores
     // TODO: handle metadata
-    store(Ptr, SI.getAlign(), Val, SI.getValueOperand()->getType());
+    store(Ptr, SI.getAlign(), Val, SI.getValueOperand()->getType(),
+          SI.isVolatile());
     if (!hasProgramExited() && !Handler.onInstructionExecuted(SI, AnyValue()))
       setFailed();
   }
@@ -2421,6 +2466,23 @@ class InstExecutor : public InstVisitor<InstExecutor, void>,
       if (Top.State == FrameState::Exit) {
         assert((Top.Func.getReturnType()->isVoidTy() || !Top.RetVal.isNone()) &&
                "Expected return value to be set on function exit.");
+        // Apply captures attributes
+        for (auto &[Prov, CI] : Top.CapturedProvenances) {
+          if (Top.RetVal.isPointer()) {
+            auto &PtrVal = Top.RetVal.asPointer();
+            if (&PtrVal.provenance() == Prov.get()) {
+              CaptureComponents Capability = PtrVal.provenance().capability();
+              if ((Capability & CI.getRetComponents()) !=
+                  (Capability & CI.getOtherComponents())) {
+                // We have to create a copy of provenance.
+                auto NewProv = Prov->clone();
+                NewProv->captureCapability(CI.getRetComponents());
+                Top.RetVal = PtrVal.getWithNewProvenance(std::move(NewProv));
+              }
+            }
+          }
+          Prov->captureCapability(CI.getOtherComponents());
+        }
         Handler.onFunctionExit(Top.Func, Top.RetVal);
         // Free stack objects allocated in this frame.
         for (auto &Obj : Top.Allocas)
diff --git a/llvm/tools/llubi/lib/Library.cpp b/llvm/tools/llubi/lib/Library.cpp
index c3452b8af92c0..d224b2e02041b 100644
--- a/llvm/tools/llubi/lib/Library.cpp
+++ b/llvm/tools/llubi/lib/Library.cpp
@@ -35,7 +35,8 @@ std::optional<std::string> Library::readStringFromMemory(const Pointer &Ptr) {
 
   while (true) {
     auto [MO, ValidOffset] = Executor.verifyMemAccess(
-        Ptr.getWithNewAddr(Address + Offset), 1, Align(1), /*IsStore=*/false);
+        Ptr.getWithNewAddr(Address + Offset), 1, Align(1), /*IsStore=*/false,
+        /*IsVolatile=*/false);
     if (!MO)
       return std::nullopt;
 
@@ -278,7 +279,8 @@ AnyValue Library::executePrintf(ArrayRef<AnyValue> Args) {
     case 'n': {
       OS.flush();
       Executor.store(Arg, Align(4), AnyValue(APInt(32, Output.size())),
-                     Type::getInt32Ty(Ctx.getContext()));
+                     Type::getInt32Ty(Ctx.getContext()),
+                     /*IsVolatile=*/false);
       break;
     }
     case 'p': {
diff --git a/llvm/tools/llubi/lib/Value.cpp b/llvm/tools/llubi/lib/Value.cpp
index 3f1c2567cc51d..840716c341d0f 100644
--- a/llvm/tools/llubi/lib/Value.cpp
+++ b/llvm/tools/llubi/lib/Value.cpp
@@ -47,7 +47,9 @@ void Pointer::print(raw_ostream &OS) const {
   } else {
     OS << (Prov->isWildcard() ? "wildcard" : "nullary");
   }
-  // TODO: print provenance
+  CaptureComponents CC = Prov->capability();
+  if (CC != CaptureComponents::All)
+    OS << ' ' << CC;
   OS << "]";
 }
 
diff --git a/llvm/tools/llubi/lib/Value.h b/llvm/tools/llubi/lib/Value.h
index 77e428813a832..accea302c053d 100644
--- a/llvm/tools/llubi/lib/Value.h
+++ b/llvm/tools/llubi/lib/Value.h
@@ -14,6 +14,7 @@
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Type.h"
+#include "llvm/Support/ModRef.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace llvm::ubi {
@@ -176,8 +177,9 @@ class Provenance : public RefCountedBase<Provenance> {
   // Null if it is concrete.
   IntrusiveRefCntPtr<WildcardProvenance> Wildcard;
 
+  CaptureComponents Capability;
+
   // TODO: modeling nofree
-  // TODO: modeling captures
   // TODO: modeling inrange(Start, End) attribute
 
   const APInt &getTag() const { return Tag; }
@@ -186,10 +188,22 @@ class Provenance : public RefCountedBase<Provenance> {
   friend class Context;
 
 public:
-  Provenance(IntrusiveRefCntPtr<MemoryObject> Obj) : Obj(std::move(Obj)) {}
+  Provenance(IntrusiveRefCntPtr<MemoryObject> Obj,
+             CaptureComponents Capability = CaptureComponents::All)
+      : Obj(std::move(Obj)), Capability(Capability) {}
   static IntrusiveRefCntPtr<Provenance> nullary();
+  IntrusiveRefCntPtr<Provenance> clone() const {
+    IntrusiveRefCntPtr<Provenance> Res =
+        makeIntrusiveRefCnt<Provenance>(Obj, Capability);
+    Res->Wildcard = Wildcard;
+    return Res;
+  }
+  void captureCapability(CaptureComponents CapturedMask) {
+    Capability &= CapturedMask;
+  }
   IntrusiveRefCntPtr<Provenance> getWithKnownMemoryObject(MemoryObject &Obj);
   MemoryObject *getMemoryObject() const { return Obj.get(); }
+  CaptureComponents capability() const { return Capability; }
   bool isWildcard() const { return Wildcard != nullptr; }
 };
 



More information about the llvm-commits mailing list