Reorganised some code.

Reorganised some code.