--- a/popups_reference.module	Mon Mar 23 12:44:22 2009 +0100
+++ b/popups_reference.module	Thu Mar 26 14:59:27 2009 +0100
@@ -28,23 +28,25 @@
     // Add the "Add New: Node Type" links.
     $node = $form['#node'];
     if ($form['type']['#value'] .'_node_form' == $form_id) {  
-      $fields = content_fields();
-      foreach ($form as $key => $item) {
-        if (is_array($item)) {
-          $type = $item['#type'];
-          if ($type == 'fieldset') { // Loop through all the subitems.
-            foreach ($form[$key] as $subkey => $subitem) {
-              popups_reference_alter_item($form[$key], $subkey, $subitem, $fields);
-            }
-          }
-          else {
-            popups_reference_alter_item($form, $key, $item, $fields);
-          }
-        }
-        
+      $form['#after_build'][] = 'popups_reference_after';
+    }
+  }
+}
+
+function popups_reference_after($form, &$form_state) {
+  $fields = content_fields();
+  foreach ($form as $key => $item) {
+    if (element_child($key) && is_array($item)) {
+      $type = $item['#type'];
+      if ($type == 'fieldset' || is_numeric($key)) { // Loop through all the subitems.
+        $form[$key] = popups_reference_after($form[$key], $form_state);
+      }
+      else {
+        popups_reference_alter_item($form, $key, $item, $fields);
       }
     }
   }
+  return $form;
 }
 
 /**
@@ -93,9 +95,8 @@
     if ($links) {
       // Put the nodereference widget and links in an wrapper.
       // Makes it easy to find for Ahah targeting, and popups_reference behavior selecting.
-      $form[$key]['#popups_reference_prefix'] = '<div id="'. $wrapper_id .'">';
-      $form[$key]['#popups_reference_suffix'] = '<div>Add New: ' . implode(', ', $links) .'</div></div>';
-      $form[$key]['#after_build'][] = 'popups_reference_after_build';
+      $form[$key]['#prefix'] = '<div id="'. $wrapper_id .'">' . $form[$key]['#prefix'];
+      $form[$key]['#suffix'] = $form[$key]['#suffix'] . '<div>Add New: ' . implode(', ', $links) .'</div></div>';
     }
   }
 }