feat: implement sticky headers in Tree component

Author: alfredxio

src/Tree.js

@@ -6,6 +6,14 @@ import {FlattenedNode} from './shapes/nodeShapes';
 import TreeState, {State} from './state/TreeState';
 
 export default class Tree extends React.Component {
+  constructor(props) {
+    super(props);
+    this.state = {
+      topStickyHeader: null, // The header that should be sticky
+    };
+    this._listRef = React.createRef();
+  }
+
   _cache = new CellMeasurerCache({
     fixedWidth: true,
     minHeight: 20,
@@ -35,8 +43,80 @@ export default class Tree extends React.Component {
       : nodes[index];
   };
 
+  // Determine if a node is a group header
+  isGroupHeader = node => {
+    // Group headers are typically parent nodes with children
+    // and deepness of 0 (root level)
+    return node.children && node.children.length > 0 && node.deepness === 0;
+  };
+
+  componentDidMount() {
+    // Initial check for headers after mounting
+    if (this._listRef.current) {
+      const list = this._listRef.current;
+      const grid = list && list.Grid;
+      if (grid) {
+        this.handleScroll({
+          scrollTop: grid.state.scrollTop,
+          scrollHeight: grid.state.scrollHeight,
+          clientHeight: grid.props.height,
+        });
+      }
+    }
+  }
+
+  // Get all headers in the current data - simplified to reduce loops
+  getAllHeaders = () => {
+    const rowCount = this.getRowCount();
+    const headers = [];
+    let cumulativeHeight = 0;
+
+    for (let i = 0; i < rowCount; i++) {
+      const node = this.getNode(i);
+      
+      if (this.isGroupHeader(node)) {
+        headers.push({
+          node,
+          index: i,
+          top: cumulativeHeight,
+        });
+      }
+      
+      cumulativeHeight += this._cache.rowHeight({index: i});
+    }
+
+    return headers;
+  };
+
+  // Handle scroll events to update sticky headers - simplified logic
+  handleScroll = ({scrollTop, scrollHeight, clientHeight}) => {
+    if (!this._listRef.current) return;
+
+    // Get all headers in the tree
+    const allHeaders = this.getAllHeaders();
+
+    // Find the header that should be sticky (last one whose top is <= scrollTop)
+    const topStickyHeader = allHeaders
+      .filter(h => h.top <= scrollTop)
+      .pop() || null;
+
+    // Only update state if something has changed
+    const currentStickyId = this.state.topStickyHeader && this.state.topStickyHeader.node && this.state.topStickyHeader.node.id;
+    const newStickyId = topStickyHeader && topStickyHeader.node && topStickyHeader.node.id;
+
+    if (currentStickyId !== newStickyId) {
+      this.setState({
+        topStickyHeader,
+      });
+    }
+  };
+
   rowRenderer = ({node, key, measure, style, NodeRenderer, index}) => {
     const {nodeMarginLeft} = this.props;
+    const isHeader = this.isGroupHeader(node);
+
+    // Add a class to identify group headers
+    const className = isHeader ? 'tree-group-header' : '';
 
     return (
       <NodeRenderer
@@ -47,10 +127,43 @@ export default class Tree extends React.Component {
           userSelect: 'none',
           cursor: 'pointer',
         }}
+        className={className}
         node={node}
         onChange={this.props.onChange}
         measure={measure}
         index={index}
+        isGroupHeader={isHeader}
+      />
+    );
+  };
+
+  // Render the sticky header - simplified render method
+  renderStickyHeader = () => {
+    const {topStickyHeader} = this.state;
+    if (!topStickyHeader) return null;
+
+    const {NodeRenderer, nodeMarginLeft} = this.props;
+    // Always use the current node from the tree to ensure we have the latest state
+    const index = topStickyHeader.index;
+    const currentNode = this.getNode(index);
+
+    return (
+      <NodeRenderer
+        key={`sticky-header-${currentNode.id}`}
+        style={{
+          marginLeft: currentNode.deepness * nodeMarginLeft,
+          userSelect: 'none',
+          cursor: 'pointer',
+          width: '100%',
+          background: '#fff',
+          zIndex: 10,
+        }}
+        className="tree-group-header tree-sticky"
+        node={currentNode}
+        onChange={this.props.onChange}
+        index={index}
+        isGroupHeader={true}
+        isSticky={true}
       />
     );
   };
@@ -66,25 +179,63 @@ export default class Tree extends React.Component {
     );
   };
 
+  componentDidUpdate(prevProps) {
+    // If nodes change, reset the cache
+    if (prevProps.nodes !== this.props.nodes) {
+      this._cache.clearAll();
+      if (this._listRef.current) {
+        this._listRef.current.recomputeRowHeights();
+      }
+      
+      // Force rerender of sticky header when nodes change
+      this.forceUpdate();
+    }
+  }
+
   render() {
     const {nodes, width, scrollToIndex, scrollToAlignment} = this.props;
+    const {topStickyHeader} = this.state;
+
+    // Calculate the height of the sticky header to properly offset the list
+    const stickyHeaderHeight = topStickyHeader ? this._cache.rowHeight({index: topStickyHeader.index}) : 0;
 
     return (
-      <AutoSizer disableWidth={Boolean(width)}>
-        {({height, width: autoWidth}) => (
-          <List
-            deferredMeasurementCache={this._cache}
-            ref={r => (this._list = r)}
-            height={height}
-            rowCount={this.getRowCount()}
-            rowHeight={this._cache.rowHeight}
-            rowRenderer={this.measureRowRenderer(nodes)}
-            width={width || autoWidth}
-            scrollToIndex={scrollToIndex}
-            scrollToAlignment={scrollToAlignment}
-          />
+      <div className="tree-container" style={{position: 'relative', height: '100%'}}>
+        {/* Sticky header container - simplified markup */}
+        {topStickyHeader && (
+          <div
+            className="tree-sticky-header-container"
+            style={{
+              position: 'absolute',
+              top: 0,
+              left: 0,
+              right: 0,
+              zIndex: 100,
+              height: `${stickyHeaderHeight}px`,
+            }}
+          >
+            {this.renderStickyHeader()}
+          </div>
         )}
-      </AutoSizer>
+
+        <AutoSizer disableWidth={Boolean(width)}>
+          {({height, width: autoWidth}) => (
+            <List
+              deferredMeasurementCache={this._cache}
+              ref={this._listRef}
+              height={height}
+              rowCount={this.getRowCount()}
+              rowHeight={this._cache.rowHeight}
+              rowRenderer={this.measureRowRenderer(nodes)}
+              width={width || autoWidth}
+              scrollToIndex={scrollToIndex}
+              scrollToAlignment={scrollToAlignment}
+              onScroll={this.handleScroll}
+              overscanRowCount={20}
+            />
+          )}
+        </AutoSizer>
+      </div>
     );
   }
 }