diff --git a/ext/rugged/rugged_tree.c b/ext/rugged/rugged_tree.c index 9d39f8cfa..fb72dc049 100644 --- a/ext/rugged/rugged_tree.c +++ b/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); } /* diff --git a/test/tree_test.rb b/test/tree_test.rb index 9a4053ba0..a0e1dc05e 100644 --- a/test/tree_test.rb +++ b/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