feat: better mount and reset lines on random and daily

Author: kawre

lua/leetcode-ui/lines/button/menu/exit.lua

@@ -8,7 +8,7 @@ function MenuExitButton:init()
     MenuExitButton.super.init(self, "Exit", {
         icon = "󰩈",
         sc = "q",
-        on_press = vim.cmd.quitall,
+        on_press = function() vim.cmd("qa!") end,
     })
 end
 

lua/leetcode-ui/question.lua

@@ -16,10 +16,11 @@ local log = require("leetcode.logger")
 ---@field console lc.ui.Console
 ---@field lang string
 ---@field cache lc.cache.Question
+---@field reset boolean
 local Question = Object("LeetQuestion")
 
 ---@param raw? boolean
-function Question:get_snippet(raw)
+function Question:snippet(raw)
     local snippets = self.q.code_snippets ~= vim.NIL and self.q.code_snippets or {}
     local snip = vim.tbl_filter(function(snip) return snip.lang_slug == self.lang end, snippets)[1]
     if not snip then return end
@@ -28,7 +29,30 @@ function Question:get_snippet(raw)
     return raw and code or self:injector(code)
 end
 
-function Question:create_file()
+---@param code? string
+function Question:set_lines(code)
+    if not vim.api.nvim_buf_is_valid(self.bufnr) then return end
+
+    pcall(vim.cmd.undojoin)
+    local s_i, e_i = self:range()
+    code = code and code or (self:snippet(true) or "")
+    vim.api.nvim_buf_set_lines(self.bufnr, s_i - 1, e_i, false, vim.split(code, "\n"))
+end
+
+function Question:reset_lines()
+    if not self.reset then return end
+
+    local new_lines = self:snippet(true) or ""
+
+    vim.schedule(function() --
+        log.info("Previous code found, resetting…\nTo undo, simply press `u`")
+    end)
+
+    self:set_lines(new_lines)
+end
+
+---@return string path, boolean? existed
+function Question:path()
     local lang = utils.get_lang(self.lang)
     local alt = lang.alt and ("." .. lang.alt) or ""
 
@@ -38,47 +62,39 @@ function Question:create_file()
     self.file = config.storage.home:joinpath(fn_legacy)
 
     if self.file:exists() then --
-        return self.file:absolute()
+        return self.file:absolute(), true
     end
 
     local fn = ("%s.%s%s.%s"):format(self.q.frontend_id, self.q.title_slug, alt, lang.ft)
     self.file = config.storage.home:joinpath(fn)
+    local existed = true
 
     if not self.file:exists() then --
-        self.file:write(self:get_snippet(), "w")
+        self.file:write(self:snippet(), "w")
+        existed = false
     end
 
-    return self.file:absolute()
+    return self.file:absolute(), existed
 end
 
----@param new_tabp? boolean
----@return boolean was_loaded
-function Question:create_buffer(new_tabp)
-    local file_name = self:create_file()
+function Question:create_buffer()
+    local path, existed = self:path()
 
-    local buf = vim.fn.bufadd(file_name)
-    assert(buf ~= 0, "Failed to create buffer")
+    vim.cmd("$tabe " .. path)
 
-    self.bufnr = buf
-    if vim.fn.bufloaded(self.bufnr) == 1 then
-        return true
-    else
-        vim.fn.bufload(self.bufnr)
-    end
+    self.bufnr = vim.api.nvim_get_current_buf()
+    self.winid = vim.api.nvim_get_current_win()
 
-    local cmd
-    if new_tabp then
-        cmd = ("$tabe %s"):format(file_name)
-    else
-        cmd = ("edit %s"):format(file_name)
-    end
+    vim.api.nvim_set_option_value("buflisted", true, { buf = self.bufnr })
 
     local i = self:fold_range()
-    if i then cmd = cmd .. (" | %d,%dfold"):format(1, i) end
-
-    vim.api.nvim_exec2(cmd, {})
+    if i then --
+        pcall(vim.cmd, ("%d,%dfold"):format(1, i))
+    end
 
-    return false
+    if existed then --
+        self:reset_lines()
+    end
 end
 
 ---@param before boolean
@@ -99,7 +115,7 @@ function Question:inject(before)
     end
 
     if res and res ~= "" then
-        return before and (res .. "\n\n") or ("\n\n" .. res)
+        return before and (res .. "\n") or ("\n" .. res)
     else
         return nil
     end
