Add persistent highlight functionality

Also refactoring