From 5ee2be25eb28af41cc88814bf8d5eee9c75458a5 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Fri, 16 Dec 2016 17:30:14 +0300
Subject: [PATCH] PY-21840 Consider surrogate pairs in Python's
 StringLiteralTextEscaper

Namely, detect a single unsplittable piece of a string returned by
getDecodedFragments() by checking whether the decoded text contains
only one code point, not code unit (as returned by String.length()),
since some escape sequences, in particular, long unicode escape
sequences might be decoded into a surrogate pair.
---
 .../psi/impl/PyStringLiteralExpressionImpl.java     | 10 ++++++----
 .../testSrc/com/jetbrains/python/PyEditingTest.java | 10 ++++++++++
 .../com/jetbrains/python/PyStringLiteralTest.java   | 13 +++++++++++++
 3 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/python/src/com/jetbrains/python/psi/impl/PyStringLiteralExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyStringLiteralExpressionImpl.java
index ea57e55c3142..26ab8638e783 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyStringLiteralExpressionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyStringLiteralExpressionImpl.java
@@ -351,7 +351,7 @@ public class PyStringLiteralExpressionImpl extends PyElementImpl implements PySt
         if (intersection != null && !intersection.isEmpty()) {
           final String value = fragment.getSecond();
           final String intersectedValue;
-          if (value.length() == 1 || value.length() == intersection.getLength()) {
+          if (value.codePointCount(0, value.length()) == 1 || value.length() == intersection.getLength()) {
             intersectedValue = value;
           }
           else {
@@ -367,7 +367,7 @@ public class PyStringLiteralExpressionImpl extends PyElementImpl implements PySt
 
     @Override
     public int getOffsetInHost(final int offsetInDecoded, @NotNull final TextRange rangeInsideHost) {
-      int offset = 0;
+      int offset = 0; // running offset in the decoded fragment
       int endOffset = -1;
       for (Pair<TextRange, String> fragment : myHost.getDecodedFragments()) {
         final TextRange encodedTextRange = fragment.getFirst();
@@ -379,13 +379,15 @@ public class PyStringLiteralExpressionImpl extends PyElementImpl implements PySt
           if (valueLength == 0) {
             return -1;
           }
-          else if (valueLength == 1) {
+          // A long unicode escape of form \U01234567 can be decoded into a surrogate pair
+          else if (value.codePointCount(0, valueLength) == 1) {
             if (offset == offsetInDecoded) {
               return intersection.getStartOffset();
             }
-            offset++;
+            offset += valueLength;
           }
           else {
+            // Literal fragment without escapes: it's safe to use intersection length instead of value length
             if (offset + intersectionLength >= offsetInDecoded) {
               final int delta = offsetInDecoded - offset;
               return intersection.getStartOffset() + delta;
diff --git a/python/testSrc/com/jetbrains/python/PyEditingTest.java b/python/testSrc/com/jetbrains/python/PyEditingTest.java
index 8c9ca82bed74..721f1f099e93 100644
--- a/python/testSrc/com/jetbrains/python/PyEditingTest.java
+++ b/python/testSrc/com/jetbrains/python/PyEditingTest.java
@@ -472,6 +472,16 @@ public class PyEditingTest extends PyTestCase {
                 ")");
   }
 
+  // PY-21840
+  public void testEditInjectedRegexpFragmentWithLongUnicodeEscape() {
+    myFixture.configureByText(PythonFileType.INSTANCE,
+                              "import re\n" +
+                              "re.compile(ur'\\U00010000<caret>')");
+    doTyping("t");
+    myFixture.checkResult("import re\n" +
+                          "re.compile(ur'\\U00010000t')");
+  }
+
   private String doTestTyping(final String text, final int offset, final char character) {
     final PsiFile file = WriteCommandAction.runWriteCommandAction(null, new Computable<PsiFile>() {
       @Override
diff --git a/python/testSrc/com/jetbrains/python/PyStringLiteralTest.java b/python/testSrc/com/jetbrains/python/PyStringLiteralTest.java
index 4efa8cf2a9f7..90b489535d43 100644
--- a/python/testSrc/com/jetbrains/python/PyStringLiteralTest.java
+++ b/python/testSrc/com/jetbrains/python/PyStringLiteralTest.java
@@ -115,6 +115,19 @@ public class PyStringLiteralTest extends PyTestCase {
     assertEquals(-1, escaper.getOffsetInHost(9, range));
   }
 
+  public void testEscaperOffsetInLongUnicodeEscape() {
+    final PyStringLiteralExpression expr = createLiteralFromText("u'XXX a\\U0001F600\\U0001F600b YYY'");
+    final LiteralTextEscaper<? extends PsiLanguageInjectionHost> escaper = expr.createLiteralTextEscaper();
+    final TextRange range = TextRange.create(6, 28);
+    assertEquals(6, escaper.getOffsetInHost(0, range));
+    assertEquals(7, escaper.getOffsetInHost(1, range));
+    // Each \\U0001F600 is represented as a surrogate pair, hence 2 characters-wide step in decoded text
+    assertEquals(17, escaper.getOffsetInHost(3, range)); 
+    assertEquals(27, escaper.getOffsetInHost(5, range));
+    assertEquals(28, escaper.getOffsetInHost(6, range));
+    assertEquals(-1, escaper.getOffsetInHost(7, range));
+  }
+
   public void testStringValue() {
     assertEquals("foo", createLiteralFromText("\"\"\"foo\"\"\"").getStringValue());
     assertEquals("foo", createLiteralFromText("u\"foo\"").getStringValue());