Refactor JS code a little