Add back some changes that somehow got lost from last commit