From 2c7c68ce9a03fa2e894b9baa27df76757c1603f0 Mon Sep 17 00:00:00 2001
From: Andrew Patton <andrew@acusti.ca>
Date: Mon, 10 Oct 2016 00:18:55 -0700
Subject: [PATCH] Adapt restoreSelection to work for all activeElements

---
 .../dom/shared/ReactInputSelection.js         | 114 +++++++++++++-----
 .../__tests__/ReactInputSelection-test.js     |   5 +-
 2 files changed, 86 insertions(+), 33 deletions(-)

diff --git a/src/renderers/dom/shared/ReactInputSelection.js b/src/renderers/dom/shared/ReactInputSelection.js
index 41bc0d11a0bc40..a934c65a13a22f 100644
--- a/src/renderers/dom/shared/ReactInputSelection.js
+++ b/src/renderers/dom/shared/ReactInputSelection.js
@@ -19,7 +19,7 @@ var focusNode = require('fbjs/lib/focusNode');
 var getActiveElement = require('fbjs/lib/getActiveElement');
 
 function isInDocument(node) {
-  return containsNode(document.documentElement, node);
+  return node.ownerDocument && containsNode(node.ownerDocument.documentElement, node);
 }
 
 function getFocusedElement() {
@@ -36,6 +36,62 @@ function getFocusedElement() {
   return focusedElem;
 }
 
+function getElementsWithSelections(acc, win) {
+  acc = acc || [];
+  win = win || window;
+  var doc;
+  try {
+    doc = win.document;
+  } catch (e) {
+    return acc;
+  }
+  var element = null;
+  if (win.getSelection) {
+    var selection = win.getSelection();
+    var startNode = selection.anchorNode;
+    var endNode = selection.focusNode;
+    var startOffset = selection.anchorOffset;
+    var endOffset = selection.focusOffset;
+    if (startNode && startNode.childNodes.length) {
+      if (startNode.childNodes[startOffset] === endNode.childNodes[endOffset]) {
+        element = startNode.childNodes[startOffset];
+      }
+    } else {
+      element = startNode;
+    }
+  } else if (doc.selection) {
+    var range = doc.selection.createRange();
+    element = range.parentElement();
+  }
+  if (ReactInputSelection.hasSelectionCapabilities(element)) {
+    acc = acc.concat(element);
+  }
+  return Array.prototype.reduce.call(win.frames, getElementsWithSelections, acc);
+}
+
+function focusNodePreservingScroll(element) {
+  // Focusing a node can change the scroll position, which is undesirable
+  const ancestors = [];
+  let ancestor = element;
+  while ((ancestor = ancestor.parentNode)) {
+    if (ancestor.nodeType === ELEMENT_NODE) {
+      ancestors.push({
+        element: ancestor,
+        left: ancestor.scrollLeft,
+        top: ancestor.scrollTop,
+      });
+    }
+  }
+
+  focusNode(element);
+
+  for (let i = 0; i < ancestors.length; i++) {
+    const info = ancestors[i];
+    info.element.scrollLeft = info.left;
+    info.element.scrollTop = info.top;
+  }
+}
+
 /**
  * @ReactInputSelection: React input selection module. Based on Selection.js,
  * but modified to be suitable for react and has a couple of bug fixes (doesn't
@@ -54,12 +110,17 @@ var ReactInputSelection = {
   },
 
   getSelectionInformation: function() {
-    var focusedElem = getFocusedElement();
+    var focusedElement = getFocusedElement();
     return {
-      focusedElem: focusedElem,
-      selectionRange: ReactInputSelection.hasSelectionCapabilities(focusedElem)
-        ? ReactInputSelection.getSelection(focusedElem)
-        : null,
+      focusedElement: focusedElement,
+      activeElements: getElementsWithSelections().map(function(element) {
+        return {
+          element: element,
+          selectionRange: ReactInputSelection.hasSelectionCapabilities(element)
+            ? ReactInputSelection.getSelection(element)
+            : null,
+        };
+      }),
     };
   },
 
@@ -69,34 +130,25 @@ var ReactInputSelection = {
    * nodes and place them back in, resulting in focus being lost.
    */
   restoreSelection: function(priorSelectionInformation) {
-    var curFocusedElem = getFocusedElement();
-    var priorFocusedElem = priorSelectionInformation.focusedElem;
-    var priorSelectionRange = priorSelectionInformation.selectionRange;
-    if (curFocusedElem !== priorFocusedElem && isInDocument(priorFocusedElem)) {
-      if (ReactInputSelection.hasSelectionCapabilities(priorFocusedElem)) {
-        ReactInputSelection.setSelection(priorFocusedElem, priorSelectionRange);
-      }
-
-      // Focusing a node can change the scroll position, which is undesirable
-      const ancestors = [];
-      let ancestor = priorFocusedElem;
-      while ((ancestor = ancestor.parentNode)) {
-        if (ancestor.nodeType === ELEMENT_NODE) {
-          ancestors.push({
-            element: ancestor,
-            left: ancestor.scrollLeft,
-            top: ancestor.scrollTop,
-          });
+    priorSelectionInformation.activeElements.forEach(function(activeElement) {
+      var element = activeElement.element;
+      if (isInDocument(element) &&
+          getActiveElement(element.ownerDocument) !== element) {
+        if (ReactInputSelection.hasSelectionCapabilities(element)) {
+          ReactInputSelection.setSelection(
+            element,
+            activeElement.selectionRange
+          );
+          focusNodePreservingScroll(element);
         }
       }
+    });
 
-      focusNode(priorFocusedElem);
-
-      for (let i = 0; i < ancestors.length; i++) {
-        const info = ancestors[i];
-        info.element.scrollLeft = info.left;
-        info.element.scrollTop = info.top;
-      }
+    var curFocusedElement = getFocusedElement();
+    var priorFocusedElement = priorSelectionInformation.focusedElement;
+    if (curFocusedElement !== priorFocusedElement &&
+        isInDocument(priorFocusedElement)) {
+      focusNodePreservingScroll(priorFocusedElement);
     }
   },
 
diff --git a/src/renderers/dom/shared/__tests__/ReactInputSelection-test.js b/src/renderers/dom/shared/__tests__/ReactInputSelection-test.js
index 3748539c2d8b77..bb90c0dfb1061c 100644
--- a/src/renderers/dom/shared/__tests__/ReactInputSelection-test.js
+++ b/src/renderers/dom/shared/__tests__/ReactInputSelection-test.js
@@ -149,8 +149,9 @@ describe('ReactInputSelection', () => {
       input.selectionStart = 1;
       input.selectionEnd = 10;
       var selectionInfo = ReactInputSelection.getSelectionInformation();
-      expect(selectionInfo.focusedElem).toBe(input);
-      expect(selectionInfo.selectionRange).toEqual({start: 1, end: 10});
+      expect(selectionInfo.focusedElement).toBe(input);
+      expect(selectionInfo.activeElements[0].element).toBe(input);
+      expect(selectionInfo.activeElements[0].selectionRange).toEqual({start: 1, end: 10});
       expect(document.activeElement).toBe(input);
       input.setSelectionRange(0, 0);
       document.body.removeChild(input);