--- a/popups_reference.module	Wed Jan 21 14:39:42 2009 +0100
+++ b/popups_reference.module	Mon Jan 26 14:23:35 2009 +0100
@@ -92,24 +92,19 @@
     if ($links) {
       // Put the nodereference widget and links in an wrapper.
       // Makes it easy to find for Ahah targeting, and popups_reference behavior selecting.
-      $prefix = '<div id="'. $wrapper_id .'">';
-      if (isset($form[$key]['#prefix'])) {
-        $form[$key]['#prefix'] .= $prefix;
-      }
-      else {
-        $form[$key]['#prefix'] = $prefix;
-      }
-      $suffix = '<div>Add New: ' . implode(', ', $links) .'</div></div>';
-      if (isset($form[$key]['#suffix'])) {
-        $form[$key]['#suffix'] = $suffix . $form[$key]['#suffix'];
-      }
-      else {
-        $form[$key]['#suffix'] = $suffix;
-      }
+      $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';
     }
   }
 }
 
+function popups_reference_after_build($form_element) {
+  $form_element['#prefix'] = $form_element['#popups_reference_prefix'] . $form_element['#prefix'];
+  $form_element['#suffix'] .= $form_element['#popups_reference_suffix'];
+  return $form_element;
+}
+
 /**
  * Generates 'Add new...' link
  * for each allowed content type