Fix highlighting bug