Fixed up some code that was half working.