#!/bin/bash
# Copyright (c) 2017, Medicine Yeh
# MIT License

# Upstream page: https://github.com/MedicineYeh/p-clang-format
# Local modifications:
#   - https://github.com/MedicineYeh/p-clang-format/pull/2
#   - CUDA_CHEVRON_* hack to protect "<<<" and ">>>" from clang-format
#     (see https://github.com/llvm/llvm-project/issues/52881)

# Turn on this to replace any pragma
#REPLACE_ANY_PRAGMA=1

function replace_all()
{
    local prefix="$1"
    local keywords=("${!2}")
    local pattern="$3"
    local file_pathes=("${!4}")

    # echo "Keywords     : ${keywords[*]}"
    # echo "Pattern      : ${pattern}"
    # echo "File Path    : ${file_pathes[*]}"

    for file_path in "${file_pathes[@]}"; do
        [[ ! -r "$file_path" ]] && continue
        for w in "${keywords[@]}"; do
            # echo "Replace '${prefix}$w' in '$file_path'"
            sed -i "s/${prefix}${w}/${pattern}${prefix}${w}/g" "${file_path}"
        done

        # HACK: protect "<<<" and ">>>" from clang-format
        sed -i "s/<<</CUDA_CHEVRON_OPEN/g" "${file_path}"
        sed -i "s/>>>/CUDA_CHEVRON_CLOSE/g" "${file_path}"
    done
}

function re_replace_all()
{
    local prefix="$1"
    local keywords=("${!2}")
    local pattern="$3"
    local file_pathes=("${!4}")

    # echo "Keywords     : ${keywords[*]}"
    # echo "Anti-Pattern : ${pattern}"
    # echo "File Path    : ${file_pathes[*]}"

    for file_path in "${file_pathes[@]}"; do
        [[ ! -r "$file_path" ]] && continue
        for w in "${keywords[@]}"; do
            # echo "Re:replace '${prefix}$w' in '$file_path'"
            sed -i "s/${pattern}${prefix}${w}/${prefix}${w}/g" "${file_path}"
        done

        # HACK: unprotect "<<<" and ">>>"
        sed -i "s/CUDA_CHEVRON_OPEN/<<</g" "${file_path}"
        sed -i "s/CUDA_CHEVRON_CLOSE/>>>/g" "${file_path}"
    done
}

function format()
{
    clang-format "$@"
}

function main()
{
    # This is the pattern for replacing to comments
    local comment_pattern="\/\/"
    # This is the pattern for replacing comments back to original string
    local re_comment_pattern="\/\/ *"
    # The path of files
    local file_pathes=()
    # Define the keywords of pragma to be escaped from formatting
    local pragma_prefix="#pragma "
    local pragma_key_words=()
    # Loop specific
    pragma_key_words+=(omp simd loop unroll ivdep vector)
    # Memory specific
    pragma_key_words+=(alloc_section)
    # Misc.
    pragma_key_words+=("cilk grainsize" offload)

    # Turn on the following line to indent any pragma
    [[ "$REPLACE_ANY_PRAGMA" == "1" ]] && pragma_key_words=("")

    # Find all files in the arguments
    for var in "$@"; do
        if [[ ! "$var" =~ ^-.* ]]; then
            file_pathes+=("$var")
        fi
    done

    # Create and trap a temporary file
    tmpfile=$(mktemp --tmpdir="${TMPDIR:-/tmp}" p-clang-format.XXXXXXXXXX)
    trap 'rm -f -- "$tmpfile"' INT TERM HUP EXIT
    # Append the temporary file to file_pathes so pragmas will be replaced in it
    # (this handles the case when the user did not specify -i to edit files in-place)
    file_pathes+=("$tmpfile")

    # Replace "#pragma" -> "//#pragma" in the files
    replace_all "$pragma_prefix" "pragma_key_words[@]" "$comment_pattern" "file_pathes[@]"
    # Run clang-format and redirect its output to the $tmpfile
    format "$@" > "$tmpfile"
    # Replace "//#pragma" -> "#pragma" in the files
    re_replace_all "$pragma_prefix" "pragma_key_words[@]" "$re_comment_pattern" "file_pathes[@]"
    # Show the clang-format's stdout after replacement
    cat "$tmpfile"
}

main "$@"
