From 1b5f2838099f283857729e820cc05e2b19df7a2c Mon Sep 17 00:00:00 2001
From: "Lucas de A. Vasco" <67849851+LucasAVasco@users.noreply.github.com>
Date: Fri, 12 Jul 2024 13:46:23 -0300
Subject: [PATCH] feat(provider): marker provider (#218)

* feat(provider): marker provider

Add a provider that can fold generic markers. The default ones are:

* 'foldmarker' option
* '#region' and '#endregion' from VS Code
---
 lua/ufo/model/foldingrange.lua |  1 +
 lua/ufo/provider/init.lua      |  2 +-
 lua/ufo/provider/marker.lua    | 81 ++++++++++++++++++++++++++++++++++
 3 files changed, 83 insertions(+), 1 deletion(-)
 create mode 100644 lua/ufo/provider/marker.lua

diff --git a/lua/ufo/model/foldingrange.lua b/lua/ufo/model/foldingrange.lua
index e436434..89d60b9 100644
--- a/lua/ufo/model/foldingrange.lua
+++ b/lua/ufo/model/foldingrange.lua
@@ -2,6 +2,7 @@
 ---| 'comment'
 ---| 'imports'
 ---| 'region'
+---| 'marker'
 
 ---@class UfoFoldingRange
 ---@field startLine number
diff --git a/lua/ufo/provider/init.lua b/lua/ufo/provider/init.lua
index f5302ae..cacf265 100644
--- a/lua/ufo/provider/init.lua
+++ b/lua/ufo/provider/init.lua
@@ -9,7 +9,7 @@ local log = require('ufo.lib.log')
 ---@field modules table
 local Provider = {
     modulePathPrefix = 'ufo.provider.',
-    innerProviders = {'lsp', 'treesitter', 'indent'}
+    innerProviders = {'lsp', 'treesitter', 'indent', 'marker'}
 }
 
 local function needFallback(reason)
diff --git a/lua/ufo/provider/marker.lua b/lua/ufo/provider/marker.lua
new file mode 100644
index 0000000..e8d6401
--- /dev/null
+++ b/lua/ufo/provider/marker.lua
@@ -0,0 +1,81 @@
+local foldingrange = require('ufo.model.foldingrange')
+local bufmanager = require('ufo.bufmanager')
+local utils = require('ufo.utils')
+
+
+-- Provider implementation
+
+local Marker = {}
+
+
+-- Data necessary to query folding ranges from VS Code region folding
+local vs_code_marker = {
+    '#region',     -- Start of marker
+    '#endregion',  -- End of marker
+    'region',      -- Kind to be applied to a VS Code region folding
+}
+
+
+--- Function that returns folds for the provided buffer based in the markers
+--- @param bufnr number Vim buffer number
+--- @return UfoFoldingRange[]|nil Folds List of marker folds in the buffer, or `nil` if they can not be queried
+function Marker.getFolds(bufnr)
+    local buf = bufmanager:get(bufnr)
+    local winid = utils.getWinByBuf(bufnr)
+
+    -- Does not work with buffers or windows that are not managed by UFO
+    if not buf or winid < 0 then
+        return
+    end
+
+    -- Defines the 'start' and 'end' markers that the provider will search, and the kind to apply
+    -- to these markers. Each element of the `markers` list is a list of the 'start', 'end' markers
+    -- and kind applied, in this order. Example: `local markers = { { 'start marker', 'end marker', 'marker kind' } }`
+    -- The search is done by marker pair. One marker pair does not affect the other. So the end marker of `markers[0]`
+    -- will not close the start marker of `markers[1]`, by example.
+    local markers = {
+        vim.fn.split(vim.wo[winid].foldmarker .. ',marker', ','),  -- Configured Vim marker
+        vs_code_marker
+    }
+
+    -- Query the markers, generate the folding ranges and save in the `folds` variable
+    local lines = buf:lines(1, -1)
+    local folds = {}
+
+    for _, marker in ipairs(markers) do
+        local openMarkerLines = {}
+
+        for lineNum, line in ipairs(lines) do
+            -- Open marker
+            local start_column, end_column = line:find(marker[1], 1, true)
+
+            if start_column then
+                table.insert(openMarkerLines, lineNum)
+            end
+
+            -- Close marker
+            start_column = line:find(marker[2], end_column or 1, true)
+
+            if start_column then
+                local relatedOpenMarkerLine = table.remove(openMarkerLines)
+
+                if relatedOpenMarkerLine then
+                    table.insert(
+                    folds,
+                    foldingrange.new(relatedOpenMarkerLine - 1, lineNum - 1, nil, nil, marker[3])
+                    )
+                end
+            end
+        end
+
+        -- Closes all remaining open markers (they will be open to the end of the file)
+        for _, markerStart in ipairs(openMarkerLines) do
+            table.insert(folds, foldingrange.new(markerStart - 1, #lines, nil, nil, marker[3]))
+        end
+    end
+
+    return folds
+end
+
+
+return Marker