Fixing highlighting issue with @while