Fix code review comment: repcace some ifs to one expression