tree: allow for a tree limit in Tree::count_recursive

Sometimes you find yourself with a repository that not only has a lot of blobs
but a lot of trees without many blobs in between. Such a repository can cause
`count_recursive` to take a long time to return.

Allow for a second limit for the number of trees. If unset we use the limit for
blobs which will not have any effect unless you have less than one blob per
tree, which is roughly what we're trying to protect against here.

Author: carlosmn

ext/rugged/rugged_tree.c

@@ -81,32 +81,41 @@ static VALUE rb_git_tree_entrycount(VALUE self)
 
 struct rugged_treecount_cb_payload
 {
-	int count;
-	int limit;
+	int entry_count;
+	int tree_count;
+	int entry_limit;
+	int tree_limit;
 };
 
 static int rugged__treecount_cb(const char *root, const git_tree_entry *entry, void *data)
 {
 	struct rugged_treecount_cb_payload *payload = data;
 
-	if (payload->limit >= 0 && payload->count >= payload->limit) {
+	if (payload->entry_limit >= 0 && payload->entry_count >= payload->entry_limit) {
+		return -1;
+	} else if (payload->tree_limit >= 0 && payload->tree_count >= payload->tree_limit) {
 		return -1;
 	} else if(git_tree_entry_type(entry) == GIT_OBJ_TREE) {
+		++(payload->tree_count);
 		return 0;
 	} else {
-		++(payload->count);
+		++(payload->entry_count);
 		return 1;
 	}
 }
 
 /*
  *  call-seq:
- *    tree.count_recursive(limit=nil) -> count
+ *    tree.count_recursive(entry_limit=nil, tree_limit=nil) -> count
  *
- *  `limit` - The maximum number of blobs to the count in the repository.
- *  Rugged will stop walking the tree after `limit` items to avoid long
+ *  `entry_limit` - The maximum number of blobs to the count in the repository.
+ *  Rugged will stop walking the trees after `entry_limit` items to avoid long
  *  execution times.
  *
+ *  `tree_limit` - The maximum number of trees to look in. Rugged will stop
+ *  walking the trees after `tree_limit` items to avoid long execution times. If
+ *  it is unset but `entry_limit` is set, `tree_limit` is set to the same value.
+ *
  *  Return the number of blobs (up to the limit) contained in the tree and
  *  all subtrees.
  */
@@ -115,18 +124,29 @@ static VALUE rb_git_tree_entrycount_recursive(int argc, VALUE* argv, VALUE self)
 	git_tree *tree;
 	int error;
 	struct rugged_treecount_cb_payload payload;
-	VALUE rb_limit;
+	VALUE rb_entry_limit, rb_tree_limit;
 
 	TypedData_Get_Struct(self, git_tree, &rugged_object_type, tree);
 
-	rb_scan_args(argc, argv, "01", &rb_limit);
+	rb_scan_args(argc, argv, "02", &rb_entry_limit, &rb_tree_limit);
+
+	payload.entry_limit = -1;
+	payload.tree_limit = -1;
+	payload.entry_count = 0;
+	payload.tree_count = 0;
 
-	payload.limit = -1;
-	payload.count = 0;
+	if (!NIL_P(rb_entry_limit)) {
+		Check_Type(rb_entry_limit, T_FIXNUM);
+		payload.entry_limit = FIX2INT(rb_entry_limit);
+	}
+
+	if (!NIL_P(rb_tree_limit)) {
+		Check_Type(rb_tree_limit, T_FIXNUM);
+		payload.tree_limit = FIX2INT(rb_tree_limit);
+	}
 
-	if (!NIL_P(rb_limit)) {
-		Check_Type(rb_limit, T_FIXNUM);
-		payload.limit = FIX2INT(rb_limit);
+	if (payload.tree_limit < 0 && payload.entry_limit >= 0) {
+		payload.tree_limit = payload.entry_limit;
 	}
 
 
@@ -139,7 +159,7 @@ static VALUE rb_git_tree_entrycount_recursive(int argc, VALUE* argv, VALUE self)
 
 	rugged_exception_check(error);
 
-	return INT2FIX(payload.count);
+	return INT2FIX(payload.entry_count);
 }
 
 /*

test/tree_test.rb

@@ -48,6 +48,8 @@ def test_read_tree_data
     assert_equal 6, @tree.count_recursive
     assert_equal 5, @tree.count_recursive(5)
     assert_equal 6, @tree.count_recursive(10)
+    assert_equal 2, @tree.count_recursive(10, 1)
+    assert_equal 4, @tree.count_recursive(10, 2)
     assert_raises(TypeError) do
       @tree.count_recursive("NaN")
     end