@@ -109,50 +125,62 @@ end
 function Question:injector(code)
     local lang = utils.get_lang(self.lang)
 
-    local inj_before = self:inject(true) or ""
-    local inj_after = self:inject(false) or ""
+    local parts = {
+        ("%s @leet start"):format(lang.comment),
+        code,
+        ("%s @leet end"):format(lang.comment),
+    }
 
-    return inj_before --
-        .. ("%s @leet start\n"):format(lang.comment)
-        .. code
-        .. ("\n%s @leet end"):format(lang.comment)
-        .. inj_after
-end
+    local before = self:inject(true)
+    if before then table.insert(parts, 1, before) end
 
----@param pre? boolean
-function Question:_unmount(pre)
-    self.info:unmount()
-    self.console:unmount()
-    self.description:unmount()
+    local after = self:inject(false)
+    if after then table.insert(parts, after) end
 
-    if not pre and vim.api.nvim_buf_is_valid(self.bufnr) then
-        vim.api.nvim_buf_delete(self.bufnr, { force = true })
-    end
+    return table.concat(parts, "\n")
+end
 
-    if vim.api.nvim_win_is_valid(self.winid) then --
-        vim.api.nvim_win_close(self.winid, true)
+function Question:unmount()
+    if vim.v.dying ~= 0 then --
+        return
     end
 
-    _Lc_questions = vim.tbl_filter(function(q) return q.bufnr ~= self.bufnr end, _Lc_questions)
+    vim.schedule(function()
+        self.info:unmount()
+        self.console:unmount()
+        self.description:unmount()
+
+        if vim.api.nvim_buf_is_valid(self.bufnr) then
+            vim.api.nvim_buf_delete(self.bufnr, { force = true, unload = false })
+        end
+
+        _Lc_questions = vim.tbl_filter(function(q) --
+            return q.bufnr ~= self.bufnr
+        end, _Lc_questions)
 
-    self = nil
+        self = nil
+    end)
 end
 
----@param self lc.ui.Question
----@param pre? boolean
-Question.unmount = vim.schedule_wrap(function(self, pre) self:_unmount(pre) end)
+function Question:_unmount()
+    if vim.api.nvim_win_is_valid(self.winid) then vim.api.nvim_win_close(self.winid, true) end
+end
 
-function Question:handle_mount()
-    self:create_buffer(true)
+local group = vim.api.nvim_create_augroup("leetcode_questions", { clear = true })
+function Question:autocmds()
+    vim.api.nvim_create_autocmd("WinClosed", {
+        group = group,
+        buffer = self.bufnr,
+        callback = function() self:unmount() end,
+    })
+end
 
-    self.winid = vim.api.nvim_get_current_win()
+function Question:handle_mount()
+    self:create_buffer()
 
     table.insert(_Lc_questions, self)
 
-    vim.api.nvim_create_autocmd("QuitPre", {
-        buffer = self.bufnr,
-        callback = function() self:unmount(true) end,
-    })
+    self:autocmds()
 
     self.description = Description(self):mount()
     self.console = Console(self)
@@ -173,7 +201,7 @@ function Question:mount()
     end
     self.q = q
 
-    if self:get_snippet() then
+    if self:snippet() then
         self:handle_mount()
     else
         local msg = ("Snippet for `%s` not found. Select a different language"):format(self.lang)
