// RUN: mlir-opt -allow-unregistered-dialect -verify-diagnostics -ownership-based-buffer-deallocation \
// RUN:  --buffer-deallocation-simplification -split-input-file %s | FileCheck %s
// RUN: mlir-opt -allow-unregistered-dialect -verify-diagnostics -ownership-based-buffer-deallocation=private-function-dynamic-ownership=true -split-input-file %s > /dev/null

// RUN: mlir-opt %s -buffer-deallocation-pipeline --split-input-file --verify-diagnostics > /dev/null

// Test Case: Nested regions - This test defines a BufferBasedOp inside the
// region of a RegionBufferBasedOp.
// BufferDeallocation expected behavior: The AllocOp for the BufferBasedOp
// should remain inside the region of the RegionBufferBasedOp and it should insert
// the missing DeallocOp in the same region. The missing DeallocOp should be
// inserted after CopyOp.

func.func @nested_regions_and_cond_branch(
  %arg0: i1,
  %arg1: memref<2xf32>,
  %arg2: memref<2xf32>) {
  cf.cond_br %arg0, ^bb1, ^bb2
^bb1:
  cf.br ^bb3(%arg1 : memref<2xf32>)
^bb2:
  %0 = memref.alloc() : memref<2xf32>
  test.region_buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %1 = memref.alloc() : memref<2xf32>
    test.buffer_based in(%arg1: memref<2xf32>) out(%1: memref<2xf32>)
    %tmp1 = math.exp %gen1_arg0 : f32
    test.region_yield %tmp1 : f32
  }
  cf.br ^bb3(%0 : memref<2xf32>)
^bb3(%1: memref<2xf32>):
  test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
  return
}

// CHECK-LABEL: func @nested_regions_and_cond_branch
//  CHECK-SAME: ([[ARG0:%.+]]: i1, [[ARG1:%.+]]: memref<2xf32>, [[ARG2:%.+]]: memref<2xf32>)
//       CHECK: ^bb1:
//   CHECK-NOT:   bufferization.clone
//   CHECK-NOT:   bufferization.dealloc
//       CHECK:   cf.br ^bb3([[ARG1]], %false
//       CHECK: ^bb2:
//       CHECK:   [[ALLOC0:%.+]] = memref.alloc()
//       CHECK:   test.region_buffer_based
//       CHECK:     [[ALLOC1:%.+]] = memref.alloc()
//       CHECK:     test.buffer_based
//       CHECK:     bufferization.dealloc ([[ALLOC1]] : memref<2xf32>) if (%true
//  CHECK-NEXT:     test.region_yield
//   CHECK-NOT:   bufferization.clone
//   CHECK-NOT:   bufferization.dealloc
//       CHECK:   cf.br ^bb3([[ALLOC0]], %true
//       CHECK: ^bb3([[A0:%.+]]: memref<2xf32>, [[COND0:%.+]]: i1):
//       CHECK:   test.copy
//  CHECK-NEXT:   [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[A0]]
//  CHECK-NEXT:   bufferization.dealloc ([[BASE]] : {{.*}}) if ([[COND0]])
//       CHECK:   return

// -----

// Test Case: nested region control flow
// The alloc %1 flows through both if branches until it is finally returned.
// Hence, it does not require a specific dealloc operation. However, %3
// requires a dealloc.

func.func @nested_region_control_flow(
  %arg0 : index,
  %arg1 : index) -> memref<?x?xf32> {
  %0 = arith.cmpi eq, %arg0, %arg1 : index
  %1 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
  %2 = scf.if %0 -> (memref<?x?xf32>) {
    scf.yield %1 : memref<?x?xf32>
  } else {
    %3 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
    "test.read_buffer"(%3) : (memref<?x?xf32>) -> ()
    scf.yield %1 : memref<?x?xf32>
  }
  return %2 : memref<?x?xf32>
}

// CHECK-LABEL: func @nested_region_control_flow
//       CHECK:   [[ALLOC:%.+]] = memref.alloc(
//       CHECK:   [[V0:%.+]]:2 = scf.if
//       CHECK:     scf.yield [[ALLOC]], %false
//       CHECK:     [[ALLOC1:%.+]] = memref.alloc(
//       CHECK:     bufferization.dealloc ([[ALLOC1]] :{{.*}}) if (%true{{[0-9_]*}})
//   CHECK-NOT: retain
//       CHECK:     scf.yield [[ALLOC]], %false
//       CHECK:   [[V1:%.+]] = scf.if [[V0]]#1
//       CHECK:     scf.yield [[V0]]#0
//       CHECK:     [[CLONE:%.+]] = bufferization.clone [[V0]]#0
//       CHECK:     scf.yield [[CLONE]]
//       CHECK:   [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
//       CHECK:   bufferization.dealloc ([[ALLOC]], [[BASE]] : {{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V1]] :
//       CHECK:   return [[V1]]

