diff --git a/src/main.ts b/src/main.ts
index 9bc0ba7..89682b0 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -97,7 +97,17 @@ export default class CrossbowPlugin extends Plugin {
   }
 
   public runWithoutCacheUpdate(fileHasChanged: boolean): void {
-    const targetEditor = this.app.workspace.activeEditor?.editor;
+    // Get editor of current file
+    const fileView = app.workspace
+      .getLeavesOfType('markdown')
+      .find((leaf) => leaf.view instanceof MarkdownView && leaf.view.file === this.currentFile)?.view as
+      | MarkdownView
+      | undefined;
+
+    if (!fileView) return;
+
+    const targetEditor = fileView.editor;
+
     if (!targetEditor) return;
 
     const wordLookup = this.tokenizationService.getWordLookupFromEditor(targetEditor);
diff --git a/src/model/suggestion.ts b/src/model/suggestion.ts
index 16f2655..80fb06a 100644
--- a/src/model/suggestion.ts
+++ b/src/model/suggestion.ts
@@ -56,8 +56,6 @@ export class Suggestion implements ITreeNodeData {
   public constructor(word: string, matchSet: Set<CacheMatch>, matchOccurrences: EditorPosition[]) {
     this.word = word;
     this.children = matchOccurrences.map((p) => {
-      console.log(word);
-
       const matchOccurrenceEnd = { ch: p.ch + word.length, line: p.line } as EditorPosition;
       const matches = Array.from(matchSet);
       return new Occurrence(this, p, matchOccurrenceEnd, matches);
diff --git a/src/model/trie.ts b/src/model/trie.ts
new file mode 100644
index 0000000..fc94750
--- /dev/null
+++ b/src/model/trie.ts
@@ -0,0 +1,109 @@
+// Copyright (C) 2023 - shoedler - github.com/shoedler
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+export type Trie<T> = {
+  value?: T;
+  next: { [key: string]: Trie<T> };
+};
+
+const createTrie = <T>(items: T[], keySelector: (item: T) => string): Trie<T> => {
+  const root: Trie<T> = { next: {} };
+
+  for (let i = 0; i < items.length; i++) {
+    const item = items[i];
+    addOrUpdate(root, item, keySelector, (o, n) => {
+      throw new Error(`Duplicate entry: ${JSON.stringify(n)}`);
+    });
+  }
+
+  return root;
+};
+
+const addOrUpdate = <T>(
+  root: Trie<T>,
+  item: T,
+  keySelector: (item: T) => string,
+  update: (oldItem: T, newItem: T) => T
+): void => {
+  const key = keySelector(item).toLocaleLowerCase();
+
+  let trie = root;
+
+  for (let i = 0; i < key.length; i++) {
+    const c = key[i];
+
+    trie.next[c] = trie.next[c] || { next: {} };
+
+    if (i === key.length - 1) {
+      if (trie.next[c].value) {
+        trie.next[c].value = update(trie.next[c].value as T, item);
+      } else {
+        trie.next[c].value = item;
+      }
+
+      break;
+    }
+
+    trie = trie.next[c] as Trie<T>;
+  }
+};
+
+const tryGet = <T>(root: Trie<T>, key: string): T | undefined => {
+  let trie = root;
+  let i = 0;
+
+  while (i < key.length && trie) {
+    trie = trie.next[key[i++]] as Trie<T>;
+  }
+
+  return trie.value as T | undefined;
+};
+
+const toArray = <T>(root: Trie<T>): T[] => {
+  const result: T[] = [];
+
+  for (const value of traverse(root)) {
+    result.push(value);
+  }
+
+  return result;
+};
+
+const length = <T>(root: Trie<T>): number => {
+  let result = 0;
+
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  for (const _ of traverse(root)) {
+    result++;
+  }
+
+  return result;
+};
+
+function* traverse<T>(trie: Trie<T>): Generator<T> {
+  if (trie.value) {
+    yield trie.value;
+  }
+
+  for (const key in trie.next) {
+    yield* traverse(trie.next[key]);
+  }
+}
+
+export const TrieTools = {
+  createTrie,
+  addOrUpdate,
+  tryGet,
+  toArray,
+  length,
+  traverse,
+};
diff --git a/src/view/tree/tree.ts b/src/view/tree/tree.ts
index be0a804..2149fd2 100644
--- a/src/view/tree/tree.ts
+++ b/src/view/tree/tree.ts
@@ -41,8 +41,10 @@ export class Tree extends HTMLElement implements ITreeContextProvider {
 
   public update<T extends ITreeNodeData>(data: T[]): void {
     const oldNodes = this.getChildTreeNodes();
-    const newNodes = data.map((d) => new TreeNode<T>(d, this));
-    const batch = this.treeUpdater.update(newNodes, oldNodes);
+    const batch = this.treeUpdater.update(
+      data.map((d) => new TreeNode(d, this)),
+      oldNodes
+    );
 
     requestAnimationFrame(() => {
       batch.forEach((update) => update(this));
diff --git a/src/view/tree/treeNode.ts b/src/view/tree/treeNode.ts
index 3ee4846..85b47d0 100644
--- a/src/view/tree/treeNode.ts
+++ b/src/view/tree/treeNode.ts
@@ -117,17 +117,7 @@ export class TreeNode<TData extends ITreeNodeData> extends HTMLElement {
       this.mainWrapper.addEventListener('click', () => (this.isCollapsed() ? this.expand() : this.collapse()));
 
       // lazy-create children nodes
-      const children = this.value.children as ITreeNodeData[];
-      const childrenWrapper = this.childrenWrapper as HTMLDivElement;
-
-      this.mainWrapper.addEventListener(
-        'click',
-        () => {
-          const nodes = children.map((child) => new TreeNode(child, this.manager));
-          childrenWrapper.replaceChildren(...nodes);
-        },
-        { once: true }
-      );
+      this.mainWrapper.addEventListener('click', () => this.generateChildren(), { once: true });
     }
 
     if (this.value.actions.length > 0) {
@@ -155,6 +145,14 @@ export class TreeNode<TData extends ITreeNodeData> extends HTMLElement {
     };
   }
 
+  public generateChildren(): void {
+    if (!this.childrenWrapper || !this.value.children) return;
+    if (this.childrenWrapper.children.length > 0) return;
+
+    const nodes = this.value.children.map((child) => new TreeNode(child, this.manager));
+    this.childrenWrapper.replaceChildren(...nodes);
+  }
+
   public isCollapsed() {
     if (!this.childrenWrapper || !this.iconWrapper) return true;
 
diff --git a/src/view/tree/treeUpdater.ts b/src/view/tree/treeUpdater.ts
index d4e042b..fa5e4e2 100644
--- a/src/view/tree/treeUpdater.ts
+++ b/src/view/tree/treeUpdater.ts
@@ -37,21 +37,19 @@ export class TreeUpdater {
             const child = newNode.childTreeNodes.find((child) => equals(child.value, expandedChild.value));
             if (child) {
               child.expand();
+              child.generateChildren();
             }
           });
 
         // Replace existing node with new node
         updates.push((container) => {
-          console.log(
-            'replacing',
-            existingNode.value.constructor.name,
-            existingNode.value.uid,
-            'with',
-            newNode.value.constructor.name,
-            newNode.value.uid
-          );
           container.insertAfter(newNode, existingNode);
-          existingNode.isCollapsed() ? newNode.collapse() : newNode.expand();
+
+          if (!existingNode.isCollapsed()) {
+            newNode.expand();
+            newNode.generateChildren();
+          }
+
           existingNode.remove();
         });
       } else {
@@ -59,7 +57,6 @@ export class TreeUpdater {
         const insertionIndex = oldNodes.findIndex((oldNode) => newNode.value.uid.localeCompare(oldNode.value.uid) < 0);
 
         updates.push((container) => {
-          console.log('inserting', newNode.value.constructor.name, newNode.value.uid);
           if (insertionIndex === -1) {
             container.appendChild(newNode);
           } else {
@@ -72,7 +69,6 @@ export class TreeUpdater {
     // Now, we're left with the existing suggestions that we need to remove
     updates.push((container) => {
       oldNodes.forEach((item) => {
-        console.log('removing', item.value.constructor.name, item.value.uid);
         item.remove();
       });
     });