Move shifts over flow control

Re-implement old optimization which was removed during the switch to rule
rewriting for simplicity reasons.