text: # RUN: %PYTHON %s | FileCheck %s
  
  from mlir.ir import *
  from mlir.dialects import arith
  from mlir.dialects import scf
  from mlir.dialects import std
  from mlir.dialects import builtin
  
  
  def constructAndPrintInModule(f):
    print("\nTEST:", f.__name__)
    with Context(), Location.unknown():
      module = Module.create()
      with InsertionPoint(module.body):
        f()
      print(module)
    return f
  
  
  # CHECK-LABEL: TEST: testSimpleLoop
  @constructAndPrintInModule
  def testSimpleLoop():
    index_type = IndexType.get()
  
    @builtin.FuncOp.from_py_func(index_type, index_type, index_type)
    def simple_loop(lb, ub, step):
      loop = scf.ForOp(lb, ub, step, [lb, lb])
      with InsertionPoint(loop.body):
        scf.YieldOp(loop.inner_iter_args)
      return
  
  
  # CHECK: func @simple_loop(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
  # CHECK: scf.for %{{.*}} = %[[ARG0]] to %[[ARG1]] step %[[ARG2]]
  # CHECK: iter_args(%[[I1:.*]] = %[[ARG0]], %[[I2:.*]] = %[[ARG0]])
  # CHECK: scf.yield %[[I1]], %[[I2]]
  
  
  # CHECK-LABEL: TEST: testInductionVar
  @constructAndPrintInModule
  def testInductionVar():
    index_type = IndexType.get()
  
    @builtin.FuncOp.from_py_func(index_type, index_type, index_type)
    def induction_var(lb, ub, step):
      loop = scf.ForOp(lb, ub, step, [lb])
      with InsertionPoint(loop.body):
        scf.YieldOp([loop.induction_variable])
      return
  
  
  # CHECK: func @induction_var(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
  # CHECK: scf.for %[[IV:.*]] = %[[ARG0]] to %[[ARG1]] step %[[ARG2]]
  # CHECK: scf.yield %[[IV]]
  
  
  @constructAndPrintInModule
  def testOpsAsArguments():
    index_type = IndexType.get()
    callee = builtin.FuncOp(
        "callee", ([], [index_type, index_type]), visibility="private")
    func = builtin.FuncOp("ops_as_arguments", ([], []))
    with InsertionPoint(func.add_entry_block()):
      lb = arith.ConstantOp.create_index(0)
      ub = arith.ConstantOp.create_index(42)
      step = arith.ConstantOp.create_index(2)
      iter_args = std.CallOp(callee, [])
      loop = scf.ForOp(lb, ub, step, iter_args)
      with InsertionPoint(loop.body):
        scf.YieldOp(loop.inner_iter_args)
      std.ReturnOp([])
  
  
  # CHECK-LABEL: TEST: testOpsAsArguments
  # CHECK: func private @callee() -> (index, index)
  # CHECK: func @ops_as_arguments() {
  # CHECK:   %[[LB:.*]] = arith.constant 0
  # CHECK:   %[[UB:.*]] = arith.constant 42
  # CHECK:   %[[STEP:.*]] = arith.constant 2
  # CHECK:   %[[ARGS:.*]]:2 = call @callee()
  # CHECK:   scf.for %arg0 = %c0 to %c42 step %c2
  # CHECK:   iter_args(%{{.*}} = %[[ARGS]]#0, %{{.*}} = %[[ARGS]]#1)
  # CHECK:     scf.yield %{{.*}}, %{{.*}}
  # CHECK:   return
  