#!/bin/bash
set -e
USAGE=$(cat <<"EOF"
  COMMANDS

./miri install <flags>:
Installs the miri driver and cargo-miri. <flags> are passed to `cargo
install`. Sets up the rpath such that the installed binary should work in any
working directory.

./miri build <flags>:
Just build miri. <flags> are passed to `cargo build`.

./miri check <flags>:
Just check miri. <flags> are passed to `cargo check`.

./miri test <flags>:
Build miri, set up a sysroot and then run the test suite. <flags> are passed
to the final `cargo test` invocation.

./miri run <flags>:
Build miri, set up a sysroot and then run the driver with the given <flags>.
(Also respects MIRIFLAGS environment variable.)

The commands above also exist in a "-debug" variant (e.g. "./miri run-debug
<flags>") which uses debug builds instead of release builds, for faster build
times and slower execution times.

./miri fmt <flags>:
Format all sources and tests. <flags> are passed to `rustfmt`.

./miri clippy <flags>:
Runs clippy on all sources. <flags> are passed to `cargo clippy`.

./miri many-seeds <command>:
Runs <command> over and over again with different seeds for Miri. The MIRIFLAGS
variable is set to its original value appended with ` -Zmiri-seed=$SEED` for
many different seeds.

./miri bench <benches>:
Runs the benchmarks from bench-cargo-miri in hyperfine. hyperfine needs to be installed.
<benches> can explicitly list the benchmarks to run; by default, all of them are run.

  ENVIRONMENT VARIABLES

MIRI_SYSROOT:
If already set, the "sysroot setup" step is skipped.

CARGO_EXTRA_FLAGS:
Pass extra flags to all cargo invocations.
EOF
)

## Preparation
# macOS does not have a useful readlink/realpath so we have to use Python instead...
MIRIDIR=$(python3 -c 'import os, sys; print(os.path.dirname(os.path.realpath(sys.argv[1])))' "$0")
TOOLCHAIN=$(cd "$MIRIDIR"; rustup show active-toolchain | head -n 1 | cut -d ' ' -f 1)

# Determine command.
COMMAND="$1"
[ $# -gt 0 ] && shift

## Handle some commands early, since they should *not* alter the environment.
case "$COMMAND" in
many-seeds)
    for SEED in $({ echo obase=16; seq 0 255; } | bc); do
        echo "Trying seed: $SEED"
        MIRIFLAGS="$MIRIFLAGS -Zmiri-seed=$SEED" $@ || { echo "Failing seed: $SEED"; break; }
    done
    exit 0
    ;;
bench)
    # Make sure we have an up-to-date Miri installed
    "$0" install
    # Run the requested benchmarks
    if [ -z "$@" ]; then
        BENCHES=( $(ls "$MIRIDIR/bench-cargo-miri" ) )
    else
        BENCHES=("$@")
    fi
    for BENCH in "${BENCHES[@]}"; do
        hyperfine -w 1 -m 5 --shell=none "cargo +$TOOLCHAIN miri run --manifest-path bench-cargo-miri/$BENCH/Cargo.toml"
    done
    exit 0
    ;;
esac

## Prepare the environment
# Determine some toolchain properties
TARGET=$(rustc +$TOOLCHAIN --version --verbose | grep "^host:" | cut -d ' ' -f 2)
SYSROOT=$(rustc +$TOOLCHAIN --print sysroot)
LIBDIR=$SYSROOT/lib/rustlib/$TARGET/lib
if ! test -d "$LIBDIR"; then
    echo "Something went wrong determining the library dir."
    echo "I got $LIBDIR but that does not exist."
    echo "Please report a bug at https://github.com/rust-lang/miri/issues."
    exit 2
fi

# Prepare flags for cargo and rustc.
CARGO="cargo +$TOOLCHAIN"
if [ -z "$CARGO_INCREMENTAL" ]; then
    # Default CARGO_INCREMENTAL to 1.
    export CARGO_INCREMENTAL=1
fi
if [ -z "$CARGO_TARGET_DIR" ]; then
    # Share target dir between `miri` and `cargo-miri`.
    export CARGO_TARGET_DIR="$MIRIDIR/target"
fi
# We set the rpath so that Miri finds the private rustc libraries it needs.
export RUSTFLAGS="-C link-args=-Wl,-rpath,$LIBDIR $RUSTFLAGS"
# Determine flags passed to all cargo invocations.
# This is a bit more annoying that one would hope due to
# <https://github.com/rust-lang/cargo/issues/6992>.
case "$COMMAND" in
*-debug)
    CARGO_INSTALL_FLAGS="--target $TARGET --debug $CARGO_EXTRA_FLAGS"
    CARGO_BUILD_FLAGS="--target $TARGET $CARGO_EXTRA_FLAGS"
    ;;
*)
    CARGO_INSTALL_FLAGS="--target $TARGET $CARGO_EXTRA_FLAGS"
    CARGO_BUILD_FLAGS="--target $TARGET --release $CARGO_EXTRA_FLAGS"
    ;;
esac

## Helper functions

