added missing logic