#!/bin/bash
# shellcheck disable=SC1091
#
# Utility for emitting AMD GPU target architectures in a given distribution
#
# Christian Kastner <ckk@kvr.at>
# License: MIT
set -eu

if ! [ -f "debian/changelog" ]; then
    echo "debian/changelog not found." >&2
    exit 1
fi

DATADIR=/usr/share/pkg-rocm-tools/data/build-targets
DIST="$(dpkg-parsechangelog -Sdistribution)"
SEP=" "
REDUCE=0

usage() {
    cat >&2 <<-EOF | fmt -w "${COLUMNS:-80}"

    Usage: $0 [ --dist <dist> ] [ --sep <sep> ]

    Package build helper for determining ROCm target architectures.

    These architectures come from central lists which are maintained
    per-distribution by the Debian ROCm Team. If no list exists for the current
    distribution, then the list for Debian's unstable or 'sid' distribution is
    used as a fallback.

    This utility must be called from within an unpackage package source. The
    distribution is read from debian/changelog. This can be overriden with
    the --dist option.

    If the --reduce option is used, then if any build dependency, or one of its
    dependencies, listed in debian/control has a X-ROCm-GPU-Architecture
    attribute, then the list of target architectures will be correspondingly
    reduced. This will entail a call to dpkg-checkbuilddeps.

    The default separator is the space character ' '. This can be overriden
    with --sep. The separator can be an arbitrary string.

    If the environment variable ROCM_TARGET_ARCH_FIXED is set to a space-
    separated of architectures, then that list will be emitted instead. This
    can be useful to package maintainers who might want quick local builds with
    fewer or just architecture during development, without having to modify
    d/rules. This can also be useful to Salsa CI.

    Options:
      -h             Show this help
      --dist <dist>  Emit list for <dist> instead
      --sep <sep>    When emitting the list, use <sep> as separator
      --reduce       Reduce list of supported architectures as per the B-Ds.
EOF
}

#
# Get package names from a list of relationships
#
# Given the value of a Depends or Build-Depends field, return a list of all
# package names occuring in that list.
#
# get_package_names "libfoo (<= 1.0), libbar [amd64]" -> "libfoo libbar"
#
get_package_names() {
    multiline_str="$(echo "$1" | tr "," "\n")"

    # Exploit the fact that this expansion will create a tokens split by
    # whitespace, which we can then filter for tokens that are valid package
    # names. Because all other possible tokens (parentheses, pipes, version
    # numbers, build profiles, ...) cannot be valid package names.
    for package in $multiline_str; do
        if echo "$package" | grep -q '^[a-z0-9][a-z0-9+.-]*$'; then
            echo "$package"
        fi
    done
}

#
# Extract the value of get_build_depends from debian/control
#
# Must be in a source package for this to work
#
get_build_depends() {
    # Get the value of Build-Depends, with line continuation
    raw_field=$(awk '/^Build-Depends:/ {
                            value=$0;
                            sub(/^Build-Depends:[ \t]*/, "", value);
                            while(getline > 0 && /^[ \t]/) {
                                value = value $0
                            }
                            print value;
                            exit
                        }' debian/control)
    get_package_names "$raw_field"
}

#
# Query the Depends field of a package from the dpkg status database
#
# Package must be installed for this to work, which is assumed to be the case
# for all build dependencies
#
get_depends() {
    package_name="$1"

    raw_field="$(dpkg-query -f '${Depends}' -W "$package_name")"
    get_package_names "$raw_field"
}

#
# Like get_depends(), but for the X-ROCm-GPU-Architecture field
#
# The (possibly empty) result will be a space-separated list
#
get_rocm_built_arches() {
    package_name="$1"

    dpkg-query -f '${X-ROCm-GPU-Architecture}' -W "$package_name" 2>/dev/null
}

#
# Read and output support GPU architectures from a file
#
# The architectures are returned as a space-separated list
#
get_supported_gfxarches() {
    supportfile="$1"

    gfxarches=
    while read -r gfxarch; do
        # Skip empty lines; treat lines beginning with '#' as comments
        if [ -z "$gfxarch" ] || [ "${gfxarch:0:1}" = "#" ]; then
            continue
        fi
        gfxarches="$gfxarches $(printf "%s" "$gfxarch")"
    done < "$supportfile"
    echo "${gfxarches# }"
}

