// RUN: mlir-opt %s -split-input-file \
// RUN:    -async-to-async-runtime="eliminate-blocking-await-ops=true" \
// RUN:   | FileCheck %s --dump-input=always

// CHECK-LABEL: func @simple_callee
// CHECK-SAME: (%[[ARG:.*]]: f32)
// CHECK-SAME: -> (!async.token, !async.value<f32> {builtin.foo = "bar"})
func @simple_callee(%arg0: f32) -> (f32 {builtin.foo = "bar"}) {
// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token
// CHECK: %[[RETURNED_STORAGE:.*]] = async.runtime.create : !async.value<f32>
// CHECK: %[[ID:.*]] = async.coro.id
// CHECK: %[[HDL:.*]] = async.coro.begin %[[ID]]
// CHECK: br ^[[ORIGINAL_ENTRY:.*]]
// CHECK  ^[[ORIGINAL_ENTRY]]:
// CHECK:   %[[VAL:.*]] = arith.addf %[[ARG]], %[[ARG]] : f32
  %0 = arith.addf %arg0, %arg0 : f32
// CHECK:   %[[VAL_STORAGE:.*]] = async.runtime.create : !async.value<f32>
  %1 = async.runtime.create: !async.value<f32>
// CHECK:   async.runtime.store %[[VAL]], %[[VAL_STORAGE]] : !async.value<f32>
  async.runtime.store %0, %1: !async.value<f32>
// CHECK:   async.runtime.set_available %[[VAL_STORAGE]] : !async.value<f32>
  async.runtime.set_available %1: !async.value<f32>

// CHECK:   %[[SAVED:.*]] = async.coro.save %[[HDL]]
// CHECK:   async.runtime.await_and_resume %[[VAL_STORAGE]], %[[HDL]]
// CHECK:   async.coro.suspend %[[SAVED]]
// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME:.*]], ^[[CLEANUP:.*]]
  %2 = async.await %1 : !async.value<f32>

// CHECK: ^[[RESUME]]:
// CHECK:   %[[IS_ERROR:.*]] = async.runtime.is_error %[[VAL_STORAGE]] : !async.value<f32>
// CHECK:   cond_br %[[IS_ERROR]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_OK:.*]]

// CHECK: ^[[BRANCH_OK]]:
// CHECK:   %[[LOADED:.*]] = async.runtime.load %[[VAL_STORAGE]] : !async.value<f32>
// CHECK:   %[[RETURNED:.*]] = arith.mulf %[[ARG]], %[[LOADED]] : f32
// CHECK:   async.runtime.store %[[RETURNED]], %[[RETURNED_STORAGE]] : !async.value<f32>
// CHECK:   async.runtime.set_available %[[RETURNED_STORAGE]]
// CHECK:   async.runtime.set_available %[[TOKEN]]
// CHECK:   br ^[[CLEANUP]]
  %3 = arith.mulf %arg0, %2 : f32
  return %3: f32

// CHECK: ^[[BRANCH_ERROR]]:
// CHECK:   async.runtime.set_error %[[TOKEN]]
// CHECK:   async.runtime.set_error %[[RETURNED_STORAGE]]
// CHECK:   br ^[[CLEANUP]]


// CHECK: ^[[CLEANUP]]:
// CHECK:   async.coro.free %[[ID]], %[[HDL]]
// CHECK:   br ^[[SUSPEND]]

// CHECK: ^[[SUSPEND]]:
// CHECK:   async.coro.end %[[HDL]]
// CHECK:   return %[[TOKEN]], %[[RETURNED_STORAGE]] : !async.token, !async.value<f32>
}

