Additional code cleanup