Also re-arrange code a bit for better readbility Signed-off-by: Aayush Gupta <aayushgupta219@gmail.com>