// -----

// Test Case: nested region control flow with a nested buffer allocation in a
// divergent branch.
// Buffer deallocation places a copy for both  %1 and %3, since they are
// returned in the end.

func.func @nested_region_control_flow_div(
  %arg0 : index,
  %arg1 : index) -> memref<?x?xf32> {
  %0 = arith.cmpi eq, %arg0, %arg1 : index
  %1 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
  %2 = scf.if %0 -> (memref<?x?xf32>) {
    scf.yield %1 : memref<?x?xf32>
  } else {
    %3 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
    scf.yield %3 : memref<?x?xf32>
  }
  return %2 : memref<?x?xf32>
}

// CHECK-LABEL: func @nested_region_control_flow_div
//       CHECK:   [[ALLOC:%.+]] = memref.alloc(
//       CHECK:   [[V0:%.+]]:2 = scf.if
//       CHECK:     scf.yield [[ALLOC]], %false
//       CHECK:     [[ALLOC1:%.+]] = memref.alloc(
//       CHECK:     scf.yield [[ALLOC1]], %true
//       CHECK:   [[V1:%.+]] = scf.if [[V0]]#1
//       CHECK:     scf.yield [[V0]]#0
//       CHECK:     [[CLONE:%.+]] = bufferization.clone [[V0]]#0
//       CHECK:     scf.yield [[CLONE]]
//       CHECK:   [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
//       CHECK:   bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V1]] :
//       CHECK:   return [[V1]]

// -----

// Test Case: nested region control flow within a region interface.
// No copies are required in this case since the allocation finally escapes
// the method.

func.func @inner_region_control_flow(%arg0 : index) -> memref<?x?xf32> {
  %0 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
  %1 = test.region_if %0 : memref<?x?xf32> -> (memref<?x?xf32>) then {
    ^bb0(%arg1 : memref<?x?xf32>):
      test.region_if_yield %arg1 : memref<?x?xf32>
  } else {
    ^bb0(%arg1 : memref<?x?xf32>):
      test.region_if_yield %arg1 : memref<?x?xf32>
  } join {
    ^bb0(%arg1 : memref<?x?xf32>):
      test.region_if_yield %arg1 : memref<?x?xf32>
  }
  return %1 : memref<?x?xf32>
}

// CHECK-LABEL: func.func @inner_region_control_flow
//       CHECK:   [[ALLOC:%.+]] = memref.alloc(
//       CHECK:   [[V0:%.+]]:2 = test.region_if [[ALLOC]], %false
//       CHECK:   ^bb0([[ARG1:%.+]]: memref<?x?xf32>, [[ARG2:%.+]]: i1):
//       CHECK:     test.region_if_yield [[ARG1]], [[ARG2]]
//       CHECK:   ^bb0([[ARG1:%.+]]: memref<?x?xf32>, [[ARG2:%.+]]: i1):
//       CHECK:     test.region_if_yield [[ARG1]], [[ARG2]]
//       CHECK:   ^bb0([[ARG1:%.+]]: memref<?x?xf32>, [[ARG2:%.+]]: i1):
//       CHECK:     test.region_if_yield [[ARG1]], [[ARG2]]
//       CHECK:   [[V1:%.+]] = scf.if [[V0]]#1
//       CHECK:     scf.yield [[V0]]#0
//       CHECK:     [[CLONE:%.+]] = bufferization.clone [[V0]]#0
//       CHECK:     scf.yield [[CLONE]]
//       CHECK:   [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
//       CHECK:   bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V1]] :
//       CHECK:   return [[V1]]

// -----

func.func @nestedRegionsAndCondBranchAlloca(
  %arg0: i1,
  %arg1: memref<2xf32>,
  %arg2: memref<2xf32>) {
  cf.cond_br %arg0, ^bb1, ^bb2
^bb1:
  cf.br ^bb3(%arg1 : memref<2xf32>)
^bb2:
  %0 = memref.alloc() : memref<2xf32>
  test.region_buffer_based in(%arg1: memref<2xf32>) out(%0: memref<2xf32>) {
  ^bb0(%gen1_arg0: f32, %gen1_arg1: f32):
    %1 = memref.alloca() : memref<2xf32>
    test.buffer_based in(%arg1: memref<2xf32>) out(%1: memref<2xf32>)
    %tmp1 = math.exp %gen1_arg0 : f32
    test.region_yield %tmp1 : f32
  }
  cf.br ^bb3(%0 : memref<2xf32>)
^bb3(%1: memref<2xf32>):
  test.copy(%1, %arg2) : (memref<2xf32>, memref<2xf32>)
  return
}

