Merge pull request #71 from barrett-ruth/feat/derive

:CP deriving

doc/cp.txt

@@ -25,6 +25,14 @@ COMMANDS                                            *cp-commands*
                                                             *:CP*
 cp.nvim uses a single :CP command with intelligent argument parsing:
 
+State Restoration ~
+
+:CP                             Restore contest context from current file.
+                                Automatically detects platform, contest, problem,
+                                and language from cached state. Use this after
+                                switching files to restore your CP environment.
+                                Requires previous setup with full :CP command.
+
 Setup Commands ~
 
 :CP {platform} {contest_id} {problem_id} [--lang={language}]
@@ -92,12 +100,12 @@ Here's an example configuration with lazy.nvim: >
           },
           debug = false,
           scrapers = { ... }, -- all scrapers enabled by default
-          filename = default_filename,
+          filename = default_filename,  -- <contest id> + <problem id>
           run_panel = {
             diff_mode = 'vim',
             next_test_key = '<c-n>',
             prev_test_key = '<c-p>',
-            toggle_diff_key = '<c-t>',
+            toggle_diff_key = 't',
             max_output_lines = 50,
           },
           diff = {
@@ -152,7 +160,7 @@ Here's an example configuration with lazy.nvim: >
                         Git provides character-level precision, vim uses built-in diff.
       - {next_test_key} (`string`, default: `"<c-n>"`) Key to navigate to next test case.
       - {prev_test_key} (`string`, default: `"<c-p>"`) Key to navigate to previous test case.
-      - {toggle_diff_key} (`string`, default: `"<c-t>"`) Key to toggle diff mode between vim and git.
+      - {toggle_diff_key} (`string`, default: `"t"`) Key to toggle diff mode between vim and git.
       - {max_output_lines} (`number`, default: `50`) Maximum lines of test output to display.
 
 *cp.DiffConfig*
@@ -269,7 +277,12 @@ Example: Setting up and solving AtCoder contest ABC324
 <   This automatically sets up problem B
 
 6. Continue solving problems with :CP next/:CP prev navigation
-7. Submit solutions on AtCoder website
+7. Switch to another file (e.g., previous contest): >
+    :e ~/contests/abc323/a.cpp
+    :CP
+<   Automatically restores abc323 contest context
+
+8. Submit solutions on AtCoder website
 <
 
 RUN PANEL                                          *cp-run*
@@ -341,8 +354,8 @@ Keymaps ~
                                                 *cp-test-keys*
 <c-n>                           Navigate to next test case (configurable via run_panel.next_test_key)
 <c-p>                           Navigate to previous test case (configurable via run_panel.prev_test_key)
-<c-t>                           Toggle diff mode between vim and git (configurable via run_panel.toggle_diff_key)
-<c-q>                           Exit test panel and restore layout
+t                           Toggle diff mode between vim and git (configurable via run_panel.toggle_diff_key)
+q                           Exit test panel and restore layout
 
 Diff Modes ~
 

lua/cp/cache.lua

@@ -1,5 +1,12 @@
+---@class FileState
+---@field platform string
+---@field contest_id string
+---@field problem_id? string
+---@field language? string
+
 ---@class CacheData
 ---@field [string] table<string, ContestData>
+---@field file_states? table<string, FileState>
 
 ---@class ContestData
 ---@field problems Problem[]
@@ -228,4 +235,46 @@ function M.get_constraints(platform, contest_id, problem_id)
   return problem_data.timeout_ms, problem_data.memory_mb
 end
 
+---@param file_path string
+---@return FileState?
+function M.get_file_state(file_path)
+  vim.validate({
+    file_path = { file_path, 'string' },
+  })
+
+  if not cache_data.file_states then
+    return nil
+  end
+
+  return cache_data.file_states[file_path]
+end
+
+---@param file_path string
+---@param platform string
+---@param contest_id string
+---@param problem_id? string
+---@param language? string
+function M.set_file_state(file_path, platform, contest_id, problem_id, language)
+  vim.validate({
+    file_path = { file_path, 'string' },
+    platform = { platform, 'string' },
+    contest_id = { contest_id, 'string' },
+    problem_id = { problem_id, { 'string', 'nil' }, true },
+    language = { language, { 'string', 'nil' }, true },
+  })
+
+  if not cache_data.file_states then
+    cache_data.file_states = {}
+  end
+
+  cache_data.file_states[file_path] = {
+    platform = platform,
+    contest_id = contest_id,
+    problem_id = problem_id,
+    language = language,
+  }
+
+  M.save()
+end
+
 return M

lua/cp/config.lua

@@ -100,7 +100,7 @@ M.defaults = {
     diff_mode = 'vim',
     next_test_key = '<c-n>',
     prev_test_key = '<c-p>',
-    toggle_diff_key = '<c-t>',
+    toggle_diff_key = 't',
     max_output_lines = 50,
   },
   diff = {
@@ -139,8 +139,6 @@ function M.setup(user_config)
                 return false
               end
 
-              -- Allow any language and extension configurations
-
               return true
             end,
             'contest configuration',

lua/cp/init.lua

@@ -140,6 +140,8 @@ local function setup_problem(contest_id, problem_id, language)
     config.hooks.setup_code(ctx)
   end
 
+  cache.set_file_state(vim.fn.expand('%:p'), state.platform, contest_id, problem_id, language)
+
   logger.log(('switched to problem %s'):format(ctx.problem_name))
 end
 
@@ -276,6 +278,9 @@ local function toggle_run_panel(is_debug)
     vim.api.nvim_win_call(actual_win, function()
       vim.cmd.diffthis()
     end)
+    -- NOTE: diffthis() sets foldcolumn, so override it after
+    vim.api.nvim_set_option_value('foldcolumn', '0', { win = expected_win })
+    vim.api.nvim_set_option_value('foldcolumn', '0', { win = actual_win })
 
     return {
       buffers = { expected_buf, actual_buf },
@@ -346,7 +351,7 @@ local function toggle_run_panel(is_debug)
       actual_content = actual_content
     end
 
-    local desired_mode = should_show_diff and config.run_panel.diff_mode or 'vim'
+    local desired_mode = config.run_panel.diff_mode
 
     if current_diff_layout and current_mode ~= desired_mode then
       local saved_pos = vim.api.nvim_win_get_cursor(0)
@@ -444,7 +449,7 @@ local function toggle_run_panel(is_debug)
   end
 
   setup_keybindings_for_buffer = function(buf)
-    vim.keymap.set('n', '<c-q>', function()
+    vim.keymap.set('n', 'q', function()
       toggle_run_panel()
     end, { buffer = buf, silent = true })
     vim.keymap.set('n', config.run_panel.toggle_diff_key, function()
@@ -549,11 +554,51 @@ local function navigate_problem(delta, language)
   end
 end
 
+local function restore_from_current_file()
+  local current_file = vim.fn.expand('%:p')
+  if current_file == '' then
+    logger.log('No file is currently open', vim.log.levels.ERROR)
+    return false
+  end
+
+  cache.load()
+  local file_state = cache.get_file_state(current_file)
+  if not file_state then
+    logger.log(
+      'No cached state found for current file. Use :CP <platform> <contest> <problem> first.',
+      vim.log.levels.ERROR
+    )
+    return false
+  end
+
+  logger.log(
+    ('Restoring from cached state: %s %s %s'):format(
+      file_state.platform,
+      file_state.contest_id,
+      file_state.problem_id or 'CSES'
+    )
+  )
+
+  if not set_platform(file_state.platform) then
+    return false
+  end
+
+  state.contest_id = file_state.contest_id
+  state.problem_id = file_state.problem_id
+
+  if file_state.platform == 'cses' then
+    setup_problem(file_state.contest_id, nil, file_state.language)
+  else
+    setup_problem(file_state.contest_id, file_state.problem_id, file_state.language)
+  end
+
+  return true
+end
+
 local function parse_command(args)
   if #args == 0 then
     return {
-      type = 'error',
-      message = 'Usage: :CP <platform> <contest> [problem] [--lang=<language>] | :CP <action> | :CP <problem>',
+      type = 'restore_from_file',
     }
   end
 
@@ -649,6 +694,11 @@ function M.handle_command(opts)
     return
   end
 
+  if cmd.type == 'restore_from_file' then
+    restore_from_current_file()
+    return
+  end
+
   if cmd.type == 'action' then
     if cmd.action == 'run' then
       toggle_run_panel(cmd.debug)

lua/cp/window.lua

@@ -75,10 +75,11 @@ function M.restore_layout(state, tile_fn)
     local source_file
     if source_files ~= '' then
       local files = vim.split(source_files, '\n')
-      local valid_extensions = vim.tbl_keys(constants.filetype_to_language)
+      -- Prefer known extensions first, but accept any extension
+      local known_extensions = vim.tbl_keys(constants.filetype_to_language)
       for _, file in ipairs(files) do
         local ext = vim.fn.fnamemodify(file, ':e')
-        if vim.tbl_contains(valid_extensions, ext) then
+        if vim.tbl_contains(known_extensions, ext) then
           source_file = file
           break
         end

spec/cache_spec.lua

@@ -105,4 +105,52 @@ describe('cp.cache', function()
       assert.is_nil(result)
     end)
   end)
+
+  describe('file state', function()
+    it('stores and retrieves file state', function()
+      local file_path = '/tmp/test.cpp'
+
+      cache.set_file_state(file_path, 'atcoder', 'abc123', 'a', 'cpp')
+      local result = cache.get_file_state(file_path)
+
+      assert.is_not_nil(result)
+      assert.equals('atcoder', result.platform)
+      assert.equals('abc123', result.contest_id)
+      assert.equals('a', result.problem_id)
+      assert.equals('cpp', result.language)
+    end)
+
+    it('handles cses file state without problem_id', function()
+      local file_path = '/tmp/cses.py'
+
+      cache.set_file_state(file_path, 'cses', '1068', nil, 'python')
+      local result = cache.get_file_state(file_path)
+
+      assert.is_not_nil(result)
+      assert.equals('cses', result.platform)
+      assert.equals('1068', result.contest_id)
+      assert.is_nil(result.problem_id)
+      assert.equals('python', result.language)
+    end)
+
+    it('returns nil for missing file state', function()
+      local result = cache.get_file_state('/nonexistent/file.cpp')
+      assert.is_nil(result)
+    end)
+
+    it('overwrites existing file state', function()
+      local file_path = '/tmp/overwrite.cpp'
+
+      cache.set_file_state(file_path, 'atcoder', 'abc123', 'a', 'cpp')
+      cache.set_file_state(file_path, 'codeforces', '1934', 'b', 'python')
+
+      local result = cache.get_file_state(file_path)
+
+      assert.is_not_nil(result)
+      assert.equals('codeforces', result.platform)
+      assert.equals('1934', result.contest_id)
+      assert.equals('b', result.problem_id)
+      assert.equals('python', result.language)
+    end)
+  end)
 end)

spec/command_parsing_spec.lua

@@ -32,7 +32,7 @@ describe('cp command parsing', function()
   end)
 
   describe('empty arguments', function()
-    it('logs error for no arguments', function()
+    it('attempts file state restoration for no arguments', function()
       local opts = { fargs = {} }
 
       cp.handle_command(opts)
@@ -40,7 +40,10 @@ describe('cp command parsing', function()
       assert.is_true(#logged_messages > 0)
       local error_logged = false
       for _, log_entry in ipairs(logged_messages) do
-        if log_entry.level == vim.log.levels.ERROR and log_entry.msg:match('Usage:') then
+        if
+          log_entry.level == vim.log.levels.ERROR
+          and log_entry.msg:match('No file is currently open')
+        then
           error_logged = true
           break
         end
@@ -185,6 +188,23 @@ describe('cp command parsing', function()
     end)
   end)
 
+  describe('restore from file', function()
+    it('returns restore_from_file type for empty args', function()
+      local opts = { fargs = {} }
+      local logged_error = false
+
+      cp.handle_command(opts)
+
+      for _, log in ipairs(logged_messages) do
+        if log.level == vim.log.levels.ERROR and log.msg:match('No file is currently open') then
+          logged_error = true
+        end
+      end
+
+      assert.is_true(logged_error)
+    end)
+  end)
+
   describe('invalid commands', function()
     it('logs error for invalid platform', function()
       local opts = { fargs = { 'invalid_platform' } }