Updating code based on review comments