Refactor and adjust design a bit.