From 8b60cd7abdf2f055c1e59bafc3cf375baffab1f6 Mon Sep 17 00:00:00 2001
From: Vincent Chan <okcdz@diverse.space>
Date: Thu, 4 Aug 2022 11:22:38 +0800
Subject: [PATCH] refactor: use node iterator for selection

---
 .../lib/document/node_iterator.dart           | 10 ++++++++
 .../lib/service/selection_service.dart        | 23 ++++++++-----------
 2 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/frontend/app_flowy/packages/flowy_editor/lib/document/node_iterator.dart b/frontend/app_flowy/packages/flowy_editor/lib/document/node_iterator.dart
index 8603c043e4..1f321e937a 100644
--- a/frontend/app_flowy/packages/flowy_editor/lib/document/node_iterator.dart
+++ b/frontend/app_flowy/packages/flowy_editor/lib/document/node_iterator.dart
@@ -61,4 +61,14 @@ class NodeIterator implements Iterator<Node> {
   Node get current {
     return _currentNode!;
   }
+
+  List<Node> toList() {
+    final result = <Node>[];
+
+    while (moveNext()) {
+      result.add(current);
+    }
+
+    return result;
+  }
 }
diff --git a/frontend/app_flowy/packages/flowy_editor/lib/service/selection_service.dart b/frontend/app_flowy/packages/flowy_editor/lib/service/selection_service.dart
index a304272846..55b08f9279 100644
--- a/frontend/app_flowy/packages/flowy_editor/lib/service/selection_service.dart
+++ b/frontend/app_flowy/packages/flowy_editor/lib/service/selection_service.dart
@@ -1,5 +1,7 @@
 import 'dart:async';
 
+import 'package:flowy_editor/document/node_iterator.dart';
+import 'package:flowy_editor/document/state_tree.dart';
 import 'package:flutter/gestures.dart';
 import 'package:flutter/material.dart';
 
@@ -212,7 +214,7 @@ class _FlowySelectionState extends State<FlowySelection>
 
   @override
   List<Node> getNodesInSelection(Selection selection) =>
-      _selectedNodesInSelection(editorState.document.root, selection);
+      _selectedNodesInSelection(editorState.document, selection);
 
   @override
   void initState() {
@@ -462,8 +464,7 @@ class _FlowySelectionState extends State<FlowySelection>
   }
 
   void _updateSelection(Selection selection) {
-    final nodes =
-        _selectedNodesInSelection(editorState.document.root, selection);
+    final nodes = _selectedNodesInSelection(editorState.document, selection);
 
     currentSelection = selection;
     currentSelectedNodes.value = nodes;
@@ -572,16 +573,10 @@ class _FlowySelectionState extends State<FlowySelection>
     currentState?.show();
   }
 
-  List<Node> _selectedNodesInSelection(Node node, Selection selection) {
-    List<Node> result = [];
-    if (node.parent != null) {
-      if (node.inSelection(selection)) {
-        result.add(node);
-      }
-    }
-    for (final child in node.children) {
-      result.addAll(_selectedNodesInSelection(child, selection));
-    }
-    return result;
+  List<Node> _selectedNodesInSelection(
+      StateTree stateTree, Selection selection) {
+    final startNode = stateTree.nodeAtPath(selection.start.path)!;
+    final endNode = stateTree.nodeAtPath(selection.end.path)!;
+    return NodeIterator(stateTree, startNode, endNode).toList();
   }
 }