From 81592e724a28cb4cc1d0ea29277f15f5d3725c32 Mon Sep 17 00:00:00 2001
From: Myzel394 <50424412+Myzel394@users.noreply.github.com>
Date: Sun, 14 Apr 2024 10:06:49 +0200
Subject: [PATCH 1/6] feat: Create first WIP for LSP support

---
 lua/telescope/_extensions/jsonfly.lua | 257 +++++++++++++++++++++-----
 1 file changed, 206 insertions(+), 51 deletions(-)

diff --git a/lua/telescope/_extensions/jsonfly.lua b/lua/telescope/_extensions/jsonfly.lua
index 39f7dae..a019ec4 100644
--- a/lua/telescope/_extensions/jsonfly.lua
+++ b/lua/telescope/_extensions/jsonfly.lua
@@ -24,16 +24,30 @@ local conf = require("telescope.config").values
 local make_entry = require "telescope.make_entry"
 local entry_display = require "telescope.pickers.entry_display"
 
-local function get_recursive_keys(t)
+---@class InputEntry
+---@field value InputEntry|table|number|string|boolean|nil
+---@field newlines number
+---@field key_start number
+---@field value_start number
+--
+---@class Entry
+---@field key string
+---@field entry InputEntry
+--
+---@param t table
+---@return Entry[]
+local function get_entries_from_lua_json(t)
     local keys = {}
 
+    --@type k string
+    --@type raw_value InputEntry
     for k, raw_value in pairs(t) do
         table.insert(keys, {key = k, entry = raw_value})
 
         local v = raw_value.value
 
         if type(v) == "table" then
-            local sub_keys = get_recursive_keys(v)
+            local sub_keys = get_entries_from_lua_json(v)
             for _, sub_key in ipairs(sub_keys) do
                 table.insert(keys, {key = k .. "." .. sub_key.key, entry = sub_key.entry})
             end
@@ -43,6 +57,84 @@ local function get_recursive_keys(t)
     return keys
 end
 
+---@param result Symbol
+---@return string|number|table|boolean|nil
+local function parse_lsp_value(result)
+    if result.kind == 2 then
+        local value = {}
+
+        for _, child in ipairs(result.children) do
+            value[child.name] = parse_lsp_value(child)
+        end
+
+        return value
+    elseif result.kind == 16 then
+        return tonumber(result.detail)
+    elseif result.kind == 15 then
+        return result.detail
+    elseif result.kind == 18 then
+        local value = {}
+
+        for i, child in ipairs(result.children) do
+            value[i] = parse_lsp_value(child)
+        end
+
+        return value
+    elseif result.kind == 13 then
+        return nil
+    elseif result.kind == 17 then
+        return result.detail == "true"
+    end
+end
+
+
+---@class Symbol
+---@field name string
+---@field kind number 2 = Object, 16 = Number, 15 = String, 18 = Array, 13 = Null, 17 = Boolean
+---@field range Range
+---@field selectionRange Range
+---@field detail string
+---@field children Symbol[]
+--
+---@class Range
+---@field start Position
+---@field ["end"] Position
+--
+---@class Position
+---@field line number
+---@field character number
+--
+---@param symbols Symbol[]
+---@return Entry[]
+local function get_entries_from_lsp_symbols(symbols)
+    local keys = {}
+
+    for _, symbol in ipairs(symbols) do
+        local key = symbol.name
+
+        if symbol.kind == 2 then
+            local sub_keys = get_entries_from_lsp_symbols(symbol.children)
+            for _, sub_key in ipairs(sub_keys) do
+                table.insert(keys, {key = key .. "." .. sub_key.key, entry = sub_key.entry})
+            end
+        end
+
+        ---@type Entry
+        local entry = {
+            key = key,
+            entry = {
+                value = parse_lsp_value(symbol),
+                newlines = symbol.range["end"].line,
+                key_start = symbol.range.start.character,
+                value_start = symbol.selectionRange.start.character,
+            }
+        }
+        table.insert(keys, entry)
+    end
+
+    return keys
+end
+
 local function truncate_overflow(value, max_length, overflow_marker)
     if vim.fn.strdisplaywidth(value) > max_length then
         return value:sub(1, max_length - vim.fn.strdisplaywidth(overflow_marker)) .. overflow_marker
@@ -67,7 +159,7 @@ local function create_display_preview(value, opts)
         local preview_table = {}
 
         for k, v in pairs(value) do
-            table.insert(preview_table, k .. ": " .. create_display_preview(v.value, opts))
+            table.insert(preview_table, k .. ": " .. create_display_preview(v, opts))
         end
 
         return "{ " .. table.concat(preview_table, ", ") .. " }", "other"
