#!/usr/bin/env bash
# Copyright (c) 2022, the Dart project authors.  Please see the AUTHORS file
# for details. All rights reserved. Use of this source code is governed by a
# BSD-style license that can be found in the LICENSE file.

# Script to compile a benchmark using dart2wasm. Assumes the Dart repo's
# directory structure.

set -e

function follow_links() {
  file="$1"
  while [ -h "$file" ]; do
    # On Mac OS, readlink -f doesn't work.
    file="$(readlink "$file")"
  done
  echo "$file"
}

function host_arch() {
  # Use uname to determine the host architecture.
  case `uname -m` in
    x86_64)
    echo "X64"
    ;;
    aarch64 | arm64 | armv8*)
    echo "ARM64"
    ;;
    *)
    echo "Unknown host architecture" `uname -m` >&2
    exit 1
    ;;
  esac
}

# Unlike $0, $BASH_SOURCE points to the absolute path of this file.
PROG_NAME="$(follow_links "$BASH_SOURCE")"

# Handle the case where dart-sdk/bin has been symlinked to.
PROG_DIR="$(cd "${PROG_NAME%/*}" ; pwd -P)"
SDK_DIR="$(cd "${PROG_DIR}/../../.." ; pwd -P)"

# Locate build directory, containing executables, snapshots and platform dill.
if [[ `uname` == 'Darwin' ]]; then
  OUT_DIR="$SDK_DIR/xcodebuild"
else
  OUT_DIR="$SDK_DIR/out"
fi
HOST_ARCH="$(host_arch)"
DART_CONFIGURATION=${DART_CONFIGURATION:-Release$HOST_ARCH}
BIN_DIR="$OUT_DIR/$DART_CONFIGURATION"

BINARYEN="$BIN_DIR/wasm-opt"
DART="$BIN_DIR/dart"
DART_AOT_RUNTIME="$BIN_DIR/dartaotruntime"
LIBRARIES_JSON_ARG="--libraries-spec=$SDK_DIR/sdk/lib/libraries.json"

function find_flags {
  echo -en "$(sed -n "/$1 =/,/end of $1/ p" $SDK_DIR/pkg/dartdev/lib/src/commands/compile.dart  | sed  '1d' | sed '$d' | tr '\n' ' ' | sed 's#\s\+# #g' | sed 's#^\s\+##' | sed 's#\s\+$##')"
}
# Use same flags as `dart compile exe`
BINARYEN_FLAGS=($(find_flags 'binaryenFlags'))
BINARYEN_FLAGS_DEFERRED_LOADING=($(find_flags 'binaryenFlagsDeferredLoading'))
OPT_FLAGS_L0=($(find_flags 'optimizationLevel0Flags'))
OPT_FLAGS_L1=($(find_flags 'optimizationLevel1Flags'))
OPT_FLAGS_L2=($(find_flags 'optimizationLevel2Flags'))
OPT_FLAGS_L3=($(find_flags 'optimizationLevel3Flags'))
OPT_FLAGS_L4=($(find_flags 'optimizationLevel4Flags'))

RUN_BINARYEN=1
MULTI_MODULE=0
RUN_SRC=0
GENERATE_SOURCE_MAP=1
COMPILE_BENCHMARK_BASE_NAME=""
PLATFORM_FILENAME="$BIN_DIR/dart2wasm_platform.dill"
SNAPSHOT_NAME="dart2wasm"

# All arguments will be passed along to dart2wasm except specially recognized
# flags.
VM_ARGS=()
DART2WASM_ARGS=("--require-js-string-builtin")
DART_FILE=""
WASM_FILE=""
while [ $# -gt 0 ]; do
    case "$1" in
      --compile-benchmark=*)
      COMPILE_BENCHMARK_BASE_NAME=$(echo $1 | sed 's/--compile-benchmark=//g')
      shift
      ;;

    --src)
      RUN_SRC=1
      shift
      ;;

    --js-compatibility)
      PLATFORM_FILENAME="$BIN_DIR/dart2wasm_js_compatibility_platform.dill"
      DART2WASM_ARGS+=("--js-compatibility")
      shift
      ;;

    -g | --no-strip-wasm)
      BINARYEN_FLAGS+=("-g")
      shift
      ;;

    --compiler-asserts)
      SNAPSHOT_NAME="dart2wasm_asserts"
      VM_ARGS+=("--enable-asserts")
      shift
      ;;

    -O0 | --optimization-level=0)
      DART2WASM_ARGS+=(${OPT_FLAGS_L0[@]})
      RUN_BINARYEN=0
      shift
      ;;

    -O1 | --optimization-level=1)
      DART2WASM_ARGS+=(${OPT_FLAGS_L1[@]})
      RUN_BINARYEN=1
      shift
      ;;

    -O2 | --optimization-level=2)
      DART2WASM_ARGS+=(${OPT_FLAGS_L2[@]})
      RUN_BINARYEN=1
      shift
      ;;

    -O3 | --optimization-level=3)
      DART2WASM_ARGS+=(${OPT_FLAGS_L3[@]})
      RUN_BINARYEN=1
      shift
      ;;

    -O4 | --optimization-level=4)
      DART2WASM_ARGS+=(${OPT_FLAGS_L4[@]})
      RUN_BINARYEN=1
      shift
      ;;

    --extra-compiler-option=--platform=*)
      PLATFORM_FILENAME="${1#--extra-compiler-option=--platform=}"
      shift
      ;;

    --extra-compiler-option=--enable-deferred-loading)
      MULTI_MODULE=1
      DART2WASM_ARGS+=("--enable-deferred-loading")
      shift
      ;;

    --extra-compiler-option=--enable-multi-module-stress-test-mode)
      MULTI_MODULE=1
      DART2WASM_ARGS+=("--enable-multi-module-stress-test-mode")
      shift
      ;;

    --extra-compiler-option=--no-js-compatibility)
      PLATFORM_FILENAME="$BIN_DIR/dart2wasm_platform.dill"
      DART2WASM_ARGS+=("--no-js-compatibility")
      shift
      ;;

    --extra-compiler-option=*)
      DART2WASM_ARGS+=("${1#--extra-compiler-option=}")
      shift
      ;;

    --no-source-maps)
      GENERATE_SOURCE_MAP=0
      shift
      ;;

    -o)
      shift
      WASM_FILE="$1"
      shift
      ;;

    --* | -*)
      DART2WASM_ARGS+=("$1")
      shift
      ;;

    *)
      if [ -z "$DART_FILE" ]; then
        DART_FILE="$1"
        shift
      else
        if [ -z "$WASM_FILE" ]; then
          WASM_FILE="$1"
          shift
        else
          echo "Unexpected argument $1"
          exit 1
        fi
      fi
      ;;
  esac
done

if [ $GENERATE_SOURCE_MAP -eq 1 ]; then
  BINARYEN_FLAGS+=("-ism" "${WASM_FILE}.map" "-osm" "${WASM_FILE}.map")
fi

if [ -z "$DART_FILE" -o -z "$WASM_FILE" ]; then
  echo "Expected <file.dart> <file.wasm>"
  exit 1
fi


PLATFORM_ARG="--platform=$PLATFORM_FILENAME"
DART2WASM_AOT_SNAPSHOT="$BIN_DIR/$SNAPSHOT_NAME.snapshot"
DART2WASM_SRC="$SDK_DIR/pkg/dart2wasm/bin/dart2wasm.dart"

function measure() {
  set +e
  RESULT=$( { /usr/bin/time --format="\nMemory: %M KB, Time: %e s" $@; } 2>&1 )
  EXIT=$?
  if [ $EXIT -ne 0 ]; then
    echo "Running $@ resulted in exit code $EXIT"
    echo "$RESULT"
    exit 1
  fi
  set -e
  MEMORY=$(echo "$RESULT" | tail -n1 | sed 's/Memory: \([0-9.]\+\) KB, Time: \([0-9.]\+\) s.*/\1/g')
  MEMORY=$(($MEMORY * 1024))
  TIME=$(echo "$RESULT" | tail -n1 | sed 's/Memory: \([0-9.]\+\) KB, Time: \([0-9.]\+\) s.*/\2/g')
}

