# RUN: %PYTHON %s | FileCheck %s

from mlir.ir import *
import mlir.dialects.builtin as builtin
import mlir.dialects.std as std


def run(f):
  print("\nTEST:", f.__name__)
  f()
  return f


# CHECK-LABEL: TEST: testFromPyFunc
@run
def testFromPyFunc():
  with Context() as ctx, Location.unknown() as loc:
    ctx.allow_unregistered_dialects = True
    m = builtin.ModuleOp()
    f32 = F32Type.get()
    f64 = F64Type.get()
    with InsertionPoint(m.body):
      # CHECK-LABEL: func @unary_return(%arg0: f64) -> f64
      # CHECK: return %arg0 : f64
      @builtin.FuncOp.from_py_func(f64)
      def unary_return(a):
        return a

      # CHECK-LABEL: func @binary_return(%arg0: f32, %arg1: f64) -> (f32, f64)
      # CHECK: return %arg0, %arg1 : f32, f64
      @builtin.FuncOp.from_py_func(f32, f64)
      def binary_return(a, b):
        return a, b

      # CHECK-LABEL: func @none_return(%arg0: f32, %arg1: f64)
      # CHECK: return
      @builtin.FuncOp.from_py_func(f32, f64)
      def none_return(a, b):
        pass

      # CHECK-LABEL: func @call_unary
      # CHECK: %0 = call @unary_return(%arg0) : (f64) -> f64
      # CHECK: return %0 : f64
      @builtin.FuncOp.from_py_func(f64)
      def call_unary(a):
        return unary_return(a)

      # CHECK-LABEL: func @call_binary
      # CHECK: %0:2 = call @binary_return(%arg0, %arg1) : (f32, f64) -> (f32, f64)
      # CHECK: return %0#0, %0#1 : f32, f64
      @builtin.FuncOp.from_py_func(f32, f64)
      def call_binary(a, b):
        return binary_return(a, b)

      # We expect coercion of a single result operation to a returned value.
      # CHECK-LABEL: func @single_result_op
      # CHECK: %0 = "custom.op1"() : () -> f32
      # CHECK: return %0 : f32
      @builtin.FuncOp.from_py_func()
      def single_result_op():
        return Operation.create("custom.op1", results=[f32])

      # CHECK-LABEL: func @call_none
      # CHECK: call @none_return(%arg0, %arg1) : (f32, f64) -> ()
      # CHECK: return
      @builtin.FuncOp.from_py_func(f32, f64)
      def call_none(a, b):
        return none_return(a, b)

      ## Variants and optional feature tests.
      # CHECK-LABEL: func @from_name_arg
      @builtin.FuncOp.from_py_func(f32, f64, name="from_name_arg")
      def explicit_name(a, b):
        return b

      @builtin.FuncOp.from_py_func(f32, f64)
      def positional_func_op(a, b, func_op):
        assert isinstance(func_op, builtin.FuncOp)
        return b

      @builtin.FuncOp.from_py_func(f32, f64)
      def kw_func_op(a, b=None, func_op=None):
        assert isinstance(func_op, builtin.FuncOp)
        return b

      @builtin.FuncOp.from_py_func(f32, f64)
      def kwargs_func_op(a, b=None, **kwargs):
        assert isinstance(kwargs["func_op"], builtin.FuncOp)
        return b

      # CHECK-LABEL: func @explicit_results(%arg0: f32, %arg1: f64) -> f64
      # CHECK: return %arg1 : f64
      @builtin.FuncOp.from_py_func(f32, f64, results=[f64])
      def explicit_results(a, b):
        std.ReturnOp([b])

  print(m)


# CHECK-LABEL: TEST: testFromPyFuncErrors
@run
def testFromPyFuncErrors():
  with Context() as ctx, Location.unknown() as loc:
    m = builtin.ModuleOp()
    f32 = F32Type.get()
    f64 = F64Type.get()
    with InsertionPoint(m.body):
      try:

        @builtin.FuncOp.from_py_func(f64, results=[f64])
        def unary_return(a):
          return a
      except AssertionError as e:
        # CHECK: Capturing a python function with explicit `results=` requires that the wrapped function returns None.
        print(e)