@@ -128,14 +220,14 @@ return require"telescope".register_extension {
         opts = vim.tbl_deep_extend("force", opts, extension_config or {})
     end,
     exports = {
-        jsonfly = function()
+        jsonfly = function(xopts)
             local current_buf = vim.api.nvim_get_current_buf()
             local filename = vim.api.nvim_buf_get_name(current_buf)
             local content_lines = vim.api.nvim_buf_get_lines(current_buf, 0, -1, false)
             local content = table.concat(content_lines, "\n")
 
             local parsed = json:decode(content)
-            local keys = get_recursive_keys(parsed)
+            local keys = get_entries_from_lua_json(parsed)
 
             local displayer = entry_display.create {
                 separator = " ",
@@ -146,57 +238,120 @@ return require"telescope".register_extension {
                 },
             }
 
-            pickers.new(opts, {
-                prompt_title = opts.prompt_title,
-                finder = finders.new_table {
-                    results = keys,
-                    entry_maker = function(entry)
-                        local _, raw_depth = entry.key:gsub("%.", ".")
-                        local depth = (raw_depth or 0) + 1
+            local params = vim.lsp.util.make_position_params(xopts.winnr)
+            local result = vim.lsp.buf_request(
+                current_buf,
+                "textDocument/documentSymbol",
+                params,
+                function(_, result)
+                    local keys = get_entries_from_lsp_symbols(result)
 
-                        return make_entry.set_default_entry_mt({
-                            value = current_buf,
-                            ordinal = entry.key,
-                            display = function(_)
-                                local preview, hl_group_key = create_display_preview(entry.entry.value, opts)
+                    pickers.new(opts, {
+                        prompt_title = opts.prompt_title,
+                        finder = finders.new_table {
+                            results = keys,
+                            ---@param entry Entry
+                            entry_maker = function(entry)
+                                local _, raw_depth = entry.key:gsub("%.", ".")
+                                local depth = (raw_depth or 0) + 1
 
-                                local key = opts.subkeys_display == "normal" and entry.key or replace_previous_keys(entry.key, " ")
+                                return make_entry.set_default_entry_mt({
+                                    value = current_buf,
+                                    ordinal = entry.key,
+                                    display = function(_)
+                                        local preview, hl_group_key = create_display_preview(entry.entry.value, opts)
 
-                                return displayer {
-                                    { depth, "TelescopeResultsNumber"},
-                                    {
-                                         truncate_overflow(
-                                            key,
-                                            opts.key_max_length,
-                                            opts.overflow_marker
-                                        ),
-                                        "@property.json",
-                                    },
-                                    {
-                                        truncate_overflow(
-                                            preview,
-                                            opts.max_length,
-                                            opts.overflow_marker
-                                        ),
-                                        opts.highlights[hl_group_key] or "TelescopeResultsString",
-                                    },
-                                }
+                                        local key = opts.subkeys_display == "normal" and entry.key or replace_previous_keys(entry.key, " ")
+
+                                        return displayer {
+                                            { depth, "TelescopeResultsNumber"},
+                                            {
+                                                 truncate_overflow(
+                                                    key,
+                                                    opts.key_max_length,
+                                                    opts.overflow_marker
+                                                ),
+                                                "@property.json",
+                                            },
+                                            {
+                                                truncate_overflow(
+                                                    preview,
+                                                    opts.max_length,
+                                                    opts.overflow_marker
+                                                ),
+                                                opts.highlights[hl_group_key] or "TelescopeResultsString",
+                                            },
+                                        }
+                                    end,
+
+                                    bufnr = current_buf,
+                                    filename = filename,
+                                    lnum = entry.entry.newlines + 1,
+                                    col = opts.jump_behavior == "key_start"
+                                            and entry.entry.key_start
+                                            -- Use length ("#" operator) as vim jumps to the bytes, not characters
+                                            or entry.entry.value_start
+                                }, opts)
                             end,
+                        },
+                        previewer = conf.grep_previewer(opts),
+                        sorter = conf.generic_sorter(opts),
+                        sorting_strategy = "ascending",
+                    }):find()
+                end
+            )
 
-                            bufnr = current_buf,
-                            filename = filename,
-                            lnum = entry.entry.newlines + 1,
-                            col = opts.jump_behavior == "key_start"
-                                    and entry.entry.key_start
-                                    -- Use length ("#" operator) as vim jumps to the bytes, not characters
-                                    or entry.entry.value_start
-                        }, opts)
-                    end,
-                },
-                previewer = conf.grep_previewer(opts),
-                sorter = conf.generic_sorter(opts),
-                sorting_strategy = "ascending",
-            }):find()
+            -- pickers.new(opts, {
+            --     prompt_title = opts.prompt_title,
+            --     finder = finders.new_table {
+            --         results = keys,
+            --         entry_maker = function(entry)
+            --             local _, raw_depth = entry.key:gsub("%.", ".")
+            --             local depth = (raw_depth or 0) + 1
+            --
+            --             return make_entry.set_default_entry_mt({
+            --                 value = current_buf,
+            --                 ordinal = entry.key,
+            --                 display = function(_)
+            --                     local preview, hl_group_key = create_display_preview(entry.entry.value, opts)
+            --
+            --                     local key = opts.subkeys_display == "normal" and entry.key or replace_previous_keys(entry.key, " ")
+            --
+            --                     return displayer {
+            --                         { depth, "TelescopeResultsNumber"},
+            --                         {
+            --                              truncate_overflow(
+            --                                 key,
+            --                                 opts.key_max_length,
+            --                                 opts.overflow_marker
+            --                             ),
+            --                             "@property.json",
+            --                         },
+            --                         {
+            --                             truncate_overflow(
+            --                                 preview,
+            --                                 opts.max_length,
+            --                                 opts.overflow_marker
+            --                             ),
+            --                             opts.highlights[hl_group_key] or "TelescopeResultsString",
+            --                         },
+            --                     }
+            --                 end,
+            --
+            --                 bufnr = current_buf,
+            --                 filename = filename,
+            --                 lnum = entry.entry.newlines + 1,
+            --                 col = opts.jump_behavior == "key_start"
+            --                         and entry.entry.key_start
+            --                         -- Use length ("#" operator) as vim jumps to the bytes, not characters
+            --                         or entry.entry.value_start
+            --             }, opts)
+            --         end,
+            --     },
+            --     previewer = conf.grep_previewer(opts),
+            --     sorter = conf.generic_sorter(opts),
+            --     sorting_strategy = "ascending",
+            -- }):find()
         end
     }
 }

From c405ae33012e8173e85ca0a3f47cc90411ed5832 Mon Sep 17 00:00:00 2001
From: Myzel394 <50424412+Myzel394@users.noreply.github.com>
Date: Sun, 14 Apr 2024 10:19:38 +0200
Subject: [PATCH 2/6] chore: Improve types

---
 lua/jsonfly/json.lua                  |  4 +-
 lua/telescope/_extensions/jsonfly.lua | 71 +++++++++++++++++++--------
 2 files changed, 53 insertions(+), 22 deletions(-)

diff --git a/lua/jsonfly/json.lua b/lua/jsonfly/json.lua
index fb58ddd..145b055 100644
--- a/lua/jsonfly/json.lua
+++ b/lua/jsonfly/json.lua
@@ -387,7 +387,7 @@ local function grok_object(self, text, start, options)
       ---- Add start position so we can quickly jump to it
       VALUE[key] = {
           value = new_val,
-          newlines = newlines or 0,
+          line_number = (newlines or 0) + 1,
           key_start = relative_start + 1,
           value_start = get_relative_i(text, i),
       }