function measure_size() {
  SIZE=$(cat $1 | wc -c)
  GZIP_SIZE=$(cat $1 | gzip -c6 | wc -c)
}

function run_if_binaryen_enabled() {
  if [ $RUN_BINARYEN -eq 1 ]; then
    $@
  fi
}

COMPILER_SIZE=0
COMPILER_GZIP_SIZE=0

function run_compiler() {
  if [ $RUN_SRC -eq 1 ]; then
    dart2wasm_command=("$DART" "${VM_ARGS[@]}" "$DART2WASM_SRC" "$LIBRARIES_JSON_ARG" "${DART2WASM_ARGS[@]}" "$DART_FILE" "$WASM_FILE")
  else
    dart2wasm_command=("$DART_AOT_RUNTIME" "${VM_ARGS[@]}" "$DART2WASM_AOT_SNAPSHOT" "$PLATFORM_ARG" "${DART2WASM_ARGS[@]}" "$DART_FILE" "$WASM_FILE")
  fi

  if [ -n "$COMPILE_BENCHMARK_BASE_NAME" ]; then
    measure "${dart2wasm_command[@]}"
    COMPILER_TIME=$TIME
    COMPILER_MEMORY=$MEMORY

    measure_size ${WASM_FILE%.wasm}.mjs
    MJS_SIZE=$SIZE
    MJS_GZIP_SIZE=$GZIP_SIZE

    if [ $MULTI_MODULE -eq 1 ]; then
      for WASM_FILE in "${WASM_FILE%.wasm}"*.wasm; do
        measure_size $WASM_FILE
        (( COMPILER_SIZE+=$SIZE ))
        (( COMPILER_GZIP_SIZE+=$GZIP_SIZE ))
      done
    else
      measure_size $WASM_FILE
      COMPILER_SIZE=$SIZE
      COMPILER_GZIP_SIZE=$GZIP_SIZE
    fi
  else
    "${dart2wasm_command[@]}"
  fi
}

