From 4aa16d285889dbfa9991dda6f274253f40b47681 Mon Sep 17 00:00:00 2001 From: Barrett Ruth Date: Mon, 15 Sep 2025 11:35:22 -0400 Subject: [PATCH] feat: flesh out language support --- after/ftplugin/cpout.lua | 2 +- doc/cp.txt | 71 ++++++++++++++++++------ lua/cp/config.lua | 70 ++++++++++++++++++++--- lua/cp/init.lua | 52 +++++++++++------ lua/cp/problem.lua | 5 +- lua/cp/snippets.lua | 117 +++++++++++++++++++++------------------ plugin/cp.lua | 36 +++++++++++- 7 files changed, 254 insertions(+), 99 deletions(-) diff --git a/after/ftplugin/cpout.lua b/after/ftplugin/cpout.lua index 3f0bcdc..857a799 100644 --- a/after/ftplugin/cpout.lua +++ b/after/ftplugin/cpout.lua @@ -4,4 +4,4 @@ vim.opt_local.statuscolumn = "" vim.opt_local.signcolumn = "no" vim.opt_local.wrap = true vim.opt_local.linebreak = true -vim.opt_local.modifiable = false +vim.opt_local.modifiable = true diff --git a/doc/cp.txt b/doc/cp.txt index 306c70a..375adbd 100644 --- a/doc/cp.txt +++ b/doc/cp.txt @@ -27,11 +27,12 @@ cp.nvim uses a single :CP command with intelligent argument parsing: Setup Commands ~ -:CP {platform} {contest_id} {problem_id} +:CP {platform} {contest_id} {problem_id} [--lang={language}] Full setup: set platform, load contest metadata, and set up specific problem. Scrapes test cases and creates source file. Example: :CP codeforces 1933 a + Example: :CP codeforces 1933 a --lang=python :CP {platform} {contest_id} Contest setup: set platform and load contest metadata for navigation. Caches problem list. @@ -40,9 +41,11 @@ Setup Commands ~ :CP {platform} Platform setup: set platform only. Example: :CP cses -:CP {problem_id} Problem switch: switch to different problem +:CP {problem_id} [--lang={language}] + Problem switch: switch to different problem within current contest context. Example: :CP b (switch to problem b) + Example: :CP b --lang=python Action Commands ~ @@ -75,21 +78,36 @@ Optional configuration with lazy.nvim: > debug = false, contests = { default = { - cpp_version = 20, - compile_flags = { "-O2", "-DLOCAL", "-Wall", "-Wextra" }, - debug_flags = { "-g3", "-fsanitize=address,undefined", "-DLOCAL" }, + cpp = { + compile = { + 'g++', '-std=c++{version}', '-O2', '-Wall', '-Wextra', + '-DLOCAL', '{source}', '-o', '{binary}', + }, + run = { '{binary}' }, + debug = { + 'g++', '-std=c++{version}', '-g3', + '-fsanitize=address,undefined', '-DLOCAL', + '{source}', '-o', '{binary}', + }, + version = 20, + extension = "cc", + }, + python = { + run = { 'python3', '{source}' }, + debug = { 'python3', '{source}' }, + extension = "py", + }, timeout_ms = 2000, }, - atcoder = { cpp_version = 23 }, + codeforces = { cpp = { version = 23 } }, }, hooks = { before_run = function(problem_id) vim.cmd.w() end, before_debug = function(problem_id) ... end, }, - tile = function(source_buf, input_buf, output_buf) - end, - filename = function(contest, problem_id, problem_letter) - end, + snippets = { ... }, -- LuaSnip snippets + tile = function(source_buf, input_buf, output_buf) ... end, + filename = function(contest, problem_id, problem_letter) ... end, } } < @@ -98,10 +116,21 @@ Configuration options: contests Dictionary of contest configurations - each contest inherits from 'default'. - cpp_version c++ standard version (e.g. 20, 23) - compile_flags compiler flags for run builds - debug_flags compiler flags for debug builds - timeout_ms duration (ms) to run/debug before timeout + cpp C++ language configuration + compile Compile command template with {version}, {source}, {binary} placeholders + run Run command template with {binary} placeholder + debug Debug compile command template + version C++ standard version (e.g. 20, 23) + extension File extension for C++ files (default: "cc") + + python Python language configuration + run Run command template with {source} placeholder + debug Debug run command template + extension File extension for Python files (default: "py") + + default_language Default language when --lang not specified (default: "cpp") + + timeout_ms Duration (ms) to run/debug before timeout snippets LuaSnip snippets by contest type @@ -221,8 +250,8 @@ cp.nvim creates the following file structure upon problem setup: build/ {contest_id}{problem_id}.run " Compiled binary io/ - {contest_id}{problem_id}.in " Test input - {contest_id}{problem_id}.out " Program output + {contest_id}{problem_id}.cpin " Test input + {contest_id}{problem_id}.cpout " Program output {contest_id}{problem_id}.expected " Expected output The plugin automatically manages this structure and navigation between problems @@ -233,9 +262,17 @@ SNIPPETS *cp-snippets* cp.nvim integrates with LuaSnip for automatic template expansion. When you open a new problem file, type the contest name and press to expand. -Built-in snippets include basic C++ templates for each contest type. +Built-in snippets include basic C++ and Python templates for each contest type. Custom snippets can be added via configuration. +IMPORTANT: Snippet trigger names must exactly match the contest/platform names: +- "codeforces" for Codeforces problems +- "atcoder" for AtCoder problems +- "cses" for CSES problems + +The plugin automatically selects the appropriate template based on the file +extension (e.g., .cc files get C++ templates, .py files get Python templates). + HEALTH CHECK *cp-health* Run |:checkhealth| cp to verify your setup. diff --git a/lua/cp/config.lua b/lua/cp/config.lua index 48a245e..3ae18fe 100644 --- a/lua/cp/config.lua +++ b/lua/cp/config.lua @@ -1,13 +1,34 @@ +---@class LanguageConfig +---@field compile? string[] Compile command template +---@field run string[] Run command template +---@field debug? string[] Debug command template +---@field executable? string Executable name +---@field version? number Language version +---@field extension string File extension + +---@class ContestConfig +---@field cpp LanguageConfig +---@field python LanguageConfig +---@field default_language string +---@field timeout_ms number + ---@class cp.Config ----@field contests table ----@field snippets table +---@field contests table +---@field snippets table[] ---@field hooks table ---@field debug boolean ---@field tile? fun(source_buf: number, input_buf: number, output_buf: number) ----@field filename? fun(contest: string, problem_id: string, problem_letter?: string): string +---@field filename? fun(contest: string, contest_id: string, problem_id?: string, config: cp.Config, language?: string): string local M = {} +local filetype_to_language = { + cc = "cpp", + c = "cpp", + py = "python", + py3 = "python", +} + ---@type cp.Config M.defaults = { contests = { @@ -37,13 +58,16 @@ M.defaults = { }, executable = nil, version = 20, + extension = "cc", }, python = { compile = nil, run = { "{source}" }, debug = { "{source}" }, executable = "python3", + extension = "py", }, + default_language = "cpp", timeout_ms = 2000, }, atcoder = { @@ -74,8 +98,8 @@ local function extend_contest_config(base_config, contest_config) return result end ----@param user_config table|nil ----@return table +---@param user_config cp.Config|nil +---@return cp.Config function M.setup(user_config) vim.validate({ user_config = { user_config, { "table", "nil" }, true }, @@ -97,6 +121,20 @@ function M.setup(user_config) before_debug = { user_config.hooks.before_debug, { "function", "nil" }, true }, }) end + + if user_config.contests then + for contest_name, contest_config in pairs(user_config.contests) do + for lang_name, lang_config in pairs(contest_config) do + if type(lang_config) == "table" and lang_config.extension then + if not vim.tbl_contains(vim.tbl_keys(filetype_to_language), lang_config.extension) then + error(("Invalid extension '%s' for language '%s' in contest '%s'. Valid extensions: %s"):format( + lang_config.extension, lang_name, contest_name, table.concat(vim.tbl_keys(filetype_to_language), ", ") + )) + end + end + end + end + end end local config = vim.tbl_deep_extend("force", M.defaults, user_config or {}) @@ -111,14 +149,32 @@ function M.setup(user_config) return config end -local function default_filename(contest, contest_id, problem_id) +---@param contest string +---@param contest_id string +---@param problem_id? string +---@param config cp.Config +---@param language? string +---@return string +local function default_filename(contest, contest_id, problem_id, config, language) + vim.validate({ + contest = { contest, "string" }, + contest_id = { contest_id, "string" }, + problem_id = { problem_id, { "string", "nil" }, true }, + config = { config, "table" }, + language = { language, { "string", "nil" }, true }, + }) + local full_problem_id = contest_id:lower() if contest == "atcoder" or contest == "codeforces" then if problem_id then full_problem_id = full_problem_id .. problem_id:lower() end end - return full_problem_id .. ".cc" + + local contest_config = config.contests[contest] or config.contests.default + local target_language = language or contest_config.default_language + local language_config = contest_config[target_language] + return full_problem_id .. "." .. language_config.extension end M.default_filename = default_filename diff --git a/lua/cp/init.lua b/lua/cp/init.lua index 01a046a..50792b3 100644 --- a/lua/cp/init.lua +++ b/lua/cp/init.lua @@ -48,7 +48,8 @@ end ---@param contest_id string ---@param problem_id? string -local function setup_problem(contest_id, problem_id) +---@param language? string +local function setup_problem(contest_id, problem_id, language) if not state.platform then logger.log("no platform set. run :CP first", vim.log.levels.ERROR) return @@ -88,7 +89,7 @@ local function setup_problem(contest_id, problem_id) state.test_cases = cached_test_cases end - local ctx = problem.create_context(state.platform, contest_id, problem_id, config) + local ctx = problem.create_context(state.platform, contest_id, problem_id, config, language) local scrape_result = scrape.scrape_problem(ctx) @@ -305,33 +306,52 @@ end local function parse_command(args) if #args == 0 then - return { type = "error", message = "Usage: :CP [problem] | :CP | :CP " } + return { type = "error", message = "Usage: :CP [problem] [--lang=] | :CP | :CP " } end - local first = args[1] + local language = nil + + for i, arg in ipairs(args) do + local lang_match = arg:match("^--lang=(.+)$") + if lang_match then + language = lang_match + elseif arg == "--lang" then + if i + 1 <= #args then + language = args[i + 1] + else + return { type = "error", message = "--lang requires a value" } + end + end + end + + local filtered_args = vim.tbl_filter(function(arg) + return not (arg:match("^--lang") or arg == language) + end, args) + + local first = filtered_args[1] if vim.tbl_contains(actions, first) then return { type = "action", action = first } end if vim.tbl_contains(platforms, first) then - if #args == 1 then - return { type = "platform_only", platform = first } - elseif #args == 2 then + if #filtered_args == 1 then + return { type = "platform_only", platform = first, language = language } + elseif #filtered_args == 2 then if first == "cses" then - return { type = "cses_problem", platform = first, problem = args[2] } + return { type = "cses_problem", platform = first, problem = filtered_args[2], language = language } else - return { type = "contest_setup", platform = first, contest = args[2] } + return { type = "contest_setup", platform = first, contest = filtered_args[2], language = language } end - elseif #args == 3 then - return { type = "full_setup", platform = first, contest = args[2], problem = args[3] } + elseif #filtered_args == 3 then + return { type = "full_setup", platform = first, contest = filtered_args[2], problem = filtered_args[3], language = language } else return { type = "error", message = "Too many arguments" } end end if state.platform and state.contest_id then - return { type = "problem_switch", problem = first } + return { type = "problem_switch", problem = first, language = language } end return { type = "error", message = "Unknown command or no contest context" } @@ -398,7 +418,7 @@ function M.handle_command(opts) ) end - setup_problem(cmd.contest, cmd.problem) + setup_problem(cmd.contest, cmd.problem, cmd.language) end return end @@ -412,16 +432,16 @@ function M.handle_command(opts) vim.log.levels.WARN ) end - setup_problem(cmd.problem) + setup_problem(cmd.problem, nil, cmd.language) end return end if cmd.type == "problem_switch" then if state.platform == "cses" then - setup_problem(cmd.problem) + setup_problem(cmd.problem, nil, cmd.language) else - setup_problem(state.contest_id, cmd.problem) + setup_problem(state.contest_id, cmd.problem, cmd.language) end return end diff --git a/lua/cp/problem.lua b/lua/cp/problem.lua index 60406fd..6904d27 100644 --- a/lua/cp/problem.lua +++ b/lua/cp/problem.lua @@ -15,10 +15,11 @@ local M = {} ---@param contest_id string ---@param problem_id? string ---@param config cp.Config +---@param language? string ---@return ProblemContext -function M.create_context(contest, contest_id, problem_id, config) +function M.create_context(contest, contest_id, problem_id, config, language) local filename_fn = config.filename or require("cp.config").default_filename - local source_file = filename_fn(contest, contest_id, problem_id) + local source_file = filename_fn(contest, contest_id, problem_id, config, language) local base_name = vim.fn.fnamemodify(source_file, ":t:r") return { diff --git a/lua/cp/snippets.lua b/lua/cp/snippets.lua index fbddfff..34b056e 100644 --- a/lua/cp/snippets.lua +++ b/lua/cp/snippets.lua @@ -10,102 +10,113 @@ function M.setup(config) local s, i, fmt = ls.snippet, ls.insert_node, require("luasnip.extras.fmt").fmt - local default_snippets = { - s( - "codeforces", - fmt( - [[#include + local filetype_to_language = { + cc = "cpp", + c = "cpp", + py = "python", + py3 = "python", + } + + local language_to_filetype = {} + for ext, lang in pairs(filetype_to_language) do + language_to_filetype[lang] = ext + end + + local template_definitions = { + cpp = { + codeforces = [[#include using namespace std; -void solve() {{ +void solve() { {} -}} +} -int main() {{ +int main() { std::cin.tie(nullptr)->sync_with_stdio(false); int tc = 1; std::cin >> tc; - for (int t = 0; t < tc; ++t) {{ + for (int t = 0; t < tc; ++t) { solve(); - }} + } return 0; -}}]], - { i(1) } - ) - ), +}]], - s( - "atcoder", - fmt( - [[#include + atcoder = [[#include using namespace std; -void solve() {{ +void solve() { {} -}} +} -int main() {{ +int main() { std::cin.tie(nullptr)->sync_with_stdio(false); #ifdef LOCAL int tc; std::cin >> tc; - for (int t = 0; t < tc; ++t) {{ + for (int t = 0; t < tc; ++t) { solve(); - }} + } #else solve(); #endif return 0; -}}]], - { i(1) } - ) - ), +}]], - s( - "cses", - fmt( - [[#include + cses = [[#include using namespace std; -int main() {{ +int main() { std::cin.tie(nullptr)->sync_with_stdio(false); {} return 0; -}}]], - { i(1) } - ) - ), +}]], + }, + + python = { + codeforces = [[def solve(): + {} + +if __name__ == "__main__": + tc = int(input()) + for _ in range(tc): + solve()]], + + atcoder = [[def solve(): + {} + +if __name__ == "__main__": + solve()]], + + cses = [[{}]], + }, } - local default_map = {} - for _, snippet in pairs(default_snippets) do - default_map[snippet.trigger] = snippet + for language, filetype in pairs(language_to_filetype) do + local snippets = {} + + for contest, template in pairs(template_definitions[language] or {}) do + table.insert(snippets, s(contest, fmt(template, { i(1) }))) + end + + for _, snippet in ipairs(config.snippets or {}) do + if snippet.filetype == filetype then + table.insert(snippets, snippet) + end + end + + ls.add_snippets(filetype, snippets) end - - local user_map = {} - for _, snippet in pairs(config.snippets or {}) do - user_map[snippet.trigger] = snippet - end - - local merged_map = vim.tbl_extend("force", default_map, user_map) - - local all_snippets = {} - for _, snippet in pairs(merged_map) do - table.insert(all_snippets, snippet) - end - - ls.add_snippets("cpp", all_snippets) end return M diff --git a/plugin/cp.lua b/plugin/cp.lua index 83a2816..754fa09 100644 --- a/plugin/cp.lua +++ b/plugin/cp.lua @@ -11,15 +11,41 @@ vim.api.nvim_create_user_command("CP", function(opts) cp.handle_command(opts) end, { nargs = "*", + desc = "Competitive programming helper", complete = function(ArgLead, CmdLine, _) + local filetype_to_language = { + cc = "cpp", + c = "cpp", + py = "python", + py3 = "python", + } + + local languages = vim.tbl_keys(vim.tbl_add_reverse_lookup(filetype_to_language)) + + if ArgLead:match("^--lang=") then + local lang_completions = {} + for _, lang in ipairs(languages) do + table.insert(lang_completions, "--lang=" .. lang) + end + return vim.tbl_filter(function(completion) + return completion:find(ArgLead, 1, true) == 1 + end, lang_completions) + end + + if ArgLead == "--lang" then + return { "--lang" } + end + local args = vim.split(vim.trim(CmdLine), "%s+") local num_args = #args if CmdLine:sub(-1) == " " then num_args = num_args + 1 end + local lang_flag_present = vim.tbl_contains(args, "--lang") or vim.iter(args):any(function(arg) return arg:match("^--lang=") end) + if num_args == 2 then - local candidates = {} + local candidates = { "--lang" } vim.list_extend(candidates, platforms) vim.list_extend(candidates, actions) local cp = require("cp") @@ -37,13 +63,17 @@ end, { return vim.tbl_filter(function(cmd) return cmd:find(ArgLead, 1, true) == 1 end, candidates) - elseif num_args == 4 then + elseif args[#args-1] == "--lang" then + return vim.tbl_filter(function(lang) + return lang:find(ArgLead, 1, true) == 1 + end, languages) + elseif num_args == 4 and not lang_flag_present then if vim.tbl_contains(platforms, args[2]) then local cache = require("cp.cache") cache.load() local contest_data = cache.get_contest_data(args[2], args[3]) if contest_data and contest_data.problems then - local candidates = {} + local candidates = { "--lang" } for _, problem in ipairs(contest_data.problems) do table.insert(candidates, problem.id) end