// CHECK-LABEL: func @nestedRegionsAndCondBranchAlloca
//  CHECK-SAME: ([[ARG0:%.+]]: i1, [[ARG1:%.+]]: memref<2xf32>, [[ARG2:%.+]]: memref<2xf32>)
//       CHECK: ^bb1:
//       CHECK:   cf.br ^bb3([[ARG1]], %false
//       CHECK: ^bb2:
//       CHECK:   [[ALLOC:%.+]] = memref.alloc()
//       CHECK:   test.region_buffer_based
//       CHECK:     memref.alloca()
//       CHECK:     test.buffer_based
//   CHECK-NOT:     bufferization.dealloc
//   CHECK-NOT:     bufferization.clone
//       CHECK:     test.region_yield
//       CHECK:   }
//       CHECK:   cf.br ^bb3([[ALLOC]], %true
//       CHECK: ^bb3([[A0:%.+]]: memref<2xf32>, [[COND:%.+]]: i1):
//       CHECK:   test.copy
//       CHECK:   [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[A0]]
//       CHECK:   bufferization.dealloc ([[BASE]] :{{.*}}) if ([[COND]])

// -----

func.func @nestedRegionControlFlowAlloca(
  %arg0 : index, %arg1 : index, %arg2: f32) -> memref<?x?xf32> {
  %0 = arith.cmpi eq, %arg0, %arg1 : index
  %1 = memref.alloc(%arg0, %arg0) : memref<?x?xf32>
  %2 = scf.if %0 -> (memref<?x?xf32>) {
    scf.yield %1 : memref<?x?xf32>
  } else {
    %3 = memref.alloca(%arg0, %arg1) : memref<?x?xf32>
    %c0 = arith.constant 0 : index
    memref.store %arg2, %3[%c0, %c0] : memref<?x?xf32>
    scf.yield %1 : memref<?x?xf32>
  }
  return %2 : memref<?x?xf32>
}

// CHECK-LABEL: func @nestedRegionControlFlowAlloca
//       CHECK: [[ALLOC:%.+]] = memref.alloc(
//       CHECK: [[V0:%.+]]:2 = scf.if
//       CHECK:   scf.yield [[ALLOC]], %false
//       CHECK:   memref.alloca(
//       CHECK:   scf.yield [[ALLOC]], %false
//       CHECK: [[V1:%.+]] = scf.if [[V0]]#1
//       CHECK:   scf.yield [[V0]]#0
//       CHECK:   [[CLONE:%.+]] = bufferization.clone [[V0]]#0
//       CHECK:   scf.yield [[CLONE]]
//       CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
//       CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V1]] :
//       CHECK: return [[V1]]

// -----

// Test Case: structured control-flow loop using a nested alloc.
// The iteration argument %iterBuf has to be freed before yielding %3 to avoid
// memory leaks.

func.func @loop_alloc(
  %lb: index,
  %ub: index,
  %step: index,
  %buf: memref<2xf32>,
  %res: memref<2xf32>) {
  %0 = memref.alloc() : memref<2xf32>
  "test.read_buffer"(%0) : (memref<2xf32>) -> ()
  %1 = scf.for %i = %lb to %ub step %step
    iter_args(%iterBuf = %buf) -> memref<2xf32> {
    %2 = arith.cmpi eq, %i, %ub : index
    %3 = memref.alloc() : memref<2xf32>
    scf.yield %3 : memref<2xf32>
  }
  test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
  return
}

// CHECK-LABEL: func @loop_alloc
//  CHECK-SAME: ([[ARG0:%.+]]: index, [[ARG1:%.+]]: index, [[ARG2:%.+]]: index, [[ARG3:%.+]]: memref<2xf32>, [[ARG4:%.+]]: memref<2xf32>)
//       CHECK: [[ALLOC:%.+]] = memref.alloc()
//       CHECK: [[V0:%.+]]:2 = scf.for {{.*}} iter_args([[ARG6:%.+]] = [[ARG3]], [[ARG7:%.+]] = %false
//       CHECK:   [[ALLOC1:%.+]] = memref.alloc()
//       CHECK:   [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[ARG6]]
//       CHECK:   bufferization.dealloc ([[BASE]] :{{.*}}) if ([[ARG7]])
//   CHECK-NOT:       retain
//       CHECK:   scf.yield [[ALLOC1]], %true
//       CHECK: test.copy
//       CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
//       CHECK: bufferization.dealloc ([[ALLOC]] :{{.*}}) if (%true
//   CHECK-NOT: retain
//       CHECK: bufferization.dealloc ([[BASE]] :{{.*}}) if ([[V0]]#1)
//   CHECK-NOT: retain

