2 files changed, 44 insertions, 25 deletions

diff --git a/src/client/util/UndoManager.ts b/src/client/util/UndoManager.ts
index 34910bac3..8e9e11a11 100644
--- a/src/client/util/UndoManager.ts
+++ b/src/client/util/UndoManager.ts
@@ -1,36 +1,49 @@
 import { observable, action } from "mobx";
 import { Opt } from "../../fields/Field";
 
-export function undoBatch(target: any, key: string | symbol, descriptor?: TypedPropertyDescriptor<any>): any {
-    let fn: (...args: any[]) => any;
-    let patchedFn: Opt<(...args: any[]) => any>;
-
-    if (descriptor) {
-        fn = descriptor.value;
-    }
-
-    return {
+function propertyDecorator(target: any, key: string | symbol) {
+    Object.defineProperty(target, key, {
         configurable: true,
         enumerable: false,
-        get() {
-            if (!patchedFn) {
-                patchedFn = (...args: any[]) => {
+        get: function () {
+            return 5;
+        },
+        set: function (value: any) {
+            Object.defineProperty(this, key, {
+                enumerable: false,
+                writable: true,
+                configurable: true,
+                value: function (...args: any[]) {
                     try {
-                        UndoManager.StartBatch()
-                        return fn.call(this, ...args)
+                        UndoManager.StartBatch();
+                        return value.apply(this, args);
                     } finally {
-                        UndoManager.EndBatch()
+                        UndoManager.EndBatch();
                     }
-                };
-            }
-            return patchedFn;
-        },
-        set(newFn: any) {
-            patchedFn = undefined;
-            fn = newFn;
+                }
+            })
+        }
+    })
+}
+export function undoBatch(target: any, key: string | symbol, descriptor: TypedPropertyDescriptor<any>): any {
+    if (!descriptor) {
+        propertyDecorator(target, key);
+        return;
+    }
+    const oldFunction = descriptor.value;
+
+    descriptor.value = function (...args: any[]) {
+        try {
+            UndoManager.StartBatch()
+            return oldFunction.apply(this, args)
+        } finally {
+            UndoManager.EndBatch()
         }
     }
+
+    return descriptor;
 }
+
 export namespace UndoManager {
     export interface UndoEvent {
         undo: () => void;
diff --git a/src/client/views/collections/CollectionFreeFormView.tsx b/src/client/views/collections/CollectionFreeFormView.tsx
index 3a66ebb75..07e9c0899 100644
--- a/src/client/views/collections/CollectionFreeFormView.tsx
+++ b/src/client/views/collections/CollectionFreeFormView.tsx
@@ -4,7 +4,7 @@ import { action, computed } from "mobx";
 import { CollectionFreeFormDocumentView } from "../nodes/CollectionFreeFormDocumentView";
 import { DragManager } from "../../util/DragManager";
 import "./CollectionFreeFormView.scss";
-import { CollectionViewBase, COLLECTION_BORDER_WIDTH } from "./CollectionViewBase";
+import { CollectionViewBase, COLLECTION_BORDER_WIDTH, CollectionViewProps } from "./CollectionViewBase";
 import { KeyStore } from "../../../fields/KeyStore";
 import { Document } from "../../../fields/Document";
 import { ListField } from "../../../fields/ListField";
@@ -38,9 +38,15 @@ export class CollectionFreeFormView extends CollectionViewBase {
     @computed
     get resizeScaling() { return this.isAnnotationOverlay ? this.props.Document.GetNumber(KeyStore.Width, 0) / this.nativeWidth : 1; }
 
-    @undoBatch
+    constructor(props: CollectionViewProps) {
+        super(props);
+
+        this.drop = this.drop.bind(this);
+    }
+
     @action
-    drop = (e: Event, de: DragManager.DropEvent) => {
+    @undoBatch
+    drop(e: Event, de: DragManager.DropEvent) {
         const doc: DocumentView = de.data["document"];
         if (doc.props.ContainingCollectionView && doc.props.ContainingCollectionView !== this) {
             doc.props.ContainingCollectionView.removeDocument(doc.props.Document);