diff --git a/src/main/java/org/aya/intellij/actions/lsp/AyaLsp.java b/src/main/java/org/aya/intellij/actions/lsp/AyaLsp.java
index 74a0bc7..4799c27 100644
--- a/src/main/java/org/aya/intellij/actions/lsp/AyaLsp.java
+++ b/src/main/java/org/aya/intellij/actions/lsp/AyaLsp.java
@@ -66,7 +66,7 @@ public final class AyaLsp extends InMemoryCompilerAdvisor implements AyaLanguage
   private static final @NotNull Logger LOG = Logger.getInstance(AyaLsp.class);
   private final @NotNull AyaLanguageServer server;
   private final @NotNull Project project;
-  private final @NotNull MutableSet<VirtualFile> libraryPathCache = MutableSet.create();
+  private final @NotNull MutableSet<VirtualFile> librarySrcPathCache = MutableSet.create();
   private final @NotNull MutableMap<Path, ImmutableSeq<Problem>> problemCache = MutableMap.create();
   private final @NotNull ExecutorService compilerPool = Executors.newFixedThreadPool(1);
 
@@ -216,14 +216,17 @@ void recompile(@NotNull Runnable compile, @Nullable Runnable callback) {
 
   public void registerLibrary(@NotNull VirtualFile library) {
     if (JB.fileSupported(library)) {
-      libraryPathCache.add(library);
-      server.registerLibrary(JB.canonicalize(library));
+      var root = JB.canonicalize(library);
+      server.registerLibrary(root).forEach(registeredLibrary ->
+        registeredLibrary.modulePath()
+          .mapNotNull(path -> library.findFileByRelativePath(root.relativize(path).toString()))
+          .forEach(librarySrcPathCache::add));
     }
   }
 
   public boolean isInLibrary(@Nullable VirtualFile file) {
     while (file != null && file.isValid() && JB.fileSupported(file)) {
-      if (libraryPathCache.contains(file)) return true;
+      if (librarySrcPathCache.contains(file)) return true;
       file = file.getParent();
     }
     return false;
diff --git a/src/main/java/org/aya/intellij/ui/AyaTreeView.java b/src/main/java/org/aya/intellij/ui/AyaTreeView.java
index 1cf26f3..715092b 100644
--- a/src/main/java/org/aya/intellij/ui/AyaTreeView.java
+++ b/src/main/java/org/aya/intellij/ui/AyaTreeView.java
@@ -28,7 +28,6 @@
 import org.aya.intellij.psi.utils.AyaPsiUtils;
 import org.aya.intellij.service.AyaSettingService;
 import org.aya.intellij.ui.toolwindow.AyaToolWindow;
-import org.aya.util.TreeBuilder;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
diff --git a/src/main/java/org/aya/intellij/ui/TreeBuilder.java b/src/main/java/org/aya/intellij/ui/TreeBuilder.java
new file mode 100644
index 0000000..873eb72
--- /dev/null
+++ b/src/main/java/org/aya/intellij/ui/TreeBuilder.java
@@ -0,0 +1,46 @@
+// Copyright (c) 2020-2023 Tesla (Yinsen) Zhang.
+// Use of this source code is governed by the MIT license that can be found in the LICENSE.md file.
+package org.aya.intellij.ui;
+
+import kala.collection.mutable.MutableList;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.ArrayDeque;
+import java.util.Deque;
+import java.util.Objects;
+
+/** @param <T> the tree, no longer used in Aya. */
+public abstract class TreeBuilder<T extends TreeBuilder.Tree<T>> {
+  public interface Tree<T extends Tree<T>> {
+    @NotNull MutableList<T> children();
+  }
+
+  protected final Deque<@NotNull MutableList<@NotNull T>> tops = new ArrayDeque<>();
+
+  public @NotNull MutableList<@NotNull T> root() {
+    return tops.getFirst();
+  }
+
+  {
+    tops.addLast(MutableList.create());
+  }
+
+  public void append(@NotNull T trace) {
+    shift(trace);
+    reduce();
+  }
+
+  public void shift(@NotNull T trace) {
+    Objects.requireNonNull(tops.getLast()).append(trace);
+    tops.addLast(trace.children());
+  }
+
+  public void unshift() {
+    var buffer = Objects.requireNonNull(tops.getLast());
+    buffer.removeAt(buffer.size() - 1);
+  }
+  public void reduce() { tops.removeLast(); }
+  public void reduceAndUnshift() {
+    tops.removeLast();
+  }
+}