# RUN: %PYTHON %s | FileCheck %s

import gc
from mlir.ir import *


def run(f):
  print("\nTEST:", f.__name__)
  f()
  gc.collect()
  assert Context._get_live_count() == 0
  return f


# CHECK-LABEL: TEST: testDialectDescriptor
@run
def testDialectDescriptor():
  ctx = Context()
  d = ctx.get_dialect_descriptor("func")
  # CHECK: <DialectDescriptor func>
  print(d)
  # CHECK: func
  print(d.namespace)
  try:
    _ = ctx.get_dialect_descriptor("not_existing")
  except ValueError:
    pass
  else:
    assert False, "Expected exception"


# CHECK-LABEL: TEST: testUserDialectClass
@run
def testUserDialectClass():
  ctx = Context()
  # Access using attribute.
  d = ctx.dialects.func
  # CHECK: <Dialect func (class mlir.dialects._func_ops_gen._Dialect)>
  print(d)
  try:
    _ = ctx.dialects.not_existing
  except AttributeError:
    pass
  else:
    assert False, "Expected exception"

  # Access using index.
  d = ctx.dialects["func"]
  # CHECK: <Dialect func (class mlir.dialects._func_ops_gen._Dialect)>
  print(d)
  try:
    _ = ctx.dialects["not_existing"]
  except IndexError:
    pass
  else:
    assert False, "Expected exception"

  # Using the 'd' alias.
  d = ctx.d["func"]
  # CHECK: <Dialect func (class mlir.dialects._func_ops_gen._Dialect)>
  print(d)


# CHECK-LABEL: TEST: testCustomOpView
# This test uses the standard dialect AddFOp as an example of a user op.
# TODO: Op creation and access is still quite verbose: simplify this test as
# additional capabilities come online.
@run
def testCustomOpView():

  def createInput():
    op = Operation.create("pytest_dummy.intinput", results=[f32])
    # TODO: Auto result cast from operation
    return op.results[0]

  with Context() as ctx, Location.unknown():
    ctx.allow_unregistered_dialects = True
    m = Module.create()

    with InsertionPoint(m.body):
      f32 = F32Type.get()
      # Create via dialects context collection.
      input1 = createInput()
      input2 = createInput()
      op1 = ctx.dialects.arith.AddFOp(input1, input2)

      # Create via an import
      from mlir.dialects.arith import AddFOp
      AddFOp(input1, op1.result)

  # CHECK: %[[INPUT0:.*]] = "pytest_dummy.intinput"
  # CHECK: %[[INPUT1:.*]] = "pytest_dummy.intinput"
  # CHECK: %[[R0:.*]] = arith.addf %[[INPUT0]], %[[INPUT1]] : f32
  # CHECK: %[[R1:.*]] = arith.addf %[[INPUT0]], %[[R0]] : f32
  m.operation.print()


# CHECK-LABEL: TEST: testIsRegisteredOperation
@run
def testIsRegisteredOperation():
  ctx = Context()

  # CHECK: cf.cond_br: True
  print(f"cf.cond_br: {ctx.is_registered_operation('cf.cond_br')}")
  # CHECK: func.not_existing: False
  print(f"func.not_existing: {ctx.is_registered_operation('func.not_existing')}")