// CHECK-LABEL: func @simple_caller()
// CHECK-SAME: -> (!async.token, !async.value<f32>)
func @simple_caller() -> f32 {
// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token
// CHECK: %[[RETURNED_STORAGE:.*]] = async.runtime.create : !async.value<f32>
// CHECK: %[[ID:.*]] = async.coro.id
// CHECK: %[[HDL:.*]] = async.coro.begin %[[ID]]
// CHECK: br ^[[ORIGINAL_ENTRY:.*]]
// CHECK  ^[[ORIGINAL_ENTRY]]:

// CHECK:   %[[CONSTANT:.*]] = arith.constant
  %c = arith.constant 1.0 : f32
// CHECK:   %[[RETURNED_TO_CALLER:.*]]:2 = call @simple_callee(%[[CONSTANT]]) : (f32) -> (!async.token, !async.value<f32>)
// CHECK:   %[[SAVED:.*]] = async.coro.save %[[HDL]]
// CHECK:   async.runtime.await_and_resume %[[RETURNED_TO_CALLER]]#0, %[[HDL]]
// CHECK:   async.coro.suspend %[[SAVED]]
// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME:.*]], ^[[CLEANUP:.*]]
  %r = call @simple_callee(%c): (f32) -> f32

// CHECK: ^[[RESUME]]:
// CHECK:   %[[IS_TOKEN_ERROR:.*]] = async.runtime.is_error %[[RETURNED_TO_CALLER]]#0 : !async.token
// CHECK:   cond_br %[[IS_TOKEN_ERROR]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_TOKEN_OK:.*]]

// CHECK: ^[[BRANCH_TOKEN_OK]]:
// CHECK:   %[[IS_VALUE_ERROR:.*]] = async.runtime.is_error %[[RETURNED_TO_CALLER]]#1 : !async.value<f32>
// CHECK:   cond_br %[[IS_VALUE_ERROR]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_VALUE_OK:.*]]

// CHECK: ^[[BRANCH_VALUE_OK]]:
// CHECK:   %[[LOADED:.*]] = async.runtime.load %[[RETURNED_TO_CALLER]]#1 : !async.value<f32>
// CHECK:   async.runtime.store %[[LOADED]], %[[RETURNED_STORAGE]] : !async.value<f32>
// CHECK:   async.runtime.set_available %[[RETURNED_STORAGE]]
// CHECK:   async.runtime.set_available %[[TOKEN]]
// CHECK:   br ^[[CLEANUP]]
  return %r: f32
// CHECK: ^[[BRANCH_ERROR]]:
// CHECK:   async.runtime.set_error %[[TOKEN]]
// CHECK:   async.runtime.set_error %[[RETURNED_STORAGE]]
// CHECK:   br ^[[CLEANUP]]


// CHECK: ^[[CLEANUP]]:
// CHECK:   async.coro.free %[[ID]], %[[HDL]]
// CHECK:   br ^[[SUSPEND]]

// CHECK: ^[[SUSPEND]]:
// CHECK:   async.coro.end %[[HDL]]
// CHECK:   return %[[TOKEN]], %[[RETURNED_STORAGE]] : !async.token, !async.value<f32>
}

