Removing duplicate code, small fixes.