Restore accidentally removed code