diff --git a/role_expire.js b/role_expire.js
index 582df8b..24027c5 100644
--- a/role_expire.js
+++ b/role_expire.js
@@ -14,8 +14,11 @@
     attach: function (context, settings) {
       $('input.role-expire-role-expiry', context).parent().hide();
 
-      $('#edit-roles input.form-checkbox', context).each(function() {
-        var textfieldId = this.id.replace("roles", "role-expire");
+      var rolesKey = $('#edit-roles-change').length == 0 ? 'roles' : 'roles-change';
+      var checkBoxes = $('#edit-' + rolesKey + ' input.form-checkbox', context);
+      
+      checkBoxes.each(function() {
+        var textfieldId = this.id.replace(rolesKey, "role-expire");
 
         // Move all expiry date fields under corresponding checkboxes
         $(this).parent().after($('#'+textfieldId).parent());
@@ -26,8 +29,8 @@
         }
       });
 
-      $('#edit-roles input.form-checkbox', context).click(function() {
-        var textfieldId = this.id.replace("roles", "role-expire");
+      checkBoxes.click(function() {
+        var textfieldId = this.id.replace(rolesKey, "role-expire");
 
         // Toggle expiry date fields
         if ($(this).attr("checked")) {
diff --git a/role_expire.module b/role_expire.module
index 4557f39..06db781 100644
--- a/role_expire.module
+++ b/role_expire.module
@@ -216,14 +216,14 @@ function role_expire_permission() {
  * Implements hook_form_FORM-ID_alter().
  */
 function role_expire_form_user_register_form_alter(&$form, $form_state) {
-  $form = array_merge_recursive($form, role_expire_add_expiration_input());
+  role_expire_add_expiration_input($form);
 }
 
 /**
  * Implements hook_form_FORM-ID_alter().
  */
 function role_expire_form_user_profile_form_alter(&$form, $form_state) {
-  $form = array_merge_recursive($form, role_expire_add_expiration_input($form['#user']));
+  role_expire_add_expiration_input($form);
 }
 
 /**
@@ -452,22 +452,22 @@ function role_expire_cron() {
 
 
 /**
- * Add form element that accepts the role expiration time.
+ * Add form elements to the given form that accept the role expiration times.
  *
- * @param $account
- *   The user object.
- * @return
- *   Form element.
+ * @param $form array
+ *   The user edit form to add role expire input elements to.
  */
-function role_expire_add_expiration_input($account = NULL) {
-  $form = array();
-  if (user_access('administer users') || user_access('administer role expire')) {
+function role_expire_add_expiration_input(array &$form) {
+  $roles_key = user_access('administer permissions') ? 'roles' : 'roles_change';
+  if (user_access('administer role expire') && !empty($form['account'][$roles_key])) {
+    $uid = (!empty($form['#user']) && is_object($form['#user']) && !empty($form['#user']->uid)) ? $form['#user']->uid : NULL;
+    
     drupal_add_js(drupal_get_path('module', 'role_expire') . '/role_expire.js');
-    $form['roles']['#attributes'] = array('class' => array('role-expire-roles'));
-
-    foreach (_role_expire_get_role() as $rid => $role) {
-      if (is_object($account) and array_key_exists('uid', $account)) {
-        $expiry_timestamp = role_expire_get_user_role_expiry_time($account->uid, $rid);
+    $form['account'][$roles_key]['#attributes'] = array('class' => array('role-expire-roles'));
+    
+    foreach ($form['account'][$roles_key]['#options'] as $rid => $role) {
+      if ($uid !== NULL) {
+        $expiry_timestamp = role_expire_get_user_role_expiry_time($uid, $rid);
       }
       else {
         $expiry_timestamp = '';
@@ -481,10 +481,9 @@ function role_expire_add_expiration_input($account = NULL) {
         '#description' => t("Leave blank for default role expiry (never, or the duration you have set for the role), enter date and time in format: <em>yyyy-mm-dd hh:mm:ss</em> or use relative time i.e. 1 day, 2 months, 1 year, 3 years.")
       );
     }
-
+    
     $form['#validate'][] = '_role_expire_validate_role_expires';
   }
-  return $form;
 }
 
 /*******************************************************************************