nvim-treesitter/lua/nvim-treesitter/fold.lua

local api = vim.api
local tsutils = require "nvim-treesitter.ts_utils"
local query = require "nvim-treesitter.query"
local parsers = require "nvim-treesitter.parsers"

local M = {}

-- This is cached on buf tick to avoid computing that multiple times
-- Especially not for every line in the file when `zx` is hit
local folds_levels = tsutils.memoize_by_buf_tick(function(bufnr)
  local max_fold_level = api.nvim_win_get_option(0, "foldnestmax")
  local parser = parsers.get_parser(bufnr)

  if not parser then
    return {}
  end

  local matches = query.get_capture_matches_recursively(bufnr, function(lang)
    if query.has_folds(lang) then
      return "@fold", "folds"
    elseif query.has_locals(lang) then
      return "@scope", "locals"
    end
  end)

  local levels_tmp = {}

  for _, node in ipairs(matches) do
    local start, _, stop, stop_col = node.node:range()

    if stop_col > 0 then
      stop = stop + 1
    end

    local should_fold = start + 1 < stop -- Only fold for 2+ lines

    -- This can be folded
    -- Fold only multiline nodes that are not exactly the same as previously met folds
    if should_fold and not (levels_tmp[start] and levels_tmp[stop]) then
      levels_tmp[start] = (levels_tmp[start] or 0) + 1
      levels_tmp[stop] = (levels_tmp[stop] or 0) - 1
    end
  end

  local levels = {}
  local current_level = 0

  -- We now have the list of fold opening and closing, fill the gaps and mark where fold start
  for lnum = 0, api.nvim_buf_line_count(bufnr) do
    local prefix = ""
    local shift = levels_tmp[lnum] or 0

    -- Determine if it's the start of a fold
    if levels_tmp[lnum] and shift >= 0 then
      prefix = ">"
    end

    current_level = current_level + shift

    -- Ignore folds greater than max_fold_level
    if current_level > max_fold_level then
      levels[lnum + 1] = max_fold_level
    else
      levels[lnum + 1] = prefix .. tostring(current_level)
    end
  end

  return levels
end)

function M.get_fold_indic(lnum)
  if not parsers.has_parser() or not lnum then
    return "0"
  end

  local buf = api.nvim_get_current_buf()

  local levels = folds_levels(buf) or {}

  return levels[lnum] or "0"
end

return M