Refine the code according to Jon's comments.