// -----

// Test Case: structured control-flow loop with a nested if operation.
// The loop yields buffers that have been defined outside of the loop and the
// backedges only use the iteration arguments (or one of its aliases).
// Therefore, we do not have to (and are not allowed to) free any buffers
// that are passed via the backedges.

func.func @loop_nested_if_no_alloc(
  %lb: index,
  %ub: index,
  %step: index,
  %buf: memref<2xf32>,
  %res: memref<2xf32>) {
  %0 = memref.alloc() : memref<2xf32>
  %1 = scf.for %i = %lb to %ub step %step
    iter_args(%iterBuf = %buf) -> memref<2xf32> {
    %2 = arith.cmpi eq, %i, %ub : index
    %3 = scf.if %2 -> (memref<2xf32>) {
      scf.yield %0 : memref<2xf32>
    } else {
      scf.yield %iterBuf : memref<2xf32>
    }
    scf.yield %3 : memref<2xf32>
  }
  test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
  return
}

// CHECK-LABEL: func @loop_nested_if_no_alloc
//  CHECK-SAME: ({{.*}}, [[ARG3:%.+]]: memref<2xf32>, [[ARG4:%.+]]: memref<2xf32>)
//       CHECK: [[ALLOC:%.+]] = memref.alloc()
//       CHECK: [[V0:%.+]]:2 = scf.for {{.*}} iter_args([[ARG6:%.+]] = [[ARG3]], [[ARG7:%.+]] = %false
//       CHECK:   [[V1:%.+]]:2 = scf.if
//       CHECK:     scf.yield [[ALLOC]], %false
//       CHECK:     scf.yield [[ARG6]], %false
//       CHECK:   [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[ARG6]]
//       CHECK:   [[OWN:%.+]] = bufferization.dealloc ([[BASE]] :{{.*}}) if ([[ARG7]]) retain ([[V1]]#0 :
//       CHECK:   [[OWN_AGG:%.+]] = arith.ori [[OWN]], [[V1]]#1
//       CHECK:   scf.yield [[V1]]#0, [[OWN_AGG]]
//       CHECK: test.copy
//       CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
//       CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1)

// TODO: we know statically that the inner dealloc will never deallocate
//       anything, i.e., we can optimize it away

// -----

// Test Case: structured control-flow loop with a nested if operation using
// a deeply nested buffer allocation.

func.func @loop_nested_if_alloc(
  %lb: index,
  %ub: index,
  %step: index,
  %buf: memref<2xf32>) -> memref<2xf32> {
  %0 = memref.alloc() : memref<2xf32>
  %1 = scf.for %i = %lb to %ub step %step
    iter_args(%iterBuf = %buf) -> memref<2xf32> {
    %2 = arith.cmpi eq, %i, %ub : index
    %3 = scf.if %2 -> (memref<2xf32>) {
      %4 = memref.alloc() : memref<2xf32>
      scf.yield %4 : memref<2xf32>
    } else {
      scf.yield %0 : memref<2xf32>
    }
    scf.yield %3 : memref<2xf32>
  }
  return %1 : memref<2xf32>
}

// CHECK-LABEL: func @loop_nested_if_alloc
//  CHECK-SAME: ({{.*}}, [[ARG3:%.+]]: memref<2xf32>)
//       CHECK: [[ALLOC:%.+]] = memref.alloc()
//       CHECK: [[V0:%.+]]:2 = scf.for {{.*}} iter_args([[ARG5:%.+]] = [[ARG3]], [[ARG6:%.+]] = %false
//       CHECK:   [[V1:%.+]]:2 = scf.if
//       CHECK:     [[ALLOC1:%.+]] = memref.alloc()
//       CHECK:     scf.yield [[ALLOC1]], %true
//       CHECK:     scf.yield [[ALLOC]], %false
//       CHECK:   [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[ARG5]]
//       CHECK:   [[OWN:%.+]] = bufferization.dealloc ([[BASE]] :{{.*}}) if ([[ARG6]]) retain ([[V1]]#0 :
//       CHECK:   [[OWN_AGG:%.+]] = arith.ori [[OWN]], [[V1]]#1
//       CHECK:   scf.yield [[V1]]#0, [[OWN_AGG]]
//       CHECK: }
//       CHECK: [[V2:%.+]] = scf.if [[V0]]#1
//       CHECK:   scf.yield [[V0]]#0
//       CHECK:   [[CLONE:%.+]] = bufferization.clone [[V0]]#0
//       CHECK:   scf.yield [[CLONE]]
//       CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
//       CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V2]] :
//       CHECK: return [[V2]]

// -----