# Build a sysroot and set MIRI_SYSROOT to use it. Arguments are passed to `cargo miri setup`.
build_sysroot() {
    # Build once, for the user to see.
    $CARGO run $CARGO_BUILD_FLAGS --manifest-path "$MIRIDIR"/cargo-miri/Cargo.toml -- miri setup "$@"
    # Call again, to just set env var.
    export MIRI_SYSROOT="$($CARGO run $CARGO_BUILD_FLAGS --manifest-path "$MIRIDIR"/cargo-miri/Cargo.toml -q -- miri setup --print-sysroot "$@")"
}

# Prepare and set MIRI_SYSROOT. Respects `MIRI_TEST_TARGET` and takes into account
# locally built vs. distributed rustc.
find_sysroot() {
    if [ -n "$MIRI_SYSROOT" ]; then
        # Sysroot already set, use that.
        return 0
    fi
    # We need to build a sysroot.
    if [ -n "$MIRI_TEST_TARGET" ]; then
        build_sysroot --target "$MIRI_TEST_TARGET"
    else
        build_sysroot
    fi
}

## Main

# Run command.
case "$COMMAND" in
install|install-debug)
    # "--locked" to respect the Cargo.lock file if it exists,
    # "--offline" to avoid querying the registry (for yanked packages).
    $CARGO install $CARGO_INSTALL_FLAGS --path "$MIRIDIR" --force --locked --offline "$@"
    $CARGO install $CARGO_INSTALL_FLAGS --path "$MIRIDIR"/cargo-miri --force --locked --offline "$@"
    ;;
check|check-debug)
    # Check, and let caller control flags.
    $CARGO check $CARGO_BUILD_FLAGS --manifest-path "$MIRIDIR"/Cargo.toml --all-targets "$@"
    $CARGO check $CARGO_BUILD_FLAGS --manifest-path "$MIRIDIR"/cargo-miri/Cargo.toml "$@"
    ;;
build|build-debug)
    # Build, and let caller control flags.
    $CARGO build $CARGO_BUILD_FLAGS --manifest-path "$MIRIDIR"/Cargo.toml "$@"
    $CARGO build $CARGO_BUILD_FLAGS --manifest-path "$MIRIDIR"/cargo-miri/Cargo.toml "$@"
    ;;
test|test-debug|bless|bless-debug)
    # First build and get a sysroot.
    $CARGO build $CARGO_BUILD_FLAGS --manifest-path "$MIRIDIR"/Cargo.toml
    find_sysroot
    case "$COMMAND" in
    bless|bless-debug)
        export MIRI_BLESS="Gesundheit"
        ;;
    esac
    # Then test, and let caller control flags.
    # Only in root project and ui_test as `cargo-miri` has no tests.
    $CARGO test $CARGO_BUILD_FLAGS --manifest-path "$MIRIDIR"/Cargo.toml "$@"
    $CARGO test $CARGO_BUILD_FLAGS --manifest-path "$MIRIDIR"/ui_test/Cargo.toml "$@"
    ;;
run|run-debug)
    # Scan for "--target" to overwrite the "MIRI_TEST_TARGET" env var so
    # that we set the MIRI_SYSROOT up the right way.
    FOUND_TARGET_OPT=0
    for ARG in "$@"; do
        if [ "$LAST_ARG" = "--target" ]; then
            # Found it!
            export MIRI_TEST_TARGET="$ARG"
            FOUND_TARGET_OPT=1
            break
        fi
        LAST_ARG="$ARG"
    done
    if [ "$FOUND_TARGET_OPT" = "0" ] && [ -n "$MIRI_TEST_TARGET" ]; then
        # Make sure Miri actually uses this target.
        MIRIFLAGS="$MIRIFLAGS --target $MIRI_TEST_TARGET"
    fi
    # First build and get a sysroot.
    $CARGO build $CARGO_BUILD_FLAGS --manifest-path "$MIRIDIR"/Cargo.toml
    find_sysroot
    # Then run the actual command.
    exec $CARGO run $CARGO_BUILD_FLAGS --manifest-path "$MIRIDIR"/Cargo.toml -- --sysroot "$MIRI_SYSROOT" $MIRIFLAGS "$@"
    ;;
fmt)
    find "$MIRIDIR" -not \( -name target -prune \) -name '*.rs' \
        | xargs rustfmt +$TOOLCHAIN --edition=2021 --config-path "$MIRIDIR/rustfmt.toml" "$@"
    ;;
clippy)
    $CARGO clippy $CARGO_BUILD_FLAGS --manifest-path "$MIRIDIR"/Cargo.toml --all-targets "$@"
    $CARGO clippy $CARGO_BUILD_FLAGS --manifest-path "$MIRIDIR"/ui_test/Cargo.toml --all-targets "$@"
    $CARGO clippy $CARGO_BUILD_FLAGS --manifest-path "$MIRIDIR"/cargo-miri/Cargo.toml "$@"
    ;;
*)
    if [ -n "$COMMAND" ]; then
      echo "Unknown command: $COMMAND"
      echo
    fi
    echo "$USAGE"
    exit 1
esac