#!/bin/bash
# chmod +x avx2

terminal_output=true
found_avx1=false
found_avx2=false

AVX1_PATTERN='vaddpd|vaddps|vsubpd|vsubps|vmovapd|vmovaps|vmovupd|vmovups|vbroadcastsd|vbroadcastss|vperm2f128|vshufpd|vshufps|vcmppd|vcmpps|vpcmpeqb|vpcmpeqw|vpcmpeqd|vpcmpeqq|vgatherdpd|vgatherqpd|vgatherdps|vgatherqps'
AVX2_PATTERN='vpaddd|vpaddq|vpsubd|vpsubq|vpmulld|vpmullq|vandpd|vandps|vandnpd|vandnps|vorpd|vorps|vxorpd|vxorps|vpbroadcastb|vpbroadcastw|vpbroadcastd|vpbroadcastq|vpermd|vpblendd|vpblendvb|vpslldq|vpsrldq|vpsllw|vpslld|vpsllq|vpsrlw|vpsrld|vpsrlq|vpminsb|vpminsw|vpminsd|vpminsq|vpmaxsb|vpmaxsw|vpmaxsd|vpmaxsq'


##### functions block #####

function echo_red () {
    $terminal_output && printf "%s \n" "$(tput setaf 1)$1$(tput sgr0)" || echo "$1"
}


function echo_green () {
    $terminal_output && printf "%s \n" "$(tput setaf 2)$1$(tput sgr0)" || echo "$1"
}


function print_and_overwrite() {
    local message="$1"
    # Get the terminal width
    local term_width
    term_width=$(tput cols)
    
    printf "%s" "$message"
    
    # Calculate the number of lines the message occupies
    local message_length=${#message}
    cursor_lines=$(( (message_length + term_width - 1) / term_width ))
}


function clear_message() {
    # Move the cursor back to the start of the line
    printf "\r"
    
    # Clear all lines occupied by the previous message
    for ((i = 0; i < cursor_lines; i++)); do
        printf "%*s\r" "$(tput cols)" " "
        # Move up one line if not on the last line
        (( i < cursor_lines - 1 )) && printf "\033[A"
    done
    
    cursor_lines=0
}


function is_x86() {
    [[ $(file -b "$1") == *x86* ]] && return 0 || return 1
}  


# Function to check for AVX and AVX2 instructions in executable code sections
function has_avx12() {
    [ -e /tmp/otool.txt ] && rm /tmp/otool.txt
    print_and_overwrite "..disassembling $1"
    otool -tv "$1" 2>/dev/null > /tmp/otool.txt
    clear_message

    local report_avx1=false
    local report_avx2=false


    # Check for AVX instructions
    if $check_avx1; then
        print_and_overwrite "..checking for AVX in $1"
        if grep -q -m1 -E "$AVX1_PATTERN" /tmp/otool.txt; then
            found_avx1=true
            report_avx1=true
        fi
        clear_message
    fi

    # Check for AVX2 instructions
    if $check_avx2; then
        print_and_overwrite "..checking for AVX2 in $1"
        if grep -q -m1 -E "$AVX2_PATTERN" /tmp/otool.txt; then
            found_avx2=true
            report_avx2=true
        fi
        clear_message
    fi


    if $count_commands; then
        count=0
        if $check_avx1 && $report_avx1; then
            print_and_overwrite "..counting AVX commands in $1"
            count=$((count + $(grep -oE "$AVX1_PATTERN" /tmp/otool.txt | wc -l)))
            clear_message
        fi

        if $check_avx2 && $report_avx2; then
            print_and_overwrite "..counting AVX2 commands in $1"
            count=$((count + $(grep -oE "$AVX2_PATTERN" /tmp/otool.txt | wc -l)))
            clear_message
        fi
        
        count="$count "
        commands_string=" commands"
    else
        count=""
        commands_string=""
    fi


    [ -e /tmp/otool.txt ] && rm /tmp/otool.txt

    # preset, check avx and avx2
    if $check_avx1 && $check_avx2; then
        $report_avx1 && $report_avx2 && echo_red "Executable with ${count}AVX and AVX2${commands_string}: $1" && return 0
        $report_avx1 && echo_red "Executable with ${count}AVX${commands_string}: $1" && return 0
        $report_avx2 && echo_red "Executable with ${count}AVX2${commands_string}: $1" && return 0

    # check avx2, only by -avx2 argument
    elif $check_avx2 && ! $check_avx1; then
        $report_avx2 && echo_red "Executable with ${count}AVX2${commands_string}: $1" && return 0

    # check avx1, only by -avx1 argument
    elif $check_avx1 && ! $check_avx2; then
        $report_avx1 && echo_red "Executable with ${count}AVX${commands_string}: $1" && return 0
    fi
        
    # If AVX1 or 2 is not found
    echo_green "Executable without $avx_string: $1"
    return 1
     
}

##### end of functions block #####




echo "Macschrauber's AVX/AVX2 Scanner 15-6-2025"

# Check if command line tools / otool are available
# also triggers Mac Os installing command line tools
if ! otool --version  2>/dev/null; then
    echo_red "command line tools are missing"
    exit 1
fi

# Check if input is provided
if [ $# -eq 0 ]; then
    echo "Usage: $0 <directory> or <file>"
    echo "options: -c[ount]     count AVX, AVX2 or both"
    echo "         -avx1        check only for AVX commands"
    echo "         -avx2        check only for AVX2 commands"
    exit 1
fi


count_commands=false
check_avx1=true
check_avx2=true

given_file_or_directory="$1"
shift


# Parse flags
for arg in "$@"; do
  case "$arg" in
    -c*) count_commands=true ;;
    -avx1) check_avx2=false ;;
    -avx2) check_avx1=false ;;
    -*) echo "Unknown option: $arg" >&2; exit 1 ;;
  esac