// Test Case: several nested structured control-flow loops with a deeply nested
// buffer allocation inside an if operation.

func.func @loop_nested_alloc(
  %lb: index,
  %ub: index,
  %step: index,
  %buf: memref<2xf32>,
  %res: memref<2xf32>) {
  %0 = memref.alloc() : memref<2xf32>
  "test.read_buffer"(%0) : (memref<2xf32>) -> ()
  %1 = scf.for %i = %lb to %ub step %step
    iter_args(%iterBuf = %buf) -> memref<2xf32> {
    %2 = scf.for %i2 = %lb to %ub step %step
      iter_args(%iterBuf2 = %iterBuf) -> memref<2xf32> {
      %3 = scf.for %i3 = %lb to %ub step %step
        iter_args(%iterBuf3 = %iterBuf2) -> memref<2xf32> {
        %4 = memref.alloc() : memref<2xf32>
        "test.read_buffer"(%4) : (memref<2xf32>) -> ()
        %5 = arith.cmpi eq, %i, %ub : index
        %6 = scf.if %5 -> (memref<2xf32>) {
          %7 = memref.alloc() : memref<2xf32>
          scf.yield %7 : memref<2xf32>
        } else {
          scf.yield %iterBuf3 : memref<2xf32>
        }
        scf.yield %6 : memref<2xf32>
      }
      scf.yield %3 : memref<2xf32>
    }
    scf.yield %2 : memref<2xf32>
  }
  test.copy(%1, %res) : (memref<2xf32>, memref<2xf32>)
  return
}

// CHECK-LABEL: func @loop_nested_alloc
//       CHECK: ({{.*}}, [[ARG3:%.+]]: memref<2xf32>, {{.*}}: memref<2xf32>)
//       CHECK: [[ALLOC:%.+]] = memref.alloc()
//       CHECK: [[V0:%.+]]:2 = scf.for {{.*}} iter_args([[ARG6:%.+]] = [[ARG3]], [[ARG7:%.+]] = %false
//       CHECK:   [[V1:%.+]]:2 = scf.for {{.*}} iter_args([[ARG9:%.+]] = [[ARG6]], [[ARG10:%.+]] = %false
//       CHECK:     [[V2:%.+]]:2 = scf.for {{.*}} iter_args([[ARG12:%.+]] = [[ARG9]], [[ARG13:%.+]] = %false
//       CHECK:       [[ALLOC1:%.+]] = memref.alloc()
//       CHECK:       [[V3:%.+]]:2 = scf.if
//       CHECK:         [[ALLOC2:%.+]] = memref.alloc()
//       CHECK:         scf.yield [[ALLOC2]], %true
//       CHECK:       } else {
//       CHECK:         scf.yield [[ARG12]], %false
//       CHECK:       }
//       CHECK:       [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[ARG12]]
//       CHECK:       [[OWN:%.+]] = bufferization.dealloc ([[BASE]] :{{.*}}) if ([[ARG13]]) retain ([[V3]]#0 :
//       CHECK:       bufferization.dealloc ([[ALLOC1]] :{{.*}}) if (%true{{[0-9_]*}})
//   CHECK-NOT: retain
//       CHECK:       [[OWN_AGG:%.+]] = arith.ori [[OWN]], [[V3]]#1
//       CHECK:       scf.yield [[V3]]#0, [[OWN_AGG]]
//       CHECK:     }
//       CHECK:     [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[ARG9]]
//       CHECK:     [[OWN:%.+]] = bufferization.dealloc ([[BASE]] :{{.*}}) if ([[ARG10]]) retain ([[V2]]#0 :
//       CHECK:     [[OWN_AGG:%.+]] = arith.ori [[OWN]], [[V2]]#1
//       CHECK:     scf.yield [[V2]]#0, [[OWN_AGG]]
//       CHECK:   }
//       CHECK:   [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[ARG6]]
//       CHECK:   [[OWN:%.+]] = bufferization.dealloc ([[BASE]] :{{.*}}) if ([[ARG7]]) retain ([[V1]]#0 :
//       CHECK:   [[OWN_AGG:%.+]] = arith.ori [[OWN]], [[V1]]#1
//       CHECK:   scf.yield [[V1]]#0, [[OWN_AGG]]
//       CHECK: }
//       CHECK: test.copy
//       CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
//       CHECK: bufferization.dealloc ([[ALLOC]] :{{.*}}) if (%true
//       CHECK: bufferization.dealloc ([[BASE]] :{{.*}}) if ([[V0]]#1)

// TODO: all the retain operands could be removed by doing some more thorough analysis

// -----

