Added in some missing code