=== accounttypes.module
==================================================================
--- accounttypes.module	(revision 14971)
+++ accounttypes.module	(local)
@@ -208,7 +208,7 @@
     $all_users[$result->uid] = $result->uid;
   }
   foreach ($all_users as $user => $uid) {
-    if (!_accounttypes_user_exists($uid)) {
+    if (!accounttypes_user_accounttype($uid)) {
       $result1 = db_query("INSERT INTO {accounttypes_users} VALUES (%d, %d)", $uid, variable_get('accounttypes_default', STD_ACCOUNT_TYPE));
       if (!$result1) {
         $error = 1;
@@ -1154,13 +1154,16 @@
 } //function accounttypes_user_filter_form_submit()
 
 /**
- * See if an entry exists in accounttypes_users for the given uid
+ * Checks for the account type of the given UID. If none is set, a value of NULL
+ * is returned.
  * 
- * @return object A database result object for one database record.
+ * @param $uid the UID of the user to check
+ * @return the numeric account type ID assigned to the user with the given uid
+ * or FALSE
  */
-function _accounttypes_user_exists($uid) {
-  return (int)db_result(db_query('SELECT COUNT(*) FROM {accounttypes_users} WHERE uid = %d', $uid));
-} //function _accounttypes_user_exists()
+function accounttypes_user_accounttype($uid) {
+  return db_result(db_query('SELECT atid FROM {accounttypes_users} WHERE uid = %d', $uid));
+} //function accounttypes_user_accounttype()
 
 /**
  * implementation of hook_user
@@ -1172,25 +1175,30 @@
  * @param mixed &$account
  * @param string $category
  */
-function accounttypes_user($op, $edit, &$account, $category = NULL) {//TODO cleanup/add better error checking
+function accounttypes_user($op, &$edit, &$account, $category = NULL) {//TODO cleanup/add better error checking
   $administer = user_access('administer accounttypes');
   $assign = user_access('assign accounttypes');
   $allacctypes = get_accounttypes();
-  if (!$edit['selectAT']) {
-    if (isset($account->selectAT)) {
-      $selectAT = $account->selectAT;
-    }
+
+  if ($edit['selectAT']) {
+    $selectAT = $edit['selectAT'];
   } 
   else {
-    $selectAT = $edit['selectAT'];
+    $selectAT = accounttypes_user_accounttype($account->uid);
   }
+  
   switch ($op) {
     case 'insert':
     case 'update': 
-      $problem = 0;
+
+      // Stop now if this is update does not explicitly update selectAT
+      if ($op == 'update' && !$edit['selectAT']) {
+        break;
+      }
+
+      // Which roles should be initially assigned to this account type?
       $new_valid_rids = array();
       $new_initial_rids = array();
-      $existing_user_rids = array();
       $result1 = db_query('SELECT ar.rid, ar.initial FROM {accounttypes_roles} ar WHERE ar.atid = %d', $selectAT);
       while ($theentry = db_fetch_object($result1)) {
         $new_valid_rids[$theentry->rid] = $theentry->rid;
@@ -1198,80 +1206,66 @@
           $new_initial_rids[$theentry->rid] = $theentry->rid;
         }
       }
-      $result2 = db_query('SELECT ur.rid, ur.uid FROM {users_roles} ur WHERE ur.uid = %d', $account->uid);
-      while ($theentry = db_fetch_object($result2)) {
-        $existing_user_rids[$theentry->rid] = $theentry->uid;
-      }
-      //--------------------first of three cases------------------------------
-      if (_accounttypes_user_exists($account->uid) && ($op == 'update')) { //this is just so we don't duplicate error checking; maybe we should.
-        foreach ($existing_user_rids as $eur => $value) {
-          if (!array_key_exists($eur, $new_valid_rids)) {
-            $result2a = db_query('DELETE FROM {users_roles} WHERE rid = %d AND uid = %d', $eur, $value);
-            if (!$result2a) {
-              $problem = 1;
-            }
+
+      // If roles have been given, make sure only allowed roles are set
+      if ($edit['roles']) {
+
+        // If the user has made a naughty role selection, remove it.
+        $picked_rids = array_keys($edit['roles']);
+        foreach ($picked_rids as $picked_rid) {
+          if (!$new_valid_rids[ $picked_rid ]) {
+            unset($edit['roles'][$picked_rid]);
           }
         }
-        foreach($new_initial_rids as $key => $value) {
-          if (!array_key_exists($key, $existing_user_rids)) {
-            $result2b = db_query("INSERT INTO {users_roles} (rid, uid) VALUES (%d, %d)", $key, $account->uid);
-          }
-        }
-        $result3 = db_query("UPDATE {accounttypes_users} SET atid = %d WHERE uid = %d", $selectAT, $account->uid);
       }
-      //--------------------second of three cases------------------------------
-      elseif ($op == 'insert') { //this is just so we don't duplicate error checking; maybe we should.
-        if (isset($edit['selectAT'])) {
-          $result3 = db_query("INSERT INTO {accounttypes_users} VALUES (%d, %d)", $account->uid, $edit['selectAT']);
-          //$result4 = db_query("INSERT INTO {users_roles} (rid, uid) VALUES (%d, %d)", 2, $account->uid);
-          foreach($new_initial_rids as $key => $value) {
-            $account->roles[$key] = $value;
-            $result2b = db_query("INSERT INTO {users_roles} (rid, uid) VALUES (%d, %d)", $key, $account->uid);
-          }
+
+      // If roles are missing, populate it with the initial roles
+      else {
+        $edit['roles'] = $new_initial_rids;
+      }
+
+      // In case of update...
+      if ($op == 'update') {
+        $current_accounttype = accounttypes_user_accounttype($account->uid);
+
+        if ($current_accounttype !== FALSE) {
+          $result_at = db_query("UPDATE {accounttypes_users} SET atid = %d WHERE uid = %d", $selectAT, $account->uid);
         }
-        else{
-          $result3 = db_query("INSERT INTO {accounttypes_users} VALUES (%d, %d)", $account->uid, variable_get('accounttypes_default', STD_ACCOUNT_TYPE));
-          //$result4 = db_query("INSERT INTO {users_roles} (rid, uid) VALUES (%d, %d)", 2, $account->uid);
-          foreach($new_initial_rids as $key => $value) {
-            $account->roles[$key] = $value;
-            $result2b = db_query("INSERT INTO {users_roles} (rid, uid) VALUES (%d, %d)", $key, $account->uid);
-          }
+
+        else {
+          $result_at = db_query("INSERT INTO {accounttypes_users} (uid, atid) VALUES (%d, %d)", $account->uid, $selectAT);
         }
       }
-      //--------------------third of three cases------------------------------
-      elseif (!_accounttypes_user_exists($account->uid) && ($op == 'update')) { //this is for the case where old users exist and we want to assign account types one by one
-        if (isset($edit['selectAT'])) {
-          foreach ($existing_user_rids as $eur => $value) {
-            if (!array_key_exists($eur, $new_valid_rids)) {
-              $result2a = db_query('DELETE FROM {users_roles} WHERE rid = %d AND uid = %d', $eur, $value);
-              if (!$result2a) {
-                $problem = 1;
-              }
-            }
-          }
-          foreach($new_initial_rids as $key => $value) {
-            if (!array_key_exists($key, $existing_user_rids)) {
-              $result2b = db_query("INSERT INTO {users_roles} (rid, uid) VALUES (%d, %d)", $key, $account->uid);
-            }
-          }
-          $result3 = db_query("INSERT INTO {accounttypes_users} VALUES (%d, %d)", $account->uid, $edit['selectAT']);
+
+      // In case of insert...
+      else {
+
+        // Assign the standard if none given (weird)
+        if (!$selectAT) {
+          $selectAT = STD_ACCOUNT_TYPE;
         }
+
+        // insert the new account type
+        $result_at = db_query("INSERT INTO {accounttypes_users} VALUES (%d, %d)", $account->uid, $edit['selectAT']);
       }
-      if (!$result3) {
-        drupal_set_message('Account Type update was unsuccessful', 'error');
+
+      // Check for errors
+      if (!$result_at) {
+        drupal_set_message(t('Account Type update was unsuccessful.'), 'error');
       }
+
+      // Go ahead and drop this into the account
       $account->selectAT = $selectAT;
       $edit['selectAT'] = NULL;
       break;
     case 'delete':
-      if (_accounttypes_user_exists($account->uid)) {
+      if (accounttypes_user_accounttype($account->uid)) {
         $result = db_query("DELETE FROM {accounttypes_users} WHERE uid = %d", $account->uid);
         unset($account->selectAT);
       }
       break;
     case 'load':
-      $result = db_result(db_query('SELECT atid FROM {accounttypes_users} WHERE uid = %d', $account->uid));
-      $account->selectAT = $result;
+      $account->selectAT = accounttypes_user_accounttype($account->uid);
       break;
     case 'validate':
       if (!array_key_exists($edit['selectAT'], $allacctypes)) {