Refactor for unification bug fix