Some code refactoring and cleanup; one small bug fix