@@ -232,37 +260,55 @@ end
 ---@param self lc.ui.Question
 ---@param lang lc.lang
 Question.change_lang = vim.schedule_wrap(function(self, lang)
-    if vim.api.nvim_get_current_win() ~= self.winid then
-        vim.api.nvim_set_current_win(self.winid)
-    end
-
     local old_lang, old_bufnr = self.lang, self.bufnr
-    self.lang = lang
 
-    local ok, was_loaded = pcall(Question.create_buffer, self)
-    if ok then
-        vim.api.nvim_buf_set_option(old_bufnr, "buflisted", false)
-        vim.api.nvim_buf_set_option(self.bufnr, "buflisted", true)
+    local ok, err = pcall(function()
+        self.lang = lang
+        local path, existed = self:path()
+
+        self.bufnr = vim.fn.bufadd(path)
+        assert(self.bufnr ~= 0, "Failed to create buffer " .. path)
 
-        if was_loaded then
-            vim.api.nvim_win_set_buf(self.winid, self.bufnr)
-        else
+        local loaded = vim.api.nvim_buf_is_loaded(self.bufnr)
+        vim.fn.bufload(self.bufnr)
+
+        vim.api.nvim_win_set_buf(self.winid, self.bufnr)
+
+        vim.api.nvim_set_option_value("buflisted", false, { buf = old_bufnr })
+        vim.api.nvim_set_option_value("buflisted", true, { buf = self.bufnr })
+
+        local i = self:fold_range()
+        if i then --
+            pcall(vim.cmd, ("%d,%dfold"):format(1, i))
+        end
+
+        if existed then --
+            self:reset_lines()
+        end
+
+        if not loaded then --
             utils.exec_hook("question_enter", self)
         end
-    else
-        log.error("Changing language failed")
+
+        self:autocmds()
+    end)
+
+    if not ok then
+        log.error("Failed to change language\n" .. err)
         self.lang = old_lang
         self.bufnr = old_bufnr
     end
 end)
 
 ---@param problem lc.cache.Question
-function Question:init(problem)
+---@param reset boolean
+function Question:init(problem, reset)
     self.cache = problem
     self.lang = config.lang
+    self.reset = reset and true or false
 end
 
----@type fun(question: lc.cache.Question): lc.ui.Question
+---@type fun(question: lc.cache.Question, reset?: boolean): lc.ui.Question
 local LeetQuestion = Question
 
 return LeetQuestion

lua/leetcode/command/init.lua

@@ -90,14 +90,14 @@ function cmd.delete_cookie()
     cookie.delete()
 end
 
-cmd.q_close_all = vim.schedule_wrap(function()
+cmd.q_close_all = function()
     local utils = require("leetcode.utils")
     local qs = utils.question_tabs()
 
     for _, tabp in ipairs(qs) do
-        tabp.question:unmount()
+        tabp.question:_unmount()
     end
-end)
+end
 
 cmd.expire = vim.schedule_wrap(function()
     local tabp = api.nvim_get_current_tabpage()
@@ -127,7 +127,7 @@ function cmd.qot()
     problems.question_of_today(function(qot, err)
         if err then return log.err(err) end
         local problemlist = require("leetcode.cache.problemlist")
-        Question(problemlist.get_by_title_slug(qot.title_slug)):mount()
+        Question(problemlist.get_by_title_slug(qot.title_slug), true):mount()
     end)
 end
 
@@ -151,7 +151,7 @@ function cmd.random_question(opts)
 
     local item = problems.get_by_title_slug(q.title_slug) or {}
     local Question = require("leetcode-ui.question")
-    Question(item):mount()
+    Question(item, true):mount()
 end
 
 function cmd.start_with_cmd()
@@ -281,8 +281,7 @@ function cmd.reset()
     local q = utils.curr_question()
     if not q then return end
 
-    local snip = q:get_snippet(true)
-    utils.set_question_lines(q, snip)
+    q:set_lines()
 end
 
 function cmd.last_submit()
@@ -303,8 +302,8 @@ function cmd.last_submit()
             return
         end
 
-        if type(res) == "table" and res.code and api.nvim_buf_is_valid(q.bufnr) then
-            utils.set_question_lines(q, res.code)
+        if type(res) == "table" and res.code then
+            q:set_lines(res.code)
         else
             log.error("Something went wrong")
         end
@@ -359,7 +358,7 @@ function cmd.inject()
         end
 
         if end_i == nil then
-            log.error("`@leet start` not found")
+            log.error("`@leet end` not found")
         else
             local after = q:inject(false)
             if after then

lua/leetcode/pickers/question.lua

@@ -124,7 +124,7 @@ return {
                 }),
                 sorter = conf.generic_sorter(theme),
                 attach_mappings = function(prompt_bufnr, map)
-                    actions.select_default:replace(function()
+                    local function mount_question(reset)
                         local selection = action_state.get_selected_entry()
                         if not selection then return end
 
@@ -134,8 +134,12 @@ return {
                         end
 
                         actions.close(prompt_bufnr)
-                        Question(q):mount()
-                    end)
+                        Question(q, reset):mount()
+                    end
+
+                    actions.select_default:replace(function() mount_question() end)
+                    map({ "n", "i" }, "<C-Enter>", function() mount_question(true) end)
+
                     return true
                 end,
             })