refactor and fix bug I introduced