Add some minor code changes (polishing)