Minor code reorganization for clarity