feat: 🎸 added TreeState support to Tree component

Author: diogofcunha

src/Tree.js

@@ -3,14 +3,39 @@ import PropTypes from 'prop-types';
 import {AutoSizer, List, CellMeasurerCache, CellMeasurer} from 'react-virtualized';
 
 import {FlattenedNode} from './shapes/nodeShapes';
+import TreeState, {State} from './state/TreeState';
 
 export default class Tree extends React.Component {
   _cache = new CellMeasurerCache({
     fixedWidth: true,
     minHeight: 20,
   });
 
-  rowRenderer = ({node, key, measure, style, NodeRenderer}) => {
+  getRowCount = () => {
+    const {nodes} = this.props;
+
+    return nodes instanceof State ? nodes.flattenedTree.length : nodes.length;
+  };
+
+  getNodeDeepness = (node, index) => {
+    const {nodes} = this.props;
+
+    if (nodes instanceof State) {
+      TreeState.getNodeDeepness(nodes, index);
+    }
+
+    return nodes instanceof State ? TreeState.getNodeDeepness(nodes, index) : node.deepness;
+  };
+
+  getNode = index => {
+    const {nodes} = this.props;
+
+    return nodes instanceof State
+      ? {...TreeState.getNodeAt(nodes, index), deepness: this.getNodeDeepness({}, index)}
+      : nodes[index];
+  };
+
+  rowRenderer = ({node, key, measure, style, NodeRenderer, index}) => {
     const {nodeMarginLeft} = this.props;
 
     return (
@@ -25,17 +50,18 @@ export default class Tree extends React.Component {
         node={node}
         onChange={this.props.onChange}
         measure={measure}
+        index={index}
       />
     );
   };
 
   measureRowRenderer = nodes => ({key, index, style, parent}) => {
     const {NodeRenderer} = this.props;
-    const node = nodes[index];
+    const node = this.getNode(index);
 
     return (
       <CellMeasurer cache={this._cache} columnIndex={0} key={key} rowIndex={index} parent={parent}>
-        {m => this.rowRenderer({...m, node, key, style, NodeRenderer})}
+        {m => this.rowRenderer({...m, index, node, key, style, NodeRenderer})}
       </CellMeasurer>
     );
   };
@@ -50,7 +76,7 @@ export default class Tree extends React.Component {
             deferredMeasurementCache={this._cache}
             ref={r => (this._list = r)}
             height={height}
-            rowCount={nodes.length}
+            rowCount={this.getRowCount()}
             rowHeight={this._cache.rowHeight}
             rowRenderer={this.measureRowRenderer(nodes)}
             width={width || autoWidth}