Fix bugs and improve performance (I hope)