Based on changes from refactor branch for code parity Signed-off-by: Aayush Gupta <aayushgupta219@gmail.com>