Update and enhance code