1 files changed, 35 insertions, 29 deletions

diff --git a/src/fields/Doc.ts b/src/fields/Doc.ts
index c10751698..c6cabe269 100644
--- a/src/fields/Doc.ts
+++ b/src/fields/Doc.ts
@@ -94,6 +94,8 @@ export function DocListCastOrNull(field: FieldResult) {
 
 export const WidthSym = Symbol('Width');
 export const HeightSym = Symbol('Height');
+export const AnimationSym = Symbol('Animation');
+export const HighlightSym = Symbol('Highlight');
 export const DataSym = Symbol('Data');
 export const LayoutSym = Symbol('Layout');
 export const FieldsSym = Symbol('Fields');
@@ -329,6 +331,8 @@ export class Doc extends RefField {
     @observable private ___fieldKeys: any = {};
     @observable public [AclSym]: { [key: string]: symbol } = {};
     @observable public [DirectLinksSym]: Set<Doc> = new Set();
+    @observable public [AnimationSym]: Opt<Doc>;
+    @observable public [HighlightSym]: boolean = false;
 
     private [UpdatingFromServer]: boolean = false;
     private [ForceServerWrite]: boolean = false;
@@ -1261,53 +1265,55 @@ export namespace Doc {
     }
 
     export function linkFollowUnhighlight() {
-        Doc.UnhighlightAll();
+        UnhighlightWatchers.length = 0;
+        highlightedDocs.forEach(doc => Doc.UnHighlightDoc(doc));
         document.removeEventListener('pointerdown', linkFollowUnhighlight);
-        runInAction(() => (HighlightBrush.linkFollowEffect = undefined));
     }
 
-    let _lastDate = 0;
+    let UnhighlightWatchers: (() => void)[] = [];
+    let UnhighlightTimer: any;
+    export function AddUnlightWatcher(watcher: () => void) {
+        if (UnhighlightTimer) {
+            UnhighlightWatchers.push(watcher);
+        } else watcher();
+    }
     export function linkFollowHighlight(destDoc: Doc | Doc[], dataAndDisplayDocs = true, presEffect?: Doc) {
-        //linkFollowUnhighlight();
-        // runInAction(() => presEffect && (HighlightBrush.linkFollowEffect = undefined));
-        // setTimeout(() => runInAction(() => presEffect && (HighlightBrush.linkFollowEffect = presEffect)));
-        runInAction(() => presEffect && (HighlightBrush.linkFollowEffect = presEffect));
-        (destDoc instanceof Doc ? [destDoc] : destDoc).forEach(doc => Doc.HighlightDoc(doc, dataAndDisplayDocs));
+        linkFollowUnhighlight();
+        (destDoc instanceof Doc ? [destDoc] : destDoc).forEach(doc => Doc.HighlightDoc(doc, dataAndDisplayDocs, presEffect));
         document.removeEventListener('pointerdown', linkFollowUnhighlight);
         document.addEventListener('pointerdown', linkFollowUnhighlight);
-        const lastDate = (_lastDate = Date.now());
-        window.setTimeout(() => _lastDate === lastDate && linkFollowUnhighlight(), 5000);
+        if (UnhighlightTimer) clearTimeout(UnhighlightTimer);
+        UnhighlightTimer = window.setTimeout(() => {
+            UnhighlightWatchers.forEach(watcher => watcher());
+            linkFollowUnhighlight();
+            UnhighlightTimer = 0;
+        }, 5000);
     }
 
-    export class HighlightBrush {
-        @observable HighlightedDoc: Map<Doc, boolean> = new Map();
-        @observable static linkFollowEffect: Doc | undefined;
-    }
-    const highlightManager = new HighlightBrush();
+    var highlightedDocs = new Set<Doc>();
     export function IsHighlighted(doc: Doc) {
         if (!doc || GetEffectiveAcl(doc) === AclPrivate || GetEffectiveAcl(Doc.GetProto(doc)) === AclPrivate || doc.opacity === 0) return false;
-        return highlightManager.HighlightedDoc.get(doc) || highlightManager.HighlightedDoc.get(Doc.GetProto(doc));
+        return doc[HighlightSym] || Doc.GetProto(doc)[HighlightSym];
     }
-    export function HighlightDoc(doc: Doc, dataAndDisplayDocs = true) {
+    export function HighlightDoc(doc: Doc, dataAndDisplayDocs = true, presEffect?: Doc) {
         runInAction(() => {
-            highlightManager.HighlightedDoc.set(doc, true);
-            dataAndDisplayDocs && highlightManager.HighlightedDoc.set(Doc.GetProto(doc), true);
+            doc[AnimationSym] = presEffect;
+            highlightedDocs.add(doc);
+            doc[HighlightSym] = true;
+            if (dataAndDisplayDocs) {
+                highlightedDocs.add(Doc.GetProto(doc));
+                Doc.GetProto(doc)[HighlightSym] = true;
+            }
         });
     }
     export function UnHighlightDoc(doc: Doc) {
         runInAction(() => {
-            highlightManager.HighlightedDoc.set(doc, false);
-            highlightManager.HighlightedDoc.set(Doc.GetProto(doc), false);
+            highlightedDocs.delete(doc);
+            highlightedDocs.delete(Doc.GetProto(doc));
+            doc[HighlightSym] = Doc.GetProto(doc)[HighlightSym] = false;
+            doc[AnimationSym] = undefined;
         });
     }
-    export function UnhighlightAll() {
-        const mapEntries = highlightManager.HighlightedDoc.keys();
-        let docEntry: IteratorResult<Doc>;
-        while (!(docEntry = mapEntries.next()).done) {
-            const targetDoc = docEntry.value;
-            targetDoc && Doc.UnHighlightDoc(targetDoc);
-        }
-    }
     export function UnBrushAllDocs() {
         brushManager.BrushedDoc.clear();
     }