// CHECK-LABEL: func @double_caller()
// CHECK-SAME: -> (!async.token, !async.value<f32>)
func @double_caller() -> f32 {
// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token
// CHECK: %[[RETURNED_STORAGE:.*]] = async.runtime.create : !async.value<f32>
// CHECK: %[[ID:.*]] = async.coro.id
// CHECK: %[[HDL:.*]] = async.coro.begin %[[ID]]
// CHECK: br ^[[ORIGINAL_ENTRY:.*]]
// CHECK  ^[[ORIGINAL_ENTRY]]:

// CHECK:   %[[CONSTANT:.*]] = arith.constant
  %c = arith.constant 1.0 : f32
// CHECK:   %[[RETURNED_TO_CALLER_1:.*]]:2 = call @simple_callee(%[[CONSTANT]]) : (f32) -> (!async.token, !async.value<f32>)
// CHECK:   %[[SAVED_1:.*]] = async.coro.save %[[HDL]]
// CHECK:   async.runtime.await_and_resume %[[RETURNED_TO_CALLER_1]]#0, %[[HDL]]
// CHECK:   async.coro.suspend %[[SAVED_1]]
// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_1:.*]], ^[[CLEANUP:.*]]
  %r = call @simple_callee(%c): (f32) -> f32

// CHECK: ^[[RESUME_1]]:
// CHECK:   %[[IS_TOKEN_ERROR_1:.*]] = async.runtime.is_error %[[RETURNED_TO_CALLER_1]]#0 : !async.token
// CHECK:   cond_br %[[IS_TOKEN_ERROR_1]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_TOKEN_OK_1:.*]]

// CHECK: ^[[BRANCH_TOKEN_OK_1]]:
// CHECK:   %[[IS_VALUE_ERROR_1:.*]] = async.runtime.is_error %[[RETURNED_TO_CALLER_1]]#1 : !async.value<f32>
// CHECK:   cond_br %[[IS_VALUE_ERROR_1]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_VALUE_OK_1:.*]]

// CHECK: ^[[BRANCH_VALUE_OK_1]]:
// CHECK:   %[[LOADED_1:.*]] = async.runtime.load %[[RETURNED_TO_CALLER_1]]#1 : !async.value<f32>
// CHECK:   %[[RETURNED_TO_CALLER_2:.*]]:2 = call @simple_callee(%[[LOADED_1]]) : (f32) -> (!async.token, !async.value<f32>)
// CHECK:   %[[SAVED_2:.*]] = async.coro.save %[[HDL]]
// CHECK:   async.runtime.await_and_resume %[[RETURNED_TO_CALLER_2]]#0, %[[HDL]]
// CHECK:   async.coro.suspend %[[SAVED_2]]
// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_2:.*]], ^[[CLEANUP:.*]]
  %s = call @simple_callee(%r): (f32) -> f32

// CHECK: ^[[RESUME_2]]:
// CHECK:   %[[IS_TOKEN_ERROR_2:.*]] = async.runtime.is_error %[[RETURNED_TO_CALLER_2]]#0 : !async.token
// CHECK:   cond_br %[[IS_TOKEN_ERROR_2]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_TOKEN_OK_2:.*]]

// CHECK: ^[[BRANCH_TOKEN_OK_2]]:
// CHECK:   %[[IS_VALUE_ERROR_2:.*]] = async.runtime.is_error %[[RETURNED_TO_CALLER_2]]#1 : !async.value<f32>
// CHECK:   cond_br %[[IS_VALUE_ERROR_2]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_VALUE_OK_2:.*]]

// CHECK: ^[[BRANCH_VALUE_OK_2]]:
// CHECK:   %[[LOADED_2:.*]] = async.runtime.load %[[RETURNED_TO_CALLER_2]]#1 : !async.value<f32>
// CHECK:   async.runtime.store %[[LOADED_2]], %[[RETURNED_STORAGE]] : !async.value<f32>
// CHECK:   async.runtime.set_available %[[RETURNED_STORAGE]]
// CHECK:   async.runtime.set_available %[[TOKEN]]
// CHECK:   br ^[[CLEANUP]]
  return %s: f32
// CHECK: ^[[BRANCH_ERROR]]:
// CHECK:   async.runtime.set_error %[[TOKEN]]
// CHECK:   async.runtime.set_error %[[RETURNED_STORAGE]]
// CHECK:   br ^[[CLEANUP]]

// CHECK: ^[[CLEANUP]]:
// CHECK:   async.coro.free %[[ID]], %[[HDL]]
// CHECK:   br ^[[SUSPEND]]

// CHECK: ^[[SUSPEND]]:
// CHECK:   async.coro.end %[[HDL]]
// CHECK:   return %[[TOKEN]], %[[RETURNED_STORAGE]] : !async.token, !async.value<f32>
}

