[MS] Refactor code