From 382e018856a884a96a94ad551dbc1d7b0317b2e5 Mon Sep 17 00:00:00 2001
From: Yuan Fu <casouri@gmail.com>
Date: Sun, 29 Jan 2023 00:06:09 -0800
Subject: [PATCH] Add treesit-subtree-stat

* src/treesit.c (Ftreesit_subtree_stat): New function.
* lisp/treesit.el (treesit): Add to shortdoc.
---
 lisp/treesit.el | 11 ++++++---
 src/treesit.c   | 63 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 3 deletions(-)

diff --git a/lisp/treesit.el b/lisp/treesit.el
index 363692eabdf..b3029707376 100644
--- a/lisp/treesit.el
+++ b/lisp/treesit.el
@@ -2972,10 +2972,10 @@ function signals an error."
    :no-value (treesit-parser-set-included-ranges parser '((1 . 4) (5 . 8))))
   (treesit-parser-included-ranges
    :no-eval (treesit-parser-included-ranges parser)
-   :eg-result '((1 . 4) (5 . 8)))
+   :eg-result ((1 . 4) (5 . 8)))
   (treesit-query-range
    :no-eval (treesit-query-range node '((script_element) @cap))
-   :eg-result-string '((1 . 4) (5 . 8)))
+   :eg-result ((1 . 4) (5 . 8)))
 
 
   "Retrieving a node"
@@ -3121,7 +3121,12 @@ function signals an error."
    :eg-result-string "#<treesit-node (translation_unit) in 1-11>")
   (treesit-query-string
    :no-eval (treesit-query-string "int c = 0;" '((identifier) @id) 'c)
-   :eg-result-string "((id . #<treesit-node (identifier) in 5-6>))"))
+   :eg-result-string "((id . #<treesit-node (identifier) in 5-6>))")
+
+  "Misc"
+  (treesit-subtree-stat
+   :no-eval (treesit-subtree-stat node)
+   :eg-result (6 33 487)))
 
 (provide 'treesit)
 
diff --git a/src/treesit.c b/src/treesit.c
index 917db582676..b210ec0923a 100644
--- a/src/treesit.c
+++ b/src/treesit.c
@@ -3312,6 +3312,68 @@ a regexp.  */)
     return parent;
 }
 
+DEFUN ("treesit-subtree-stat",
+       Ftreesit_subtree_stat,
+       Streesit_subtree_stat, 1, 1, 0,
+       doc: /* Return information about the subtree of NODE.
+
+Return a list (MAX-DEPTH MAX-WIDTH COUNT), where MAX-DEPTH is the
+maximum depth of the subtree, MAX-WIDTH is the maximum number of
+direct children of nodes in the subtree, and COUNT is the number of
+nodes in the subtree, including NODE.  */)
+  (Lisp_Object node)
+{
+  /* Having a limit on the depth to traverse doesn't have much impact
+     on the time it takes, so I left that out.  */
+  CHECK_TS_NODE (node);
+
+  treesit_initialize ();
+
+  TSTreeCursor cursor = ts_tree_cursor_new (XTS_NODE (node)->node);
+  ptrdiff_t max_depth = 1;
+  ptrdiff_t max_width = 0;
+  ptrdiff_t count = 0;
+  ptrdiff_t current_depth = 0;
+
+  /* Traverse the subtree depth-first.  */
+  while (true)
+    {
+      count++;
+
+      /* Go down depth-first.  */
+      while (ts_tree_cursor_goto_first_child (&cursor))
+	{
+	  current_depth++;
+	  count++;
+	  /* While we're at here, measure the number of siblings.  */
+	  ptrdiff_t width_count = 1;
+	  while (ts_tree_cursor_goto_next_sibling (&cursor))
+	    width_count++;
+	  max_width = max (max_width, width_count);
+	  /* Go back to the first sibling.  */
+	  treesit_assume_true (ts_tree_cursor_goto_parent (&cursor));
+	  treesit_assume_true (ts_tree_cursor_goto_first_child (&cursor));
+	}
+      max_depth = max (max_depth, current_depth);
+
+      /* Go to next sibling.  If there is no next sibling, go to
+         parent's next sibling, and so on.  If there is no more
+         parent, we've traversed the whole subtree, stop.  */
+      while (!ts_tree_cursor_goto_next_sibling (&cursor))
+	{
+	  if (ts_tree_cursor_goto_parent (&cursor))
+	    current_depth--;
+	  else
+	    {
+	      ts_tree_cursor_delete (&cursor);
+	      return list3 (make_fixnum (max_depth),
+			    make_fixnum (max_width),
+			    make_fixnum (count));
+	    }
+	}
+    }
+}
+
 #endif	/* HAVE_TREE_SITTER */
 
 DEFUN ("treesit-available-p", Ftreesit_available_p,
@@ -3511,6 +3573,7 @@ then in the system default locations for dynamic libraries, in that order.  */);
   defsubr (&Streesit_search_subtree);
   defsubr (&Streesit_search_forward);
   defsubr (&Streesit_induce_sparse_tree);
+  defsubr (&Streesit_subtree_stat);
 #endif /* HAVE_TREE_SITTER */
   defsubr (&Streesit_available_p);
 }
-- 
2.39.2