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

DATADIR=/usr/share/pkg-rocm-tools/data/build-targets
DIST=
SEP=
FOR_BUILD=0

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

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

	Emit the list of ROCm target ISAs for a distribution.

	The lists for all distributions are maintained by the Debian ROCm Team.

	The distribution is read from /etc/os-release, and the default separator
	is the space character ' '.

	The --for-build option can only be used from within a Debian package source
	directory. It will infer the target distribution from debian/changelog. If
	any binary dependency was built with the X-ROCm-Built-For attribute, then
	the list of target ISAs will be correspondingly reduced, if necessary.

	If the environment variable ROCM_TARGET_ISA_FIXED is set to a
	space-separated of ISAs, 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.

	Options:
	  -h              Show this help
	  --dist <dist>   Emit for <dist>, rather than for the host distribution
	  --sep <sep>     When emitting the list, separate ISAs by <sep>
	  --for-build     See above
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-Built-For field
#
# The (possibly empty) result will be a space-separated list
#
get_rocm_built_for() {
    package_name="$1"

    dpkg-query -f '${X-ROCm-Built-For}' -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
        ;;
     --for-build)
        FOR_BUILD=1
        shift
        ;;
    -h|--help)
        usage
        exit 0
        ;;
    *)
        echo "Error: unknown option $1" >&2
        usage
        exit 1
        ;;
esac
done

#
# Initialize defaults
#
if [ "$FOR_BUILD" -eq 1 ]; then
    [ -n "$SEP" ] || SEP=';'
    if [ -n "$DIST" ]; then
        echo "Cannot use --dist with --for-build" >&2
        exit 1
    fi
    DIST="$(dpkg-parsechangelog -Sdistribution)"
else
    [ -n "$SEP" ] || SEP=' '
    if [ -z "$DIST" ]; then
        DIST="$(. /etc/os-release && echo "$VERSION_CODENAME")"
    fi
fi

#
# Load the default arch list
#
if ! [ -f "$DATADIR/$DIST" ]; then
    echo "Error: no AMD GPU support data for distribution '$DIST'" >&2
    exit 1
fi
GFXARCHES=$(get_supported_gfxarches "$DATADIR/$DIST")

#
# Finally... output
#
if [ -n "${ROCM_TARGET_ISA_FIXED:-}" ]; then
    printf "%s" "$ROCM_TARGET_ISA_FIXED" | tr ' ' "$SEP"
elif [ "$FOR_BUILD" -eq 1 ]; then
    if ! [ -f "debian/control" ]; then
        echo "--for-build can only be used in an unpacked source." >&2
        exit 1
    fi
    # 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_for="$(get_rocm_built_for "$binary_dep")"
            if [ -n "$built_for" ]; then
                GFXARCHES="$(intersect_gfxarches "$GFXARCHES" "$built_for")"
            fi
        done
    done
    printf "%s" "$GFXARCHES" | tr ' ' "$SEP"
else
    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
