Added benchmarking code and updated TODO