1 files changed, 97 insertions, 34 deletions

diff --git a/taxonomy_node_tree.module b/taxonomy_node_tree.module
index cbd9b26..619d5d7 100644
--- a/taxonomy_node_tree.module
+++ b/taxonomy_node_tree.module
@@ -64,6 +64,10 @@ function taxonomy_node_tree_index($nodes, $terms) {
       $tree[$term->tid] = $term;
     }
   }
+
+  // Rewind result so it can be used again.
+  $nodes->data_seek(0);
+
   return $tree;
 }
 
@@ -118,76 +122,125 @@ function taxonomy_node_tree_sort(&$tree, $term) {
  * @param $tree
  *   Tree with term objects.
  *
- * @param $clean
- *   Set to TRUE to remove branches without nodes.
- *
  * @return
  *   Hierarchical term tree.
  */
-function taxonomy_node_tree_hierarchy($tree, $clean = FALSE) {
+function taxonomy_node_tree_hierarchy($tree) {
   if (is_array($tree) && $tree != NULL) {
     foreach ($tree as &$term) {
       taxonomy_node_tree_sort($tree, $term);
-
-      if ($clean) {
-        taxonomy_node_tree_clean($term);
-      }
     }
   }
   return $tree;
 }
 
 /**
- * @todo
+ * Count descendant nodes.
+ *
+ * @param $term
+ *   Term object.
+ *
+ * @param $count
+ *   Optional initial count offset. Also used for recursion.
+ *
+ * @return
+ *   Number of descendant nodes.
  */
-function taxonomy_node_tree_clean(&$term) {
-  $list = taxonomy_node_tree_get_empty($term);
-  taxonomy_node_tree_remove($term, $list);
-}
-
-function taxonomy_node_tree_remove(&$term, $list) {
-  if (in_array($term->tid, $list)) {
-    $term = NULL;
-    return;
-  }
-
+function taxonomy_node_tree_count($term, &$count = 0) {
   if (isset($term->below)) {
     foreach ($term->below as &$below) {
-      taxonomy_node_tree_remove($below, $list);
+      taxonomy_node_tree_count($below, $count);
     }
   }
+
+  if (isset($term->nodes)) {
+    $count += count($term->nodes);
+  }
+
+  return $count;
 }
 
 /**
- * TODO: skip parents should be controlled by a parameter.
+ * Build a list of empty terms.
+ *
+ * @param $term
+ *   Term object.
+ *
+ * @param $clean
+ *   Set to true to cleanup term before adding to list.
+ *
+ * @param $list
+ *   Array to store term list.
+ *
+ * @return
+ *   Array with terms without descendant nodes.
+ *
+ * @todo
+ *   Skip parents should be controlled by a parameter.
  */
-function taxonomy_node_tree_get_empty($term, &$list = array()) {
+function taxonomy_node_tree_get_empty($term, $clean = FALSE, &$list = array()) {
   if (taxonomy_node_tree_count($term) == 0) {
     // Skip parents.
     if ($term->parents[0] != 0) {
-      $list[] = $term->tid;
+      $save = drupal_clone($term);
+
+      if ($clean == TRUE) {
+        unset($save->nodes);
+        unset($save->below);
+        unset($save->children);
+      }
+
+      $list[] = $save;
     }
   }
   else if (isset($term->below)) {
     foreach ($term->below as $below) {
-      taxonomy_node_tree_get_empty($below, $list);
+      taxonomy_node_tree_get_empty($below, $clean, $list);
     }
   }
 
   return $list;
 }
 
-function taxonomy_node_tree_count($term, &$count = 0) {
-  if (isset($term->below)) {
-    foreach ($term->below as &$below) {
-      taxonomy_node_tree_count($below, $count);
+/**
+ * Build a list of no-nempty terms.
+ *
+ * @param $term
+ *   Term object.
+ *
+ * @param $clean
+ *   Set to true to cleanup term before adding to list.
+ *
+ * @param $list
+ *   Array to store term list.
+ *
+ * @return
+ *   Array with terms with descendant nodes.
+ *
+ * @todo
+ *   Skip parents should be controlled by a parameter.
+ */
+function taxonomy_node_tree_get_non_empty($term, $clean = FALSE, &$list = array()) {
+  // Add parents and terms with descendant nodes.
+  if ($term->parents[0] == 0 || taxonomy_node_tree_count($term) > 0) {
+    $save = drupal_clone($term);
+
+    if ($clean == TRUE) {
+      unset($save->nodes);
+      unset($save->below);
+      unset($save->children);
     }
+
+    $list[] = $save;
   }
-  else if (isset($term->nodes)) {
-    $count += count($term->nodes);
+
+  if (is_array($term->below)) {
+    foreach ($term->below as $below) {
+      taxonomy_node_tree_get_non_empty($below, $clean, $list);
+    }
   }
 
-  return $count;
+  return $list;
 }
 
 /**
@@ -203,7 +256,7 @@ function taxonomy_node_tree_count($term, &$count = 0) {
  *   Array with term objects.
  *
  * @param $clean
- *   Set to TRUE to remove branches without nodes.
+ *   Set to TRUE to remove branches without descendant nodes.
  *
  * @return
  *  Hierarchical tree.
@@ -211,7 +264,17 @@ function taxonomy_node_tree_count($term, &$count = 0) {
 function taxonomy_node_tree_build($nodes, $terms, $clean = FALSE) {
   $tree = taxonomy_node_tree_index($nodes, $terms);
   $tree = taxonomy_node_tree_relation($tree);
-  $tree = taxonomy_node_tree_hierarchy($tree, $clean);
+  $tree = taxonomy_node_tree_hierarchy($tree);
+
+  if ($clean == TRUE && is_array($tree)) {
+    $terms = array();
+    foreach ($tree as $term) {
+      $list  = taxonomy_node_tree_get_non_empty($term, TRUE);
+      $terms = array_merge($terms, $list);
+    }
+
+    $tree = taxonomy_node_tree_build($nodes, $terms, FALSE);
+  }
 
   return $tree;
 }