Cleaned up some code after merge