irongit

A Neovim plugin enabling you to solve LeetCode problems.

462 lines12 KBLua
1local Description = require("leetcode-ui.split.description")
2local Console = require("leetcode-ui.layout.console")
3local Info = require("leetcode-ui.popup.info")
4local Object = require("nui.object")
5
6local api_question = require("leetcode.api.question")
7local utils = require("leetcode.utils")
8local ui_utils = require("leetcode-ui.utils")
9local config = require("leetcode.config")
10local log = require("leetcode.logger")
11
12---@alias lc.editor.section "imports" | "code"
13
14---@class lc.ui.Question
15---@field file Path
16---@field q lc.question_res
17---@field description lc.ui.Description
18---@field bufnr integer
19---@field console lc.ui.Console
20---@field lang string
21---@field cache lc.cache.Question
22---@field reset boolean
23local Question = Object("LeetQuestion")
24
25---@param raw? boolean
26function Question:snippet(raw)
27 local snippets = self.q.code_snippets ~= vim.NIL and self.q.code_snippets or {}
28 local snip = vim.tbl_filter(function(snip)
29 return snip.lang_slug == self.lang
30 end, snippets)[1]
31 if not snip then
32 return
33 end
34
35 local code = snip.code:gsub("\r\n", "\n")
36 return raw and code or self:injector(code)
37end
38
39---@param start_i integer
40---@param end_i integer
41---@param lines string[]|string
42function Question:editor_set_lines(start_i, end_i, lines)
43 if not (self.bufnr and vim.api.nvim_buf_is_valid(self.bufnr)) then
44 return
45 end
46
47 lines = type(lines) == "string" and vim.split(lines, "\n") or lines ---@cast lines string[]
48 vim.api.nvim_buf_set_lines(self.bufnr, start_i, end_i, false, lines)
49end
50
51function Question:editor_reset()
52 local new_lines = self:snippet() or ""
53 self:editor_set_lines(0, -1, new_lines)
54end
55
56function Question:editor_reset_code()
57 local new_lines = self:snippet(true) or ""
58 self:editor_section_replace(new_lines, "code")
59end
60
61function Question:reset_previous_code()
62 self:editor_reset_code()
63 vim.schedule(function()
64 log.info("Previous code found and reset. To undo, simply press `u` or use `:undo`.")
65 end)
66end
67
68---@return string path, boolean existed
69function Question:path()
70 local lang = utils.get_lang(self.lang)
71 local alt = lang.alt and ("." .. lang.alt) or ""
72
73 -- handle legacy file names first
74 local fn_legacy = --
75 ("%s.%s-%s.%s"):format(self.q.frontend_id, self.q.title_slug, lang.slug, lang.ft)
76 self.file = config.storage.home:joinpath(fn_legacy)
77
78 if self.file:exists() then
79 return self.file:absolute(), true
80 end
81
82 local fn = ("%s.%s%s.%s"):format(self.q.frontend_id, self.q.title_slug, alt, lang.ft)
83 self.file = config.storage.home:joinpath(fn)
84 local existed = self.file:exists()
85
86 if not existed then
87 self.file:write(self:snippet(), "w")
88 end
89
90 return self.file:absolute(), existed
91end
92
93function Question:create_buffer()
94 local path, existed = self:path()
95
96 vim.cmd("$tabe " .. path)
97 self.bufnr = vim.api.nvim_get_current_buf()
98 self.winid = vim.api.nvim_get_current_win()
99 ui_utils.win_set_winfixbuf(self.winid)
100
101 self:open_buffer(existed)
102end
103
104---@param strict? boolean
105function Question:editor_fold_imports(strict)
106 if not (self.bufnr and vim.api.nvim_buf_is_valid(self.bufnr)) then
107 return
108 end
109
110 local range = self:editor_section_range("imports", true)
111 if range.complete or range.end_i then
112 vim.api.nvim_buf_call(self.bufnr, function()
113 pcall(vim.cmd, ("%d,%dfold"):format(range.start_i or 1, range.end_i)) ---@diagnostic disable-line: param-type-mismatch
114 end)
115 elseif strict then
116 range.log_not_found()
117 end
118end
119
120function Question:editor_yank_code()
121 if not (self.bufnr and vim.api.nvim_buf_is_valid(self.bufnr)) then
122 return
123 end
124
125 local range = self:editor_section_range("code")
126 if range:is_valid_or_log() then
127 vim.api.nvim_buf_call(self.bufnr, function()
128 vim.cmd(("%d,%dyank"):format(range.start_i, range.end_i))
129 end)
130 end
131end
132
133---@param existed boolean
134function Question:open_buffer(existed)
135 ui_utils.buf_set_opts(self.bufnr, { buflisted = true })
136 ui_utils.win_set_buf(self.winid, self.bufnr, true)
137
138 vim.cmd([[match DiagnosticHint /@leet/]])
139
140 if config.user.editor.fold_imports then
141 self:editor_fold_imports(false)
142 end
143
144 if config.user.editor.reset_previous_code and (existed and self.cache.status == "ac") then
145 self:reset_previous_code()
146 end
147end
148
149---@return string[]?
150function Question:inject_imports()
151 local inject = config.user.injector[self.lang] or {}
152
153 local imports = inject.imports
154 local default_imports = config.imports[self.lang]
155
156 local function valid_imports(tbl)
157 if vim.tbl_isempty(tbl or {}) then
158 return false
159 end
160 if not vim.islist(tbl) then
161 log.error("Invalid imports format for language: " .. self.lang)
162 return false
163 end
164 return true
165 end
166
167 if not imports then
168 return default_imports
169 end
170
171 if type(imports) == "function" then
172 local overriden = imports(vim.deepcopy(default_imports or {}))
173 if not valid_imports(overriden) then
174 return default_imports
175 end
176 return overriden
177 end
178
179 local resolved = type(imports) == "string" and { imports } or imports
180 if not valid_imports(resolved) then
181 return
182 end
183
184 local merged, seen = {}, {}
185 local function add(list)
186 for _, v in ipairs(list) do
187 if not seen[v] then
188 table.insert(merged, v)
189 seen[v] = true
190 end
191 end
192 end
193
194 add(default_imports)
195 add(resolved)
196
197 return not vim.tbl_isempty(merged) and merged or nil
198end
199
200---@param before boolean
201---@return string?
202function Question:inject(before)
203 local inject = config.user.injector[self.lang] or {}
204 local inj = before and (inject.before or {}) or (inject.after or {})
205
206 local res
207 if type(inj) == "table" then
208 res = table.concat(inj, "\n")
209 elseif type(inj) == "string" then
210 res = inj
211 end
212
213 if res and res ~= "" then
214 return res
215 else
216 return nil
217 end
218end
219
220---@param lines string|string[]
221---@param name lc.editor.section
222---@return string
223function Question:editor_section(lines, name)
224 local comment = utils.get_lang(self.lang).comment
225
226 local start_tag = comment .. " " .. utils.section_tag(name, true)
227 local end_tag = comment .. " " .. utils.section_tag(name, false)
228
229 local str = type(lines) ~= "string" and table.concat(lines, "\n") or lines
230 return table.concat({ start_tag, str, end_tag }, "\n")
231end
232
233---@param lines string[]|string
234---@param name string
235function Question:editor_section_replace(lines, name)
236 local range = self:editor_section_range(name)
237
238 if range:is_valid_or_log() then
239 self:editor_set_lines(range.start_i - 1, range.end_i, lines)
240 end
241end
242
243---@param code string
244function Question:injector(code)
245 local inject = config.user.injector[self.lang] or {}
246
247 local parts = { self:editor_section(code, "code") }
248
249 local before = self:inject(true)
250 if before then
251 table.insert(parts, 1, before)
252 end
253
254 local imports = self:inject_imports()
255 if imports then
256 table.insert(parts, 1, self:editor_section(table.concat(imports, "\n"), "imports"))
257 end
258
259 local after = self:inject(false)
260 if after then
261 table.insert(parts, after)
262 end
263
264 local gap = (inject.gap or 1) + 1
265 return table.concat(parts, ("\n"):rep(gap))
266end
267
268function Question:_unmount()
269 if vim.v.dying ~= 0 then
270 return
271 end
272
273 vim.schedule(function()
274 self.info:unmount()
275 self.console:unmount()
276 self.description:unmount()
277
278 if self.bufnr and vim.api.nvim_buf_is_valid(self.bufnr) then
279 vim.api.nvim_buf_delete(self.bufnr, { force = true, unload = false })
280 end
281
282 _Lc_state.questions = vim.tbl_filter(function(q)
283 return q.bufnr ~= self.bufnr
284 end, _Lc_state.questions)
285
286 self = nil
287 end)
288end
289
290function Question:unmount()
291 if self.winid and vim.api.nvim_win_is_valid(self.winid) then
292 vim.api.nvim_win_close(self.winid, true)
293 end
294end
295
296local group = vim.api.nvim_create_augroup("leetcode_questions", { clear = true })
297function Question:autocmds()
298 vim.api.nvim_create_autocmd("WinClosed", {
299 group = group,
300 pattern = tostring(self.winid),
301 callback = function()
302 self:_unmount()
303 end,
304 })
305end
306
307function Question:handle_mount()
308 self:create_buffer()
309
310 self.description = Description(self):mount()
311 self.console = Console(self)
312 self.info = Info(self)
313
314 table.insert(_Lc_state.questions, self)
315
316 self:autocmds()
317 utils.exec_hooks("question_enter", self)
318
319 return self
320end
321
322function Question:mount()
323 local tabp = utils.detect_duplicate_question(self.cache.title_slug, config.lang)
324 if tabp then
325 return pcall(vim.api.nvim_set_current_tabpage, tabp)
326 end
327
328 local q = api_question.by_title_slug(self.cache.title_slug)
329 if not q or q.is_paid_only and not config.auth.is_premium then
330 return log.warn("Question is for premium users only")
331 end
332 self.q = q
333
334 if self:snippet() then
335 self:handle_mount()
336 else
337 local msg = ("Snippet for `%s` not found. Select a different language"):format(self.lang)
338 log.warn(msg)
339
340 local picker = require("leetcode.picker")
341 picker.language(self, function(slug)
342 self.lang = slug
343 self:handle_mount()
344 end)
345 end
346
347 return self
348end
349
350---@class lc.Question.Editor.Range
351---@field start_i? number
352---@field end_i? number
353---@field lines string[]
354---@field sliced string[]
355---@field complete boolean
356
357---@param name string
358---@param inclusive? boolean
359function Question:editor_section_range(name, inclusive)
360 local lines = vim.api.nvim_buf_get_lines(self.bufnr, 0, -1, false)
361 local start_i, end_i
362
363 local start_tag = utils.section_tag(name, true)
364 local end_tag = utils.section_tag(name, false)
365
366 for i, line in ipairs(lines) do
367 if line:match(start_tag) then
368 start_i = i + (inclusive and 0 or 1)
369 elseif line:match(end_tag) then
370 end_i = i - (inclusive and 0 or 1)
371 end
372 end
373
374 local res = {
375 start_i = start_i,
376 end_i = end_i,
377 lines = lines,
378 complete = start_i and end_i,
379 }
380
381 res.log_not_found = function()
382 if res.complete then
383 return
384 end
385
386 local missing = {}
387 if not res.start_i then
388 table.insert(missing, ("`%s`"):format(start_tag))
389 end
390 if not res.end_i then
391 table.insert(missing, ("`%s`"):format(end_tag))
392 end
393
394 log.warn(table.concat(missing, " and ") .. " not found.")
395 end
396
397 res.is_partial = function()
398 return res.complete or (res.start_i or res.end_i)
399 end
400
401 res.is_valid_or_log = function()
402 if res.complete then
403 return true
404 else
405 res.log_not_found()
406 return false
407 end
408 end
409
410 return res
411end
412
413---@param submit boolean
414---@return string
415function Question:editor_submit_lines(submit)
416 local range = self:editor_section_range("code")
417 assert(range.complete, "Code section not found")
418
419 local prefix = not submit and ("\n"):rep(range.start_i - 1) or ""
420 return prefix .. table.concat(range.lines, "\n", range.start_i, range.end_i)
421end
422
423---@param self lc.ui.Question
424---@param lang lc.lang
425Question.change_lang = vim.schedule_wrap(function(self, lang)
426 local old_lang, old_bufnr = self.lang, self.bufnr
427
428 local ok, err = pcall(function()
429 self.lang = lang
430 local path, existed = self:path()
431
432 self.bufnr = vim.fn.bufadd(path)
433 assert(self.bufnr ~= 0, "Failed to create buffer " .. path)
434
435 local loaded = vim.api.nvim_buf_is_loaded(self.bufnr)
436 vim.fn.bufload(self.bufnr)
437
438 vim.api.nvim_set_option_value("buflisted", false, { buf = old_bufnr })
439 self:open_buffer(existed)
440
441 if not loaded then
442 utils.exec_hooks("question_enter", self)
443 end
444 end)
445
446 if not ok then
447 log.error("Failed to change language\n" .. err)
448 self.lang = old_lang
449 self.bufnr = old_bufnr
450 end
451end)
452
453---@param problem lc.cache.Question
454function Question:init(problem)
455 self.cache = problem
456 self.lang = config.lang
457end
458
459---@type fun(question: lc.cache.Question): lc.ui.Question
460local LeetQuestion = Question ---@diagnostic disable-line: assign-type-mismatch
461
462return LeetQuestion