Some refactoring to reflect last changes