diff --git a/src/components/TreeSelect/__stories__/components/WithDndListExample.tsx b/src/components/TreeSelect/__stories__/components/WithDndListExample.tsx
index 51da701515..c9c6253581 100644
--- a/src/components/TreeSelect/__stories__/components/WithDndListExample.tsx
+++ b/src/components/TreeSelect/__stories__/components/WithDndListExample.tsx
@@ -18,7 +18,7 @@ import {reorderArray} from '../../../useList/__stories__/utils/reorderArray';
 import {TreeSelect} from '../../TreeSelect';
 import {TreeSelectItem} from '../../TreeSelectItem';
 import type {TreeSelectItemProps} from '../../TreeSelectItem';
-import type {TreeSelectProps} from '../../types';
+import type {TreeSelectProps, TreeSelectRenderContainer, TreeSelectRenderItem} from '../../types';
 
 const DraggableListItem = ({
     provided,
@@ -52,11 +52,87 @@ export const WithDndListExample = (storyProps: WithDndListExampleProps) => {
     const [items, setItems] = React.useState(randomItems);
     const [value, setValue] = React.useState<string[]>([]);
 
-    const handleDrugEnd: OnDragEndResponder = ({destination, source}) => {
-        if (typeof destination?.index === 'number' && destination.index !== source.index) {
-            setItems((items) => reorderArray(items, source.index, destination.index));
-        }
-    };
+    const renderContainer: TreeSelectRenderContainer<CustomDataType> = React.useCallback(
+        ({renderItem, visibleFlattenIds, containerRef, id}) => {
+            const handleDrugEnd: OnDragEndResponder = ({destination, source}) => {
+                if (typeof destination?.index === 'number' && destination.index !== source.index) {
+                    setItems((currentItems) =>
+                        reorderArray(currentItems, source.index, destination.index),
+                    );
+                }
+            };
+
+            return (
+                <DragDropContext onDragEnd={handleDrugEnd}>
+                    <Droppable
+                        droppableId="droppable"
+                        renderClone={(
+                            provided: DraggableProvided,
+                            snapshot: DraggableStateSnapshot,
+                            rubric: DraggableRubric,
+                        ) => {
+                            return renderItem(
+                                visibleFlattenIds[rubric.source.index],
+                                rubric.source.index,
+                                {
+                                    provided,
+                                    active: snapshot.isDragging,
+                                },
+                            );
+                        }}
+                    >
+                        {(droppableProvided: DroppableProvided) => (
+                            <ListContainerView ref={containerRef} id={id}>
+                                <div
+                                    {...droppableProvided.droppableProps}
+                                    ref={droppableProvided.innerRef}
+                                >
+                                    {visibleFlattenIds.map((listItemId, index) =>
+                                        renderItem(listItemId, index),
+                                    )}
+                                    {droppableProvided.placeholder}
+                                </div>
+                            </ListContainerView>
+                        )}
+                    </Droppable>
+                </DragDropContext>
+            );
+        },
+        [setItems],
+    );
+
+    const renderItem: TreeSelectRenderItem<CustomDataType> = React.useCallback(
+        ({data, props, index, renderContext: renderContextProps}) => {
+            const commonProps = {
+                ...props,
+                title: data.someRandomKey,
+                endSlot: <Icon data={Grip} size={16} />,
+            };
+
+            // here passed props from `renderContainer` method.
+            if (renderContextProps) {
+                return (
+                    <DraggableListItem
+                        key={`item-key-${index}`}
+                        {...commonProps}
+                        {...renderContextProps}
+                    />
+                );
+            }
+            return (
+                <Draggable draggableId={String(index)} index={index} key={`item-key-${index}`}>
+                    {(provided: DraggableProvided, snapshot: DraggableStateSnapshot) => (
+                        <DraggableListItem
+                            provided={provided}
+                            {...commonProps}
+                            active={snapshot.isDragging}
+                        />
+                    )}
+                </Draggable>
+            );
+        },
+        [],
+    );
 
     return (
         <Flex>
@@ -74,76 +150,8 @@ export const WithDndListExample = (storyProps: WithDndListExampleProps) => {
                         setValue([id]);
                     }
                 }}
-                renderContainer={({renderItem, visibleFlattenIds, containerRef, id}) => {
-                    return (
-                        <DragDropContext onDragEnd={handleDrugEnd}>
-                            <Droppable
-                                droppableId="droppable"
-                                renderClone={(
-                                    provided: DraggableProvided,
-                                    snapshot: DraggableStateSnapshot,
-                                    rubric: DraggableRubric,
-                                ) => {
-                                    return renderItem(
-                                        visibleFlattenIds[rubric.source.index],
-                                        rubric.source.index,
-                                        {
-                                            provided,
-                                            active: snapshot.isDragging,
-                                        },
-                                    );
-                                }}
-                            >
-                                {(droppableProvided: DroppableProvided) => (
-                                    <ListContainerView ref={containerRef} id={id}>
-                                        <div
-                                            {...droppableProvided.droppableProps}
-                                            ref={droppableProvided.innerRef}
-                                        >
-                                            {visibleFlattenIds.map((id, index) =>
-                                                renderItem(id, index),
-                                            )}
-                                            {droppableProvided.placeholder}
-                                        </div>
-                                    </ListContainerView>
-                                )}
-                            </Droppable>
-                        </DragDropContext>
-                    );
-                }}
-                renderItem={({data, props, index, renderContext: renderContextProps}) => {
-                    const commonProps = {
-                        ...props,
-                        title: data.someRandomKey,
-                        endSlot: <Icon data={Grip} size={16} />,
-                    };
-
-                    // here passed props from `renderContainer` method.
-                    if (renderContextProps) {
-                        return (
-                            <DraggableListItem
-                                key={`item-key-${index}`}
-                                {...commonProps}
-                                {...renderContextProps}
-                            />
-                        );
-                    }
-                    return (
-                        <Draggable
-                            draggableId={String(index)}
-                            index={index}
-                            key={`item-key-${index}`}
-                        >
-                            {(provided: DraggableProvided, snapshot: DraggableStateSnapshot) => (
-                                <DraggableListItem
-                                    provided={provided}
-                                    {...commonProps}
-                                    active={snapshot.isDragging}
-                                />
-                            )}
-                        </Draggable>
-                    );
-                }}
+                renderContainer={renderContainer}
+                renderItem={renderItem}
             />
         </Flex>
     );