Cleaning up and implementing persistence