#!/bin/bash

# Runs valgrind and reports results.
#
# Since jemalloc dropped support for valgrind
# (https://github.com/jemalloc/jemalloc/issues/369), and both Rust and TensorFlow
# use jemalloc by default, we need to compile both without it.  Unfortunately,
# compiling TensorFlow from source is expensive, so this script takes a long
# time to run.

set -e

cd $(dirname $(readlink -f "$0"))

function run {
    echo "----------------------------------------------------------------------"
    echo "Running: $@"
    "$@"
    echo
}

tensorflow_version=2.13.0

valgrind_log=valgrind.log
truncate --size=0 "$valgrind_log"

rel_log=$(readlink -f "$PWD"/"$valgrind_log")
echo "Writing valgrind output to $rel_log."

echo "Building libtensorflow.so"

# This is the very expensive step.
run cargo build -p tensorflow-sys -vvv -j 1

# Build examples and grab the path to the shared libraries.
# The filename parsing is a hack because we're not really parsing JSON.
export LD_LIBRARY_PATH="$(cargo build --features='examples_system_alloc tensorflow_unstable' --example=addition --message-format=json \
    2>/dev/null | \
sed -En 's|^.*"([^"]*/target/debug/build/tensorflow-sys-[^"/]*/out)".*$|\1|p')"
echo "LD_LIBRARY_PATH=$LD_LIBRARY_PATH"

# Run valgrind against all the things.
echo "Testing tensorflow examples"
for example in addition regression expressions regression_checkpoint regression_savedmodel; do
    run cargo build --features='examples_system_alloc tensorflow_unstable' --example="$example"
    run valgrind --leak-check=full target/debug/examples/"$example" >> "$valgrind_log" 2>&1
done

echo "Testing tensorflow eager examples"
for example in mobilenetv3; do
    run cargo build --features='examples_system_alloc tensorflow_unstable eager' --example="$example"
    run valgrind --leak-check=full target/debug/examples/"$example" >> "$valgrind_log" 2>&1
done

echo "Testing tensorflow-sys examples"
for example in multiplication; do
    cargo build --features=examples_system_alloc --example="$example" -p tensorflow-sys
    (cd tensorflow-sys && run valgrind --leak-check=full ../target/debug/examples/"$example") >> "$valgrind_log" 2>&1
done

# Build tests and grab the path to the test binary, because we can't properly run valgrind through cargo.
# The filename parsing is a hack because we're not really parsing JSON.
test_binary=$(cargo test --features='examples_system_alloc tensorflow_unstable' --message-format=json --no-run \
    2>/dev/null | \
sed -En 's|^.*"([^"]*/target/debug/deps/tensorflow-[^"]*)".*$|\1|p')

# Run valgrind against the tests.
echo "Testing tests"
run valgrind --leak-check=full "$test_binary" >> "$valgrind_log" 2>&1

# Aggregate results.
printf "Definitely lost bytes: %9d\n" $(awk '/definitely lost:/{gsub(",","",$4);sum+=$4}END{print sum}' < "$valgrind_log")
printf "Indirectly lost bytes: %9d\n" $(awk '/indirectly lost:/{gsub(",","",$4);sum+=$4}END{print sum}' < "$valgrind_log")
printf "Possibly lost bytes:   %9d\n" $(awk '/possibly lost:/{gsub(",","",$4);sum+=$4}END{print sum}' < "$valgrind_log")
printf "Still reachable bytes: %9d\n" $(awk '/still reachable:/{gsub(",","",$4);sum+=$4}END{print sum}' < "$valgrind_log")

if egrep -i -q "invalid read|invalid write" "$valgrind_log"; then
    echo "Invalid operations detected"
    exit 1
fi