func.func @affine_loop() -> f32 {
  %buffer = memref.alloc() : memref<1024xf32>
  %sum_init_0 = arith.constant 0.0 : f32
  %res = affine.for %i = 0 to 10 step 2 iter_args(%sum_iter = %sum_init_0) -> f32 {
    %t = affine.load %buffer[%i] : memref<1024xf32>
    %sum_next = arith.addf %sum_iter, %t : f32
    affine.yield %sum_next : f32
  }
  return %res : f32
}

// CHECK-LABEL: func @affine_loop
//       CHECK: [[ALLOC:%.+]] = memref.alloc()
//       CHECK: affine.for {{.*}} iter_args(%arg1 = %cst)
//       CHECK:   affine.yield
//       CHECK: bufferization.dealloc ([[ALLOC]] :{{.*}}) if (%true

// -----

func.func @assumingOp(
  %arg0: !shape.witness,
  %arg2: memref<2xf32>,
  %arg3: memref<2xf32>) {
  // Confirm the alloc will be dealloc'ed in the block.
  %1 = shape.assuming %arg0 -> memref<2xf32> {
    %0 = memref.alloc() : memref<2xf32>
    "test.read_buffer"(%0) : (memref<2xf32>) -> ()
    shape.assuming_yield %arg2 : memref<2xf32>
  }
  // Confirm the alloc will be returned and dealloc'ed after its use.
  %3 = shape.assuming %arg0 -> memref<2xf32> {
    %2 = memref.alloc() : memref<2xf32>
    shape.assuming_yield %2 : memref<2xf32>
  }
  test.copy(%3, %arg3) : (memref<2xf32>, memref<2xf32>)
  return
}

// CHECK-LABEL: func @assumingOp
//       CHECK: ({{.*}}, [[ARG1:%.+]]: memref<2xf32>, {{.*}}: memref<2xf32>)
//       CHECK: [[V0:%.+]]:2 = shape.assuming
//       CHECK:   [[ALLOC:%.+]] = memref.alloc()
//       CHECK:   bufferization.dealloc ([[ALLOC]] :{{.*}}) if (%true{{[0-9_]*}})
//   CHECK-NOT: retain
//       CHECK:   shape.assuming_yield [[ARG1]], %false
//       CHECK: }
//       CHECK: [[V1:%.+]]:2 = shape.assuming
//       CHECK:   [[ALLOC:%.+]] = memref.alloc()
//       CHECK:   shape.assuming_yield [[ALLOC]], %true
//       CHECK: }
//       CHECK: test.copy
//       CHECK: [[BASE0:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
//       CHECK: [[BASE1:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V1]]#0
//       CHECK: bufferization.dealloc ([[BASE1]] :{{.*}}) if ([[V1]]#1)
//   CHECK-NOT: retain
//       CHECK: bufferization.dealloc ([[BASE0]] :{{.*}}) if ([[V0]]#1)
//   CHECK-NOT: retain
//       CHECK: return

// -----

// Test Case: The op "test.one_region_with_recursive_memory_effects" does not
// implement the RegionBranchOpInterface. This is allowed during buffer
// deallocation because the operation's region does not deal with any MemRef
// values.

func.func @noRegionBranchOpInterface() {
  %0 = "test.one_region_with_recursive_memory_effects"() ({
    %1 = "test.one_region_with_recursive_memory_effects"() ({
      %2 = memref.alloc() : memref<2xi32>
      "test.read_buffer"(%2) : (memref<2xi32>) -> ()
      "test.return"() : () -> ()
    }) : () -> (i32)
    "test.return"() : () -> ()
  }) : () -> (i32)
  "test.return"() : () -> ()
}

// -----

// Test Case: The second op "test.one_region_with_recursive_memory_effects" does
// not implement the RegionBranchOpInterface but has buffer semantics. This is
// not allowed during buffer deallocation.

func.func @noRegionBranchOpInterface() {
  %0 = "test.one_region_with_recursive_memory_effects"() ({
    // expected-error@+1 {{All operations with attached regions need to implement the RegionBranchOpInterface.}}
    %1 = "test.one_region_with_recursive_memory_effects"() ({
      %2 = memref.alloc() : memref<2xi32>
      "test.read_buffer"(%2) : (memref<2xi32>) -> ()
      "test.return"(%2) : (memref<2xi32>) -> ()
    }) : () -> (memref<2xi32>)
    "test.return"() : () -> ()
  }) : () -> (i32)
  "test.return"() : () -> ()
}

// -----

