small code change

Small code change accoording to @adamlundrigan suggestion. It make the
code smaller and more logical, and the result is the same.