; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
; RUN: opt < %s -passes='inferattrs,alloc-token<mode=increment>' -alloc-token-extended -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"

declare ptr @malloc(i64)
declare ptr @custom_malloc(i64)
declare ptr @kmalloc(i64, i64)

define ptr @test_libcall() sanitize_alloc_token {
; CHECK-LABEL: define ptr @test_libcall(
; CHECK-SAME: ) #[[ATTR1:[0-9]+]] {
; CHECK-NEXT:  [[ENTRY:.*:]]
; CHECK-NEXT:    [[TMP0:%.*]] = call ptr @__alloc_token_malloc(i64 64, i64 0)
; CHECK-NEXT:    ret ptr [[TMP0]]
;
entry:
  %ptr1 = call ptr @malloc(i64 64)
  ret ptr %ptr1
}

define ptr @test_libcall_hint() sanitize_alloc_token {
; CHECK-LABEL: define ptr @test_libcall_hint(
; CHECK-SAME: ) #[[ATTR1]] {
; CHECK-NEXT:  [[ENTRY:.*:]]
; CHECK-NEXT:    [[TMP0:%.*]] = call ptr @__alloc_token_malloc(i64 64, i64 1), !alloc_token [[META0:![0-9]+]]
; CHECK-NEXT:    ret ptr [[TMP0]]
;
entry:
  %ptr1 = call ptr @malloc(i64 64), !alloc_token !0
  ret ptr %ptr1
}

define ptr @test_nonlibcall_nohint() sanitize_alloc_token {
; CHECK-LABEL: define ptr @test_nonlibcall_nohint(
; CHECK-SAME: ) #[[ATTR1]] {
; CHECK-NEXT:  [[ENTRY:.*:]]
; CHECK-NEXT:    [[PTR1:%.*]] = call ptr @custom_malloc(i64 8)
; CHECK-NEXT:    [[PTR2:%.*]] = call ptr @kmalloc(i64 32, i64 0)
; CHECK-NEXT:    ret ptr [[PTR1]]
;
entry:
  %ptr1 = call ptr @custom_malloc(i64 8)
  %ptr2 = call ptr @kmalloc(i64 32, i64 0)
  ret ptr %ptr1
}

define ptr @test_nonlibcall_hint() sanitize_alloc_token {
; CHECK-LABEL: define ptr @test_nonlibcall_hint(
; CHECK-SAME: ) #[[ATTR1]] {
; CHECK-NEXT:  [[ENTRY:.*:]]
; CHECK-NEXT:    [[TMP0:%.*]] = call ptr @__alloc_token_custom_malloc(i64 8, i64 2), !alloc_token [[META0]]
; CHECK-NEXT:    [[TMP1:%.*]] = call ptr @__alloc_token_kmalloc(i64 32, i64 0, i64 3), !alloc_token [[META0]]
; CHECK-NEXT:    [[TMP2:%.*]] = call ptr @__alloc_token_custom_malloc(i64 64, i64 4), !alloc_token [[META0]]
; CHECK-NEXT:    [[TMP3:%.*]] = call ptr @__alloc_token_kmalloc(i64 128, i64 2, i64 5), !alloc_token [[META0]]
; CHECK-NEXT:    ret ptr [[TMP0]]
;
entry:
  %ptr1 = call ptr @custom_malloc(i64 8), !alloc_token !0
  %ptr2 = call ptr @kmalloc(i64 32, i64 0), !alloc_token !0
  %ptr3 = call ptr @custom_malloc(i64 64), !alloc_token !0
  %ptr4 = call ptr @kmalloc(i64 128, i64 2), !alloc_token !0
  ret ptr %ptr1
}

; Functions without sanitize_alloc_token do not get instrumented
define ptr @without_attribute() {
; CHECK-LABEL: define ptr @without_attribute() {
; CHECK-NEXT:  [[ENTRY:.*:]]
; CHECK-NEXT:    [[PTR1:%.*]] = call ptr @malloc(i64 64), !alloc_token [[META0]]
; CHECK-NEXT:    [[PTR2:%.*]] = call ptr @custom_malloc(i64 8), !alloc_token [[META0]]
; CHECK-NEXT:    [[PTR3:%.*]] = call ptr @kmalloc(i64 32, i64 0), !alloc_token [[META0]]
; CHECK-NEXT:    ret ptr [[PTR1]]
;
entry:
  %ptr1 = call ptr @malloc(i64 64), !alloc_token !0
  %ptr2 = call ptr @custom_malloc(i64 8), !alloc_token !0
  %ptr3 = call ptr @kmalloc(i64 32, i64 0), !alloc_token !0
  ret ptr %ptr1
}

!0 = !{!"int", i1 0}
;.
; CHECK: [[META0]] = !{!"int", i1 false}
;.