func.func @while_two_arg(%arg0: index) {
  %a = memref.alloc(%arg0) : memref<?xf32>
  scf.while (%arg1 = %a, %arg2 = %a) : (memref<?xf32>, memref<?xf32>) -> (memref<?xf32>, memref<?xf32>) {
    // This op has a side effect, but it's not an allocate/free side effect.
    %0 = "test.side_effect_op"() {effects = [{effect="read"}]} : () -> i1
    scf.condition(%0) %arg1, %arg2 : memref<?xf32>, memref<?xf32>
  } do {
  ^bb0(%arg1: memref<?xf32>, %arg2: memref<?xf32>):
    %b = memref.alloc(%arg0) : memref<?xf32>
    scf.yield %arg1, %b : memref<?xf32>, memref<?xf32>
  }
  return
}

// CHECK-LABEL: func @while_two_arg
//       CHECK: [[ALLOC:%.+]] = memref.alloc(
//       CHECK: [[V0:%.+]]:4 = scf.while ({{.*}} = [[ALLOC]], {{.*}} = [[ALLOC]], {{.*}} = %false{{[0-9_]*}}, {{.*}} = %false{{[0-9_]*}})
//       CHECK:   scf.condition
//       CHECK: ^bb0([[ARG1:%.+]]: memref<?xf32>, [[ARG2:%.+]]: memref<?xf32>, [[ARG3:%.+]]: i1, [[ARG4:%.+]]: i1):
//       CHECK:   [[ALLOC1:%.+]] = memref.alloc(
//       CHECK:   [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[ARG2]]
//       CHECK:   [[OWN:%.+]] = bufferization.dealloc ([[BASE]] :{{.*}}) if ([[ARG4]]) retain ([[ARG1]] :
//       CHECK:   [[OWN_AGG:%.+]] = arith.ori [[OWN]], [[ARG3]]
//       CHECK:   scf.yield [[ARG1]], [[ALLOC1]], [[OWN_AGG]], %true
//       CHECK: [[BASE0:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
//       CHECK: [[BASE1:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#1
//       CHECK: bufferization.dealloc ([[ALLOC]], [[BASE0]], [[BASE1]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#2, [[V0]]#3)

// -----

func.func @while_three_arg(%arg0: index) {
  %a = memref.alloc(%arg0) : memref<?xf32>
  scf.while (%arg1 = %a, %arg2 = %a, %arg3 = %a) : (memref<?xf32>, memref<?xf32>, memref<?xf32>) -> (memref<?xf32>, memref<?xf32>, memref<?xf32>) {
    // This op has a side effect, but it's not an allocate/free side effect.
    %0 = "test.side_effect_op"() {effects = [{effect="read"}]} : () -> i1
    scf.condition(%0) %arg1, %arg2, %arg3 : memref<?xf32>, memref<?xf32>, memref<?xf32>
  } do {
  ^bb0(%arg1: memref<?xf32>, %arg2: memref<?xf32>, %arg3: memref<?xf32>):
    %b = memref.alloc(%arg0) : memref<?xf32>
    %q = memref.alloc(%arg0) : memref<?xf32>
    scf.yield %q, %b, %arg2: memref<?xf32>, memref<?xf32>, memref<?xf32>
  }
  return
}

// CHECK-LABEL: func @while_three_arg
//       CHECK: [[ALLOC:%.+]] = memref.alloc(
//       CHECK: [[V0:%.+]]:6 = scf.while ({{.*}} = [[ALLOC]], {{.*}} = [[ALLOC]], {{.*}} = [[ALLOC]], {{.*}} = %false{{[0-9_]*}}, {{.*}} = %false{{[0-9_]*}}, {{.*}} = %false
//       CHECK:   scf.condition
//       CHECK: ^bb0([[ARG1:%.+]]: memref<?xf32>, [[ARG2:%.+]]: memref<?xf32>, [[ARG3:%.+]]: memref<?xf32>, [[ARG4:%.+]]: i1, [[ARG5:%.+]]: i1, [[ARG6:%.+]]: i1):
//       CHECK:   [[ALLOC1:%.+]] = memref.alloc(
//       CHECK:   [[ALLOC2:%.+]] = memref.alloc(
//       CHECK:   [[BASE0:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[ARG1]]
//       CHECK:   [[BASE2:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[ARG3]]
//       CHECK:   [[OWN:%.+]] = bufferization.dealloc ([[BASE0]], [[BASE2]] :{{.*}}) if ([[ARG4]], [[ARG6]]) retain ([[ARG2]] :
//       CHECK:   [[OWN_AGG:%.+]] = arith.ori [[OWN]], [[ARG5]]
//       CHECK:   scf.yield [[ALLOC2]], [[ALLOC1]], [[ARG2]], %true{{[0-9_]*}}, %true{{[0-9_]*}}, [[OWN_AGG]] :
//       CHECK: }
//       CHECK: [[BASE0:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
//       CHECK: [[BASE1:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#1
//       CHECK: [[BASE2:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#2
//       CHECK: bufferization.dealloc ([[ALLOC]], [[BASE0]], [[BASE1]], [[BASE2]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#3, [[V0]]#4, [[V0]]#5)