@@ -440,7 +440,7 @@ local function grok_array(self, text, start, options)
       -- can't table.insert(VALUE, val) here because it's a no-op if val is nil
       VALUE[VALUE_INDEX] = {
           value = val,
-          newlines = newlines or 0,
+          line_number = (newlines or 0) + 1,
           value_start = relative_start,
           key_start = relative_start,
       }
diff --git a/lua/telescope/_extensions/jsonfly.lua b/lua/telescope/_extensions/jsonfly.lua
index a019ec4..bac8ac8 100644
--- a/lua/telescope/_extensions/jsonfly.lua
+++ b/lua/telescope/_extensions/jsonfly.lua
@@ -1,3 +1,4 @@
+---- Documentation for jsonfly ----
 --- Type definitions
 ---@class Options
 ---@field key_max_length number - Length for the key column, 0 for no column-like display, Default: 50
@@ -17,6 +18,18 @@
 ---@field null string - Highlight group for null values, Default: "@constant.builtin.json"
 ---@field other string - Highlight group for other types, Default: "@label.json"
 
+---- Types below are for internal use only ----
+--
+---@class EntryPosition
+---@field line_number number
+---@field key_start number
+---@field value_start number
+--
+---@class Entry
+---@field key string
+---@field value Entry|table|number|string|boolean|nil
+---@field position EntryPosition
+
 local json = require"jsonfly.json"
 local finders = require "telescope.finders"
 local pickers = require "telescope.pickers"
@@ -24,16 +37,6 @@ local conf = require("telescope.config").values
 local make_entry = require "telescope.make_entry"
 local entry_display = require "telescope.pickers.entry_display"
 
----@class InputEntry
----@field value InputEntry|table|number|string|boolean|nil
----@field newlines number
----@field key_start number
----@field value_start number
---
----@class Entry
----@field key string
----@field entry InputEntry
---
 ---@param t table
 ---@return Entry[]
 local function get_entries_from_lua_json(t)
@@ -42,14 +45,32 @@ local function get_entries_from_lua_json(t)
     --@type k string
     --@type raw_value InputEntry
     for k, raw_value in pairs(t) do
-        table.insert(keys, {key = k, entry = raw_value})
+        ---@type Entry
+        local entry = {
+            key = k,
+            value = raw_value,
+            position = {
+                line_number = raw_value.newlines,
+                key_start = raw_value.key_start,
+                value_start = raw_value.value_start,
+            }
+        }
+        table.insert(keys, entry)
 
         local v = raw_value.value
 
         if type(v) == "table" then
             local sub_keys = get_entries_from_lua_json(v)
+
             for _, sub_key in ipairs(sub_keys) do
-                table.insert(keys, {key = k .. "." .. sub_key.key, entry = sub_key.entry})
+                ---@type Entry
+                local entry = {
+                    key = k .. "." .. sub_key.key,
+                    value = sub_key,
+                    position = sub_key.position,
+                }
+
+                table.insert(keys, entry)
             end
         end
     end
@@ -114,17 +135,25 @@ local function get_entries_from_lsp_symbols(symbols)
 
         if symbol.kind == 2 then
             local sub_keys = get_entries_from_lsp_symbols(symbol.children)
+
             for _, sub_key in ipairs(sub_keys) do
