Fixing compilation issue, making code more elegant