diff --git a/nodeconnect.form.inc b/nodeconnect.form.inc
index 4cb6e20..b2cabb8 100644
--- a/nodeconnect.form.inc
+++ b/nodeconnect.form.inc
@@ -11,52 +11,69 @@
  */
 function _nodeconnect_field_attach_form($entity_type, $entity, &$form, &$form_state, $langcode=LANGUAGE_NONE, $field_name, $field) {
   $language = $form[$field_name]['#language'];
-  foreach ($form[$field_name][$language] as $key => $field) {
+  // Add a class for single value fields, this class will be referenced in js
+  $extra = $form[$field_name][$language]['#cardinality'] == 1 ? 'single-value' : '';
+  $field_title = $form[$field_name][$language]['#title'];
+
+  foreach ($form[$field_name][$language] as $key => &$widget_element) {
     if (is_numeric($key)) {
-      // we have to add some extra js for single value fields so we add a class to let the
-      // js know
-      $extra = $form[$field_name][$language]['#cardinality'] == 1 ? 'single-value' : '';
-      $form[$field_name][$language][$key]["add_nodeconnect__{$field_name}_$key"] = array(
-        '#type' => 'submit',
-        '#limit_validation_errors' => array(array($field_name)),
-        '#value' => t('Add New @field_title content', array('@field_title' => $form[$field_name][$language]['#title'])),
-        '#name' => "add_nodeconnect__{$field_name}_$key",
-        '#prefix' => "<div class = 'nodeconnect-add $extra'>",
-        '#suffix' => '</div>',
-        '#key' => $key,
-        '#field' => $field_name,
-        '#add_child' => TRUE,
-        '#language' => $language,
-        '#submit' => array('nodeconnect_include_form', 'nodeconnect_add_edit_button_submit'),
-        '#weight' => -2,
-        '#attached' => array(
-          'js' => array(
-            drupal_get_path('module', 'nodeconnect') . "/nodeconnect.js",
-          ),
-          'css' => array(
-            drupal_get_path('module', 'nodeconnect') . "/nodeconnect.css",
-          ),
-        ),
-      );
-      $form[$field_name][$language][$key]["edit_nodeconnect__{$field_name}_$key"] = array(
-        '#type' => 'submit',
-        '#limit_validation_errors' => array(array($field_name)),
-        '#value' => t('Edit referenced @field_title', array('@field_title' => $form[$field_name][$language]['#title'])),
-        '#name' => "edit_nodeconnect__{$field_name}_$key",
-        '#prefix' => "<div class = 'nodeconnect-edit $extra'>",
-        '#suffix' => '</div>',
-        '#key' => $key,
-        '#field' => $field_name,
-        '#add_child' => FALSE,
-        '#language' => 'und',
-        '#submit' => array('nodeconnect_include_form', 'nodeconnect_add_edit_button_submit'),
-        '#weight' => -2,
-      );
+      // Add button
+      _nodeconnect_form_element_new($widget_element, $language, $field_name, $field_title, $key, $extra);
+      // Edit button
+      _nodeconnect_form_element_edit($widget_element, $language, $field_name, $field_title, $key, $extra);
     }
   }
 }
 
 /*
+ * Creates a form submit element for the noderef field that goes to the child node add form
+ */
+function _nodeconnect_form_element_new(&$widget_element, $language, $field_name, $field_title, $key, $extra = '') {
+  $widget_element["add_nodeconnect__{$field_name}_$key"] = array(
+    '#type' => 'submit',
+    '#limit_validation_errors' => array(array($field_name)),
+    '#value' => t('Add New @field_title content', array('@field_title' => $field_title)),
+    '#name' => "add_nodeconnect__{$field_name}_$key",
+    '#prefix' => "<div class = 'nodeconnect-add $extra'>",
+    '#suffix' => '</div>',
+    '#key' => $key,
+    '#field' => $field_name,
+    '#add_child' => TRUE,
+    '#language' => $language,
+    '#submit' => array('nodeconnect_include_form', 'nodeconnect_add_edit_button_submit'),
+    '#weight' => -2,
+    '#attached' => array(
+      'js' => array(
+        drupal_get_path('module', 'nodeconnect') . "/nodeconnect.js",
+      ),
+      'css' => array(
+        drupal_get_path('module', 'nodeconnect') . "/nodeconnect.css",
+      ),
+    ),
+  );
+}
+
+/*
+ * Creates a form submit element for the noderef field that goes to the referenced node's edit form
+ */
+function _nodeconnect_form_element_edit(&$widget_element, $language, $field_name, $field_title, $key, $extra = '') {
+  $widget_element["edit_nodeconnect__{$field_name}_$key"] = array(
+    '#type' => 'submit',
+    '#limit_validation_errors' => array(array($field_name)),
+    '#value' => t('Edit referenced @field_title', array('@field_title' => $field_title)),
+    '#name' => "edit_nodeconnect__{$field_name}_$key",
+    '#prefix' => "<div class = 'nodeconnect-edit $extra'>",
+    '#suffix' => '</div>',
+    '#key' => $key,
+    '#field' => $field_name,
+    '#add_child' => FALSE,
+    '#language' => 'und',
+    '#submit' => array('nodeconnect_include_form', 'nodeconnect_add_edit_button_submit'),
+    '#weight' => -2,
+  );
+}
+
+/*
  * nodeconnect_add_edit_button_submit() is called when a new node is to be added
  *
  * we cache the current state and form