Adding Arrow button to revert a single block changes in cell

src/diff/cell.ts

@@ -1,5 +1,5 @@
 import { python } from '@codemirror/lang-python';
-import { MergeView } from '@codemirror/merge';
+import { MergeView, getChunks } from '@codemirror/merge';
 import { EditorView } from '@codemirror/view';
 import { jupyterTheme } from '@jupyterlab/codemirror';
 import { Message } from '@lumino/messaging';
@@ -60,29 +60,107 @@ class CodeMirrorSplitDiffWidget extends BaseDiffWidget {
         extensions: [
           basicSetup,
           python(),
-          EditorView.editable.of(false),
-          jupyterTheme
+          EditorView.editable.of(true),
+          jupyterTheme,
+          EditorView.updateListener.of(update => {
+            if (update.docChanged) {
+              const newText = update.state.doc.toString();
+
+              this._modifiedCode = newText;
+              this._newSource = newText;
+
+              this._renderArrowButtons();
+            }
+          })
         ]
       },
-      parent: this.node
+      parent: this.node,
+      gutter: true,
+      highlightChanges: true
+    });
+
+    const container = this._splitView.dom;
+    const overlay = document.createElement('div');
+    overlay.className = 'jp-DiffArrowOverlay';
+    container.appendChild(overlay);
+    this._arrowOverlay = overlay;
+
+    this._addScrollSync();
+    setTimeout(() => this._renderArrowButtons(), 50);
+  }
+
+  /**
+   * Render "merge change" buttons in the diff on left editor.
+   */
+  private _renderArrowButtons(): void {
+    if (!this._splitView) {
+      return;
+    }
+
+    const paneA = this._splitView.a;
+    const paneB = this._splitView.b;
+    const result = getChunks(paneB.state);
+    const chunks = result?.chunks ?? [];
+
+    this._arrowOverlay.innerHTML = '';
+
+    chunks.forEach(chunk => {
+      const { fromA, toA, fromB, toB } = chunk;
+      const lineBlockA = paneA.lineBlockAt(fromA);
+      const lineBlockB = paneB.lineBlockAt(fromB);
+      const midTop = (lineBlockA.top + lineBlockB.top) / 2;
+
+      const connector = document.createElement('div');
+      connector.className = 'jp-DiffConnectorLine';
+      connector.style.top = `${midTop}px`;
+
+      const arrowBtn = document.createElement('button');
+      arrowBtn.textContent = '🡪';
+      arrowBtn.className = 'jp-DiffArrow';
+      arrowBtn.title = 'Revert Block';
+
+      arrowBtn.onclick = () => {
+        const origText = paneA.state.doc.sliceString(fromA, toA);
+        paneB.dispatch({
+          changes: { from: fromB, to: toB, insert: origText }
+        });
+        this._renderArrowButtons();
+      };
+
+      connector.appendChild(arrowBtn);
+      this._arrowOverlay.appendChild(connector);
     });
   }
 
   /**
-   * Destroy the split view and clean up resources.
+   * Keep arrow overlay in sync with editor scroll.
    */
+  private _addScrollSync(): void {
+    const paneA = this._splitView.a;
+    const paneB = this._splitView.b;
+    const sync = () => this._renderArrowButtons();
+    paneA.scrollDOM.addEventListener('scroll', sync);
+    paneB.scrollDOM.addEventListener('scroll', sync);
+  }
+
   private _destroySplitView(): void {
     if (this._splitView) {
       this._splitView.destroy();
-      this._splitView = null;
+      this._splitView = null!;
+    }
+    if (this._arrowOverlay) {
+      this._arrowOverlay.remove();
     }
   }
 
   private _originalCode: string;
   private _modifiedCode: string;
-  private _splitView: MergeView | null = null;
+  private _arrowOverlay!: HTMLDivElement;
+  private _splitView!: MergeView & {
+    a: EditorView;
+    b: EditorView;
+  };
 }
-
 export async function createCodeMirrorSplitDiffWidget(
   options: IDiffWidgetOptions
 ): Promise<Widget> {

src/widget.ts

@@ -100,7 +100,7 @@ export abstract class BaseDiffWidget extends Widget {
    */
   public onAcceptClick(): void {
     if (this._cell) {
-      this._cell.sharedModel.setSource(this._newSource);
+      this._cell.sharedModel.setSource(this._newSource || this._originalSource);
       this._closeDiffView();
     }
   }
@@ -191,9 +191,9 @@ export abstract class BaseDiffWidget extends Widget {
   private _cell: ICellModel;
   private _cellFooterTracker: ICellFooterTracker;
   private _originalSource: string;
-  private _newSource: string;
   private _showActionButtons: boolean;
   private _openDiff: boolean;
   private _toggleButton: ToolbarButton | null = null;
   private _trans: TranslationBundle;
+  public _newSource: string;
 }

style/base.css

@@ -77,3 +77,35 @@
   background-color: var(--jp-layout-color3);
   border-color: var(--jp-border-color1);
 }
+
+.jp-DiffArrowOverlay {
+  position: absolute;
+  inset: 0 22px 0 0;
+  pointer-events: none;
+  z-index: 3;
+}
+
+.jp-DiffConnectorLine {
+  position: absolute;
+  width: 100%;
+  height: 30px;
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  pointer-events: none;
+  z-index: 4;
+}
+
+.jp-DiffArrow {
+  pointer-events: all;
+  background: var(--jp-layout-color1);
+  padding: 0 6px;
+  cursor: pointer;
+  border: none;
+  font-size: 12px;
+  transition: background-color 0.15s;
+}
+
+.jp-DiffArrow:hover {
+  background-color: var(--jp-layout-color3);
+}