Move main block into function