option to free_root() to not my_free() the blocks

fixed bug/updated count_distinct2 test
changed reset in count distinct to avoid calls to my_free()

include/my_sys.h

@@ -108,7 +108,11 @@ extern int NEAR my_errno;		/* Last error in mysys */
 
 	/* root_alloc flags */
 #define MY_KEEP_PREALLOC	1
-
+#define MY_MARK_BLOCKS_FREE     2 /* do not my_free() blocks,
+				     just move used into free list
+				     and mark all blocks as fully free
+				  */
+  
 	/* defines when allocating data */
 
 #ifdef SAFEMALLOC

include/my_tree.h

@@ -62,6 +62,9 @@ void init_tree(TREE *tree,uint default_alloc_size, int element_size,
 	       qsort_cmp2 compare, my_bool with_delete,
 	       void (*free_element)(void*));
 void delete_tree(TREE*);
+void reset_tree(TREE*);
+  /* similar to delete tree, except we do not my_free() blocks in mem_root
+   */
 #define is_tree_inited(tree) ((tree)->root != 0)
 
 	/* Functions on leafs */

mysql-test/r/count_distinct2.result

@@ -70,3 +70,7 @@ count(distinct s,t)
 5
 count(distinct n1)	count(distinct n2)
 2	3
+count(distinct n2)	n1
+1	NULL
+1	1
+3	2

mysql-test/t/count_distinct2.test

@@ -40,3 +40,6 @@ select distinct s,t from t1;
 select count(distinct s,t) from t1;
 
 select count(distinct n1), count(distinct n2) from t1;
+
+select count(distinct n2), n1 from t1 group by n1;
+drop table t1;

mysys/my_alloc.c

@@ -100,7 +100,41 @@ gptr alloc_root(MEM_ROOT *mem_root,unsigned int Size)
 #endif
 }
 
-	/* deallocate everything used by alloc_root */
+static inline void mark_blocks_free(MEM_ROOT* root)
+{
+  reg1 USED_MEM *next,*last = 0;
+
+  /* iterate through (partially) free blocks, mark them fully free */
+  for(next = root->free; next; next = next->next )
+    {
+      last = next;
+      next->left = next->size - ALIGN_SIZE(sizeof(USED_MEM));
+    }
+  /* if free block list was not empty, point the next of the
+     last free block to the beginning of the used list  */
+  next = root->used; /* a little optimization to avoid dereferencing root
+		      twice - we will shortly start iterating through used
+		      list */
+  if(last)
+    last->next = next;
+  else /* if free list is empty, just point it to the current used*/
+    root->free = next;
+
+  /* now go through the current used list, and mark each block
+     as fully free. Note that because of our optimization, we do not
+     need to initialize next here - see above
+  */
+  for(;next; next = next->next)
+      next->left = next->size - ALIGN_SIZE(sizeof(USED_MEM));
+
+  /* Now everything is set - we just need to indicate that nothing is used
+     anymore
+  */
+  root->used = 0;
+}
+
+	/* deallocate everything used by alloc_root or just move
+	 used blocks to free list if called with MY_USED_TO_FREE */
 
 void free_root(MEM_ROOT *root, myf MyFlags)
 {
@@ -109,6 +143,11 @@ void free_root(MEM_ROOT *root, myf MyFlags)
 
   if (!root)
     DBUG_VOID_RETURN; /* purecov: inspected */
+  if(MyFlags & MY_MARK_BLOCKS_FREE)
+    {
+      mark_blocks_free(root);
+      DBUG_VOID_RETURN;
+    }
   if (!(MyFlags & MY_KEEP_PREALLOC))
     root->pre_alloc=0;
 

mysys/tree.c

@@ -103,9 +103,9 @@ void init_tree(TREE *tree, uint default_alloc_size, int size,
   DBUG_VOID_RETURN;
 }
 
-void delete_tree(TREE *tree)
+static void free_tree(TREE *tree, myf free_flags)
 {
-  DBUG_ENTER("delete_tree");
+  DBUG_ENTER("free_tree");
   DBUG_PRINT("enter",("tree: %lx",tree));
 
   if (tree->root)				/* If initialized */
@@ -116,7 +116,7 @@ void delete_tree(TREE *tree)
     {
       if (tree->free)
 	delete_tree_element(tree,tree->root);
-      free_root(&tree->mem_root,MYF(0));
+      free_root(&tree->mem_root, free_flags);
     }
   }
   tree->root= &tree->null_element;
@@ -125,6 +125,18 @@ void delete_tree(TREE *tree)
   DBUG_VOID_RETURN;
 }
 
+void delete_tree(TREE* tree)
+{
+  free_tree(tree, MYF(0)); /* my_free() mem_root if applicable */
+}
+
+void reset_tree(TREE* tree)
+{
+  free_tree(tree, MYF(MY_MARK_BLOCKS_FREE));
+  /* do not my_free() mem_root if applicable, just mark blocks as free */
+}
+
+
 static void delete_tree_element(TREE *tree, TREE_ELEMENT *element)
 {
   if (element != &tree->null_element)

sql/item_sum.cc

@@ -924,7 +924,7 @@ bool Item_sum_count_distinct::setup(THD *thd)
 void Item_sum_count_distinct::reset()
 {
   if(use_tree)
-    delete_tree(&tree);
+    reset_tree(&tree);
   else
     {
       table->file->extra(HA_EXTRA_NO_CACHE);