Code refractoring

Refractored some code