# CHECK-LABEL: TEST: testBuildFuncOp
@run
def testBuildFuncOp():
  ctx = Context()
  with Location.unknown(ctx) as loc:
    m = builtin.ModuleOp()

    f32 = F32Type.get()
    tensor_type = RankedTensorType.get((2, 3, 4), f32)
    with InsertionPoint.at_block_begin(m.body):
      func = builtin.FuncOp(name="some_func",
                            type=FunctionType.get(
                                inputs=[tensor_type, tensor_type],
                                results=[tensor_type]),
                            visibility="nested")
      # CHECK: Name is: "some_func"
      print("Name is: ", func.name)

      # CHECK: Type is: (tensor<2x3x4xf32>, tensor<2x3x4xf32>) -> tensor<2x3x4xf32>
      print("Type is: ", func.type)

      # CHECK: Visibility is: "nested"
      print("Visibility is: ", func.visibility)

      try:
        entry_block = func.entry_block
      except IndexError as e:
        # CHECK: External function does not have a body
        print(e)

      with InsertionPoint(func.add_entry_block()):
        std.ReturnOp([func.entry_block.arguments[0]])
        pass

      try:
        func.add_entry_block()
      except IndexError as e:
        # CHECK: The function already has an entry block!
        print(e)

      # Try the callback builder and passing type as tuple.
      func = builtin.FuncOp(name="some_other_func",
                            type=([tensor_type, tensor_type], [tensor_type]),
                            visibility="nested",
                            body_builder=lambda func: std.ReturnOp(
                                [func.entry_block.arguments[0]]))

  # CHECK: module  {
  # CHECK:  func nested @some_func(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
  # CHECK:   return %arg0 : tensor<2x3x4xf32>
  # CHECK:  }
  # CHECK:  func nested @some_other_func(%arg0: tensor<2x3x4xf32>, %arg1: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
  # CHECK:   return %arg0 : tensor<2x3x4xf32>
  # CHECK:  }
  print(m)


# CHECK-LABEL: TEST: testFuncArgumentAccess
@run
def testFuncArgumentAccess():
  with Context() as ctx, Location.unknown():
    ctx.allow_unregistered_dialects = True
    module = Module.create()
    f32 = F32Type.get()
    f64 = F64Type.get()
    with InsertionPoint(module.body):
      func = builtin.FuncOp("some_func", ([f32, f32], [f32, f32]))
      with InsertionPoint(func.add_entry_block()):
        std.ReturnOp(func.arguments)
      func.arg_attrs = ArrayAttr.get([
          DictAttr.get({
              "custom_dialect.foo": StringAttr.get("bar"),
              "custom_dialect.baz": UnitAttr.get()
          }),
          DictAttr.get({"custom_dialect.qux": ArrayAttr.get([])})
      ])
      func.result_attrs = ArrayAttr.get([
          DictAttr.get({"custom_dialect.res1": FloatAttr.get(f32, 42.0)}),
          DictAttr.get({"custom_dialect.res2": FloatAttr.get(f64, 256.0)})
      ])

      other = builtin.FuncOp("other_func", ([f32, f32], []))
      with InsertionPoint(other.add_entry_block()):
        std.ReturnOp([])
      other.arg_attrs = [
          DictAttr.get({"custom_dialect.foo": StringAttr.get("qux")}),
          DictAttr.get()
      ]

  # CHECK: [{custom_dialect.baz, custom_dialect.foo = "bar"}, {custom_dialect.qux = []}]
  print(func.arg_attrs)

  # CHECK: [{custom_dialect.res1 = 4.200000e+01 : f32}, {custom_dialect.res2 = 2.560000e+02 : f64}]
  print(func.result_attrs)

  # CHECK: func @some_func(
  # CHECK: %[[ARG0:.*]]: f32 {custom_dialect.baz, custom_dialect.foo = "bar"},
  # CHECK: %[[ARG1:.*]]: f32 {custom_dialect.qux = []}) ->
  # CHECK: f32 {custom_dialect.res1 = 4.200000e+01 : f32},
  # CHECK: f32 {custom_dialect.res2 = 2.560000e+02 : f64})
  # CHECK: return %[[ARG0]], %[[ARG1]] : f32, f32
  #
  # CHECK: func @other_func(
  # CHECK: %{{.*}}: f32 {custom_dialect.foo = "qux"},
  # CHECK: %{{.*}}: f32)
  print(module)