From 0bb2810bd8543ea82870c9d4f36b902ef24e7ca2 Mon Sep 17 00:00:00 2001
From: Anthony Fu <anthonyfu117@hotmail.com>
Date: Mon, 5 Feb 2024 16:03:48 +0100
Subject: [PATCH] fix: splitLines extra new lines, fix #576

---
 packages/core/src/transformer-decorations.ts |  4 +-
 packages/core/src/utils.ts                   | 17 +++++--
 packages/core/test/utils.test.ts             | 47 ++++++++++++++++++++
 3 files changed, 63 insertions(+), 5 deletions(-)
 create mode 100644 packages/core/test/utils.test.ts

diff --git a/packages/core/src/transformer-decorations.ts b/packages/core/src/transformer-decorations.ts
index 65f4296ff..76379fb92 100644
--- a/packages/core/src/transformer-decorations.ts
+++ b/packages/core/src/transformer-decorations.ts
@@ -93,8 +93,8 @@ export function transformerDecorations(): ShikiTransformer {
 
       const lines = Array.from(codeEl.children).filter(i => i.type === 'element' && i.tagName === 'span') as Element[]
 
-      // if (lines.length !== ctx.converter.lines.length)
-      //   throw new ShikiError(`Number of lines in code element (${lines.length}) does not match the number of lines in the source (${ctx.converter.lines.length}). Failed to apply decorations.`)
+      if (lines.length !== ctx.converter.lines.length)
+        throw new ShikiError(`Number of lines in code element (${lines.length}) does not match the number of lines in the source (${ctx.converter.lines.length}). Failed to apply decorations.`)
 
       function applyLineSection(line: number, start: number, end: number, decoration: DecorationItem) {
         const lineEl = lines[line]
diff --git a/packages/core/src/utils.ts b/packages/core/src/utils.ts
index 2b734de2b..ef53917d2 100644
--- a/packages/core/src/utils.ts
+++ b/packages/core/src/utils.ts
@@ -9,8 +9,19 @@ export function toArray<T>(x: MaybeArray<T>): T[] {
 /**
  * Slipt a string into lines, each line preserves the line ending.
  */
-export function splitLines(str: string) {
-  return Array.from(str.matchAll(/^.*$/mg)).map(x => [x[0], x.index!] as [code: string, offset: number])
+export function splitLines(code: string, preserveEnding = false): [string, number][] {
+  const parts = code.split(/(\r?\n)/g)
+  let index = 0
+  const lines: [string, number][] = []
+  for (let i = 0; i < parts.length; i += 2) {
+    const line = preserveEnding
+      ? parts[i] + (parts[i + 1] || '')
+      : parts[i]
+    lines.push([line, index])
+    index += parts[i].length
+    index += parts[i + 1]?.length || 0
+  }
+  return lines
 }
 
 /**
@@ -169,7 +180,7 @@ export function stringifyTokenStyle(token: Record<string, string>) {
  * Creates a converter between index and position in a code block.
  */
 export function createPositionConverter(code: string) {
-  const lines = Array.from(code.matchAll(/.*?($|\n)/g)).map(match => match[0])
+  const lines = splitLines(code, true).map(([line]) => line)
 
   function indexToPos(index: number): Position {
     let character = index
diff --git a/packages/core/test/utils.test.ts b/packages/core/test/utils.test.ts
new file mode 100644
index 000000000..0c9da3b16
--- /dev/null
+++ b/packages/core/test/utils.test.ts
@@ -0,0 +1,47 @@
+/* eslint-disable style/no-tabs */
+import { describe, expect, it } from 'vitest'
+import { splitLines } from '../src/utils'
+
+describe('utils', () => {
+  it('splitLines', () => {
+    const lines = [
+      '\t*/\r',
+      '\tpublic void setTestingRefNum(long l) {\r',
+      '\ttestingRefNum = l;\r',
+      '\t}\r',
+    ]
+    const code = lines.join('\n')
+
+    const resultWithEnding = splitLines(code, true)
+    const resultWithoutEnding = splitLines(code, false)
+    const reconstructed = resultWithEnding.map(([line]) => line).join('')
+    expect(reconstructed).toBe(code)
+    // the offset should be the same
+    expect(resultWithoutEnding.map(i => i[1]))
+      .toEqual(resultWithEnding.map(i => i[1]))
+    expect(resultWithEnding).toMatchInlineSnapshot(`
+      [
+        [
+          "	*/
+      ",
+          0,
+        ],
+        [
+          "	public void setTestingRefNum(long l) {
+      ",
+          5,
+        ],
+        [
+          "	testingRefNum = l;
+      ",
+          46,
+        ],
+        [
+          "	}
+      ",
+          67,
+        ],
+      ]
+    `)
+  })
+})