# Add the following code before all include to avoid compilation failure.
set(UNITY_CC_BEFORE_CODE [[
#ifndef NOMINMAX
#define NOMINMAX
#endif
#ifndef _USE_MATH_DEFINES
#define _USE_MATH_DEFINES
#endif]])
set(UNITY_CU_BEFORE_CODE [[
#ifndef __CUDACC_VER_MAJOR__
#define __CUDACC_VER_MAJOR__ CUDA_COMPILER_MAJOR_VERSION
#endif
#ifndef __CUDACC_VER_MINOR__
#define __CUDACC_VER_MINOR__ CUDA_COMPILER_MINOR_VERSION
#endif]])
if(WITH_GPU)
    string(REPLACE "." ";" CUDA_COMPILER_VERSION ${CMAKE_CUDA_COMPILER_VERSION})
    list(GET CUDA_COMPILER_VERSION 0 CUDA_COMPILER_MAJOR_VERSION)
    list(GET CUDA_COMPILER_VERSION 1 CUDA_COMPILER_MINOR_VERSION)
    string(REPLACE
        "CUDA_COMPILER_MAJOR_VERSION" ${CUDA_COMPILER_MAJOR_VERSION}
        UNITY_CU_BEFORE_CODE ${UNITY_CU_BEFORE_CODE})
    string(REPLACE
        "CUDA_COMPILER_MINOR_VERSION" ${CUDA_COMPILER_MINOR_VERSION}
        UNITY_CU_BEFORE_CODE ${UNITY_CU_BEFORE_CODE})
endif()

# Group a list of source files that can be included together.
# This combination is just a guiding rule, and the source file of group
# do not have to exist.
# Here you need to specify the source type which belongs to cc or cu.
function(register_unity_group TYPE)
    # Get UNITY_TARGET from CMAKE_CURRENT_SOURCE_DIR.
    string(REPLACE "${PADDLE_SOURCE_DIR}/paddle/fluid/" "" UNITY_TARGET ${CMAKE_CURRENT_SOURCE_DIR})
    string(REPLACE "/" "_" UNITY_TARGET ${UNITY_TARGET})
    set(UNITY_TARGET "paddle_${UNITY_TARGET}_unity")

    # Variable unity_group_index is used to record the number of UNITY_TARGET groups.
    get_property(unity_group_index GLOBAL PROPERTY ${UNITY_TARGET}_${TYPE}_group_index)
    if("${unity_group_index}" STREQUAL "")
        set(unity_group_index 0)
    endif()

    # Variable unity_group_sources is used to record the sources of one group.
    set(unity_group_sources ${UNITY_TARGET}_${TYPE}_group_${unity_group_index}_sources)
    set_property(GLOBAL PROPERTY ${unity_group_sources} "")
    foreach(src ${ARGN})
        # UB use absolute path of source.
        if(NOT IS_ABSOLUTE ${src})
            set(src ${CMAKE_CURRENT_SOURCE_DIR}/${src})
        endif()
        set_property(GLOBAL APPEND PROPERTY ${unity_group_sources} ${src})
    endforeach()

    # If unity_file does not exists, nv_library or cc_library will use
    # dummy_file. Touch unity_file to avoid to use dummy file.
    set(unity_file ${CMAKE_CURRENT_BINARY_DIR}/${UNITY_TARGET}_${unity_group_index}_${TYPE}.${TYPE})
    if(NOT EXISTS ${unity_file})
        file(TOUCH ${unity_file})
    endif()

    math(EXPR unity_group_index "${unity_group_index} + 1")
    set_property(GLOBAL PROPERTY ${UNITY_TARGET}_${TYPE}_group_index ${unity_group_index})
endfunction(register_unity_group)

# Combine the original source files used by `TARGET`, then use
# `unity_target_${TYPE}_sources` to get the combined source files.
# If the source file does not hit any registed groups, use itself.
# This function put the actual combination relationship in variables instead of
# writing the unity source file. The reason is that writing unity source file
# will change the timestampe and affect the effect of retaining the build
# directory on Windows.
# Here you need to specify the source type which belongs to cc or cu.
function(compose_unity_target_sources TARGET TYPE)
    # Variable unity_target_sources represents the source file used in TARGET
    set(unity_target_sources "")
    get_property(unity_group_index_max GLOBAL PROPERTY ${TARGET}_${TYPE}_group_index)
    foreach(src ${ARGN})
        set(unity_file "")
        # UB use absolute path of source.
        if(IS_ABSOLUTE ${src})
            set(src_absolute_path ${src})
        else()
            set(src_absolute_path ${CMAKE_CURRENT_SOURCE_DIR}/${src})
        endif()
        # If `unity_group_index_max` is empty, there is no combination
        # relationship.
        # TODO(Avin0323): Whether use target property `UNITY_BUILD` of CMAKE to
        # combine source files.
        if(NOT "${unity_group_index_max}" STREQUAL "")
            # Search in each registed group.
            foreach(unity_group_index RANGE ${unity_group_index_max})
                if(${unity_group_index} GREATER_EQUAL ${unity_group_index_max})
                    break()
                endif()
                get_property(unity_group_sources GLOBAL PROPERTY ${TARGET}_${TYPE}_group_${unity_group_index}_sources)
                if(${src_absolute_path} IN_LIST unity_group_sources)
                    set(unity_file ${CMAKE_CURRENT_BINARY_DIR}/${TARGET}_${unity_group_index}_${TYPE}.${TYPE})
                    set(unity_file_sources ${TARGET}_${TYPE}_file_${unity_group_index}_sources)
                    get_property(set_unity_file_sources GLOBAL PROPERTY ${unity_file_sources} SET)
                    if(NOT ${set_unity_file_sources})
                        # Add macro before include source files.
                        set_property(GLOBAL PROPERTY ${unity_file_sources} "// Generate by Unity Build")
                        set_property(GLOBAL APPEND PROPERTY ${unity_file_sources} ${UNITY_CC_BEFORE_CODE})
                        if(WITH_GPU AND "${TYPE}" STREQUAL "cu")
                            set_property(GLOBAL APPEND PROPERTY ${unity_file_sources} ${UNITY_CU_BEFORE_CODE})
                        endif()
                    endif()
                    set_property(GLOBAL APPEND PROPERTY ${unity_file_sources} "#include \"${src_absolute_path}\"")
                    set(unity_target_sources ${unity_target_sources} ${unity_file})
                    break()
                endif()
            endforeach()
        endif()
        # Use original source file.
        if("${unity_file}" STREQUAL "")
            set(unity_target_sources ${unity_target_sources} ${src})
        endif()
    endforeach()

    set(unity_target_${TYPE}_sources ${unity_target_sources} PARENT_SCOPE)
endfunction(compose_unity_target_sources)

# Write the unity files used by `UNITY_TARGET`.
# Write dependent on whether the contents of the unity file have changed, which
# protects incremental compilation speed.
function(finish_unity_target TYPE)
    # Get UNITY_TARGET from CMAKE_CURRENT_SOURCE_DIR.
    string(REPLACE "${PADDLE_SOURCE_DIR}/paddle/fluid/" "" UNITY_TARGET ${CMAKE_CURRENT_SOURCE_DIR})
    string(REPLACE "/" "_" UNITY_TARGET ${UNITY_TARGET})
    set(UNITY_TARGET "paddle_${UNITY_TARGET}_unity")

    get_property(unity_group_index_max GLOBAL PROPERTY ${UNITY_TARGET}_${TYPE}_group_index)
    if(NOT "${unity_group_index_max}" STREQUAL "")
        foreach(unity_group_index RANGE ${unity_group_index_max})
            if(${unity_group_index} GREATER_EQUAL ${unity_group_index_max})
                break()
            endif()
            get_property(unity_file_sources GLOBAL PROPERTY ${UNITY_TARGET}_${TYPE}_file_${unity_group_index}_sources)
            set(unity_file_read_content "")
            string(JOIN "\n" unity_file_write_content ${unity_file_sources})
            set(unity_file ${CMAKE_CURRENT_BINARY_DIR}/${UNITY_TARGET}_${unity_group_index}_${TYPE}.${TYPE})
            file(READ ${unity_file} unity_file_read_content)
            if(NOT "${unity_file_read_content}" STREQUAL "${unity_file_write_content}")
                file(WRITE ${unity_file} ${unity_file_write_content})
            endif()
        endforeach()
    endif()
endfunction(finish_unity_target)
