Little refactor and better abstraction