Code cleanup for PR