Code clean-up