// CHECK-LABEL: func @recursive
// CHECK-SAME: (%[[ARG:.*]]: !async.token) -> !async.token
func @recursive(%arg: !async.token) {
// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token
// CHECK: %[[ID:.*]] = async.coro.id
// CHECK: %[[HDL:.*]] = async.coro.begin %[[ID]]
// CHECK: %[[SAVED_1:.*]] = async.coro.save %[[HDL]]
// CHECK: async.runtime.await_and_resume %[[ARG]], %[[HDL]]
// CHECK: async.coro.suspend %[[SAVED_1]]
// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_1:.*]], ^[[CLEANUP:.*]]

  async.await %arg : !async.token
// CHECK: ^[[RESUME_1]]:
// CHECK:   %[[IS_ERROR:.*]] = async.runtime.is_error %[[ARG]] : !async.token
// CHECK:   cond_br %[[IS_ERROR]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_OK:.*]]

// CHECK: ^[[BRANCH_OK]]:
// CHECK:   %[[GIVEN:.*]] = async.runtime.create : !async.token
%r = async.runtime.create : !async.token
// CHECK:   async.runtime.set_available %[[GIVEN]]
async.runtime.set_available %r: !async.token
// CHECK:   %[[RETURNED_TO_CALLER:.*]] = call @recursive(%[[GIVEN]]) : (!async.token) -> !async.token
call @recursive(%r): (!async.token) -> ()
// CHECK: %[[SAVED_2:.*]] = async.coro.save %[[HDL]]
// CHECK: async.runtime.await_and_resume %[[RETURNED_TO_CALLER]], %[[HDL]]
// CHECK: async.coro.suspend %[[SAVED_2]]
// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_2:.*]], ^[[CLEANUP:.*]]

// CHECK: ^[[RESUME_2]]:
// CHECK:   async.runtime.set_available %[[TOKEN]]
// CHECK:   br ^[[CLEANUP]]

// CHECK: ^[[BRANCH_ERROR]]:
// CHECK:   async.runtime.set_error %[[TOKEN]]
// CHECK:   br ^[[CLEANUP]]
return

// CHECK: ^[[CLEANUP]]:
// CHECK:   async.coro.free %[[ID]], %[[HDL]]
// CHECK:   br ^[[SUSPEND]]

// CHECK: ^[[SUSPEND]]:
// CHECK:   async.coro.end %[[HDL]]
// CHECK:   return %[[TOKEN]] : !async.token
}

// CHECK-LABEL: func @corecursive1
// CHECK-SAME: (%[[ARG:.*]]: !async.token) -> !async.token
func @corecursive1(%arg: !async.token) {
// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token
// CHECK: %[[ID:.*]] = async.coro.id
// CHECK: %[[HDL:.*]] = async.coro.begin %[[ID]]
// CHECK: %[[SAVED_1:.*]] = async.coro.save %[[HDL]]
// CHECK: async.runtime.await_and_resume %[[ARG]], %[[HDL]]
// CHECK: async.coro.suspend %[[SAVED_1]]
// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_1:.*]], ^[[CLEANUP:.*]]

  async.await %arg : !async.token
// CHECK: ^[[RESUME_1]]:
// CHECK:   %[[IS_ERROR:.*]] = async.runtime.is_error %[[ARG]] : !async.token
// CHECK:   cond_br %[[IS_ERROR]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_OK:.*]]

// CHECK: ^[[BRANCH_OK]]:
// CHECK:   %[[GIVEN:.*]] = async.runtime.create : !async.token
%r = async.runtime.create : !async.token
// CHECK:   async.runtime.set_available %[[GIVEN]]
async.runtime.set_available %r: !async.token
// CHECK:   %[[RETURNED_TO_CALLER:.*]] = call @corecursive2(%[[GIVEN]]) : (!async.token) -> !async.token
call @corecursive2(%r): (!async.token) -> ()
// CHECK: %[[SAVED_2:.*]] = async.coro.save %[[HDL]]
// CHECK: async.runtime.await_and_resume %[[RETURNED_TO_CALLER]], %[[HDL]]
// CHECK: async.coro.suspend %[[SAVED_2]]
// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_2:.*]], ^[[CLEANUP:.*]]

// CHECK: ^[[RESUME_2]]:
// CHECK:   async.runtime.set_available %[[TOKEN]]
// CHECK:   br ^[[CLEANUP]]

// CHECK: ^[[BRANCH_ERROR]]:
// CHECK:   async.runtime.set_error %[[TOKEN]]
// CHECK:   br ^[[CLEANUP]]
return

// CHECK: ^[[CLEANUP]]:
// CHECK:   async.coro.free %[[ID]], %[[HDL]]
// CHECK:   br ^[[SUSPEND]]

// CHECK: ^[[SUSPEND]]:
// CHECK:   async.coro.end %[[HDL]]
// CHECK:   return %[[TOKEN]] : !async.token
}

