fix(lsp): correctly parse LSP snippets #15579

Fixes #15522

4 files changed, 567 insertions, 68 deletions

diff --git a/runtime/lua/vim/lsp/_snippet.lua b/runtime/lua/vim/lsp/_snippet.lua
new file mode 100644
index 0000000000..0140b0aee3
--- /dev/null
+++ b/runtime/lua/vim/lsp/_snippet.lua
@@ -0,0 +1,399 @@
+local P = {}
+
+---Take characters until the target characters (The escape sequence is '\' + char)
+---@param targets string[] The character list for stop consuming text.
+---@param specials string[] If the character isn't contained in targets/specials, '\' will be left.
+P.take_until = function(targets, specials)
+  targets = targets or {}
+  specials = specials or {}
+
+  return function(input, pos)
+    local new_pos = pos
+    local raw = {}
+    local esc = {}
+    while new_pos <= #input do
+      local c = string.sub(input, new_pos, new_pos)
+      if c == '\\' then
+        table.insert(raw, '\\')
+        new_pos = new_pos + 1
+        c = string.sub(input, new_pos, new_pos)
+        if not vim.tbl_contains(targets, c) and not vim.tbl_contains(specials, c) then
+          table.insert(esc, '\\')
+        end
+        table.insert(raw, c)
+        table.insert(esc, c)
+        new_pos = new_pos + 1
+      else
+        if vim.tbl_contains(targets, c) then
+          break
+        end
+        table.insert(raw, c)
+        table.insert(esc, c)
+        new_pos = new_pos + 1
+      end
+    end
+
+    if new_pos == pos then
+      return P.unmatch(pos)
+    end
+
+    return {
+      parsed = true,
+      value = {
+        raw = table.concat(raw, ''),
+        esc = table.concat(esc, '')
+      },
+      pos = new_pos,
+    }
+  end
+end
+
+P.unmatch = function(pos)
+  return {
+    parsed = false,
+    value = nil,
+    pos = pos,
+  }
+end
+
+P.map = function(parser, map)
+  return function(input, pos)
+    local result = parser(input, pos)
+    if result.parsed then
+      return {
+        parsed = true,
+        value = map(result.value),
+        pos = result.pos,
+      }
+    end
+    return P.unmatch(pos)
+  end
+end
+
+P.lazy = function(factory)
+  return function(input, pos)
+    return factory()(input, pos)
+  end
+end
+
+P.token = function(token)
+  return function(input, pos)
+    local maybe_token = string.sub(input, pos, pos + #token - 1)
+    if token == maybe_token then
+      return {
+        parsed = true,
+        value = maybe_token,
+        pos = pos + #token,
+      }
+    end
+    return P.unmatch(pos)
+  end
+end
+
+P.pattern = function(p)
+  return function(input, pos)
+    local maybe_match = string.match(string.sub(input, pos), '^' .. p)
+    if maybe_match then
+      return {
+        parsed = true,
+        value = maybe_match,
+        pos = pos + #maybe_match,
+      }
+    end
+    return P.unmatch(pos)
+  end
+end
+
+P.many = function(parser)
+  return function(input, pos)
+    local values = {}
+    local new_pos = pos
+    while new_pos <= #input do
+      local result = parser(input, new_pos)
+      if not result.parsed then
+        break
+      end
+      table.insert(values, result.value)
+      new_pos = result.pos
+    end
+    if #values > 0 then
+      return {
+        parsed = true,
+        value = values,
+        pos = new_pos,
+      }
+    end
+    return P.unmatch(pos)
+  end
+end
+
+P.any = function(...)
+  local parsers = { ... }
+  return function(input, pos)
+    for _, parser in ipairs(parsers) do
+      local result = parser(input, pos)
+      if result.parsed then
+        return result
+      end
+    end
+    return P.unmatch(pos)
+  end
+end
+
+P.opt = function(parser)
+  return function(input, pos)
+    local result = parser(input, pos)
+    return {
+      parsed = true,
+      value = result.value,
+      pos = result.pos,
+    }
+  end
+end
+
+P.seq = function(...)
+  local parsers = { ... }
+  return function(input, pos)
+    local values = {}
+    local new_pos = pos
+    for _, parser in ipairs(parsers) do
+      local result = parser(input, new_pos)
+      if result.parsed then
+        table.insert(values, result.value)
+        new_pos = result.pos
+      else
+        return P.unmatch(pos)
+      end
+    end
+    return {
+      parsed = true,
+      value = values,
+      pos = new_pos,
+    }
+  end
+end
+
+local Node = {}
+
+Node.Type = {
+  SNIPPET = 0,
+  TABSTOP = 1,
+  PLACEHOLDER = 2,
+  VARIABLE = 3,
+  CHOICE = 4,
+  TRANSFORM = 5,
+  FORMAT = 6,
+  TEXT = 7,
+}
+
+function Node:__tostring()
+  local insert_text = {}
+  if self.type == Node.Type.SNIPPET then
+    for _, c in ipairs(self.children) do
+      table.insert(insert_text, tostring(c))
+    end
+  elseif self.type == Node.Type.CHOICE then
+    table.insert(insert_text, self.items[1])
+  elseif self.type == Node.Type.PLACEHOLDER then
+    for _, c in ipairs(self.children or {}) do
+      table.insert(insert_text, tostring(c))
+    end
+  elseif self.type == Node.Type.TEXT then
+    table.insert(insert_text, self.esc)
+  end
+  return table.concat(insert_text, '')
+end
+
+--@see https://code.visualstudio.com/docs/editor/userdefinedsnippets#_grammar
+
+local S = {}
+S.dollar = P.token('$')
+S.open = P.token('{')
+S.close = P.token('}')
+S.colon = P.token(':')
+S.slash = P.token('/')
+S.comma = P.token(',')
+S.pipe = P.token('|')
+S.plus = P.token('+')
+S.minus = P.token('-')
+S.question = P.token('?')
+S.int = P.map(P.pattern('[0-9]+'), function(value)
+  return tonumber(value, 10)
+end)
+S.var = P.pattern('[%a_][%w_]+')
+S.text = function(targets, specials)
+  return P.map(P.take_until(targets, specials), function(value)
+    return setmetatable({
+      type = Node.Type.TEXT,
+      raw = value.raw,
+      esc = value.esc,
+    }, Node)
+  end)
+end
+
+S.toplevel = P.lazy(function()
+  return P.any(S.placeholder, S.tabstop, S.variable, S.choice)
+end)
+
+S.format = P.any(
+  P.map(P.seq(S.dollar, S.int), function(values)
+    return setmetatable({
+      type = Node.Type.FORMAT,
+      capture_index = values[2],
+    }, Node)
+  end),
+  P.map(P.seq(S.dollar, S.open, S.int, S.close), function(values)
+    return setmetatable({
+      type = Node.Type.FORMAT,
+      capture_index = values[3],
+    }, Node)
+  end),
+  P.map(P.seq(S.dollar, S.open, S.int, S.colon, S.slash, P.any(
+    P.token('upcase'),
+    P.token('downcase'),
+    P.token('capitalize'),
+    P.token('camelcase'),
+    P.token('pascalcase')
+  ), S.close), function(values)
+    return setmetatable({
+      type = Node.Type.FORMAT,
+      capture_index = values[3],
+      modifier = values[6],
+    }, Node)
+  end),
+  P.map(P.seq(S.dollar, S.open, S.int, S.colon, P.any(
+    P.seq(S.question, P.take_until({ ':' }, { '\\' }), S.colon, P.take_until({ '}' }, { '\\' })),
+    P.seq(S.plus, P.take_until({ '}' }, { '\\' })),
+    P.seq(S.minus, P.take_until({ '}' }, { '\\' }))
+  ),  S.close), function(values)
+    return setmetatable({
+      type = Node.Type.FORMAT,
+      capture_index = values[3],
+      if_text = values[5][2].esc,
+      else_text = (values[5][4] or {}).esc,
+    }, Node)
+  end)
+)
+
+S.transform = P.map(P.seq(
+  S.slash,
+  P.take_until({ '/' }, { '\\' }),
+  S.slash,
+  P.many(P.any(S.format, S.text({ '$', '/' }, { '\\' }))),
+  S.slash,
+  P.opt(P.pattern('[ig]+'))
+), function(values)
+  return setmetatable({
+    type = Node.Type.TRANSFORM,
+    pattern = values[2].raw,
+    format = values[4],
+    option = values[6],
+  }, Node)
+end)
+
+S.tabstop = P.any(
+  P.map(P.seq(S.dollar, S.int), function(values)
+    return setmetatable({
+      type = Node.Type.TABSTOP,
+      tabstop = values[2],
+    }, Node)
+  end),
+  P.map(P.seq(S.dollar, S.open, S.int, S.close), function(values)
+    return setmetatable({
+      type = Node.Type.TABSTOP,
+      tabstop = values[3],
+    }, Node)
+  end),
+  P.map(P.seq(S.dollar, S.open, S.int, S.transform, S.close), function(values)
+    return setmetatable({
+      type = Node.Type.TABSTOP,
+      tabstop = values[3],
+      transform = values[4],
+    }, Node)
+  end)
+)
+
+S.placeholder = P.any(
+  P.map(P.seq(S.dollar, S.open, S.int, S.colon, P.many(P.any(S.toplevel, S.text({ '$', '}' }, { '\\' }))), S.close), function(values)
+    return setmetatable({
+      type = Node.Type.PLACEHOLDER,
+      tabstop = values[3],
+      children = values[5],
+    }, Node)
+  end)
+)
+
+S.choice = P.map(P.seq(
+  S.dollar,
+  S.open,
+  S.int,
+  S.pipe,
+  P.many(
+    P.map(P.seq(S.text({ ',', '|' }), P.opt(S.comma)), function(values)
+      return values[1].esc
+    end)
+  ),
+  S.pipe,
+  S.close
+), function(values)
+  return setmetatable({
+    type = Node.Type.CHOICE,
+    tabstop = values[3],
+    items = values[5],
+  }, Node)
+end)
+
+S.variable = P.any(
+  P.map(P.seq(S.dollar, S.var), function(values)
+    return setmetatable({
+      type = Node.Type.VARIABLE,
+      name = values[2],
+    }, Node)
+  end),
+  P.map(P.seq(S.dollar, S.open, S.var, S.close), function(values)
+    return setmetatable({
+      type = Node.Type.VARIABLE,
+      name = values[3],
+    }, Node)
+  end),
+  P.map(P.seq(S.dollar, S.open, S.var, S.transform, S.close), function(values)
+    return setmetatable({
+      type = Node.Type.VARIABLE,
+      name = values[3],
+      transform = values[4],
+    }, Node)
+  end),
+  P.map(P.seq(S.dollar, S.open, S.var, S.colon, P.many(P.any(S.toplevel, S.text({ '$', '}' }, { '\\' }))), S.close), function(values)
+    return setmetatable({
+      type = Node.Type.VARIABLE,
+      name = values[3],
+      children = values[5],
+    }, Node)
+  end)
+)
+
+S.snippet = P.map(P.many(P.any(S.toplevel, S.text({ '$' }, { '}', '\\' }))), function(values)
+  return setmetatable({
+    type = Node.Type.SNIPPET,
+    children = values,
+  }, Node)
+end)
+
+local M = {}
+
+---The snippet node type enum
+---@types table<string, number>
+M.NodeType = Node.Type
+
+---Parse snippet string and returns the AST
+---@param input string
+---@return table
+function M.parse(input)
+  local result = S.snippet(input, 1)
+  if not result.parsed then
+    error('snippet parsing failed.')
+  end
+  return result.value
+end
+
+return M
diff --git a/runtime/lua/vim/lsp/util.lua b/runtime/lua/vim/lsp/util.lua
index a4c8b69f6c..5a22a311e0 100644
--- a/runtime/lua/vim/lsp/util.lua
+++ b/runtime/lua/vim/lsp/util.lua
@@ -1,4 +1,5 @@
 local protocol = require 'vim.lsp.protocol'
+local snippet = require 'vim.lsp._snippet'
 local vim = vim
 local validate = vim.validate
 local api = vim.api
@@ -523,74 +524,18 @@ function M.apply_text_document_edit(text_document_edit, index)
   M.apply_text_edits(text_document_edit.edits, bufnr)
 end
 
----@private
---- Recursively parses snippets in a completion entry.
----
----@param input (string) Snippet text to parse for snippets
----@param inner (bool) Whether this function is being called recursively
----@returns 2-tuple of strings: The first is the parsed result, the second is the
----unparsed rest of the input
-local function parse_snippet_rec(input, inner)
-  local res = ""
-
-  local close, closeend = nil, nil
-  if inner then
-    close, closeend = input:find("}", 1, true)
-    while close ~= nil and input:sub(close-1,close-1) == "\\" do
-      close, closeend = input:find("}", closeend+1, true)
-    end
-  end
-
-  local didx = input:find('$',  1, true)
-  if didx == nil and close == nil then
-    return input, ""
-  elseif close ~=nil and (didx == nil or close < didx) then
-    -- No inner placeholders
-    return input:sub(0, close-1), input:sub(closeend+1)
-  end
-
-  res = res .. input:sub(0, didx-1)
-  input = input:sub(didx+1)
-
-  local tabstop, tabstopend = input:find('^%d+')
-  local placeholder, placeholderend = input:find('^{%d+:')
-  local choice, choiceend = input:find('^{%d+|')
-
-  if tabstop then
-    input = input:sub(tabstopend+1)
-  elseif choice then
-    input = input:sub(choiceend+1)
-    close, closeend = input:find("|}", 1, true)
-
-    res = res .. input:sub(0, close-1)
-    input = input:sub(closeend+1)
-  elseif placeholder then
-    -- TODO: add support for variables
-    input = input:sub(placeholderend+1)
-
-    -- placeholders and variables are recursive
-    while input ~= "" do
-      local r, tail = parse_snippet_rec(input, true)
-      r = r:gsub("\\}", "}")
-
-      res = res .. r
-      input = tail
-    end
-  else
-    res = res .. "$"
-  end
-
-  return res, input
-end
-
 --- Parses snippets in a completion entry.
 ---
----@param input (string) unparsed snippet
----@returns (string) parsed snippet
+---@param input string unparsed snippet
+---@returns string parsed snippet
 function M.parse_snippet(input)
-  local res, _ = parse_snippet_rec(input, false)
-
-  return res
+  local ok, parsed = pcall(function()
+    return tostring(snippet.parse(input))
+  end)
+  if not ok then
+    return input
+  end
+  return parsed
 end
 
 ---@private
diff --git a/test/functional/plugin/lsp/snippet_spec.lua b/test/functional/plugin/lsp/snippet_spec.lua
new file mode 100644
index 0000000000..4e127743eb
--- /dev/null
+++ b/test/functional/plugin/lsp/snippet_spec.lua
@@ -0,0 +1,152 @@
+local helpers = require('test.functional.helpers')(after_each)
+local snippet = require('vim.lsp._snippet')
+
+local eq = helpers.eq
+local exec_lua = helpers.exec_lua
+
+describe('vim.lsp._snippet', function()
+  before_each(helpers.clear)
+  after_each(helpers.clear)
+
+  local parse = function(...)
+    return exec_lua('return require("vim.lsp._snippet").parse(...)', ...)
+  end
+
+  it('should parse only text', function()
+    eq({
+      type = snippet.NodeType.SNIPPET,
+      children = {
+        {
+          type = snippet.NodeType.TEXT,
+          raw = 'TE\\$\\}XT',
+          esc = 'TE$}XT'
+        }
+      }
+    }, parse('TE\\$\\}XT'))
+  end)
+
+  it('should parse tabstop', function()
+    eq({
+      type = snippet.NodeType.SNIPPET,
+      children = {
+        {
+          type = snippet.NodeType.TABSTOP,
+          tabstop = 1,
+        },
+        {
+          type = snippet.NodeType.TABSTOP,
+          tabstop = 2,
+        }
+      }
+    }, parse('$1${2}'))
+  end)
+
+  it('should parse placeholders', function()
+    eq({
+      type = snippet.NodeType.SNIPPET,
+      children = {
+        {
+          type = snippet.NodeType.PLACEHOLDER,
+          tabstop = 1,
+          children = {
+            {
+              type = snippet.NodeType.PLACEHOLDER,
+              tabstop = 2,
+              children = {
+                {
+                  type = snippet.NodeType.TEXT,
+                  raw = 'TE\\$\\}XT',
+                  esc = 'TE$}XT'
+                },
+                {
+                  type = snippet.NodeType.TABSTOP,
+                  tabstop = 3,
+                },
+                {
+                  type = snippet.NodeType.TABSTOP,
+                  tabstop = 1,
+                  transform = {
+                    type = snippet.NodeType.TRANSFORM,
+                    pattern = 'regex',
+                    option = 'i',
+                    format = {
+                      {
+                        type = snippet.NodeType.FORMAT,
+                        capture_index = 1,
+                        modifier = 'upcase'
+                      }
+                    }
+                  },
+                },
+                {
+                  type = snippet.NodeType.TEXT,
+                  raw = 'TE\\$\\}XT',
+                  esc = 'TE$}XT'
+                },
+              }
+            }
+          }
+        },
+      }
+    }, parse('${1:${2:TE\\$\\}XT$3${1/regex/${1:/upcase}/i}TE\\$\\}XT}}'))
+  end)
+
+  it('should parse variables', function()
+    eq({
+      type = snippet.NodeType.SNIPPET,
+      children = {
+        {
+          type = snippet.NodeType.VARIABLE,
+          name = 'VAR',
+        },
+        {
+          type = snippet.NodeType.VARIABLE,
+          name = 'VAR',
+        },
+        {
+          type = snippet.NodeType.VARIABLE,
+          name = 'VAR',
+          children = {
+            {
+              type = snippet.NodeType.TABSTOP,
+              tabstop = 1,
+            }
+          }
+        },
+        {
+          type = snippet.NodeType.VARIABLE,
+          name = 'VAR',
+          transform = {
+            type = snippet.NodeType.TRANSFORM,
+            pattern = 'regex',
+            format = {
+              {
+                type = snippet.NodeType.FORMAT,
+                capture_index = 1,
+                modifier = 'upcase',
+              }
+            }
+          }
+        },
+      }
+    }, parse('$VAR${VAR}${VAR:$1}${VAR/regex/${1:/upcase}/}'))
+  end)
+
+  it('should parse choice', function()
+    eq({
+      type = snippet.NodeType.SNIPPET,
+      children = {
+        {
+          type = snippet.NodeType.CHOICE,
+          tabstop = 1,
+          items = {
+            ',',
+            '|'
+          }
+        }
+      }
+    }, parse('${1|\\,,\\||}'))
+  end)
+
+end)
+
diff --git a/test/functional/plugin/lsp_spec.lua b/test/functional/plugin/lsp_spec.lua
index a9ea26343d..7df5eb049c 100644
--- a/test/functional/plugin/lsp_spec.lua
+++ b/test/functional/plugin/lsp_spec.lua
@@ -1444,8 +1444,10 @@ describe('LSP', function()
         { label='foocar', sortText="h", insertText='foodar(${1:var1} typ1, ${2:var2} *typ2) {$0\\}', insertTextFormat=2, textEdit={} },
         -- nested snippet tokens
         { label='foocar', sortText="i", insertText='foodar(${1:var1 ${2|typ2,typ3|} ${3:tail}}) {$0\\}', insertTextFormat=2, textEdit={} },
+        -- braced tabstop
+        { label='foocar', sortText="j", insertText='foodar()${0}', insertTextFormat=2, textEdit={} },
         -- plain text
-        { label='foocar', sortText="j", insertText='foodar(${1:var1})', insertTextFormat=1, textEdit={} },
+        { label='foocar', sortText="k", insertText='foodar(${1:var1})', insertTextFormat=1, textEdit={} },
       }
       local completion_list_items = {items=completion_list}
       local expected = {
@@ -1457,8 +1459,9 @@ describe('LSP', function()
         { abbr = 'foocar', dup = 1, empty = 1, icase = 1, info = ' ', kind = 'Unknown', menu = '', word = 'foobar', user_data = { nvim = { lsp = { completion_item = { label='foocar', sortText="f", textEdit={newText='foobar'} } } } } },
         { abbr = 'foocar', dup = 1, empty = 1, icase = 1, info = ' ', kind = 'Unknown', menu = '', word = 'foobar(place holder, more ...holder{})', user_data = { nvim = { lsp = { completion_item = { label='foocar', sortText="g", insertText='foodar', insertTextFormat=2, textEdit={newText='foobar(${1:place holder}, ${2:more ...holder{\\}})'} } } } } },
         { abbr = 'foocar', dup = 1, empty = 1, icase = 1, info = ' ', kind = 'Unknown', menu = '', word = 'foodar(var1 typ1, var2 *typ2) {}', user_data = { nvim = { lsp = { completion_item = { label='foocar', sortText="h", insertText='foodar(${1:var1} typ1, ${2:var2} *typ2) {$0\\}', insertTextFormat=2, textEdit={} } } } } },
-        { abbr = 'foocar', dup = 1, empty = 1, icase = 1, info = ' ', kind = 'Unknown', menu = '', word = 'foodar(var1 typ2,typ3 tail) {}', user_data = { nvim = { lsp = { completion_item = { label='foocar', sortText="i", insertText='foodar(${1:var1 ${2|typ2,typ3|} ${3:tail}}) {$0\\}', insertTextFormat=2, textEdit={} } } } } },
-        { abbr = 'foocar', dup = 1, empty = 1, icase = 1, info = ' ', kind = 'Unknown', menu = '', word = 'foodar(${1:var1})', user_data = { nvim = { lsp = { completion_item = { label='foocar', sortText="j", insertText='foodar(${1:var1})', insertTextFormat=1, textEdit={} } } } } },
+        { abbr = 'foocar', dup = 1, empty = 1, icase = 1, info = ' ', kind = 'Unknown', menu = '', word = 'foodar(var1 typ2 tail) {}', user_data = { nvim = { lsp = { completion_item = { label='foocar', sortText="i", insertText='foodar(${1:var1 ${2|typ2,typ3|} ${3:tail}}) {$0\\}', insertTextFormat=2, textEdit={} } } } } },
+        { abbr = 'foocar', dup = 1, empty = 1, icase = 1, info = ' ', kind = 'Unknown', menu = '', word = 'foodar()', user_data = { nvim = { lsp = { completion_item = { label='foocar', sortText="j", insertText='foodar()${0}', insertTextFormat=2, textEdit={} } } } } },
+        { abbr = 'foocar', dup = 1, empty = 1, icase = 1, info = ' ', kind = 'Unknown', menu = '', word = 'foodar(${1:var1})', user_data = { nvim = { lsp = { completion_item = { label='foocar', sortText="k", insertText='foodar(${1:var1})', insertTextFormat=1, textEdit={} } } } } },
       }
 
       eq(expected, exec_lua([[return vim.lsp.util.text_document_completion_list_to_complete_items(...)]], completion_list, prefix))