BINARYEN_TIME=0
BINARYEN_MEMORY=0
BINARYEN_SIZE=0
BINARYEN_GZIP_SIZE=0

function run_binaryen_single() {
  binaryen_command=("$BINARYEN" "$@" "$WASM_FILE" -o "$WASM_FILE")
  if [ -n "$COMPILE_BENCHMARK_BASE_NAME" ]; then
    # If we're measuring run each binaryen command sequentially.
    measure ${binaryen_command[@]}
    BINARYEN_TIME=$(echo "$BINARYEN_TIME + $TIME" | bc)
    BINARYEN_MEMORY=$(($BINARYEN_MEMORY > $MEMORY ? $BINARYEN_MEMORY : $MEMORY ))
    measure_size $WASM_FILE
    BINARYEN_SIZE=$(echo "$BINARYEN_SIZE + $SIZE" | bc)
    BINARYEN_GZIP_SIZE=$(echo "$BINARYEN_GZIP_SIZE + $GZIP_SIZE" | bc)
  else
    ${binaryen_command[@]} &
  fi
}

function run_binaryen() {
  if [ $MULTI_MODULE -eq 1 ]; then
    # Iterate over all matching wasm files and optimize them concurrently in
    # different processes.
    for WASM_FILE in "${WASM_FILE%.wasm}"*.wasm; do
      run_binaryen_single "${BINARYEN_FLAGS_DEFERRED_LOADING[@]}"
    done
  else
    run_binaryen_single "${BINARYEN_FLAGS[@]}"
  fi
  wait

  if [ -n "$COMPILE_BENCHMARK_BASE_NAME" ]; then
    MAX_MEMORY=$(($COMPILER_MEMORY > $BINARYEN_MEMORY ? $COMPILER_MEMORY : $BINARYEN_MEMORY ))
    TOTAL_TIME=$(echo "$COMPILER_TIME + $BINARYEN_TIME" | bc)
  fi
}

run_compiler
run_if_binaryen_enabled run_binaryen

if [ -n "$COMPILE_BENCHMARK_BASE_NAME" ]; then
  # CompileTime
  run_if_binaryen_enabled echo "$COMPILE_BENCHMARK_BASE_NAME.Time.Total(CompileTime): $TOTAL_TIME s"
  echo "$COMPILE_BENCHMARK_BASE_NAME.Time.Dart2Wasm(CompileTime): $COMPILER_TIME s"
  run_if_binaryen_enabled echo "$COMPILE_BENCHMARK_BASE_NAME.Time.Wasm2WasmOpt(CompileTime): $BINARYEN_TIME s"

  # CodeSize
  echo "$COMPILE_BENCHMARK_BASE_NAME.Size.mjs(CodeSize): $MJS_SIZE bytes"
  echo "$COMPILE_BENCHMARK_BASE_NAME.Size.mjs.gz(CodeSize): $MJS_GZIP_SIZE bytes"
  echo "$COMPILE_BENCHMARK_BASE_NAME.Size.wasm(CodeSize): $COMPILER_SIZE bytes"
  echo "$COMPILE_BENCHMARK_BASE_NAME.Size.wasm.gz(CodeSize): $COMPILER_GZIP_SIZE bytes"
  run_if_binaryen_enabled echo "$COMPILE_BENCHMARK_BASE_NAME.Size.wasm.opt(CodeSize): $BINARYEN_SIZE bytes"
  run_if_binaryen_enabled echo "$COMPILE_BENCHMARK_BASE_NAME.Size.wasm.opt.gz(CodeSize): $BINARYEN_GZIP_SIZE bytes"

  # MemoryUse
  run_if_binaryen_enabled echo "$COMPILE_BENCHMARK_BASE_NAME.MemoryUse.Max(MemoryUse): $MAX_MEMORY bytes"
  echo "$COMPILE_BENCHMARK_BASE_NAME.MemoryUse.Dart2Wasm(MemoryUse): $COMPILER_MEMORY bytes"
  run_if_binaryen_enabled echo "$COMPILE_BENCHMARK_BASE_NAME.MemoryUse.Wasm2WasmOpt(MemoryUse): $BINARYEN_MEMORY bytes"
fi