#
# Convert space-separated list of hex numbers to dec
#
convert_hex_list_to_dec() {
    hex_numbers="$1"

    dec_numbers=
    for hex_number in $hex_numbers; do
        dec_numbers="$dec_numbers $(printf "%d\n" "0x$hex_number")"
    done
    echo "${dec_numbers# }"
}

#
# Convert space-separated list of dec numbers to hex
#
convert_dec_list_to_hex() {
    dec_numbers="$1"

    hex_numbers=
    for dec_number in $dec_numbers; do
        hex_numbers="$hex_numbers $(printf "%x\n" "$dec_number")"
    done
    echo "${hex_numbers# }"
}

#
# Sort a space-separated list of gfx arch identifiers numerically
#
# gfx90a gfx806 gfx900 -> gfx806 gfx900 gfx90a
#
sort_gfxarches() {
    gfxarches="$1"

    # Strip "gfx" prefix
    arches_num=
    for gfxarch in $gfxarches; do
        arches_num="$arches_num ${gfxarch#gfx}"
    done
    arches_num="${arches_num# }"

    hexes="$(convert_hex_list_to_dec "$arches_num")"
    sorted="$(echo "$hexes" | tr ' ' '\n' | sort -n -u)"
    deces="$(convert_dec_list_to_hex "$sorted")"

    gfxarches=
    for dec_number in $deces; do
        gfxarches="$gfxarches gfx$dec_number"
    done
    echo "${gfxarches# }"
}

#
# Produce the intersection of two space-separated lists of gfx architectures
#
# The output will be sorted.
#
# "gfx90a gfx806 gfx900 gfx1100" "gfx1000 gfx90a gfx1100" -> "gfx90a gfx1000"
#
intersect_gfxarches() {
    left="$(sort_gfxarches "$1" | tr ' ' '\n')"
    right="$(sort_gfxarches "$2" | tr ' ' '\n')"

    common=
    for rightarch in $right; do
        if echo "$left" | grep -q "\b$rightarch\b"; then
            common="$common $rightarch"
        fi
    done
    echo "${common# }"
}

#
# Parse options
#
while [ $# -gt 0 ]; do
case "$1" in
    --dist)
        if [ -z "$2" ]; then
            echo "Error: Missing argument to --dist" >&2
            usage
            exit 1
        fi
        DIST="$2"
        shift 2
        ;;
    --dist=*)
        DIST="${1#--dist=}"
        shift 1
        ;;
    --sep)
        if [ -z "$2" ]; then
            echo "Error: Missing argument to --sep" >&2
            usage
            exit 1
        fi
        SEP="$2"
        shift 2
        ;;
    --sep=*)
        SEP="${1#--sep=}"
        shift 1
        ;;
    --reduce)
        REDUCE=1
        shift 1
        ;;
    -h|--help)
        usage
        exit 0
        ;;
    *)
        echo "Error: unknown option $1" >&2
        usage
        exit 1
        ;;
esac
done

#
# Load the default arch list
#
if ! [ -f "$DATADIR/$DIST" ]; then
    echo "WARNING: No AMD GPU support data for distribution '$DIST'" >&2
    echo "Falling back to support data for Debian unstable (sid)" >&2
    DIST=unstable
fi
GFXARCHES=$(get_supported_gfxarches "$DATADIR/$DIST")

#
# Finally... output
#
if [ -n "${ROCM_TARGET_ARCH_FIXED:-}" ]; then
    printf "%s" "$ROCM_TARGET_ARCH_FIXED" | tr ' ' "$SEP"
else
    if [ "$REDUCE" -eq 1 ]; then
        # Cannot query B-Ds, and their dependencies, if they aren't installed
        dpkg-checkbuilddeps

        for build_dep in $(get_build_depends); do
            [ -n "$build_dep" ] || continue
            for binary_dep in $(get_depends "$build_dep"); do
                [ -n "$binary_dep" ] || continue
                built_arches="$(get_rocm_built_arches "$binary_dep")"
                if [ -n "$built_arches" ]; then
                    GFXARCHES="$(intersect_gfxarches "$GFXARCHES" "$built_arches")"
                fi
            done
        done
    fi

    printf "%s" "$GFXARCHES" | tr ' ' "$SEP"
fi

# Special case: if the separator chosen is a newline, then most probably this
# is going to stdout, so the  final list entry is expected to be newline-
# terminated as well
if [ "$SEP" = '\n' ]; then
    printf '\n'
fi
