Added back some code that was lost in previous commit.