minor update, code refactor