Small refactor and docs update.