Minor refactor and fixes.