--- invite/invite.module.orig	2010-07-23 10:12:00.000000000 -0400
+++ invite/invite.module.new	2010-07-23 12:12:17.000000000 -0400
@@ -423,7 +423,20 @@ function invite_user($op, &$edit, &$acco
         }
       }
       if ($invite) {
-        _invite_accept($invite, $account);
+        // Process the invitation and assign target roles to the user. These
+        // will be saved by user_save().
+        $roles = invite_process($invite, $account);
+
+        if ($roles) {
+          if (!isset($edit['roles']) || !is_array($edit['roles'])) {
+            $edit['roles'] = array();
+          }
+          $edit['roles'] += $roles;
+
+          // We need to notify other modules of the role change, otherwise they
+          // will be unaware of it.
+          user_module_invoke('update', $edit, $account);
+        }
 
         // Flag the inviting user, this triggers status notifications and
         // saves us some queries otherwise.
@@ -442,66 +455,67 @@ function invite_user($op, &$edit, &$acco
 }
 
 /**
- * Set an invitation's status to accepted.
+ * Process a user that accepted an invitation.
  *
  * @param $invite
  *   An invite object.
  * @param $account
- *   The user object of the invitee.
+ *   The user object that accepted the invitation.
+ * @return
+ *   Array of target roles for the invited user.
  */
-function _invite_accept($invite, $account) {
-  // Update the invitation record.
-  db_query("UPDATE {invite} SET email = '%s', invitee = %d, joined = %d WHERE reg_code = '%s'", $account->mail, $account->uid, time(), $invite->reg_code);
-  // Delete all invites to these e-mail addresses, except this one.
-  db_query("DELETE FROM {invite} WHERE (email = '%s' OR email = '%s') AND reg_code <> '%s'", $invite->email, $account->mail, $invite->reg_code);
-  // Add all users who invited this particular e-mail address to the
-  // notification queue.
-  db_query("INSERT INTO {invite_notifications} (uid, invitee) SELECT uid, %d from {invite} WHERE (email = '%s' OR email = '%s') AND canceled = 0", $account->uid, $invite->email, $account->mail);
-  // Escalate the invitee's role.
-  _invite_escalate_role($account);
-  // Unblock user account.
-  db_query("UPDATE {users} SET status = 1 WHERE uid = %d", $account->uid);
+function invite_process($invite, $account) {
+   // Update the invitation record.
+   db_query("UPDATE {invite} SET email = '%s', invitee = %d, joined = %d WHERE reg_code = '%s'", $account->mail, $account->uid, time(), $invite->reg_code);
+
+   // Delete all invites to these e-mail addresses, except this one.
+   db_query("DELETE FROM {invite} WHERE (email = '%s' OR email = '%s') AND reg_code <> '%s'", $invite->email, $account->mail, $invite->reg_code);
+
+   // Add all users who invited this particular e-mail address to the
+   // notification queue.
+   db_query("INSERT INTO {invite_notifications} (uid, invitee) SELECT uid, %d from {invite} WHERE (email = '%s' OR email = '%s') AND canceled = 0", $account->uid, $invite->email, $account->mail);
+
+  // Unblock the user account.
+   db_query("UPDATE {users} SET status = 1 WHERE uid = %d", $account->uid);
+
+  // Determine target roles for invited user.
+  $roles = invite_target_roles($invite);
+
+  // Allow other modules to act on the role escalation.
+  $args = array('invitee' => $account, 'inviter' => $invite->inviter, 'roles' => &$roles);
+  module_invoke_all('invite', 'escalate', $args);
+
+  return $roles;
 }
 
 /**
- * Escalates an invited user's role, based on the role(s) of the inviter.
+ * Determine target roles based on the roles of an inviter.
  *
- * @param $account
- *   The user object of the invitee.
+ * @param $invite
+ *   An invite object.
+ * @return
+ *   Array of target roles for an invited user.
  */
-function _invite_escalate_role($account) {
-  // Default target role.
-  $roles = array('default');
-
-  // Add roles of inviter.
-  $inviter_uid = db_result(db_query("SELECT uid FROM {invite} WHERE invitee = %d", $account->uid));
-  if ($inviter_uid && $inviter = user_load(array('uid' => $inviter_uid))) {
-    $roles = array_merge($roles, array_intersect($inviter->roles, user_roles(0, 'send invitations')));
+function invite_target_roles($invite) {
+  $targets = array();
+
+   // Add a dummy entry to retrieve the default target role setting.
+   $roles = array('default' => 'default');
+
+   // Add roles of inviter.
+  if ($invite->inviter) {
+    $roles = array_merge($roles, array_intersect($invite->inviter->roles, user_roles(FALSE, 'send invitations')));
   }
 
   // Map to configured target roles.
-  $targets = array();
-  foreach ($roles as $role) {
-    $role_no_space = str_replace(' ', '_', $role);
-    $target = variable_get('invite_target_role_'. $role_no_space, DRUPAL_AUTHENTICATED_RID);
+  foreach ($roles as $rid => $role) {
+    $target = variable_get('invite_target_role_'. str_replace(' ', '_', $role), DRUPAL_AUTHENTICATED_RID);
     if ($target != DRUPAL_AUTHENTICATED_RID) {
       $targets[$target] = $target;
     }
   }
 
-  // Notify other modules of changed user.
-  $edit = array('roles' => $targets);
-  user_module_invoke('update', $edit, $account);
-
-  // Save new user role(s).
-  foreach ($targets as $target) {
-    db_query("DELETE FROM {users_roles} WHERE uid = %d AND rid = %d", $account->uid, $target);
-    db_query("INSERT INTO {users_roles} (uid, rid) VALUES (%d, %d)", $account->uid, $target);
-  }
-
-  // Notify other modules of role escalation.
-  $args = array('invitee' => $account, 'inviter' => $inviter, 'roles' => $targets);
-  module_invoke_all('invite', 'escalate', $args);
+  return $targets;
 }
 
 /**