Code clean up part 3