-                table.insert(keys, {key = key .. "." .. sub_key.key, entry = sub_key.entry})
+                ---@type Entry
+                local entry = {
+                    key = key .. "." .. sub_key.key,
+                    value = sub_key.value,
+                    position = sub_key.position,
+                }
+
+                table.insert(keys, entry)
             end
         end
 
         ---@type Entry
         local entry = {
             key = key,
-            entry = {
-                value = parse_lsp_value(symbol),
-                newlines = symbol.range["end"].line,
+            value = parse_lsp_value(symbol),
+            position = {
+                line_number = symbol.range["end"].line,
                 key_start = symbol.range.start.character,
                 value_start = symbol.selectionRange.start.character,
             }
@@ -255,11 +284,13 @@ return require"telescope".register_extension {
                                 local _, raw_depth = entry.key:gsub("%.", ".")
                                 local depth = (raw_depth or 0) + 1
 
+                                print(vim.inspect(entry))
+
                                 return make_entry.set_default_entry_mt({
                                     value = current_buf,
                                     ordinal = entry.key,
                                     display = function(_)
-                                        local preview, hl_group_key = create_display_preview(entry.entry.value, opts)
+                                        local preview, hl_group_key = create_display_preview(entry.value, opts)
 
                                         local key = opts.subkeys_display == "normal" and entry.key or replace_previous_keys(entry.key, " ")
 
@@ -286,11 +317,11 @@ return require"telescope".register_extension {
 
                                     bufnr = current_buf,
                                     filename = filename,
-                                    lnum = entry.entry.newlines + 1,
+                                    lnum = entry.position.line_number,
                                     col = opts.jump_behavior == "key_start"
-                                            and entry.entry.key_start
+                                            and entry.position.key_start
                                             -- Use length ("#" operator) as vim jumps to the bytes, not characters
-                                            or entry.entry.value_start
+                                            or entry.position.value_start
                                 }, opts)
                             end,
                         },

From 03723b59cb48ecc92262fda25b2f211c0c550bcd Mon Sep 17 00:00:00 2001
From: Myzel394 <50424412+Myzel394@users.noreply.github.com>
Date: Sun, 14 Apr 2024 10:33:49 +0200
Subject: [PATCH 3/6] fix: Fix LSP keys behavior

---
 lua/telescope/_extensions/jsonfly.lua | 32 +++++++++++++--------------
 1 file changed, 16 insertions(+), 16 deletions(-)

diff --git a/lua/telescope/_extensions/jsonfly.lua b/lua/telescope/_extensions/jsonfly.lua
index bac8ac8..aa923ac 100644
--- a/lua/telescope/_extensions/jsonfly.lua
+++ b/lua/telescope/_extensions/jsonfly.lua
@@ -42,8 +42,6 @@ local entry_display = require "telescope.pickers.entry_display"
 local function get_entries_from_lua_json(t)
     local keys = {}
 
-    --@type k string
-    --@type raw_value InputEntry
     for k, raw_value in pairs(t) do
         ---@type Entry
         local entry = {
@@ -133,6 +131,22 @@ local function get_entries_from_lsp_symbols(symbols)
     for _, symbol in ipairs(symbols) do
         local key = symbol.name
 
+        ---@type Entry
+        local entry = {
+            key = key,
+            value = parse_lsp_value(symbol),
+            position = {
+                line_number = symbol.range.start.line + 2,
+                key_start = symbol.range.start.character + 2,
+                -- The LSP doesn't return the start of the value, so we'll just assume it's 3 characters after the key
+                -- We assume a default JSON file like:
+                -- `"my_key": "my_value"`
+                -- Since we get the end of the key, we can just add 4 to get the start of the value
+                value_start = symbol.selectionRange["end"].character + 4
+            }
+        }
+        table.insert(keys, entry)
+
         if symbol.kind == 2 then
             local sub_keys = get_entries_from_lsp_symbols(symbol.children)
 
@@ -147,18 +161,6 @@ local function get_entries_from_lsp_symbols(symbols)
                 table.insert(keys, entry)
             end
         end
-
-        ---@type Entry
-        local entry = {
-            key = key,
-            value = parse_lsp_value(symbol),
-            position = {
-                line_number = symbol.range["end"].line,
-                key_start = symbol.range.start.character,
-                value_start = symbol.selectionRange.start.character,
-            }
-        }
-        table.insert(keys, entry)
     end
 
     return keys
@@ -284,8 +286,6 @@ return require"telescope".register_extension {
                                 local _, raw_depth = entry.key:gsub("%.", ".")
                                 local depth = (raw_depth or 0) + 1
 
-                                print(vim.inspect(entry))
-
                                 return make_entry.set_default_entry_mt({
                                     value = current_buf,
                                     ordinal = entry.key,

From a923438f6781469c3d9d9a0264a0c852682d3d3d Mon Sep 17 00:00:00 2001
From: Myzel394 <50424412+Myzel394@users.noreply.github.com>
Date: Sun, 14 Apr 2024 12:01:28 +0200
Subject: [PATCH 4/6] refactor: Outsource functions into own modules

---
 lua/jsonfly/parsers.lua               | 142 +++++++++++++++++
 lua/jsonfly/utils.lua                 |  64 ++++++++
 lua/telescope/_extensions/jsonfly.lua | 220 ++------------------------
 3 files changed, 218 insertions(+), 208 deletions(-)
 create mode 100644 lua/jsonfly/parsers.lua
 create mode 100644 lua/jsonfly/utils.lua

diff --git a/lua/jsonfly/parsers.lua b/lua/jsonfly/parsers.lua
new file mode 100644
index 0000000..d86d3bc
--- /dev/null
+++ b/lua/jsonfly/parsers.lua
@@ -0,0 +1,142 @@
+---@class EntryPosition
+---@field line_number number
+---@field key_start number
+---@field value_start number
+--
+---@class Entry
+---@field key string
+---@field value Entry|table|number|string|boolean|nil
+---@field position EntryPosition
+
+local M = {}
+
+---@param t table
+---@return Entry[]
+function M:get_entries_from_lua_json(t)
+    local keys = {}
+
+    for k, raw_value in pairs(t) do
+        ---@type Entry
+        local entry = {
+            key = k,
+            value = raw_value,
+            position = {
+                line_number = raw_value.newlines,
+                key_start = raw_value.key_start,
+                value_start = raw_value.value_start,
+            }
+        }
+        table.insert(keys, entry)
+
+        local v = raw_value.value
+
+        if type(v) == "table" then
+            local sub_keys = M:get_entries_from_lua_json(v)
+
+            for _, sub_key in ipairs(sub_keys) do
+                ---@type Entry
+                local entry = {
+                    key = k .. "." .. sub_key.key,
+                    value = sub_key,
+                    position = sub_key.position,
+                }
+
+                table.insert(keys, entry)
+            end
+        end
+    end
+
+    return keys
+end
+
+---@param result Symbol
+---@return string|number|table|boolean|nil
+function M:parse_lsp_value(result)
+    if result.kind == 2 then
+        local value = {}
+
+        for _, child in ipairs(result.children) do
+            value[child.name] = M:parse_lsp_value(child)
+        end
+
+        return value
+    elseif result.kind == 16 then
+        return tonumber(result.detail)
+    elseif result.kind == 15 then
+        return result.detail
+    elseif result.kind == 18 then
+        local value = {}
+
+        for i, child in ipairs(result.children) do
+            value[i] = M:parse_lsp_value(child)
+        end
+
+        return value
+    elseif result.kind == 13 then
+        return nil
+    elseif result.kind == 17 then
+        return result.detail == "true"
+    end
+end
+
+
+---@class Symbol
+---@field name string
+---@field kind number 2 = Object, 16 = Number, 15 = String, 18 = Array, 13 = Null, 17 = Boolean
+---@field range Range
+---@field selectionRange Range
+---@field detail string
+---@field children Symbol[]
+--
+---@class Range
+---@field start Position
+---@field ["end"] Position
+--
+---@class Position
+---@field line number
+---@field character number
+--
+---@param symbols Symbol[]
+---@return Entry[]
+function M:get_entries_from_lsp_symbols(symbols)
+    local keys = {}
+
+    for _, symbol in ipairs(symbols) do
+        local key = symbol.name
+
+        ---@type Entry
+        local entry = {
+            key = key,
+            value = M:parse_lsp_value(symbol),
+            position = {
+                line_number = symbol.range.start.line + 2,
+                key_start = symbol.range.start.character + 2,
+                -- The LSP doesn't return the start of the value, so we'll just assume it's 3 characters after the key
+                -- We assume a default JSON file like:
+                -- `"my_key": "my_value"`
+                -- Since we get the end of the key, we can just add 4 to get the start of the value
+                value_start = symbol.selectionRange["end"].character + 4
+            }
+        }
+        table.insert(keys, entry)
+
+        if symbol.kind == 2 then
+            local sub_keys = M:get_entries_from_lsp_symbols(symbol.children)
+
+            for _, sub_key in ipairs(sub_keys) do
+                ---@type Entry
+                local entry = {
+                    key = key .. "." .. sub_key.key,
+                    value = sub_key.value,
+                    position = sub_key.position,
+                }
+
+                table.insert(keys, entry)
+            end
+        end
+    end
+
+    return keys
+end
+
+return M
diff --git a/lua/jsonfly/utils.lua b/lua/jsonfly/utils.lua
new file mode 100644
index 0000000..1e2ffbb
--- /dev/null
+++ b/lua/jsonfly/utils.lua
@@ -0,0 +1,64 @@
+local M = {}
+
+function M:truncate_overflow(value, max_length, overflow_marker)
+    if vim.fn.strdisplaywidth(value) > max_length then
+        return value:sub(1, max_length - vim.fn.strdisplaywidth(overflow_marker)) .. overflow_marker
+    end
+
+    return value
+end
+
+---@param value any
+---@param opts Options
+function M:create_display_preview(value, opts)
+    local t = type(value)
+    local conceal
+
+    if opts.conceal == "auto" then
+        conceal = vim.o.conceallevel > 0
+    else
+        conceal = opts.conceal
+    end
+
+    if t == "table" then
+        local preview_table = {}
+
+        for k, v in pairs(value) do
+            table.insert(preview_table, k .. ": " .. M:create_display_preview(v, opts))
+        end
+
+        return "{ " .. table.concat(preview_table, ", ") .. " }", "other"
+    elseif t == "nil" then
+        return "null", "null"
+    elseif t == "number" then
+        return tostring(value), "number"
+    elseif t == "string" then
+        if conceal then
+            return value, "string"
+        else
+            return "\"" .. value .. "\"", "string"
+        end
+    elseif t == "boolean" then
+        return value and "true" or "false", "boolean"
+    end
+end
+
+---@param key string
+---@param replacement string
+---@return string
+---Replaces all previous keys with the replacement
+---Example: replace_previous_keys("a.b.c", "x") => "xxx.c"
+function M:replace_previous_keys(key, replacement)
+    for i = #key, 1, -1 do
+        if key:sub(i, i) == "." then
+            local len = i - 1
+            local before = replacement:rep(len)
+
+            return before .. "." .. key:sub(i + 1)
+        end
+    end
+
+    return key
+end
+
+return M
diff --git a/lua/telescope/_extensions/jsonfly.lua b/lua/telescope/_extensions/jsonfly.lua
index aa923ac..657f080 100644
--- a/lua/telescope/_extensions/jsonfly.lua
+++ b/lua/telescope/_extensions/jsonfly.lua
@@ -10,6 +10,7 @@
 ---@field highlights Highlights - Highlight groups for different types
 ---@field jump_behavior "key_start"|"value_start" - Behavior for jumping to the location, "key_start" == Jump to the start of the key, "value_start" == Jump to the start of the value, Default: "key_start"
 ---@field subkeys_display "normal"|"waterfall" - Display subkeys in a normal or waterfall style, Default: "normal"
+---@field backend "lua"|"lsp" - Backend to use for parsing JSON, "lua" = Use our own Lua parser to parse the JSON, "lsp" = Use your LSP to parse the JSON (currently only https://github.com/Microsoft/vscode-json-languageservice is supported). If the "lsp" backend is selected but the LSP fails, it will fallback to the "lua" backend, Default: "lsp"
 ---
 ---@class Highlights
 ---@field number string - Highlight group for numbers, Default: "@number.json"
@@ -18,215 +19,17 @@
 ---@field null string - Highlight group for null values, Default: "@constant.builtin.json"
 ---@field other string - Highlight group for other types, Default: "@label.json"
 
----- Types below are for internal use only ----
---
----@class EntryPosition
----@field line_number number
----@field key_start number
----@field value_start number
---
----@class Entry
----@field key string
----@field value Entry|table|number|string|boolean|nil
----@field position EntryPosition
+local parsers = require"jsonfly.parsers"
+local json = require"jsonfly.json"
+local utils = require"jsonfly.utils"
 
 local json = require"jsonfly.json"
 local finders = require "telescope.finders"
 local pickers = require "telescope.pickers"
-local conf = require("telescope.config").values
+local conf = require"telescope.config".values
 local make_entry = require "telescope.make_entry"
 local entry_display = require "telescope.pickers.entry_display"
 
----@param t table
----@return Entry[]
-local function get_entries_from_lua_json(t)
-    local keys = {}
-
-    for k, raw_value in pairs(t) do
-        ---@type Entry
-        local entry = {
-            key = k,
-            value = raw_value,
-            position = {
-                line_number = raw_value.newlines,
-                key_start = raw_value.key_start,
-                value_start = raw_value.value_start,
-            }
-        }
-        table.insert(keys, entry)
-
-        local v = raw_value.value
-
-        if type(v) == "table" then
-            local sub_keys = get_entries_from_lua_json(v)
-
-            for _, sub_key in ipairs(sub_keys) do
-                ---@type Entry
-                local entry = {
-                    key = k .. "." .. sub_key.key,
-                    value = sub_key,
-                    position = sub_key.position,
-                }
-
-                table.insert(keys, entry)
-            end
-        end
-    end
-
-    return keys
-end
-
----@param result Symbol
----@return string|number|table|boolean|nil
-local function parse_lsp_value(result)
-    if result.kind == 2 then
-        local value = {}
-
-        for _, child in ipairs(result.children) do
-            value[child.name] = parse_lsp_value(child)
-        end
-
-        return value
-    elseif result.kind == 16 then
-        return tonumber(result.detail)
-    elseif result.kind == 15 then
-        return result.detail
-    elseif result.kind == 18 then
-        local value = {}
-
-        for i, child in ipairs(result.children) do
-            value[i] = parse_lsp_value(child)
-        end
-
-        return value
-    elseif result.kind == 13 then
-        return nil
-    elseif result.kind == 17 then
-        return result.detail == "true"
-    end
-end
-
-
----@class Symbol
----@field name string
----@field kind number 2 = Object, 16 = Number, 15 = String, 18 = Array, 13 = Null, 17 = Boolean
----@field range Range
----@field selectionRange Range
----@field detail string
----@field children Symbol[]
---
----@class Range
----@field start Position
----@field ["end"] Position
---
----@class Position
----@field line number
----@field character number
---
----@param symbols Symbol[]
----@return Entry[]
-local function get_entries_from_lsp_symbols(symbols)
-    local keys = {}
-
-    for _, symbol in ipairs(symbols) do
-        local key = symbol.name
-
-        ---@type Entry
-        local entry = {
-            key = key,
-            value = parse_lsp_value(symbol),
-            position = {
-                line_number = symbol.range.start.line + 2,
-                key_start = symbol.range.start.character + 2,
-                -- The LSP doesn't return the start of the value, so we'll just assume it's 3 characters after the key
-                -- We assume a default JSON file like:
-                -- `"my_key": "my_value"`
-                -- Since we get the end of the key, we can just add 4 to get the start of the value
-                value_start = symbol.selectionRange["end"].character + 4
-            }
-        }
-        table.insert(keys, entry)
-
-        if symbol.kind == 2 then
-            local sub_keys = get_entries_from_lsp_symbols(symbol.children)
-
-            for _, sub_key in ipairs(sub_keys) do
-                ---@type Entry
-                local entry = {
-                    key = key .. "." .. sub_key.key,
-                    value = sub_key.value,
-                    position = sub_key.position,
-                }
-
-                table.insert(keys, entry)
-            end
-        end
-    end
-
-    return keys
-end
-
-local function truncate_overflow(value, max_length, overflow_marker)
-    if vim.fn.strdisplaywidth(value) > max_length then
-        return value:sub(1, max_length - vim.fn.strdisplaywidth(overflow_marker)) .. overflow_marker
-    end
-
-    return value
-end
-
----@param value any
----@param opts Options
-local function create_display_preview(value, opts)
-    local t = type(value)
-    local conceal
-
-    if opts.conceal == "auto" then
-        conceal = vim.o.conceallevel > 0
-    else
-        conceal = opts.conceal
-    end
-
-    if t == "table" then
-        local preview_table = {}
-
-        for k, v in pairs(value) do
-            table.insert(preview_table, k .. ": " .. create_display_preview(v, opts))
-        end
-
-        return "{ " .. table.concat(preview_table, ", ") .. " }", "other"
-    elseif t == "nil" then
-        return "null", "null"
-    elseif t == "number" then
-        return tostring(value), "number"
-    elseif t == "string" then
-        if conceal then
-            return value, "string"
-        else
-            return "\"" .. value .. "\"", "string"
-        end
-    elseif t == "boolean" then
-        return value and "true" or "false", "boolean"
-    end
-end
-
----@param key string
----@param replacement string
----@return string
----Replaces all previous keys with the replacement
----Example: replace_previous_keys("a.b.c", "x") => "xxx.c"
-local function replace_previous_keys(key, replacement)
-    for i = #key, 1, -1 do
-        if key:sub(i, i) == "." then
-            local len = i - 1
-            local before = replacement:rep(len)
-
-            return before .. "." .. key:sub(i + 1)
-        end
-    end
-
-    return key
-end
-
 ---@type Options
 local opts = {
     key_max_length = 50,
@@ -244,6 +47,7 @@ local opts = {
     },
     jump_behavior = "key_start",
     subkeys_display = "normal",
+    backend = "lsp",
 }
 
 return require"telescope".register_extension {
@@ -258,7 +62,7 @@ return require"telescope".register_extension {
             local content = table.concat(content_lines, "\n")
 
             local parsed = json:decode(content)
-            local keys = get_entries_from_lua_json(parsed)
+            local keys = parsers:get_entries_from_lua_json(parsed)
 
             local displayer = entry_display.create {
                 separator = " ",
@@ -275,7 +79,7 @@ return require"telescope".register_extension {
                 "textDocument/documentSymbol",
                 params,
                 function(_, result)
-                    local keys = get_entries_from_lsp_symbols(result)
+                    local keys = parsers:get_entries_from_lsp_symbols(result)
 
                     pickers.new(opts, {
                         prompt_title = opts.prompt_title,
@@ -290,14 +94,14 @@ return require"telescope".register_extension {
                                     value = current_buf,
                                     ordinal = entry.key,
                                     display = function(_)
-                                        local preview, hl_group_key = create_display_preview(entry.value, opts)
+                                        local preview, hl_group_key = utils:create_display_preview(entry.value, opts)
 
-                                        local key = opts.subkeys_display == "normal" and entry.key or replace_previous_keys(entry.key, " ")
+                                        local key = opts.subkeys_display == "normal" and entry.key or utils:replace_previous_keys(entry.key, " ")
 
                                         return displayer {
                                             { depth, "TelescopeResultsNumber"},
                                             {
-                                                 truncate_overflow(
+                                                utils:truncate_overflow(
                                                     key,
                                                     opts.key_max_length,
                                                     opts.overflow_marker
@@ -305,7 +109,7 @@ return require"telescope".register_extension {
                                                 "@property.json",
                                             },
                                             {
-                                                truncate_overflow(
+                                                utils:truncate_overflow(
                                                     preview,
                                                     opts.max_length,
                                                     opts.overflow_marker

From e4bd8b0b64e78dd46d826521023fb5ba4ba53317 Mon Sep 17 00:00:00 2001
From: Myzel394 <50424412+Myzel394@users.noreply.github.com>
Date: Sun, 14 Apr 2024 12:23:21 +0200
Subject: [PATCH 5/6] fix: Fix lua parser

---
 lua/jsonfly/parsers.lua               |  37 ++++-
 lua/telescope/_extensions/jsonfly.lua | 217 +++++++++++---------------
 2 files changed, 126 insertions(+), 128 deletions(-)

diff --git a/lua/jsonfly/parsers.lua b/lua/jsonfly/parsers.lua
index d86d3bc..2ffe8ee 100644
--- a/lua/jsonfly/parsers.lua
+++ b/lua/jsonfly/parsers.lua
@@ -7,21 +7,48 @@
 ---@field key string
 ---@field value Entry|table|number|string|boolean|nil
 ---@field position EntryPosition
+--
+---@class JSONEntry
+---@field value JSONEntry|string|number|boolean|nil
+---@field line_number number
+---@field value_start number
+---@field key_start number
+
 
 local M = {}
 
+---@param entry JSONEntry
+local function get_contents_from_json_value(entry)
+    local value = entry.value
+
+    if type(value) == "table" then
+        -- Recursively get the contents of the table
+        local contents = {}
+
+        for k, v in pairs(value) do
+            contents[k] = get_contents_from_json_value(v)
+        end
+
+        return contents
+    else
+        return entry.value
+    end
+end
+
 ---@param t table
 ---@return Entry[]
 function M:get_entries_from_lua_json(t)
     local keys = {}
 
-    for k, raw_value in pairs(t) do
+    for k, _raw_value in pairs(t) do
+        ---@type JSONEntry
+        local raw_value = _raw_value
         ---@type Entry
         local entry = {
             key = k,
-            value = raw_value,
+            value = get_contents_from_json_value(raw_value),
             position = {
-                line_number = raw_value.newlines,
+                line_number = raw_value.line_number,
                 key_start = raw_value.key_start,
                 value_start = raw_value.value_start,
             }
@@ -37,7 +64,7 @@ function M:get_entries_from_lua_json(t)
                 ---@type Entry
                 local entry = {
                     key = k .. "." .. sub_key.key,
-                    value = sub_key,
+                    value = get_contents_from_json_value(sub_key),
                     position = sub_key.position,
                 }
 
@@ -109,7 +136,7 @@ function M:get_entries_from_lsp_symbols(symbols)
             key = key,
             value = M:parse_lsp_value(symbol),
             position = {
-                line_number = symbol.range.start.line + 2,
+                line_number = symbol.range.start.line + 1,
                 key_start = symbol.range.start.character + 2,
                 -- The LSP doesn't return the start of the value, so we'll just assume it's 3 characters after the key
                 -- We assume a default JSON file like:
diff --git a/lua/telescope/_extensions/jsonfly.lua b/lua/telescope/_extensions/jsonfly.lua
index 657f080..8e4cfe9 100644
--- a/lua/telescope/_extensions/jsonfly.lua
+++ b/lua/telescope/_extensions/jsonfly.lua
@@ -47,9 +47,79 @@ local opts = {
     },
     jump_behavior = "key_start",
     subkeys_display = "normal",
-    backend = "lsp",
+    backend = "lua",
 }
 
+---@param results Entry[]
+---@param buffer number
+local function show_picker(results, buffer)
+    local filename = vim.api.nvim_buf_get_name(buffer)
+
+    local displayer = entry_display.create {
+        separator = " ",
+        items = {
+            { width = 1 },
+            opts.key_exact_length and { width = opts.key_max_length } or { remaining = true },
+            { remaining = true },
+        },
+    }
+
+    pickers.new(opts, {
+        prompt_title = opts.prompt_title,
+        finder = finders.new_table {
+            results = results,
+            ---@param entry Entry
+            entry_maker = function(entry)
+                local _, raw_depth = entry.key:gsub("%.", ".")
+                local depth = (raw_depth or 0) + 1
+
+                return make_entry.set_default_entry_mt({
+                    value = buffer,
+                    ordinal = entry.key,
+                    display = function(_)
+                        local preview, hl_group_key = utils:create_display_preview(entry.value, opts)
+
+                        local key = opts.subkeys_display == "normal" and entry.key or utils:replace_previous_keys(entry.key, " ")
+
+                        print(vim.inspect(entry))
+
+                        return displayer {
+                            { depth, "TelescopeResultsNumber"},
+                            {
+                                utils:truncate_overflow(
+                                    key,
+                                    opts.key_max_length,
+                                    opts.overflow_marker
+                                ),
+                                "@property.json",
+                            },
+                            {
+                                utils:truncate_overflow(
+                                    preview,
+                                    opts.max_length,
+                                    opts.overflow_marker
+                                ),
+                                opts.highlights[hl_group_key] or "TelescopeResultsString",
+                            },
+                        }
+                    end,
+
+                    bufnr = buffer,
+                    filename = filename,
+                    lnum = entry.position.line_number,
+                    col = opts.jump_behavior == "key_start"
+                            and entry.position.key_start
+                            -- Use length ("#" operator) as vim jumps to the bytes, not characters
+                            or entry.position.value_start
+                }, opts)
+            end,
+        },
+        previewer = conf.grep_previewer(opts),
+        sorter = conf.generic_sorter(opts),
+        sorting_strategy = "ascending",
+    }):find()
+end
+
 return require"telescope".register_extension {
     setup = function(extension_config)
         opts = vim.tbl_deep_extend("force", opts, extension_config or {})
@@ -57,136 +127,37 @@ return require"telescope".register_extension {
     exports = {
         jsonfly = function(xopts)
             local current_buf = vim.api.nvim_get_current_buf()
-            local filename = vim.api.nvim_buf_get_name(current_buf)
             local content_lines = vim.api.nvim_buf_get_lines(current_buf, 0, -1, false)
             local content = table.concat(content_lines, "\n")
 
-            local parsed = json:decode(content)
-            local keys = parsers:get_entries_from_lua_json(parsed)
+            function run_lua_parser()
+                local parsed = json:decode(content)
+                local keys = parsers:get_entries_from_lua_json(parsed)
 
-            local displayer = entry_display.create {
-                separator = " ",
-                items = {
-                    { width = 1 },
-                    opts.key_exact_length and { width = opts.key_max_length } or { remaining = true },
-                    { remaining = true },
-                },
-            }
+                show_picker(keys, current_buf)
+            end
 
-            local params = vim.lsp.util.make_position_params(xopts.winnr)
-            local result = vim.lsp.buf_request(
-                current_buf,
-                "textDocument/documentSymbol",
-                params,
-                function(_, result)
-                    local keys = parsers:get_entries_from_lsp_symbols(result)
+            if opts.backend == "lsp" then
+                local params = vim.lsp.util.make_position_params(xopts.winnr)
 
-                    pickers.new(opts, {
-                        prompt_title = opts.prompt_title,
-                        finder = finders.new_table {
-                            results = keys,
-                            ---@param entry Entry
-                            entry_maker = function(entry)
-                                local _, raw_depth = entry.key:gsub("%.", ".")
-                                local depth = (raw_depth or 0) + 1
+                vim.lsp.buf_request(
+                    current_buf,
+                    "textDocument/documentSymbol",
+                    params,
+                    function(error, lsp_response)
+                        if error then
+                            run_lua_parser()
+                            return
+                        end
 
-                                return make_entry.set_default_entry_mt({
-                                    value = current_buf,
-                                    ordinal = entry.key,
-                                    display = function(_)
-                                        local preview, hl_group_key = utils:create_display_preview(entry.value, opts)
+                        local result = parsers:get_entries_from_lsp_symbols(lsp_response)
 
-                                        local key = opts.subkeys_display == "normal" and entry.key or utils:replace_previous_keys(entry.key, " ")
-
-                                        return displayer {
-                                            { depth, "TelescopeResultsNumber"},
-                                            {
-                                                utils:truncate_overflow(
-                                                    key,
-                                                    opts.key_max_length,
-                                                    opts.overflow_marker
-                                                ),
-                                                "@property.json",
-                                            },
-                                            {
-                                                utils:truncate_overflow(
-                                                    preview,
-                                                    opts.max_length,
-                                                    opts.overflow_marker
-                                                ),
-                                                opts.highlights[hl_group_key] or "TelescopeResultsString",
-                                            },
-                                        }
-                                    end,
-
-                                    bufnr = current_buf,
-                                    filename = filename,
-                                    lnum = entry.position.line_number,
-                                    col = opts.jump_behavior == "key_start"
-                                            and entry.position.key_start
-                                            -- Use length ("#" operator) as vim jumps to the bytes, not characters
-                                            or entry.position.value_start
-                                }, opts)
-                            end,
-                        },
-                        previewer = conf.grep_previewer(opts),
-                        sorter = conf.generic_sorter(opts),
-                        sorting_strategy = "ascending",
-                    }):find()
-                end
-            )
-
-            -- pickers.new(opts, {
-            --     prompt_title = opts.prompt_title,
-            --     finder = finders.new_table {
-            --         results = keys,
-            --         entry_maker = function(entry)
-            --             local _, raw_depth = entry.key:gsub("%.", ".")
-            --             local depth = (raw_depth or 0) + 1
-            --
-            --             return make_entry.set_default_entry_mt({
-            --                 value = current_buf,
-            --                 ordinal = entry.key,
-            --                 display = function(_)
-            --                     local preview, hl_group_key = create_display_preview(entry.entry.value, opts)
-            --
-            --                     local key = opts.subkeys_display == "normal" and entry.key or replace_previous_keys(entry.key, " ")
-            --
-            --                     return displayer {
-            --                         { depth, "TelescopeResultsNumber"},
-            --                         {
-            --                              truncate_overflow(
-            --                                 key,
-            --                                 opts.key_max_length,
-            --                                 opts.overflow_marker
-            --                             ),
-            --                             "@property.json",
-            --                         },
-            --                         {
-            --                             truncate_overflow(
-            --                                 preview,
-            --                                 opts.max_length,
-            --                                 opts.overflow_marker
-            --                             ),
-            --                             opts.highlights[hl_group_key] or "TelescopeResultsString",
-            --                         },
-            --                     }
-            --                 end,
-            --
-            --                 bufnr = current_buf,
-            --                 filename = filename,
-            --                 lnum = entry.entry.newlines + 1,
-            --                 col = opts.jump_behavior == "key_start"
-            --                         and entry.entry.key_start
-            --                         -- Use length ("#" operator) as vim jumps to the bytes, not characters
-            --                         or entry.entry.value_start
-            --             }, opts)
-            --         end,
-            --     },
-            --     previewer = conf.grep_previewer(opts),
-            --     sorter = conf.generic_sorter(opts),
-            --     sorting_strategy = "ascending",
-            -- }):find()
+                        show_picker(result, current_buf)
+                    end
+                )
+            else
+                run_lua_parser()
+            end
         end
     }
 }

From 7e1fefb841caa522d7f46ca14bce6e23a7328493 Mon Sep 17 00:00:00 2001
From: Myzel394 <50424412+Myzel394@users.noreply.github.com>
Date: Sun, 14 Apr 2024 12:52:06 +0200
Subject: [PATCH 6/6] fix: Fix value_start position

---
 lua/jsonfly/parsers.lua               | 2 +-
 lua/telescope/_extensions/jsonfly.lua | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/lua/jsonfly/parsers.lua b/lua/jsonfly/parsers.lua
index 2ffe8ee..45be2f6 100644
--- a/lua/jsonfly/parsers.lua
+++ b/lua/jsonfly/parsers.lua
@@ -142,7 +142,7 @@ function M:get_entries_from_lsp_symbols(symbols)
                 -- We assume a default JSON file like:
                 -- `"my_key": "my_value"`
                 -- Since we get the end of the key, we can just add 4 to get the start of the value
-                value_start = symbol.selectionRange["end"].character + 4
+                value_start = symbol.selectionRange["end"].character + 3,
             }
         }
         table.insert(keys, entry)
diff --git a/lua/telescope/_extensions/jsonfly.lua b/lua/telescope/_extensions/jsonfly.lua
index 8e4cfe9..61db420 100644
--- a/lua/telescope/_extensions/jsonfly.lua
+++ b/lua/telescope/_extensions/jsonfly.lua
@@ -47,7 +47,7 @@ local opts = {
     },
     jump_behavior = "key_start",
     subkeys_display = "normal",
-    backend = "lua",
+    backend = "lsp",
 }
 
 ---@param results Entry[]