#!/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

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 architectures 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 architectures will be correspondingly reduced, if
    necessary.

    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.

    Options:
      -h              Show this help
      --dist <dist>   Emit for <dist>, rather than for the host distribution
      --sep <sep>     When emitting the list, use <sep> as separator
      --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-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
        ;;
     --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
    if ! [ -f 'debian/changelog' ]; then
        echo "Could not read debian/changelog" >&2
        exit 1
    fi
    [ -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_ARCH_FIXED:-}" ]; then
    printf "%s" "$ROCM_TARGET_ARCH_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_arches="$(get_rocm_built_arches "$binary_dep")"
            if [ -n "$built_arches" ]; then
                GFXARCHES="$(intersect_gfxarches "$GFXARCHES" "$built_arches")"
            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
