Moved some code to improve performance