Some refactoring for clearer code grouping

No functional change. Just re-arranging of code for better readability.