2 files changed, 43 insertions, 40 deletions
diff --git a/src/utils/hooks/use-headings-tree/use-headings-tree.test.ts b/src/utils/hooks/use-headings-tree/use-headings-tree.test.ts
index ad30a4f..2c8ff2d 100644
--- a/src/utils/hooks/use-headings-tree/use-headings-tree.test.ts
+++ b/src/utils/hooks/use-headings-tree/use-headings-tree.test.ts
@@ -1,5 +1,5 @@
 import { describe, expect, it } from '@jest/globals';
-import { renderHook } from '@testing-library/react';
+import { act, renderHook } from '@testing-library/react';
 import { useHeadingsTree } from './use-headings-tree';
 
 const labels = {
@@ -9,7 +9,7 @@ const labels = {
 };
 
 describe('useHeadingsTree', () => {
-  it('returns a ref object and the headings tree', () => {
+  it('returns a ref callback and the headings tree', () => {
     const wrapper = document.createElement('div');
 
     wrapper.innerHTML = `
@@ -19,12 +19,13 @@ describe('useHeadingsTree', () => {
 <h2>${labels.secondH2}</h2>
 <p>Totam cumque aut ipsum. Necessitatibus magnam necessitatibus. Qui illo nulla non ab. Accusamus voluptatem ab fugiat voluptas aspernatur velit dolore reprehenderit. Voluptatem quod minima asperiores voluptatum distinctio cumque quo.</p>`;
 
-    const wrapperRef = { current: wrapper };
-    const { result } = renderHook(() => useHeadingsTree(wrapperRef));
+    const { result } = renderHook(() => useHeadingsTree());
 
-    expect(result.current.length).toBe(1);
-    expect(result.current[0].label).toBe(labels.h1);
-    expect(result.current[0].children.length).toBe(2);
+    act(() => result.current.ref(wrapper));
+
+    expect(result.current.tree.length).toBe(1);
+    expect(result.current.tree[0].label).toBe(labels.h1);
+    expect(result.current.tree[0].children.length).toBe(2);
   });
 
   it('can return a headings tree starting at the specified level', () => {
@@ -37,14 +38,13 @@ describe('useHeadingsTree', () => {
 <h2>${labels.secondH2}</h2>
 <p>Totam cumque aut ipsum. Necessitatibus magnam necessitatibus. Qui illo nulla non ab. Accusamus voluptatem ab fugiat voluptas aspernatur velit dolore reprehenderit. Voluptatem quod minima asperiores voluptatum distinctio cumque quo.</p>`;
 
-    const wrapperRef = { current: wrapper };
-    const { result } = renderHook(() =>
-      useHeadingsTree(wrapperRef, { fromLevel: 2 })
-    );
+    const { result } = renderHook(() => useHeadingsTree({ fromLevel: 2 }));
+
+    act(() => result.current.ref(wrapper));
 
-    expect(result.current.length).toBe(2);
-    expect(result.current[0].label).toBe(labels.firstH2);
-    expect(result.current[1].label).toBe(labels.secondH2);
+    expect(result.current.tree.length).toBe(2);
+    expect(result.current.tree[0].label).toBe(labels.firstH2);
+    expect(result.current.tree[1].label).toBe(labels.secondH2);
   });
 
   it('can return a headings tree stopping at the specified level', () => {
@@ -57,22 +57,17 @@ describe('useHeadingsTree', () => {
 <h2>${labels.secondH2}</h2>
 <p>Totam cumque aut ipsum. Necessitatibus magnam necessitatibus. Qui illo nulla non ab. Accusamus voluptatem ab fugiat voluptas aspernatur velit dolore reprehenderit. Voluptatem quod minima asperiores voluptatum distinctio cumque quo.</p>`;
 
-    const wrapperRef = { current: wrapper };
-    const { result } = renderHook(() =>
-      useHeadingsTree(wrapperRef, { toLevel: 1 })
-    );
+    const { result } = renderHook(() => useHeadingsTree({ toLevel: 1 }));
+
+    act(() => result.current.ref(wrapper));
 
-    expect(result.current.length).toBe(1);
-    expect(result.current[0].label).toBe(labels.h1);
-    expect(result.current[0].children).toStrictEqual([]);
+    expect(result.current.tree.length).toBe(1);
+    expect(result.current.tree[0].label).toBe(labels.h1);
+    expect(result.current.tree[0].children).toStrictEqual([]);
   });
 
   it('throws an error if the options are invalid', () => {
-    const wrapperRef = { current: null };
-
-    expect(() =>
-      useHeadingsTree(wrapperRef, { fromLevel: 2, toLevel: 1 })
-    ).toThrowError(
+    expect(() => useHeadingsTree({ fromLevel: 2, toLevel: 1 })).toThrowError(
       'Invalid options: `fromLevel` must be lower or equal to `toLevel`.'
     );
   });
diff --git a/src/utils/hooks/use-headings-tree/use-headings-tree.ts b/src/utils/hooks/use-headings-tree/use-headings-tree.ts
index 6a081e7..68bdde8 100644
--- a/src/utils/hooks/use-headings-tree/use-headings-tree.ts
+++ b/src/utils/hooks/use-headings-tree/use-headings-tree.ts
@@ -1,4 +1,4 @@
-import { useEffect, useState, type RefObject } from 'react';
+import { useState, useCallback, type RefCallback } from 'react';
 import type { HeadingLevel } from '../../../components';
 
 export type HeadingsTreeNode = {
@@ -111,17 +111,26 @@ const buildHeadingsTreeFrom = (
   return treeNodes;
 };
 
+export type UseHeadingsTreeReturn<T extends HTMLElement> = {
+  /**
+   * A callback function to set a ref.
+   */
+  ref: RefCallback<T>;
+  /**
+   * The headings tree.
+   */
+  tree: HeadingsTreeNode[];
+};
+
 /**
  * React hook to retrieve the headings tree in a document or in a given wrapper.
  *
- * @param {RefObject<T>} ref - A ref to the element where to look for headings.
  * @param {UseHeadingsTreeOptions} options - The headings tree config.
- * @returns {HeadingsTreeNode[]} The headings tree.
+ * @returns {UseHeadingsTreeReturn<T>} The headings tree and a ref callback.
  */
 export const useHeadingsTree = <T extends HTMLElement = HTMLElement>(
-  ref: RefObject<T>,
   options?: UseHeadingsTreeOptions
-): HeadingsTreeNode[] => {
+): UseHeadingsTreeReturn<T> => {
   if (
     options?.fromLevel &&
     options.toLevel &&
@@ -134,15 +143,14 @@ export const useHeadingsTree = <T extends HTMLElement = HTMLElement>(
   const [tree, setTree] = useState<HeadingsTreeNode[]>([]);
   const requestedHeadingTags = getHeadingTagsList(options);
   const query = requestedHeadingTags.join(', ');
+  const ref: RefCallback<T> = useCallback(
+    (el) => {
+      const headingNodes = el?.querySelectorAll<HTMLHeadingElement>(query);
 
-  useEffect(() => {
-    if (typeof window === 'undefined') return;
-
-    const headingNodes =
-      ref.current?.querySelectorAll<HTMLHeadingElement>(query);
-
-    if (headingNodes) setTree(buildHeadingsTreeFrom(headingNodes));
-  }, [query, ref]);
+      if (headingNodes) setTree(buildHeadingsTreeFrom(headingNodes));
+    },
+    [query]
+  );
 
-  return tree;
+  return { ref, tree };
 };