7 files changed, 123 insertions, 43 deletions

diff --git a/app/assets/javascripts/ide/lib/common/model.js b/app/assets/javascripts/ide/lib/common/model.js
index d47b6704176..016dcda1fa1 100644
--- a/app/assets/javascripts/ide/lib/common/model.js
+++ b/app/assets/javascripts/ide/lib/common/model.js
@@ -32,7 +32,7 @@ export default class Model {
       );
     }
 
-    this.events = new Map();
+    this.events = new Set();
 
     this.updateContent = this.updateContent.bind(this);
     this.updateNewContent = this.updateNewContent.bind(this);
@@ -76,10 +76,11 @@ export default class Model {
   }
 
   onChange(cb) {
-    this.events.set(
-      this.path,
-      this.disposable.add(this.model.onDidChangeContent(e => cb(this, e))),
-    );
+    this.events.add(this.disposable.add(this.model.onDidChangeContent(e => cb(this, e))));
+  }
+
+  onDispose(cb) {
+    this.events.add(cb);
   }
 
   updateContent({ content, changed }) {
@@ -96,6 +97,11 @@ export default class Model {
 
   dispose() {
     this.disposable.dispose();
+
+    this.events.forEach(cb => {
+      if (typeof cb === 'function') cb();
+    });
+
     this.events.clear();
 
     eventHub.$off(`editor.update.model.dispose.${this.file.key}`, this.dispose);
diff --git a/app/assets/javascripts/ide/lib/decorations/controller.js b/app/assets/javascripts/ide/lib/decorations/controller.js
index 42904774747..13d477bb2cf 100644
--- a/app/assets/javascripts/ide/lib/decorations/controller.js
+++ b/app/assets/javascripts/ide/lib/decorations/controller.js
@@ -38,6 +38,15 @@ export default class DecorationsController {
     );
   }
 
+  hasDecorations(model) {
+    return this.decorations.has(model.url);
+  }
+
+  removeDecorations(model) {
+    this.decorations.delete(model.url);
+    this.editorDecorations.delete(model.url);
+  }
+
   dispose() {
     this.decorations.clear();
     this.editorDecorations.clear();
diff --git a/app/assets/javascripts/ide/lib/diff/controller.js b/app/assets/javascripts/ide/lib/diff/controller.js
index b136545ad11..f579424cf33 100644
--- a/app/assets/javascripts/ide/lib/diff/controller.js
+++ b/app/assets/javascripts/ide/lib/diff/controller.js
@@ -3,7 +3,7 @@ import { throttle } from 'underscore';
 import DirtyDiffWorker from './diff_worker';
 import Disposable from '../common/disposable';
 
-export const getDiffChangeType = (change) => {
+export const getDiffChangeType = change => {
   if (change.modified) {
     return 'modified';
   } else if (change.added) {
@@ -16,12 +16,7 @@ export const getDiffChangeType = (change) => {
 };
 
 export const getDecorator = change => ({
-  range: new monaco.Range(
-    change.lineNumber,
-    1,
-    change.endLineNumber,
-    1,
-  ),
+  range: new monaco.Range(change.lineNumber, 1, change.endLineNumber, 1),
   options: {
     isWholeLine: true,
     linesDecorationsClassName: `dirty-diff dirty-diff-${getDiffChangeType(change)}`,
@@ -31,6 +26,7 @@ export const getDecorator = change => ({
 export default class DirtyDiffController {
   constructor(modelManager, decorationsController) {
     this.disposable = new Disposable();
+    this.models = new Map();
     this.editorSimpleWorker = null;
     this.modelManager = modelManager;
     this.decorationsController = decorationsController;
@@ -42,7 +38,15 @@ export default class DirtyDiffController {
   }
 
   attachModel(model) {
+    if (this.models.has(model.url)) return;
+
     model.onChange(() => this.throttledComputeDiff(model));
+    model.onDispose(() => {
+      this.decorationsController.removeDecorations(model);
+      this.models.delete(model.url);
+    });
+
+    this.models.set(model.url, model);
   }
 
   computeDiff(model) {
@@ -54,7 +58,11 @@ export default class DirtyDiffController {
   }
 
   reDecorate(model) {
-    this.decorationsController.decorate(model);
+    if (this.decorationsController.hasDecorations(model)) {
+      this.decorationsController.decorate(model);
+    } else {
+      this.computeDiff(model);
+    }
   }
 
   decorate({ data }) {
@@ -65,6 +73,7 @@ export default class DirtyDiffController {
 
   dispose() {
     this.disposable.dispose();
+    this.models.clear();
 
     this.dirtyDiffWorker.removeEventListener('message', this.decorate);
     this.dirtyDiffWorker.terminate();
diff --git a/spec/javascripts/ide/components/repo_editor_spec.js b/spec/javascripts/ide/components/repo_editor_spec.js
index fc585647fbd..b06a6c62a1c 100644
--- a/spec/javascripts/ide/components/repo_editor_spec.js
+++ b/spec/javascripts/ide/components/repo_editor_spec.js
@@ -222,7 +222,7 @@ describe('RepoEditor', () => {
       vm.setupEditor();
 
       expect(vm.editor.onPositionChange).toHaveBeenCalled();
-      expect(vm.model.events.size).toBe(1);
+      expect(vm.model.events.size).toBe(2);
     });
 
     it('updates state when model content changed', done => {
diff --git a/spec/javascripts/ide/lib/common/model_spec.js b/spec/javascripts/ide/lib/common/model_spec.js
index 553a7a3746b..7a6c22b6d27 100644
--- a/spec/javascripts/ide/lib/common/model_spec.js
+++ b/spec/javascripts/ide/lib/common/model_spec.js
@@ -83,13 +83,6 @@ describe('Multi-file editor library model', () => {
   });
 
   describe('onChange', () => {
-    it('caches event by path', () => {
-      model.onChange(() => {});
-
-      expect(model.events.size).toBe(1);
-      expect(model.events.keys().next().value).toBe(model.file.key);
-    });
-
     it('calls callback on change', done => {
       const spy = jasmine.createSpy();
       model.onChange(spy);
@@ -132,5 +125,15 @@ describe('Multi-file editor library model', () => {
         jasmine.anything(),
       );
     });
+
+    it('calls onDispose callback', () => {
+      const disposeSpy = jasmine.createSpy();
+
+      model.onDispose(disposeSpy);
+
+      model.dispose();
+
+      expect(disposeSpy).toHaveBeenCalled();
+    });
   });
 });
diff --git a/spec/javascripts/ide/lib/decorations/controller_spec.js b/spec/javascripts/ide/lib/decorations/controller_spec.js
index aec325e26a9..e1c4ca570b6 100644
--- a/spec/javascripts/ide/lib/decorations/controller_spec.js
+++ b/spec/javascripts/ide/lib/decorations/controller_spec.js
@@ -117,4 +117,33 @@ describe('Multi-file editor library decorations controller', () => {
       expect(controller.editorDecorations.size).toBe(0);
     });
   });
+
+  describe('hasDecorations', () => {
+    it('returns true when decorations are cached', () => {
+      controller.addDecorations(model, 'key', [{ decoration: 'decorationValue' }]);
+
+      expect(controller.hasDecorations(model)).toBe(true);
+    });
+
+    it('returns false when no model decorations exist', () => {
+      expect(controller.hasDecorations(model)).toBe(false);
+    });
+  });
+
+  describe('removeDecorations', () => {
+    beforeEach(() => {
+      controller.addDecorations(model, 'key', [{ decoration: 'decorationValue' }]);
+      controller.decorate(model);
+    });
+
+    it('removes cached decorations', () => {
+      expect(controller.decorations.size).not.toBe(0);
+      expect(controller.editorDecorations.size).not.toBe(0);
+
+      controller.removeDecorations(model);
+
+      expect(controller.decorations.size).toBe(0);
+      expect(controller.editorDecorations.size).toBe(0);
+    });
+  });
 });
diff --git a/spec/javascripts/ide/lib/diff/controller_spec.js b/spec/javascripts/ide/lib/diff/controller_spec.js
index ff73240734e..fd8ab3b4f1d 100644
--- a/spec/javascripts/ide/lib/diff/controller_spec.js
+++ b/spec/javascripts/ide/lib/diff/controller_spec.js
@@ -3,10 +3,7 @@ import monacoLoader from '~/ide/monaco_loader';
 import editor from '~/ide/lib/editor';
 import ModelManager from '~/ide/lib/common/model_manager';
 import DecorationsController from '~/ide/lib/decorations/controller';
-import DirtyDiffController, {
-  getDiffChangeType,
-  getDecorator,
-} from '~/ide/lib/diff/controller';
+import DirtyDiffController, { getDiffChangeType, getDecorator } from '~/ide/lib/diff/controller';
 import { computeDiff } from '~/ide/lib/diff/diff';
 import { file } from '../../helpers';
 
@@ -90,6 +87,14 @@ describe('Multi-file editor library dirty diff controller', () => {
       expect(model.onChange).toHaveBeenCalled();
     });
 
+    it('adds dispose event callback', () => {
+      spyOn(model, 'onDispose');
+
+      controller.attachModel(model);
+
+      expect(model.onDispose).toHaveBeenCalled();
+    });
+
     it('calls throttledComputeDiff on change', () => {
       spyOn(controller, 'throttledComputeDiff');
 
@@ -99,6 +104,12 @@ describe('Multi-file editor library dirty diff controller', () => {
 
       expect(controller.throttledComputeDiff).toHaveBeenCalled();
     });
+
+    it('caches model', () => {
+      controller.attachModel(model);
+
+      expect(controller.models.has(model.url)).toBe(true);
+    });
   });
 
   describe('computeDiff', () => {
@@ -116,14 +127,22 @@ describe('Multi-file editor library dirty diff controller', () => {
   });
 
   describe('reDecorate', () => {
-    it('calls decorations controller decorate', () => {
+    it('calls computeDiff when no decorations are cached', () => {
+      spyOn(controller, 'computeDiff');
+
+      controller.reDecorate(model);
+
+      expect(controller.computeDiff).toHaveBeenCalledWith(model);
+    });
+
+    it('calls decorate when decorations are cached', () => {
       spyOn(controller.decorationsController, 'decorate');
 
+      controller.decorationsController.decorations.set(model.url, 'test');
+
       controller.reDecorate(model);
 
-      expect(controller.decorationsController.decorate).toHaveBeenCalledWith(
-        model,
-      );
+      expect(controller.decorationsController.decorate).toHaveBeenCalledWith(model);
     });
   });
 
@@ -133,16 +152,15 @@ describe('Multi-file editor library dirty diff controller', () => {
 
       controller.decorate({ data: { changes: [], path: model.path } });
 
-      expect(
-        controller.decorationsController.addDecorations,
-      ).toHaveBeenCalledWith(model, 'dirtyDiff', jasmine.anything());
+      expect(controller.decorationsController.addDecorations).toHaveBeenCalledWith(
+        model,
+        'dirtyDiff',
+        jasmine.anything(),
+      );
     });
 
     it('adds decorations into editor', () => {
-      const spy = spyOn(
-        controller.decorationsController.editor.instance,
-        'deltaDecorations',
-      );
+      const spy = spyOn(controller.decorationsController.editor.instance, 'deltaDecorations');
 
       controller.decorate({
         data: { changes: computeDiff('123', '1234'), path: model.path },
@@ -181,16 +199,22 @@ describe('Multi-file editor library dirty diff controller', () => {
     });
 
     it('removes worker event listener', () => {
-      spyOn(
-        controller.dirtyDiffWorker,
-        'removeEventListener',
-      ).and.callThrough();
+      spyOn(controller.dirtyDiffWorker, 'removeEventListener').and.callThrough();
 
       controller.dispose();
 
-      expect(
-        controller.dirtyDiffWorker.removeEventListener,
-      ).toHaveBeenCalledWith('message', jasmine.anything());
+      expect(controller.dirtyDiffWorker.removeEventListener).toHaveBeenCalledWith(
+        'message',
+        jasmine.anything(),
+      );
+    });
+
+    it('clears cached models', () => {
+      controller.attachModel(model);
+
+      model.dispose();
+
+      expect(controller.models.size).toBe(0);
     });
   });
 });