#!/usr/bin/env bash
#' Get the x86-64 Microarchitecture Level on the Current Machine
#'
#' Queries the CPU information to infer which level of x86-64
#' microarchitecture is supported by the CPU on the current machine,
#' i.e. x86-64-v1, x86-64-v2, x86-64-v3, or x86-64-v4.
#'
#' Usage:
#' x86-64-level
#'
#' Options:
#'   --help            Show this help
#'   --version         Show the version of this tool
#'   --verbose         Explain the identified level
#'   --assert=<level>  Assert that CPU supports x86-64-v<level> with
#'                     exit code 0 is supported, otherwise 1
#'
#' Examples:
#' $ x86-64-level
#' 3
#'
#' $ level=$(x86-64-level)
#' $ echo "x86-64-v${level}"
#' x86-64-v3
#'
#' $ x86-64-level --verbose
#' Identified x86-64-v3, because x86-64-v4 requires 'avx512f', which is
#' not supported by this CPU [Intel(R) Core(TM) i7-8650U CPU @ 1.90GHz]
#' 3
#'
#' $ cat /proc/cpuinfo | x86-64-level -
#' 3
#'
#' $ x86-64-level --assert=2
#' $ echo $?
#' 0
#'
#' $ x86-64-level --assert=4
#' The CPU [Intel(R) Core(TM) i7-8650U CPU @ 1.90GHz] on this host ('dev2')
#' supports x86-64-v3, which is less than the required x86-64-v4
#' $ echo $?
#' 1
#'
#' Version: 0.2.2-9000
#' License: CC BY-SA 4.0
#' Source code: https://github.com/ucsf-wynton/wynton-tools
#'
#' Authors:
#' * Henrik Bengtsson (expanded on Gilles implementation [2])
#' * StackExchange user 'Gilles'
#'   <https://stackexchange.com/users/164368/>
#' * StackExchange user 'gioele'
#'   <https://unix.stackexchange.com/users/14861/>
#'
#' References:
#' [1] https://www.wikipedia.org/wiki/X86-64#Microarchitecture_levels
#' [2] https://unix.stackexchange.com/a/631320

#---------------------------------------------------------------------
# CLI functions
#---------------------------------------------------------------------
help() {
    grep "^#'" "$0" | sed -E "s/^#' ?//"
}

version() {
    grep "^#' Version:" "$0" | sed 's/.* //'
}


#---------------------------------------------------------------------
# CPU functions
#---------------------------------------------------------------------
data=
read_input() {
    if [[ -z ${data} ]]; then
        if ${stdin}; then
            data=$(< /dev/stdin)
        else
            data=$(< /proc/cpuinfo)
        fi
        if [[ -z ${data} ]]; then
            echo >&2 "ERROR: Input data is empty"
            exit 1
        fi
    fi
}

get_cpu_name() {
    local name
    bfr=$(grep -E "^model name[[:space:]]*:" <<< "${data}")
    name=$(echo "${bfr}" | head -n 1)
    name="${name#model name*:}"
    echo "${name## }"
}


get_cpu_flags() {
    local flags
    flags=$(grep "^flags[[:space:]]*:" <<< "${data}" | head -n 1)
    flags="${flags#*:}"
    flags="${flags## }"
    if grep -v -q -E "^[[:lower:][:digit:]_ ]+$" <<< "${flags}"; then
        echo >&2 "ERROR: Cannot reliably infer the CPU x86-64 level, because the format of the CPU flags comprise of other symbols than only lower-case letters, digits, and underscores: '${flags}'"
        exit 1
    fi
    echo "${flags}"
}


validate_cpu_flags() {
    read_input
    get_cpu_flags > /dev/null
}


has_cpu_flags() {
    local flag
    local msg
    local cpu_name
    
    for flag; do
        ## Note, it's important to keep a trailing space
        case " ${flags} " in
            *" ${flag} "*)
                :
                ;;
            *)
                if ${verbose}; then
                    msg="Identified x86-64-v${level}, because x86-64-v$((level + 1)) requires '${flag}', which is not supported by this CPU"
                    cpu_name=$(get_cpu_name)
                    [[ -n ${cpu_name} ]] && msg="${msg} [${cpu_name}]"
                    echo >&2 "${msg}"
                fi
                return 1
                ;;
        esac
    done
}


level=0

determine_cpu_version() {
    ## x86-64-v0 (can this happen?)
    level=0
    
    ## x86-64-v1
    has_cpu_flags lm cmov cx8 fpu fxsr mmx syscall sse2 || return 0
    level=1

    ## x86-64-v2
    has_cpu_flags cx16 lahf_lm popcnt sse4_1 sse4_2 ssse3 || return 0
    level=2
    
    ## x86-64-v3
    has_cpu_flags avx avx2 bmi1 bmi2 f16c fma abm movbe xsave || return 0
    level=3

    ## x86-64-v4
    has_cpu_flags avx512f avx512bw avx512cd avx512dq avx512vl || return 0
    level=4
}


report_cpu_version() {
    read_input
    flags=$(get_cpu_flags)
    if [[ -z ${flags} ]]; then
        echo >&2 "ERROR: No CPU 'flags' entry in the input data"
        exit 1
    fi
    determine_cpu_version
    echo "${level}"
}


#---------------------------------------------------------------------
# MAIN
#---------------------------------------------------------------------
verbose=false
stdin=false
assert=

# Parse command-line options
while [[ $# -gt 0 ]]; do
    ## Options (--flags):
    if [[ "$1" == "--help" ]]; then
        help
        exit 0
    elif [[ "$1" == "--version" ]]; then
        version
        exit 0
    elif [[ "$1" == "--verbose" ]]; then
        verbose=true
    elif [[ "$1" == "-" ]]; then
        stdin=true
    elif [[ "$1" == "--assert" ]]; then
        verbose=true
    ## Options (--key=value):
    elif [[ "$1" =~ ^--.*=.*$ ]]; then
        key=${1//--}
        key=${key//=*}
        value=${1//--[[:alpha:]]*=}
        if [[ -z ${value} ]]; then
            echo >&2 "ERROR: Option '--${key}' must not be empty"
            exit 2
        fi
        if [[ "${key}" == "assert" ]]; then
            assert=${value}
            if [[ ! ${assert} =~ ^-?[0-9]+$ ]]; then
                echo >&2 "ERROR: Option --assert does not specify an integer: ${assert}"
                exit 2
            elif [[ ${assert} -lt 1 ]] || [[ ${assert} -gt 4 ]]; then
                echo >&2 "ERROR: Option --assert is out of range [1,4]: ${assert}"
                exit 2
            fi
        else
            echo >&2 "ERROR: Unknown option: $1"
            exit 2
        fi
    else
        echo >&2 "ERROR: Unknown option: $1"
        exit 2
    fi
    shift
done

if [[ -n ${assert} ]]; then
    validate_cpu_flags
    version=$(report_cpu_version)
    if [[ ${version} -lt ${assert} ]]; then
        read_input
        cpu_info=$(get_cpu_name)
        [[ -n ${cpu_info} ]] && cpu_info=" [${cpu_info}]"
        >&2 echo "The CPU${cpu_info} on this host ('${HOSTNAME}') supports x86-64-v${version}, which is less than the required x86-64-v${assert}"
        exit 1
    fi
else
    validate_cpu_flags
    report_cpu_version
fi