Minor improvements and code cleanup