// TODO: better alias analysis could simplify the dealloc inside the body further

// -----

// Memref allocated in `then` region and passed back to the parent if op.
#set = affine_set<() : (0 >= 0)>
func.func @test_affine_if_1(%arg0: memref<10xf32>) -> memref<10xf32> {
  %0 = affine.if #set() -> memref<10xf32> {
    %alloc = memref.alloc() : memref<10xf32>
    affine.yield %alloc : memref<10xf32>
  } else {
    affine.yield %arg0 : memref<10xf32>
  }
  return %0 : memref<10xf32>
}

// CHECK-LABEL: func @test_affine_if_1
//  CHECK-SAME: ([[ARG0:%.*]]: memref<10xf32>)
//       CHECK: [[V0:%.+]]:2 = affine.if
//       CHECK:   [[ALLOC:%.+]] = memref.alloc()
//       CHECK:   affine.yield [[ALLOC]], %true
//       CHECK:   affine.yield [[ARG0]], %false
//       CHECK: [[V1:%.+]] = scf.if [[V0]]#1
//       CHECK:   scf.yield [[V0]]#0
//       CHECK:   [[CLONE:%.+]] = bufferization.clone [[V0]]#0
//       CHECK:   scf.yield [[CLONE]]
//       CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
//       CHECK: bufferization.dealloc ([[BASE]] :{{.*}}) if ([[V0]]#1) retain ([[V1]] :
//       CHECK: return [[V1]]

// TODO: the dealloc could be optimized away since the memref to be deallocated
//       either aliases with V1 or the condition is false

// -----

// Memref allocated before parent IfOp and used in `then` region.
// Expected result: deallocation should happen after affine.if op.
#set = affine_set<() : (0 >= 0)>
func.func @test_affine_if_2() -> memref<10xf32> {
  %alloc0 = memref.alloc() : memref<10xf32>
  %0 = affine.if #set() -> memref<10xf32> {
    affine.yield %alloc0 : memref<10xf32>
  } else {
    %alloc = memref.alloc() : memref<10xf32>
    affine.yield %alloc : memref<10xf32>
  }
  return %0 : memref<10xf32>
}
// CHECK-LABEL: func @test_affine_if_2
//       CHECK: [[ALLOC:%.+]] = memref.alloc()
//       CHECK: [[V0:%.+]]:2 = affine.if
//       CHECK:   affine.yield [[ALLOC]], %false
//       CHECK:   [[ALLOC1:%.+]] = memref.alloc()
//       CHECK:   affine.yield [[ALLOC1]], %true
//       CHECK: [[V1:%.+]] = scf.if [[V0]]#1
//       CHECK:   scf.yield [[V0]]#0
//       CHECK:   [[CLONE:%.+]] = bufferization.clone [[V0]]#0
//       CHECK:   scf.yield [[CLONE]]
//       CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
//       CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V1]] :
//       CHECK: return [[V1]]

// -----

// Memref allocated before parent IfOp and used in `else` region.
// Expected result: deallocation should happen after affine.if op.
#set = affine_set<() : (0 >= 0)>
func.func @test_affine_if_3() -> memref<10xf32> {
  %alloc0 = memref.alloc() : memref<10xf32>
  %0 = affine.if #set() -> memref<10xf32> {
    %alloc = memref.alloc() : memref<10xf32>
    affine.yield %alloc : memref<10xf32>
  } else {
    affine.yield %alloc0 : memref<10xf32>
  }
  return %0 : memref<10xf32>
}

// CHECK-LABEL: func @test_affine_if_3
//       CHECK: [[ALLOC:%.+]] = memref.alloc()
//       CHECK: [[V0:%.+]]:2 = affine.if
//       CHECK:   [[ALLOC1:%.+]] = memref.alloc()
//       CHECK:   affine.yield [[ALLOC1]], %true
//       CHECK:   affine.yield [[ALLOC]], %false
//       CHECK: [[V1:%.+]] = scf.if [[V0]]#1
//       CHECK:   scf.yield [[V0]]#0
//       CHECK:   [[CLONE:%.+]] = bufferization.clone [[V0]]#0
//       CHECK:   scf.yield [[CLONE]]
//       CHECK: [[BASE:%[a-zA-Z0-9_]+]],{{.*}} = memref.extract_strided_metadata [[V0]]#0
//       CHECK: bufferization.dealloc ([[ALLOC]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, [[V0]]#1) retain ([[V1]]
//       CHECK: return [[V1]]