From 445f307115efcbe03beafb8ee76d91144bb4c8b1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Carlos=20Mart=C3=ADn=20Nieto?= <cmn@dwim.me>
Date: Fri, 20 Feb 2026 20:49:39 +0100
Subject: [PATCH] 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.
---
 ext/rugged/rugged_tree.c | 50 ++++++++++++++++++++++++++++------------
 test/tree_test.rb        |  2 ++
 2 files changed, 37 insertions(+), 15 deletions(-)

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