1 files changed, 123 insertions, 0 deletions
diff --git a/llvm/test/Instrumentation/AllocToken/invoke.ll b/llvm/test/Instrumentation/AllocToken/invoke.ll
new file mode 100644
index 0000000..347c99a
--- /dev/null
+++ b/llvm/test/Instrumentation/AllocToken/invoke.ll
@@ -0,0 +1,123 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt < %s -passes=inferattrs,alloc-token -alloc-token-mode=increment -S | FileCheck %s
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+
+define ptr @test_invoke_malloc() sanitize_alloc_token personality ptr @__gxx_personality_v0 {
+; CHECK-LABEL: define ptr @test_invoke_malloc(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] personality ptr @__gxx_personality_v0 {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = invoke ptr @__alloc_token_malloc(i64 64, i64 0)
+; CHECK-NEXT:            to label %[[NORMAL:.*]] unwind label %[[CLEANUP:.*]]
+; CHECK:       [[NORMAL]]:
+; CHECK-NEXT:    ret ptr [[TMP0]]
+; CHECK:       [[CLEANUP]]:
+; CHECK-NEXT:    [[LP:%.*]] = landingpad { ptr, i32 }
+; CHECK-NEXT:            cleanup
+; CHECK-NEXT:    ret ptr null
+;
+entry:
+  %ptr = invoke ptr @malloc(i64 64) to label %normal unwind label %cleanup
+
+normal:
+  ret ptr %ptr
+
+cleanup:
+  %lp = landingpad { ptr, i32 } cleanup
+  ret ptr null
+}
+
+define ptr @test_invoke_operator_new() sanitize_alloc_token personality ptr @__gxx_personality_v0 {
+; CHECK-LABEL: define ptr @test_invoke_operator_new(
+; CHECK-SAME: ) #[[ATTR0]] personality ptr @__gxx_personality_v0 {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = invoke ptr @__alloc_token__Znwm(i64 32, i64 1)
+; CHECK-NEXT:            to label %[[NORMAL:.*]] unwind label %[[CLEANUP:.*]]
+; CHECK:       [[NORMAL]]:
+; CHECK-NEXT:    ret ptr [[TMP0]]
+; CHECK:       [[CLEANUP]]:
+; CHECK-NEXT:    [[LP:%.*]] = landingpad { ptr, i32 }
+; CHECK-NEXT:            cleanup
+; CHECK-NEXT:    ret ptr null
+;
+entry:
+  %ptr = invoke ptr @_Znwm(i64 32) to label %normal unwind label %cleanup
+
+normal:
+  ret ptr %ptr
+
+cleanup:
+  %lp = landingpad { ptr, i32 } cleanup
+  ret ptr null
+}
+
+; Test complex exception flow with multiple invoke allocations
+define ptr @test_complex_invoke_flow() sanitize_alloc_token personality ptr @__gxx_personality_v0 {
+; CHECK-LABEL: define ptr @test_complex_invoke_flow(
+; CHECK-SAME: ) #[[ATTR0]] personality ptr @__gxx_personality_v0 {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = invoke ptr @__alloc_token_malloc(i64 16, i64 2)
+; CHECK-NEXT:            to label %[[FIRST_OK:.*]] unwind label %[[CLEANUP1:.*]]
+; CHECK:       [[FIRST_OK]]:
+; CHECK-NEXT:    [[TMP1:%.*]] = invoke ptr @__alloc_token__Znwm(i64 32, i64 3)
+; CHECK-NEXT:            to label %[[SECOND_OK:.*]] unwind label %[[CLEANUP2:.*]]
+; CHECK:       [[SECOND_OK]]:
+; CHECK-NEXT:    ret ptr [[TMP0]]
+; CHECK:       [[CLEANUP1]]:
+; CHECK-NEXT:    [[LP1:%.*]] = landingpad { ptr, i32 }
+; CHECK-NEXT:            cleanup
+; CHECK-NEXT:    ret ptr null
+; CHECK:       [[CLEANUP2]]:
+; CHECK-NEXT:    [[LP2:%.*]] = landingpad { ptr, i32 }
+; CHECK-NEXT:            cleanup
+; CHECK-NEXT:    ret ptr null
+;
+entry:
+  %ptr1 = invoke ptr @malloc(i64 16) to label %first_ok unwind label %cleanup1
+
+first_ok:
+  %ptr2 = invoke ptr @_Znwm(i64 32) to label %second_ok unwind label %cleanup2
+
+second_ok:
+  ret ptr %ptr1
+
+cleanup1:
+  %lp1 = landingpad { ptr, i32 } cleanup
+  ret ptr null
+
+cleanup2:
+  %lp2 = landingpad { ptr, i32 } cleanup
+  ret ptr null
+}
+
+; Test mixed call/invoke
+define ptr @test_mixed_call_invoke() sanitize_alloc_token personality ptr @__gxx_personality_v0 {
+; CHECK-LABEL: define ptr @test_mixed_call_invoke(
+; CHECK-SAME: ) #[[ATTR0]] personality ptr @__gxx_personality_v0 {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = call ptr @__alloc_token_malloc(i64 8, i64 4)
+; CHECK-NEXT:    [[TMP1:%.*]] = invoke ptr @__alloc_token_malloc(i64 16, i64 5)
+; CHECK-NEXT:            to label %[[NORMAL:.*]] unwind label %[[CLEANUP:.*]]
+; CHECK:       [[NORMAL]]:
+; CHECK-NEXT:    ret ptr [[TMP0]]
+; CHECK:       [[CLEANUP]]:
+; CHECK-NEXT:    [[LP:%.*]] = landingpad { ptr, i32 }
+; CHECK-NEXT:            cleanup
+; CHECK-NEXT:    ret ptr null
+;
+entry:
+  %ptr1 = call ptr @malloc(i64 8)
+
+  %ptr2 = invoke ptr @malloc(i64 16) to label %normal unwind label %cleanup
+
+normal:
+  ret ptr %ptr1
+
+cleanup:
+  %lp = landingpad { ptr, i32 } cleanup
+  ret ptr null
+}
+
+declare ptr @malloc(i64)
+declare ptr @_Znwm(i64)
+declare i32 @__gxx_personality_v0(...)