Optimize code a bit