Refactor the unification code