done

# gave both -avx1 and -avx2 arguments
if ! $check_avx1 && ! $check_avx2; then
    check_avx1=true
    check_avx2=true
fi

$check_avx1 && avx_string="AVX"
$check_avx2 && avx_string="AVX2"
$check_avx1 && $check_avx2 && avx_string="AVX/AVX2"



# Handle single file input
if [ -f "$given_file_or_directory" ]; then
    echo "scanning ${given_file_or_directory} for ${avx_string} code"
    $count_commands && echo "...and count commands (what takes some time)"


    if is_x86 "$given_file_or_directory"; then # Check if the file is executable
        has_avx12 "$given_file_or_directory"
    else
        echo_red "Error: The file '$given_file_or_directory' is not a x86 executable."
        exit 1
    fi

# Handle directory input
elif [ -d "$given_file_or_directory" ]; then
    echo "Scanning directory: ${given_file_or_directory} for ${avx_string} code"

    # Track whether any executable files were found in the directory
    found_executables=false

    # Use find to locate all files and process them one by one without subshells
    while IFS= read -r -d '' file; do
        if is_x86 "$file"; then # Check if the file is a x86 executable
            found_executables=true
            has_avx12 "$file"
        fi
    done < <(find "$given_file_or_directory" -type f -print0)

    # If no executables were found in the directory, notify the user.
    if ! $found_executables; then

        # maybe restrictions, scan /Contents

        if [ -d "$given_file_or_directory/Contents" ]; then
            while IFS= read -r -d '' file; do
                if is_x86 "$file"; then # Check if the file is a x86 executable
                    found_executables=true
                    has_avx12 "$file"
                fi
            done < <(find "$given_file_or_directory/Contents" -type f -print0)
        fi

            if ! $found_executables; then
                echo_red "No x86 executable files found in the directory."
                exit 1
            fi

    fi

# Handle invalid input (neither file nor directory)
else
    echo_red "Error: The input '$given_file_or_directory' is neither a valid file nor a directory."
    exit 1
fi

# Final report based on whether AVX, AVX2 code was found or not.


# preset, check avx and avx2
if $check_avx1 && $check_avx2; then
    $found_avx1 && $found_avx2 && echo_red "AVX and AVX2 code found." && exit 0
    $found_avx1 && echo_red "AVX code found." && exit 0
    $found_avx2 && echo_red "AVX2 code found." && exit 0

# check avx2, only by -avx2 argument
elif $check_avx2 && ! $check_avx1; then
    $found_avx2 && echo_red "AVX2 code found." && exit 0

# check avx1, only by -avx1 argument
elif $check_avx1 && ! $check_avx2; then
    $found_avx1 && echo_red "AVX code found." && exit 0
fi


# If AVX1 or 2 is not found    
echo_green "no ${avx_string} code found"
exit 1