Forgot to put stuff back :P
Minor changes.