################################################################################
# Python modules
# MLIR's Python modules are both directly used by the core project and are
# available for use and embedding into external projects (in their own
# namespace and with their own deps). In order to facilitate this, python
# artifacts are split between declarations, which make a subset of
# things available to be built and "add", which in line with the normal LLVM
# nomenclature, adds libraries.
################################################################################

# Function: declare_mlir_python_sources
# Declares pure python sources as part of a named grouping that can be built
# later.
# Arguments:
#   ROOT_DIR: Directory where the python namespace begins (defaults to
#     CMAKE_CURRENT_SOURCE_DIR). For non-relocatable sources, this will
#     typically just be the root of the python source tree (current directory).
#     For relocatable sources, this will point deeper into the directory that
#     can be relocated. For generated sources, can be relative to
#     CMAKE_CURRENT_BINARY_DIR. Generated and non generated sources cannot be
#     mixed.
#   ADD_TO_PARENT: Adds this source grouping to a previously declared source
#     grouping. Source groupings form a DAG.
#   SOURCES: List of specific source files relative to ROOT_DIR to include.
#   SOURCES_GLOB: List of glob patterns relative to ROOT_DIR to include.
function(declare_mlir_python_sources name)
  cmake_parse_arguments(ARG
    ""
    "ROOT_DIR;ADD_TO_PARENT"
    "SOURCES;SOURCES_GLOB"
    ${ARGN})

  if(NOT ARG_ROOT_DIR)
    set(ARG_ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
  endif()
  set(_install_destination "src/python/${name}")

  # Process the glob.
  set(_glob_sources)
  if(ARG_SOURCES_GLOB)
    set(_glob_spec ${ARG_SOURCES_GLOB})
    list(TRANSFORM _glob_spec PREPEND "${ARG_ROOT_DIR}/")
    file(GLOB_RECURSE _glob_sources
      RELATIVE "${ARG_ROOT_DIR}"
      ${_glob_spec}
    )
    list(APPEND ARG_SOURCES ${_glob_sources})
  endif()

  # We create a custom target to carry properties and dependencies for
  # generated sources.
  add_library(${name} INTERFACE)
  set_target_properties(${name} PROPERTIES
    # Yes: Leading-lowercase property names are load bearing and the recommended
    # way to do this: https://gitlab.kitware.com/cmake/cmake/-/issues/19261
    EXPORT_PROPERTIES "mlir_python_SOURCES_TYPE;mlir_python_DEPENDS"
    mlir_python_SOURCES_TYPE pure
    mlir_python_DEPENDS ""
  )

  # Use the interface include directories and sources on the target to carry the
  # properties we would like to export. These support generator expressions and
  # allow us to properly specify paths in both the local build and install scenarios.
  # The one caveat here is that because we don't directly build against the interface
  # library, we need to specify the INCLUDE_DIRECTORIES and SOURCES properties as well
  # via private properties because the evaluation would happen at configuration time
  # instead of build time.
  # Eventually this could be done using a FILE_SET simplifying the logic below.
  # FILE_SET is available in cmake 3.23+, so it is not an option at the moment.
  target_include_directories(${name} INTERFACE
    "$<BUILD_INTERFACE:${ARG_ROOT_DIR}>"
    "$<INSTALL_INTERFACE:${_install_destination}>"
  )
  set_property(TARGET ${name} PROPERTY INCLUDE_DIRECTORIES ${ARG_ROOT_DIR})

  if(ARG_SOURCES)
    list(TRANSFORM ARG_SOURCES PREPEND "${ARG_ROOT_DIR}/" OUTPUT_VARIABLE _build_sources)
    list(TRANSFORM ARG_SOURCES PREPEND "${_install_destination}/" OUTPUT_VARIABLE _install_sources)
    target_sources(${name}
      INTERFACE
        "$<INSTALL_INTERFACE:${_install_sources}>"
        "$<BUILD_INTERFACE:${_build_sources}>"
      PRIVATE ${_build_sources}
    )
  endif()

  # Add to parent.
  if(ARG_ADD_TO_PARENT)
    set_property(TARGET ${ARG_ADD_TO_PARENT} APPEND PROPERTY mlir_python_DEPENDS ${name})
  endif()

  # Install.
  set_property(GLOBAL APPEND PROPERTY MLIR_EXPORTS ${name})
  if(NOT LLVM_INSTALL_TOOLCHAIN_ONLY)
    _mlir_python_install_sources(
      ${name} "${ARG_ROOT_DIR}" "${_install_destination}"
      ${ARG_SOURCES}
    )
  endif()
endfunction()

# Function: mlir_generate_type_stubs
# Turns on automatic type stub generation for extension modules.
# Specifically, performs add_custom_command to run nanobind's stubgen on an extension module.
#
# Arguments:
#   MODULE_NAME: The fully-qualified name of the extension module (used for importing in python).
#   DEPENDS_TARGETS: List of targets these type stubs depend on being built; usually corresponding to the
#     specific extension module (e.g., something like StandalonePythonModules.extension._standaloneDialectsNanobind.dso)
#     and the core bindings extension module (e.g., something like StandalonePythonModules.extension._mlir.dso).
#   OUTPUT_DIR: The root output directory to emit the type stubs into.
#   OUTPUTS: List of expected outputs.
#   DEPENDS_TARGET_SRC_DEPS: List of cpp sources for extension library (for generating a DEPFILE).
#   IMPORT_PATHS: List of paths to add to PYTHONPATH for stubgen.
#   PATTERN_FILE: (Optional) Pattern file (see https://nanobind.readthedocs.io/en/latest/typing.html#pattern-files).
#   VERBOSE: Emit logging/status messages during stub generation (default: OFF).
# Outputs:
#   NB_STUBGEN_CUSTOM_TARGET: The target corresponding to generation which other targets can depend on.
function(mlir_generate_type_stubs)
  cmake_parse_arguments(ARG
    "VERBOSE"
    "MODULE_NAME;OUTPUT_DIR;PATTERN_FILE"
    "IMPORT_PATHS;DEPENDS_TARGETS;OUTPUTS;DEPENDS_TARGET_SRC_DEPS"
    ${ARGN})

  # Allow overriding the stubgen.py path or fetching a specific version
  # from the nanobind repository, independent of the nanobind used for
  # building. This is useful when a newer stubgen has bug fixes or features
  # not yet available in the nanobind version used for compilation.
  if(MLIR_NB_STUBGEN)
    set(NB_STUBGEN "${MLIR_NB_STUBGEN}")
  elseif(MLIR_NB_STUBGEN_VERSION)
    set(_stubgen_path "${MLIR_BINARY_DIR}/stubgen/${MLIR_NB_STUBGEN_VERSION}/stubgen.py")
    if(NOT EXISTS "${_stubgen_path}")
      message(STATUS "Downloading stubgen.py from nanobind ${MLIR_NB_STUBGEN_VERSION}...")
      file(MAKE_DIRECTORY "${MLIR_BINARY_DIR}/stubgen/${MLIR_NB_STUBGEN_VERSION}" RESULT _created_dir)
      if(NOT _created_dir EQUAL 0)
        list(GET _created_dir 1 _created_dir_error)
        message(FATAL_ERROR "Failed to create parent dir for stubgen.py: ${_created_dir_error}")
      endif()
      file(DOWNLOAD
        "https://raw.githubusercontent.com/wjakob/nanobind/${MLIR_NB_STUBGEN_VERSION}/src/stubgen.py"
        "${_stubgen_path}"
        STATUS _download_status
      )
      list(GET _download_status 0 _download_code)
      if(NOT _download_code EQUAL 0)
        list(GET _download_status 1 _download_error)
        message(FATAL_ERROR "Failed to download stubgen.py: ${_download_error}")
      endif()
    endif()
    set(NB_STUBGEN "${_stubgen_path}")
  # for people installing a distro (e.g., pip install) of nanobind
  elseif(EXISTS ${nanobind_DIR}/../src/stubgen.py)
    set(NB_STUBGEN "${nanobind_DIR}/../src/stubgen.py")
  elseif(EXISTS ${nanobind_DIR}/../stubgen.py)
    set(NB_STUBGEN "${nanobind_DIR}/../stubgen.py")
  # for people using nanobind git source tree (e.g., FetchContent_Declare and FetchContent_MakeAvailable)
  elseif(EXISTS ${nanobind_SOURCE_DIR}/src/stubgen.py)
    set(NB_STUBGEN "${nanobind_SOURCE_DIR}/src/stubgen.py")
  elseif(EXISTS ${nanobind_SOURCE_DIR}/stubgen.py)
    set(NB_STUBGEN "${nanobind_SOURCE_DIR}/stubgen.py")
  else()
    message(FATAL_ERROR "mlir_generate_type_stubs(): could not locate 'stubgen.py'!")
  endif()

  file(REAL_PATH "${NB_STUBGEN}" NB_STUBGEN)
  set(_import_paths)
  foreach(_import_path IN LISTS ARG_IMPORT_PATHS)
    file(REAL_PATH "${_import_path}" _import_path)
    list(APPEND _import_paths "-i;${_import_path}")
  endforeach()
  set(_nb_stubgen_cmd
      "${Python_EXECUTABLE}"
      "${NB_STUBGEN}"
      --module
      "${ARG_MODULE_NAME}"
      "${_import_paths}"
      --recursive
      --include-private
      --output-dir
      "${ARG_OUTPUT_DIR}")
  if(NOT ARG_VERBOSE)
    list(APPEND _nb_stubgen_cmd "--quiet")
  endif()
  if(ARG_PATTERN_FILE)
    list(APPEND _nb_stubgen_cmd "-p;${ARG_PATTERN_FILE}")
    list(APPEND ARG_DEPENDS_TARGETS "${ARG_PATTERN_FILE}")
  endif()

  list(TRANSFORM ARG_OUTPUTS PREPEND "${ARG_OUTPUT_DIR}/" OUTPUT_VARIABLE _generated_type_stubs)
  set(_depfile "${ARG_OUTPUT_DIR}/${ARG_MODULE_NAME}.d")
  if ((NOT EXISTS ${_depfile}) AND ARG_DEPENDS_TARGET_SRC_DEPS)
    list(JOIN ARG_DEPENDS_TARGET_SRC_DEPS " " _depfiles)
    list(TRANSFORM _generated_type_stubs APPEND ": ${_depfiles}" OUTPUT_VARIABLE _depfiles)
    list(JOIN _depfiles "\n" _depfiles)
    file(GENERATE OUTPUT "${_depfile}" CONTENT "${_depfiles}")
  endif()

  if(ARG_VERBOSE)
    message(STATUS "Generating type-stubs outputs ${_generated_type_stubs}")
  endif()

  # If PYTHONPATH is set and points to the build location of the python package then when stubgen runs, _mlir will get
  # imported twice and bad things will happen (e.g., Assertion `!instance && “PyGlobals already constructed”’).
  # This happens because mlir is a namespace package and the importer/loader can't distinguish between
  # mlir._mlir_libs._mlir and _mlir_libs._mlir imported from CWD.
  # So try to filter out any entries in PYTHONPATH that end in "MLIR_BINDINGS_PYTHON_INSTALL_PREFIX/.."
  # (e.g., python_packages/mlir_core/).
  set(_pythonpath "$ENV{PYTHONPATH}")
  cmake_path(CONVERT "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/.." TO_NATIVE_PATH_LIST _install_prefix NORMALIZE)
  if(WIN32)
    set(_path_sep ";")
    set(_trailing_sep "\\")
  else()
    set(_path_sep ":")
    set(_trailing_sep "/")
    # `;` is the CMake list delimiter so Windows paths are automatically lists
    # and Unix paths can be made into lists by replacing `:` with `;`
    string(REPLACE "${_path_sep}" ";" _pythonpath "${_pythonpath}")
  endif()
  string(REGEX REPLACE "${_trailing_sep}$" "" _install_prefix "${_install_prefix}")
  list(FILTER _pythonpath EXCLUDE REGEX "(${_install_prefix}|${_install_prefix}${_trailing_sep})$")
  # Note, ${_pythonpath} is a list but "${_pythonpath}" is not a list - it's a string with ";" chars in it.
  string(JOIN "${_path_sep}" _pythonpath ${_pythonpath})
  add_custom_command(
    OUTPUT ${_generated_type_stubs}
    COMMAND ${CMAKE_COMMAND} -E env PYTHONPATH="${_pythonpath}" ${_nb_stubgen_cmd}
    WORKING_DIRECTORY "${CMAKE_CURRENT_FUNCTION_LIST_DIR}"
    DEPENDS "${ARG_DEPENDS_TARGETS}"
    DEPFILE "${_depfile}"
    COMMENT "Generating type stubs using: ${_nb_stubgen_cmd}"
  )

  list(JOIN ARG_DEPENDS_TARGETS "." _name)
  set(_name "${_name}.${ARG_MODULE_NAME}.type_stubs")
  add_custom_target("${_name}" DEPENDS ${_generated_type_stubs})
  set(NB_STUBGEN_CUSTOM_TARGET "${_name}" PARENT_SCOPE)
endfunction()

# Function: declare_mlir_python_extension
# Declares a buildable python extension from C++ source files. The built
# module is considered a python source file and included as everything else.
# Arguments:
#   ROOT_DIR: Root directory where sources are interpreted relative to.
#     Defaults to CMAKE_CURRENT_SOURCE_DIR.
#   MODULE_NAME: Local import name of the module (i.e. "_mlir").
#   ADD_TO_PARENT: Same as for declare_mlir_python_sources.
#   SOURCES: C++ sources making up the module.
#   PRIVATE_LINK_LIBS: List of libraries to link in privately to the module
#     regardless of how it is included in the project (generally should be
#     static libraries that can be included with hidden visibility).
#   EMBED_CAPI_LINK_LIBS: Dependent CAPI libraries that this extension depends
#     on. These will be collected for all extensions and put into an
#     aggregate dylib that is linked against.
function(declare_mlir_python_extension name)
  cmake_parse_arguments(ARG
    "_PRIVATE_SUPPORT_LIB"
    "ROOT_DIR;MODULE_NAME;ADD_TO_PARENT"
    "SOURCES;PRIVATE_LINK_LIBS;EMBED_CAPI_LINK_LIBS"
    ${ARGN})

  if(NOT ARG_ROOT_DIR)
    set(ARG_ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
  endif()
  if(ARG__PRIVATE_SUPPORT_LIB)
    set(SOURCES_TYPE "support")
  else()
    set(SOURCES_TYPE "extension")
  endif()
  set(_install_destination "src/python/${name}")

  add_library(${name} INTERFACE)
  set_target_properties(${name} PROPERTIES
    # Yes: Leading-lowercase property names are load bearing and the recommended
    # way to do this: https://gitlab.kitware.com/cmake/cmake/-/issues/19261
    EXPORT_PROPERTIES "mlir_python_SOURCES_TYPE;mlir_python_EXTENSION_MODULE_NAME;mlir_python_EMBED_CAPI_LINK_LIBS;mlir_python_DEPENDS"
    mlir_python_SOURCES_TYPE "${SOURCES_TYPE}"
    mlir_python_EXTENSION_MODULE_NAME "${ARG_MODULE_NAME}"
    mlir_python_EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
    mlir_python_DEPENDS ""
  )

  # Set the interface source and link_libs properties of the target
  # These properties support generator expressions and are automatically exported
  list(TRANSFORM ARG_SOURCES PREPEND "${ARG_ROOT_DIR}/" OUTPUT_VARIABLE _build_sources)
  list(TRANSFORM ARG_SOURCES PREPEND "${_install_destination}/" OUTPUT_VARIABLE _install_sources)
  target_sources(${name} INTERFACE
    "$<BUILD_INTERFACE:${_build_sources}>"
    "$<INSTALL_INTERFACE:${_install_sources}>"
  )
  target_link_libraries(${name} INTERFACE
    ${ARG_PRIVATE_LINK_LIBS}
  )

  # Add to parent.
  if(ARG_ADD_TO_PARENT)
    set_property(TARGET ${ARG_ADD_TO_PARENT} APPEND PROPERTY mlir_python_DEPENDS ${name})
  endif()

  # Install.
  set_property(GLOBAL APPEND PROPERTY MLIR_EXPORTS ${name})
  if(NOT LLVM_INSTALL_TOOLCHAIN_ONLY)
    _mlir_python_install_sources(
      ${name} "${ARG_ROOT_DIR}" "${_install_destination}"
      ${ARG_SOURCES}
    )
  endif()
endfunction()

function(_mlir_python_install_sources name source_root_dir destination)
  foreach(source_relative_path ${ARGN})
    # Transform "a/b/c.py" -> "${install_prefix}/a/b" for installation.
    get_filename_component(
      dest_relative_dir "${source_relative_path}" DIRECTORY
      BASE_DIR "${source_root_dir}"
    )
    install(
      FILES "${source_root_dir}/${source_relative_path}"
      DESTINATION "${destination}/${dest_relative_dir}"
      COMPONENT mlir-python-sources
    )
  endforeach()
  get_target_export_arg(${name} MLIR export_to_mlirtargets
    UMBRELLA mlir-python-sources)
  install(TARGETS ${name}
    COMPONENT mlir-python-sources
    ${export_to_mlirtargets}
  )
endfunction()

# Define exception handling and RTTI flags.
macro(_mlir_python_eh_rtti_flags)
  set(eh_rtti_enable)
  if(MSVC)
    set(eh_rtti_enable /EHsc /GR)
  elseif(LLVM_COMPILER_IS_GCC_COMPATIBLE OR CLANG_CL)
    set(eh_rtti_enable -frtti -fexceptions)
  endif()
endmacro()

function(build_nanobind_lib)
  cmake_parse_arguments(ARG
    ""
    "INSTALL_COMPONENT;INSTALL_DESTINATION;OUTPUT_DIRECTORY;MLIR_BINDINGS_PYTHON_NB_DOMAIN"
    "EXTRA_COMPILE_OPTIONS;EXTRA_LINK_LIBS"
    ${ARGN})

  # Only build in free-threaded mode if the Python ABI supports it.
  # See https://github.com/wjakob/nanobind/blob/4ba51fcf795971c5d603d875ae4184bc0c9bd8e6/cmake/nanobind-config.cmake#L363-L371.
  if(NB_ABI MATCHES "[0-9]t")
    set(_ft "-ft")
    set(_abi3 "")
  else()
    set(_ft "")
    # Match nanobind_add_module's naming: with STABLE_ABI, the shared library
    # name includes "-abi3" (e.g., "nanobind-abi3-mlir").
    if(MLIR_ENABLE_PYTHON_STABLE_ABI)
      set(_abi3 "-abi3")
    else()
      set(_abi3 "")
    endif()
  endif()
  # nanobind does a string match on the suffix to figure out whether to build
  # the lib with free threading...
  set(NB_LIBRARY_TARGET_NAME "nanobind${_ft}${_abi3}-${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}")
  set(NB_LIBRARY_TARGET_NAME "${NB_LIBRARY_TARGET_NAME}" PARENT_SCOPE)
  nanobind_build_library(${NB_LIBRARY_TARGET_NAME} AS_SYSINCLUDE)
  target_compile_definitions(${NB_LIBRARY_TARGET_NAME}
    PRIVATE
    NB_DOMAIN=${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}
  )

  if(NOT MLIR_DISABLE_CONFIGURE_PYTHON_DEV_PACKAGES
     AND (LLVM_COMPILER_IS_GCC_COMPATIBLE OR CLANG_CL))
    # Avoid some warnings from upstream nanobind.
    # If a superproject set MLIR_DISABLE_CONFIGURE_PYTHON_DEV_PACKAGES, let
    # the super project handle compile options as it wishes.
    _mlir_python_eh_rtti_flags()
    target_compile_options(${NB_LIBRARY_TARGET_NAME}
      PRIVATE
        -Wno-c++98-compat-extra-semi
        -Wno-cast-qual
        -Wno-covered-switch-default
        -Wno-deprecated-literal-operator
        -Wno-nested-anon-types
        -Wno-unused-parameter
        -Wno-zero-length-array
        -Wno-missing-field-initializers
        ${eh_rtti_enable})
  endif()

  # Apply caller-provided extra options last so they have higher precedence.
  target_compile_options(${NB_LIBRARY_TARGET_NAME}
    PRIVATE
      ${ARG_EXTRA_COMPILE_OPTIONS}
  )
  target_link_libraries(${NB_LIBRARY_TARGET_NAME}
    PRIVATE
      ${ARG_EXTRA_LINK_LIBS}
  )
  # Propagate stable ABI to the shared nanobind library. nanobind internally
  # skips this when the interpreter is free-threaded.
  if(MLIR_ENABLE_PYTHON_STABLE_ABI AND NOT (NB_ABI MATCHES "[0-9]t"))
    target_compile_definitions(${NB_LIBRARY_TARGET_NAME}
      PUBLIC
      Py_LIMITED_API=0x030C0000
    )
    target_link_libraries(${NB_LIBRARY_TARGET_NAME} PRIVATE Python::SABIModule)
  endif()
  if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
    # nanobind handles this correctly for MacOS by explicitly setting -U for all the necessary Python symbols
    # (see https://github.com/wjakob/nanobind/blob/master/cmake/darwin-ld-cpython.sym)
    # but since we set -z,defs in llvm/cmake/modules/HandleLLVMOptions.cmake:340 for all Linux shlibs
    # we need to negate it here (we could have our own linux-ld-cpython.sym but that would be too much
    # maintenance).
    target_link_options(${NB_LIBRARY_TARGET_NAME} PRIVATE "LINKER:-z,undefs")
  endif()
  # nanobind configures with LTO for shared build which doesn't work everywhere
  # (see https://github.com/llvm/llvm-project/issues/139602).
  if(NOT LLVM_ENABLE_LTO)
    set_target_properties(${NB_LIBRARY_TARGET_NAME} PROPERTIES
      INTERPROCEDURAL_OPTIMIZATION_RELEASE OFF
      INTERPROCEDURAL_OPTIMIZATION_MINSIZEREL OFF
    )
  endif()
  set_target_properties(${NB_LIBRARY_TARGET_NAME} PROPERTIES
    LIBRARY_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
    BINARY_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
    # Needed for windows (and doesn't hurt others).
    RUNTIME_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
    ARCHIVE_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
  )
  mlir_python_setup_extension_rpath(${NB_LIBRARY_TARGET_NAME})
  install(TARGETS ${NB_LIBRARY_TARGET_NAME}
    COMPONENT ${ARG_INSTALL_COMPONENT}
    LIBRARY DESTINATION "${ARG_INSTALL_DESTINATION}"
    RUNTIME DESTINATION "${ARG_INSTALL_DESTINATION}"
  )
endfunction()

# Function: add_mlir_python_modules
# Adds python modules to a project, building them from a list of declared
# source groupings (see declare_mlir_python_sources and
# declare_mlir_python_extension). One of these must be called for each
# packaging root in use.
# Arguments:
#   ROOT_PREFIX: The directory in the build tree to emit sources. This will
#     typically be something like ${MY_BINARY_DIR}/python_packages/foobar
#     for non-relocatable modules or a deeper directory tree for relocatable.
#   INSTALL_PREFIX: Prefix into the install tree for installing the package.
#     Typically mirrors the path above but without an absolute path.
#   MLIR_BINDINGS_PYTHON_NB_DOMAIN: nanobind (and MLIR) domain within which
#     extensions will be compiled. This determines whether this package
#     will share nanobind types with other bindings packages. Expected to be unique
#     per project (and per specific set of bindings, for projects with multiple
#     bindings packages).
#   DECLARED_SOURCES: List of declared source groups to include. The entire
#     DAG of source modules is included.
#   COMMON_CAPI_LINK_LIBS: List of dylibs (typically one) to make every
#     extension depend on (see mlir_python_add_common_capi_library).
#   MLIR_BINDINGS_PYTHON_EXTRA_NANOBIND_COMPILE_OPTIONS: Extra compile options
#     to add to the shared nanobind runtime library.
#   MLIR_BINDINGS_PYTHON_EXTRA_NANOBIND_LINK_LIBS: Extra libraries to link into
#     the shared nanobind runtime library.
#   MLIR_BINDINGS_PYTHON_EXTRA_EXTENSION_COMPILE_OPTIONS: Extra compile options
#     to add to every generated extension DSO.
#   MLIR_BINDINGS_PYTHON_EXTRA_EXTENSION_LINK_LIBS: Extra libraries to link into
#     every generated extension DSO.
function(add_mlir_python_modules name)
  cmake_parse_arguments(ARG
    ""
    "ROOT_PREFIX;INSTALL_PREFIX;MLIR_BINDINGS_PYTHON_NB_DOMAIN"
    "COMMON_CAPI_LINK_LIBS;DECLARED_SOURCES;MLIR_BINDINGS_PYTHON_EXTRA_NANOBIND_COMPILE_OPTIONS;MLIR_BINDINGS_PYTHON_EXTRA_NANOBIND_LINK_LIBS;MLIR_BINDINGS_PYTHON_EXTRA_EXTENSION_COMPILE_OPTIONS;MLIR_BINDINGS_PYTHON_EXTRA_EXTENSION_LINK_LIBS"
    ${ARGN})

  # TODO(max): do the same for MLIR_PYTHON_PACKAGE_PREFIX?
  if((NOT ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN) AND MLIR_BINDINGS_PYTHON_NB_DOMAIN)
    set(ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN ${MLIR_BINDINGS_PYTHON_NB_DOMAIN})
  endif()
  if((NOT ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN) OR ("${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}" STREQUAL ""))
    message(WARNING "MLIR_BINDINGS_PYTHON_NB_DOMAIN CMake var is not set - setting to a default `mlir`.\
      It is highly recommend to set this to something unique so that your project's Python bindings do not collide with\
      others'. You also pass explicitly to `add_mlir_python_modules`.\
      See https://github.com/llvm/llvm-project/pull/171775 for more information.")
    set(ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN "mlir")
  endif()

  # Collect up all the transitive extension targets.
  _flatten_mlir_python_targets(_flat_targets ${ARG_DECLARED_SOURCES})

  # Modules with only pure python sources shouldn't build libnanobind.
  set(_needs_build_nanobind_lib OFF)
  foreach(sources_target ${_flat_targets})
    get_target_property(_source_type ${sources_target} mlir_python_SOURCES_TYPE)
    if((_source_type STREQUAL "support") OR (_source_type STREQUAL "extension"))
      set(_needs_build_nanobind_lib ON)
      break()
    endif()
  endforeach()

  # This call sets NB_LIBRARY_TARGET_NAME.
  if(_needs_build_nanobind_lib)
    build_nanobind_lib(
      INSTALL_COMPONENT ${name}
      INSTALL_DESTINATION "${ARG_INSTALL_PREFIX}/_mlir_libs"
      OUTPUT_DIRECTORY "${ARG_ROOT_PREFIX}/_mlir_libs"
      MLIR_BINDINGS_PYTHON_NB_DOMAIN ${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}
      EXTRA_COMPILE_OPTIONS
        ${ARG_MLIR_BINDINGS_PYTHON_EXTRA_NANOBIND_COMPILE_OPTIONS}
      EXTRA_LINK_LIBS
        ${ARG_MLIR_BINDINGS_PYTHON_EXTRA_NANOBIND_LINK_LIBS}
    )
  endif()

  # Helper to process an individual target.
  function(_process_target modules_target sources_target support_libs)
    get_target_property(_source_type ${sources_target} mlir_python_SOURCES_TYPE)

    if(_source_type STREQUAL "pure")
      # Pure python sources to link into the tree.
      set(_pure_sources_target "${modules_target}.sources.${sources_target}")
      add_mlir_python_sources_target(${_pure_sources_target}
        INSTALL_COMPONENT ${modules_target}
        INSTALL_DIR ${ARG_INSTALL_PREFIX}
        OUTPUT_DIRECTORY ${ARG_ROOT_PREFIX}
        SOURCES_TARGETS ${sources_target}
      )
      add_dependencies(${modules_target} ${_pure_sources_target})
    elseif(_source_type STREQUAL "extension")
      # Native CPP extension.
      get_target_property(_module_name ${sources_target} mlir_python_EXTENSION_MODULE_NAME)
      # Transform relative source to based on root dir.
      set(_extension_target "${modules_target}.extension.${_module_name}.dso")
      add_mlir_python_extension(${_extension_target} "${_module_name}" ${NB_LIBRARY_TARGET_NAME}
        INSTALL_COMPONENT ${modules_target}
        INSTALL_DIR "${ARG_INSTALL_PREFIX}/_mlir_libs"
        OUTPUT_DIRECTORY "${ARG_ROOT_PREFIX}/_mlir_libs"
        MLIR_BINDINGS_PYTHON_NB_DOMAIN ${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}
        EXTRA_COMPILE_OPTIONS
          ${ARG_MLIR_BINDINGS_PYTHON_EXTRA_EXTENSION_COMPILE_OPTIONS}
        EXTRA_LINK_LIBS
          ${ARG_MLIR_BINDINGS_PYTHON_EXTRA_EXTENSION_LINK_LIBS}
        LINK_LIBS PRIVATE
          ${sources_target}
          ${ARG_COMMON_CAPI_LINK_LIBS}
          ${support_libs}
      )
      add_dependencies(${modules_target} ${_extension_target})
      mlir_python_setup_extension_rpath(${_extension_target})
    elseif(_source_type STREQUAL "support")
      # do nothing because already built
    else()
      message(SEND_ERROR "Unrecognized source type '${_source_type}' for python source target ${sources_target}")
      return()
    endif()
  endfunction()

  # Build the modules target.
  add_custom_target(${name} ALL)

  # Build all support libs first.
  set(_mlir_python_support_libs)
  foreach(sources_target ${_flat_targets})
    get_target_property(_source_type ${sources_target} mlir_python_SOURCES_TYPE)
    if(_source_type STREQUAL "support")
      get_target_property(_module_name ${sources_target} mlir_python_EXTENSION_MODULE_NAME)
      # Use a similar mechanism as nanobind to help the runtime loader pick the correct lib.
      set(_module_name "${_module_name}-${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}")
      set(_extension_target "${name}.extension.${_module_name}.so")
      add_mlir_python_extension(${_extension_target} "${_module_name}" ${NB_LIBRARY_TARGET_NAME}
        INSTALL_COMPONENT ${name}
        INSTALL_DIR "${ARG_INSTALL_PREFIX}/_mlir_libs"
        OUTPUT_DIRECTORY "${ARG_ROOT_PREFIX}/_mlir_libs"
        MLIR_BINDINGS_PYTHON_NB_DOMAIN ${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}
        EXTRA_COMPILE_OPTIONS
          ${ARG_MLIR_BINDINGS_PYTHON_EXTRA_EXTENSION_COMPILE_OPTIONS}
        EXTRA_LINK_LIBS
          ${ARG_MLIR_BINDINGS_PYTHON_EXTRA_EXTENSION_LINK_LIBS}
        _PRIVATE_SUPPORT_LIB
        LINK_LIBS PRIVATE
        # LLVMSupport is intentionally removed to avoid introducing an LLVM dependency
        # for the mlir-python bindings. Do not add new dependencies on the C++ LLVM/MLIR
        # libraries; use the C++ standard library instead, or wrap LLVM functionality in
        # the C API first.
          ${sources_target}
          ${ARG_COMMON_CAPI_LINK_LIBS}
      )
      add_dependencies(${name} ${_extension_target})
      mlir_python_setup_extension_rpath(${_extension_target})
      list(APPEND _mlir_python_support_libs "${_extension_target}")
    endif()
  endforeach()

  # Build extensions.
  foreach(sources_target ${_flat_targets})
    _process_target(${name} ${sources_target} "${_mlir_python_support_libs}")
  endforeach()

  # Create an install target.
  if(NOT LLVM_ENABLE_IDE)
    add_llvm_install_targets(
      install-${name}
      DEPENDS ${name}
      COMPONENT ${name})
  endif()
endfunction()

# Function: declare_mlir_dialect_python_bindings
# Helper to generate source groups for dialects, including both static source
# files and a TD_FILE to generate wrappers.
#
# This will generate a source group named ${ADD_TO_PARENT}.${DIALECT_NAME}.
#
# Arguments:
#   ROOT_DIR: Same as for declare_mlir_python_sources().
#   ADD_TO_PARENT: Same as for declare_mlir_python_sources(). Unique names
#     for the subordinate source groups are derived from this.
#   TD_FILE: Tablegen file to generate source for (relative to ROOT_DIR).
#   DIALECT_NAME: Python name of the dialect.
#   SOURCES: Same as declare_mlir_python_sources().
#   SOURCES_GLOB: Same as declare_mlir_python_sources().
#   DEPENDS: Additional dependency targets.
#   GEN_ENUM_BINDINGS: Generate enum bindings.
#   GEN_ENUM_BINDINGS_TD_FILE: Optional Tablegen file to generate enums for (relative to ROOT_DIR).
#     This file is where the *EnumAttrs are defined, not where the *Enums are defined.
#     **WARNING**: This arg will shortly be removed when the just-below TODO is satisfied. Use at your
#     risk.
#
# TODO: Right now `TD_FILE` can't be the actual dialect tablegen file, since we
#       use its path to determine where to place the generated python file. If
#       we made the output path an additional argument here we could remove the
#       need for the separate "wrapper" .td files
function(declare_mlir_dialect_python_bindings)
  cmake_parse_arguments(ARG
    "GEN_ENUM_BINDINGS"
    "ROOT_DIR;ADD_TO_PARENT;TD_FILE;DIALECT_NAME"
    "SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE"
    ${ARGN})
  # Sources.
  set(_dialect_target "${ARG_ADD_TO_PARENT}.${ARG_DIALECT_NAME}")
  declare_mlir_python_sources(${_dialect_target}
    ROOT_DIR "${ARG_ROOT_DIR}"
    ADD_TO_PARENT "${ARG_ADD_TO_PARENT}"
    SOURCES "${ARG_SOURCES}"
    SOURCES_GLOB "${ARG_SOURCES_GLOB}"
  )

  # Tablegen
  if(ARG_TD_FILE)
    set(tblgen_target "${_dialect_target}.tablegen")
    set(td_file "${ARG_ROOT_DIR}/${ARG_TD_FILE}")
    get_filename_component(relative_td_directory "${ARG_TD_FILE}" DIRECTORY)
    file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/${relative_td_directory}")
    set(dialect_filename "${relative_td_directory}/_${ARG_DIALECT_NAME}_ops_gen.py")
    set(LLVM_TARGET_DEFINITIONS ${td_file})
    mlir_tablegen("${dialect_filename}"
      -gen-python-op-bindings -bind-dialect=${ARG_DIALECT_NAME}
      DEPENDS ${ARG_DEPENDS}
    )
    add_public_tablegen_target(${tblgen_target})

    set(_sources ${dialect_filename})
    if(ARG_GEN_ENUM_BINDINGS OR ARG_GEN_ENUM_BINDINGS_TD_FILE)
      if(ARG_GEN_ENUM_BINDINGS_TD_FILE)
        set(td_file "${ARG_ROOT_DIR}/${ARG_GEN_ENUM_BINDINGS_TD_FILE}")
        set(LLVM_TARGET_DEFINITIONS ${td_file})
      endif()
      set(enum_filename "${relative_td_directory}/_${ARG_DIALECT_NAME}_enum_gen.py")
      mlir_tablegen(${enum_filename} -gen-python-enum-bindings -bind-dialect=${ARG_DIALECT_NAME})
      list(APPEND _sources ${enum_filename})
    endif()

    # Generated.
    declare_mlir_python_sources("${_dialect_target}.ops_gen"
      ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
      ADD_TO_PARENT "${_dialect_target}"
      SOURCES ${_sources}
    )
  endif()
endfunction()

# Function: declare_mlir_dialect_extension_python_bindings
# Helper to generate source groups for dialect extensions, including both
# static source files and a TD_FILE to generate wrappers.
#
# This will generate a source group named ${ADD_TO_PARENT}.${EXTENSION_NAME}.
#
# Arguments:
#   ROOT_DIR: Same as for declare_mlir_python_sources().
#   ADD_TO_PARENT: Same as for declare_mlir_python_sources(). Unique names
#     for the subordinate source groups are derived from this.
#   TD_FILE: Tablegen file to generate source for (relative to ROOT_DIR).
#   DIALECT_NAME: Python name of the dialect.
#   EXTENSION_NAME: Python name of the dialect extension.
#   SOURCES: Same as declare_mlir_python_sources().
#   SOURCES_GLOB: Same as declare_mlir_python_sources().
#   DEPENDS: Additional dependency targets.
#   GEN_ENUM_BINDINGS: Generate enum bindings.
#   GEN_ENUM_BINDINGS_TD_FILE: Optional Tablegen file to generate enums for (relative to ROOT_DIR).
#     This file is where the *Attrs are defined, not where the *Enums are defined.
#     **WARNING**: This arg will shortly be removed when the TODO for
#     declare_mlir_dialect_python_bindings is satisfied. Use at your risk.
function(declare_mlir_dialect_extension_python_bindings)
  cmake_parse_arguments(ARG
    "GEN_ENUM_BINDINGS"
    "ROOT_DIR;ADD_TO_PARENT;TD_FILE;DIALECT_NAME;EXTENSION_NAME"
    "SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE"
    ${ARGN})
  # Source files.
  set(_extension_target "${ARG_ADD_TO_PARENT}.${ARG_EXTENSION_NAME}")
  declare_mlir_python_sources(${_extension_target}
    ROOT_DIR "${ARG_ROOT_DIR}"
    ADD_TO_PARENT "${ARG_ADD_TO_PARENT}"
    SOURCES "${ARG_SOURCES}"
    SOURCES_GLOB "${ARG_SOURCES_GLOB}"
  )

  # Tablegen
  if(ARG_TD_FILE)
    set(tblgen_target "${ARG_ADD_TO_PARENT}.${ARG_EXTENSION_NAME}.tablegen")
    set(td_file "${ARG_ROOT_DIR}/${ARG_TD_FILE}")
    get_filename_component(relative_td_directory "${ARG_TD_FILE}" DIRECTORY)
    file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/${relative_td_directory}")
    set(output_filename "${relative_td_directory}/_${ARG_EXTENSION_NAME}_ops_gen.py")
    set(LLVM_TARGET_DEFINITIONS ${td_file})
    mlir_tablegen("${output_filename}" -gen-python-op-bindings
                  -bind-dialect=${ARG_DIALECT_NAME}
                  -dialect-extension=${ARG_EXTENSION_NAME})
    add_public_tablegen_target(${tblgen_target})
    if(ARG_DEPENDS)
      add_dependencies(${tblgen_target} ${ARG_DEPENDS})
    endif()

    set(_sources ${output_filename})
    if(ARG_GEN_ENUM_BINDINGS OR ARG_GEN_ENUM_BINDINGS_TD_FILE)
      if(ARG_GEN_ENUM_BINDINGS_TD_FILE)
        set(td_file "${ARG_ROOT_DIR}/${ARG_GEN_ENUM_BINDINGS_TD_FILE}")
        set(LLVM_TARGET_DEFINITIONS ${td_file})
      endif()
      set(enum_filename "${relative_td_directory}/_${ARG_EXTENSION_NAME}_enum_gen.py")
      mlir_tablegen(${enum_filename} -gen-python-enum-bindings -bind-dialect=${ARG_DIALECT_NAME})
      list(APPEND _sources ${enum_filename})
    endif()

    declare_mlir_python_sources("${_extension_target}.ops_gen"
      ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
      ADD_TO_PARENT "${_extension_target}"
      SOURCES ${_sources}
    )
  endif()
endfunction()

# Function: mlir_python_setup_extension_rpath
# Sets RPATH properties on a target, assuming that it is being output to
# an _mlir_libs directory with all other libraries. For static linkage,
# the RPATH will just be the origin. If linking dynamically, then the LLVM
# library directory will be added.
# Arguments:
#   RELATIVE_INSTALL_ROOT: If building dynamically, an RPATH entry will be
#     added to the install tree lib/ directory by first traversing this
#     path relative to the installation location. Typically a number of ".."
#     entries, one for each level of the install path.
function(mlir_python_setup_extension_rpath target)
  cmake_parse_arguments(ARG
    ""
    "RELATIVE_INSTALL_ROOT"
    ""
    ${ARGN})

  # RPATH handling.
  # For the build tree, include the LLVM lib directory and the current
  # directory for RPATH searching. For install, just the current directory
  # (assumes that needed dependencies have been installed).
  if(NOT APPLE AND NOT UNIX)
    return()
  endif()

  set(_origin_prefix "\$ORIGIN")
  if(APPLE)
    set(_origin_prefix "@loader_path")
  endif()
  set_target_properties(${target} PROPERTIES
    BUILD_WITH_INSTALL_RPATH OFF
    BUILD_RPATH "${_origin_prefix}"
    INSTALL_RPATH "${_origin_prefix}"
  )

  # For static builds, that is all that is needed: all dependencies will be in
  # the one directory. For shared builds, then we also need to add the global
  # lib directory. This will be absolute for the build tree and relative for
  # install.
  # When we have access to CMake >= 3.20, there is a helper to calculate this.
  if(BUILD_SHARED_LIBS AND ARG_RELATIVE_INSTALL_ROOT)
    get_filename_component(_real_lib_dir "${LLVM_LIBRARY_OUTPUT_INTDIR}" REALPATH)
    set_property(TARGET ${target} APPEND PROPERTY
      BUILD_RPATH "${_real_lib_dir}")
    set_property(TARGET ${target} APPEND PROPERTY
      INSTALL_RPATH "${_origin_prefix}/${ARG_RELATIVE_INSTALL_ROOT}/lib${LLVM_LIBDIR_SUFFIX}")
  endif()
endfunction()

# Function: add_mlir_python_common_capi_library
# Adds a shared library which embeds dependent CAPI libraries needed to link
# all extensions.
# Arguments:
#   INSTALL_COMPONENT: Name of the install component. Typically same as the
#     target name passed to add_mlir_python_modules().
#   INSTALL_DESTINATION: Prefix into the install tree in which to install the
#     library.
#   OUTPUT_DIRECTORY: Full path in the build tree in which to create the
#     library. Typically, this will be the common _mlir_libs directory where
#     all extensions are emitted.
#   RELATIVE_INSTALL_ROOT: See mlir_python_setup_extension_rpath().
#   DECLARED_HEADERS: Source groups from which to discover headers that belong
#     to the library and should be installed with it.
#   DECLARED_SOURCES: Source groups from which to discover dependent
#     EMBED_CAPI_LINK_LIBS.
#   EMBED_LIBS: Additional libraries to embed (must be built with OBJECTS and
#     have an "obj.${name}" object library associated).
function(add_mlir_python_common_capi_library name)
  cmake_parse_arguments(ARG
    ""
    "INSTALL_COMPONENT;INSTALL_DESTINATION;OUTPUT_DIRECTORY;RELATIVE_INSTALL_ROOT"
    "DECLARED_HEADERS;DECLARED_SOURCES;EMBED_LIBS"
    ${ARGN})
  # Collect all explicit and transitive embed libs.
  set(_embed_libs ${ARG_EMBED_LIBS})
  _flatten_mlir_python_targets(_all_source_targets ${ARG_DECLARED_SOURCES})
  foreach(t ${_all_source_targets})
    get_target_property(_local_embed_libs ${t} mlir_python_EMBED_CAPI_LINK_LIBS)
    if(_local_embed_libs)
      list(APPEND _embed_libs ${_local_embed_libs})
    endif()
  endforeach()
  list(REMOVE_DUPLICATES _embed_libs)

  # Generate the aggregate .so that everything depends on.
  add_mlir_aggregate(${name}
    SHARED
    DISABLE_INSTALL
    EMBED_LIBS ${_embed_libs}
  )

  # Process any headers associated with the library
  _flatten_mlir_python_targets(_flat_header_targets ${ARG_DECLARED_HEADERS})
  set(_header_sources_target "${name}.sources")
  add_mlir_python_sources_target(${_header_sources_target}
    INSTALL_COMPONENT ${ARG_INSTALL_COMPONENT}
    INSTALL_DIR "${ARG_INSTALL_DESTINATION}/include"
    OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}/include"
    SOURCES_TARGETS ${_flat_header_targets}
  )
  add_dependencies(${name} ${_header_sources_target})

  if(WIN32)
    set_property(TARGET ${name} PROPERTY WINDOWS_EXPORT_ALL_SYMBOLS ON)
  endif()
  set_target_properties(${name} PROPERTIES
    LIBRARY_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
    BINARY_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
    # Needed for windows (and doesn't hurt others).
    RUNTIME_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
    ARCHIVE_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
  )
  mlir_python_setup_extension_rpath(${name}
    RELATIVE_INSTALL_ROOT "${ARG_RELATIVE_INSTALL_ROOT}"
  )
  install(TARGETS ${name}
    COMPONENT ${ARG_INSTALL_COMPONENT}
    LIBRARY DESTINATION "${ARG_INSTALL_DESTINATION}"
    RUNTIME DESTINATION "${ARG_INSTALL_DESTINATION}"
  )
endfunction()

function(_flatten_mlir_python_targets output_var)
  set(_flattened)
  foreach(t ${ARGN})
    get_target_property(_source_type ${t} mlir_python_SOURCES_TYPE)
    get_target_property(_depends ${t} mlir_python_DEPENDS)
    if(_source_type)
      list(APPEND _flattened "${t}")
      if(_depends)
        _flatten_mlir_python_targets(_local_flattened ${_depends})
        list(APPEND _flattened ${_local_flattened})
      endif()
    endif()
  endforeach()
  list(REMOVE_DUPLICATES _flattened)
  set(${output_var} "${_flattened}" PARENT_SCOPE)
endfunction()

# Function: add_mlir_python_sources_target
# Adds a target corresponding to an interface target that carries source
# information. This target is responsible for "building" the sources by
# placing them in the correct locations in the build and install trees.
# Arguments:
#   INSTALL_COMPONENT: Name of the install component. Typically same as the
#     target name passed to add_mlir_python_modules().
#   INSTALL_DESTINATION: Prefix into the install tree in which to install the
#     library.
#   OUTPUT_DIRECTORY: Full path in the build tree in which to create the
#     library. Typically, this will be the common _mlir_libs directory where
#     all extensions are emitted.
#   SOURCES_TARGETS: List of interface libraries that carry source information.
function(add_mlir_python_sources_target name)
  cmake_parse_arguments(ARG
  ""
  "INSTALL_COMPONENT;INSTALL_DIR;OUTPUT_DIRECTORY"
  "SOURCES_TARGETS"
  ${ARGN})

  if(ARG_UNPARSED_ARGUMENTS)
    message(FATAL_ERROR "Unhandled arguments to add_mlir_python_sources_target(${name}, ... : ${ARG_UNPARSED_ARGUMENTS}")
  endif()

  # On Windows create_symlink requires special permissions. Use copy_if_different instead.
  if(CMAKE_HOST_WIN32)
    set(_link_or_copy copy_if_different)
  else()
    set(_link_or_copy create_symlink)
  endif()

  foreach(_sources_target ${ARG_SOURCES_TARGETS})
    get_target_property(_src_paths ${_sources_target} SOURCES)
    if(NOT _src_paths)
      get_target_property(_src_paths ${_sources_target} INTERFACE_SOURCES)
      if(NOT _src_paths)
        break()
      endif()
    endif()

    get_target_property(_root_dir ${_sources_target} INCLUDE_DIRECTORIES)
    if(NOT _root_dir)
      get_target_property(_root_dir ${_sources_target} INTERFACE_INCLUDE_DIRECTORIES)
    endif()

    # Initialize an empty list of all Python source destination paths.
    set(all_dest_paths "")
    foreach(_src_path ${_src_paths})
      file(RELATIVE_PATH _source_relative_path "${_root_dir}" "${_src_path}")
      set(_dest_path "${ARG_OUTPUT_DIRECTORY}/${_source_relative_path}")

      get_filename_component(_dest_dir "${_dest_path}" DIRECTORY)
      file(MAKE_DIRECTORY "${_dest_dir}")

      add_custom_command(
        OUTPUT "${_dest_path}"
        COMMENT "Copying python source ${_src_path} -> ${_dest_path}"
        DEPENDS "${_src_path}"
        COMMAND "${CMAKE_COMMAND}" -E ${_link_or_copy}
            "${_src_path}" "${_dest_path}"
      )

      # Track the symlink or copy command output.
      list(APPEND all_dest_paths "${_dest_path}")

      if(ARG_INSTALL_DIR)
        # We have to install each file individually because we need to preserve
        # the relative directory structure in the install destination.
        # As an example, ${_source_relative_path} may be dialects/math.py
        # which would be transformed to ${ARG_INSTALL_DIR}/dialects
        # here. This could be moved outside of the loop and cleaned up
        # if we used FILE_SETS (introduced in CMake 3.23).
        get_filename_component(_install_destination "${ARG_INSTALL_DIR}/${_source_relative_path}" DIRECTORY)
        install(
          FILES ${_src_path}
          DESTINATION "${_install_destination}"
          COMPONENT ${ARG_INSTALL_COMPONENT}
        )
      endif()
    endforeach()
  endforeach()

  # Create a new custom target that depends on all symlinked or copied sources.
  add_custom_target("${name}" DEPENDS ${all_dest_paths})
endfunction()

################################################################################
# Build python extension
################################################################################
function(add_mlir_python_extension libname extname nb_library_target_name)
  cmake_parse_arguments(ARG
    "_PRIVATE_SUPPORT_LIB"
    "INSTALL_COMPONENT;INSTALL_DIR;OUTPUT_DIRECTORY;MLIR_BINDINGS_PYTHON_NB_DOMAIN"
    "SOURCES;LINK_LIBS;EXTRA_COMPILE_OPTIONS;EXTRA_LINK_LIBS"
    ${ARGN})
  if(ARG_UNPARSED_ARGUMENTS)
    message(FATAL_ERROR "Unhandled arguments to add_mlir_python_extension(${libname}, ... : ${ARG_UNPARSED_ARGUMENTS}")
  endif()

  # The extension itself must be compiled with RTTI and exceptions enabled.
  # Also, some warning classes triggered by nanobind are disabled.
  _mlir_python_eh_rtti_flags()

  if(ARG__PRIVATE_SUPPORT_LIB)
    add_library(${libname} SHARED ${ARG_SOURCES})
    if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
      # nanobind handles this correctly for MacOS by explicitly setting -U for all the necessary Python symbols
      # (see https://github.com/wjakob/nanobind/blob/master/cmake/darwin-ld-cpython.sym)
      # but since we set -z,defs in llvm/cmake/modules/HandleLLVMOptions.cmake:340 for all Linux shlibs
      # we need to negate it here (we could have our own linux-ld-cpython.sym but that would be too much
      # maintenance).
      target_link_options(${libname} PRIVATE "LINKER:-z,undefs")
    endif()
    nanobind_link_options(${libname})
    target_compile_definitions(${libname}
      PRIVATE
      NB_DOMAIN=${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}
      MLIR_CAPI_BUILDING_LIBRARY=1
    )
    if(MSVC)
      set_property(TARGET ${libname} PROPERTY WINDOWS_EXPORT_ALL_SYMBOLS ON)
    endif()
  else()
    if(MLIR_ENABLE_PYTHON_STABLE_ABI)
      set(_stable_abi_flag STABLE_ABI)
    else()
      set(_stable_abi_flag "")
    endif()
    nanobind_add_module(${libname}
      NB_DOMAIN ${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}
      FREE_THREADED
      ${_stable_abi_flag}
      NB_SHARED
      ${ARG_SOURCES}
    )
    target_compile_definitions(${libname}
      PRIVATE
      MLIR_BINDINGS_PYTHON_DOMAIN=${MLIR_BINDINGS_PYTHON_NB_DOMAIN}
    )
  endif()
  target_link_libraries(${libname} PRIVATE ${nb_library_target_name})
  target_compile_definitions(${libname}
    PRIVATE
    MLIR_BINDINGS_PYTHON_DOMAIN=${ARG_MLIR_BINDINGS_PYTHON_NB_DOMAIN}
  )
  if(APPLE)
    # In llvm/cmake/modules/HandleLLVMOptions.cmake:268 we set -Wl,-flat_namespace which breaks
    # the default name spacing on MacOS and causes "cross-wired" symbol resolution when multiple
    # bindings packages are loaded.
    target_link_options(${libname} PRIVATE "LINKER:-twolevel_namespace")
  endif()

  if (NOT MLIR_DISABLE_CONFIGURE_PYTHON_DEV_PACKAGES
      AND (LLVM_COMPILER_IS_GCC_COMPATIBLE OR CLANG_CL))
    # Avoid some warnings from upstream nanobind.
    # If a superproject set MLIR_DISABLE_CONFIGURE_PYTHON_DEV_PACKAGES, let
    # the super project handle compile options as it wishes.
    target_compile_options(${libname}
      PRIVATE
        -Wno-c++98-compat-extra-semi
        -Wno-cast-qual
        -Wno-covered-switch-default
        -Wno-deprecated-literal-operator
        -Wno-nested-anon-types
        -Wno-unused-parameter
        -Wno-zero-length-array
        -Wno-missing-field-initializers)
  endif()

  if(APPLE)
    # NanobindAdaptors.h uses PyClassMethod_New to build `pure_subclass`es but nanobind
    # doesn't declare this API as undefined in its linker flags. So we need to declare it as such
    # for downstream users that do not do something like `-undefined dynamic_lookup`.
    # Same for the rest.
    target_link_options(${libname} PUBLIC
      "LINKER:-U,_PyClassMethod_New"
    )
    if(NOT MLIR_ENABLE_PYTHON_STABLE_ABI)
      # PyCode_Addr2Location and PyFrame_GetLasti are not part of the stable
      # ABI and are not referenced when Py_LIMITED_API is defined.
      target_link_options(${libname} PUBLIC
        "LINKER:-U,_PyCode_Addr2Location"
        "LINKER:-U,_PyFrame_GetLasti"
      )
    endif()
  endif()

  target_compile_options(${libname} PRIVATE ${eh_rtti_enable})

  # Apply caller-provided extra options last so they have higher precedence.
  target_compile_options(${libname} PRIVATE ${ARG_EXTRA_COMPILE_OPTIONS})

  # Quoting CMake:
  #
  # "If you use it on normal shared libraries which other targets link against, on some platforms a
  # linker will insert a full path to the library (as specified at link time) into the dynamic section of the
  # dependent binary. Therefore, once installed, dynamic loader may eventually fail to locate the library
  # for the binary."
  #
  # So for support libs we do need an SO name but for extensions we do not (they're MODULEs anyway -
  # i.e., can't be linked against, only loaded).
  if (ARG__PRIVATE_SUPPORT_LIB)
    set(_no_soname OFF)
  else ()
    set(_no_soname ON)
  endif ()
  # Configure the output to match python expectations.
  set_target_properties(
    ${libname} PROPERTIES
    LIBRARY_OUTPUT_DIRECTORY ${ARG_OUTPUT_DIRECTORY}
    OUTPUT_NAME "${extname}"
    NO_SONAME ${_no_soname}
  )

  if(WIN32)
    # Need to also set the RUNTIME_OUTPUT_DIRECTORY on Windows in order to
    # control where the .dll gets written.
    set_target_properties(
      ${libname} PROPERTIES
      RUNTIME_OUTPUT_DIRECTORY ${ARG_OUTPUT_DIRECTORY}
      ARCHIVE_OUTPUT_DIRECTORY ${ARG_OUTPUT_DIRECTORY}
    )
  endif()

  target_link_libraries(${libname}
    PRIVATE
    ${ARG_LINK_LIBS}
    ${ARG_EXTRA_LINK_LIBS}
  )

  target_link_options(${libname}
    PRIVATE
      # On Linux, disable re-export of any static linked libraries that
      # came through.
      $<$<PLATFORM_ID:Linux>:LINKER:--exclude-libs,ALL>
  )

  if(WIN32)
    # On Windows, pyconfig.h (and by extension python.h) hardcode the version of the
    # python library which will be used for linkage depending on the flavor of the build.
    # nanobind has a workaround which depends on the definition of Py_DEBUG (if Py_DEBUG
    # is not passed in as a compile definition, nanobind undefs _DEBUG when including
    # python.h, so that the release python library would be used).
    # Since mlir uses nanobind, we can leverage their workaround by never directly
    # pyconfig.h or python.h and instead relying on the nanobind headers to include the
    # necessary python headers. This results in mlir always linking against the
    # release python library via the (undocumented) cmake property Python3_LIBRARY_RELEASE.
    target_link_libraries(${libname} PRIVATE ${Python3_LIBRARY_RELEASE})
  endif()

  ################################################################################
  # Install
  ################################################################################
  if(ARG_INSTALL_DIR)
    install(TARGETS ${libname}
      COMPONENT ${ARG_INSTALL_COMPONENT}
      LIBRARY DESTINATION ${ARG_INSTALL_DIR}
      ARCHIVE DESTINATION ${ARG_INSTALL_DIR}
      # NOTE: Even on DLL-platforms, extensions go in the lib directory tree.
      RUNTIME DESTINATION ${ARG_INSTALL_DIR}
    )
  endif()
endfunction()