// CHECK-LABEL: func @corecursive2
// CHECK-SAME: (%[[ARG:.*]]: !async.token) -> !async.token
func @corecursive2(%arg: !async.token) {
// CHECK: %[[TOKEN:.*]] = async.runtime.create : !async.token
// CHECK: %[[ID:.*]] = async.coro.id
// CHECK: %[[HDL:.*]] = async.coro.begin %[[ID]]
// CHECK: %[[SAVED_1:.*]] = async.coro.save %[[HDL]]
// CHECK: async.runtime.await_and_resume %[[ARG]], %[[HDL]]
// CHECK: async.coro.suspend %[[SAVED_1]]
// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_1:.*]], ^[[CLEANUP:.*]]

  async.await %arg : !async.token
// CHECK: ^[[RESUME_1]]:
// CHECK:   %[[IS_ERROR:.*]] = async.runtime.is_error %[[ARG]] : !async.token
// CHECK:   cond_br %[[IS_ERROR]], ^[[BRANCH_ERROR:.*]], ^[[BRANCH_OK:.*]]

// CHECK: ^[[BRANCH_OK]]:
// CHECK:   %[[GIVEN:.*]] = async.runtime.create : !async.token
%r = async.runtime.create : !async.token
// CHECK:   async.runtime.set_available %[[GIVEN]]
async.runtime.set_available %r: !async.token
// CHECK:   %[[RETURNED_TO_CALLER:.*]] = call @corecursive1(%[[GIVEN]]) : (!async.token) -> !async.token
call @corecursive1(%r): (!async.token) -> ()
// CHECK: %[[SAVED_2:.*]] = async.coro.save %[[HDL]]
// CHECK: async.runtime.await_and_resume %[[RETURNED_TO_CALLER]], %[[HDL]]
// CHECK: async.coro.suspend %[[SAVED_2]]
// CHECK-SAME: ^[[SUSPEND:.*]], ^[[RESUME_2:.*]], ^[[CLEANUP:.*]]

// CHECK: ^[[RESUME_2]]:
// CHECK:   async.runtime.set_available %[[TOKEN]]
// CHECK:   br ^[[CLEANUP]]

// CHECK: ^[[BRANCH_ERROR]]:
// CHECK:   async.runtime.set_error %[[TOKEN]]
// CHECK:   br ^[[CLEANUP]]
return

// CHECK: ^[[CLEANUP]]:
// CHECK:   async.coro.free %[[ID]], %[[HDL]]
// CHECK:   br ^[[SUSPEND]]

// CHECK: ^[[SUSPEND]]:
// CHECK:   async.coro.end %[[HDL]]
// CHECK:   return %[[TOKEN]] : !async.token
}

// CHECK-LABEL: func @caller_allowed_to_block
// CHECK-SAME: () -> f32
func @caller_allowed_to_block() -> f32 attributes { async.allowed_to_block } {
// CHECK: %[[CONSTANT:.*]] = arith.constant
  %c = arith.constant 1.0 : f32
// CHECK: %[[RETURNED_TO_CALLER:.*]]:2 = call @simple_callee(%[[CONSTANT]]) : (f32) -> (!async.token, !async.value<f32>)
// CHECK: async.runtime.await %[[RETURNED_TO_CALLER]]#0
// CHECK: async.runtime.await %[[RETURNED_TO_CALLER]]#1
// CHECK: %[[RETURNED:.*]] = async.runtime.load %[[RETURNED_TO_CALLER]]#1
  %r = call @simple_callee(%c): (f32) -> f32

// CHECK:   return %[[RETURNED]] : f32
  return %r: f32
}