Some refactoring to put things in a proper place