feat(io): multi-test case view

lua/cp/cache.lua

@@ -20,6 +20,7 @@
 ---@field id string
 ---@field name? string
 ---@field interactive? boolean
+---@field multi_test? boolean
 ---@field memory_mb? number
 ---@field timeout_ms? number
 ---@field test_cases TestCase[]
@@ -187,6 +188,7 @@ end
 ---@param timeout_ms number
 ---@param memory_mb number
 ---@param interactive boolean
+---@param multi_test boolean
 function M.set_test_cases(
   platform,
   contest_id,
@@ -194,7 +196,8 @@ function M.set_test_cases(
   test_cases,
   timeout_ms,
   memory_mb,
-  interactive
+  interactive,
+  multi_test
 )
   vim.validate({
     platform = { platform, 'string' },
@@ -204,6 +207,7 @@ function M.set_test_cases(
     timeout_ms = { timeout_ms, { 'number', 'nil' }, true },
     memory_mb = { memory_mb, { 'number', 'nil' }, true },
     interactive = { interactive, { 'boolean', 'nil' }, true },
+    multi_test = { multi_test, { 'boolean', 'nil' }, true },
   })
 
   local index = cache_data[platform][contest_id].index_map[problem_id]
@@ -212,6 +216,7 @@ function M.set_test_cases(
   cache_data[platform][contest_id].problems[index].timeout_ms = timeout_ms
   cache_data[platform][contest_id].problems[index].memory_mb = memory_mb
   cache_data[platform][contest_id].problems[index].interactive = interactive
+  cache_data[platform][contest_id].problems[index].multi_test = multi_test
 
   M.save()
 end

lua/cp/scraper.lua

@@ -198,6 +198,7 @@ function M.scrape_all_tests(platform, contest_id, callback)
             timeout_ms = ev.timeout_ms or 0,
             memory_mb = ev.memory_mb or 0,
             interactive = ev.interactive or false,
+            multi_test = ev.multi_test or false,
             problem_id = ev.problem_id,
           })
         end

lua/cp/setup.lua

@@ -98,18 +98,35 @@ local function start_tests(platform, contest_id, problems)
         cached_tests,
         ev.timeout_ms or 0,
         ev.memory_mb or 0,
-        ev.interactive
+        ev.interactive,
+        ev.multi_test
       )
 
       local io_state = state.get_io_view_state()
       if io_state then
-        local test_cases = cache.get_test_cases(platform, contest_id, state.get_problem_id())
+        local problem_id = state.get_problem_id()
+        local test_cases = cache.get_test_cases(platform, contest_id, problem_id)
         local input_lines = {}
-        for _, tc in ipairs(test_cases) do
-          for _, line in ipairs(vim.split(tc.input, '\n')) do
-            table.insert(input_lines, line)
+
+        local contest_data = cache.get_contest_data(platform, contest_id)
+        local is_multi_test = contest_data.problems[contest_data.index_map[problem_id]].multi_test
+
+        if is_multi_test and #test_cases > 1 then
+          table.insert(input_lines, tostring(#test_cases))
+          for _, tc in ipairs(test_cases) do
+            local stripped = tc.input:gsub('^1\n', '')
+            for _, line in ipairs(vim.split(stripped, '\n')) do
+              table.insert(input_lines, line)
+            end
+          end
+        else
+          for _, tc in ipairs(test_cases) do
+            for _, line in ipairs(vim.split(tc.input, '\n')) do
+              table.insert(input_lines, line)
+            end
           end
         end
+
         require('cp.utils').update_buffer_content(io_state.input_buf, input_lines, nil, nil)
       end
     end)

lua/cp/ui/views.lua

@@ -315,9 +315,21 @@ function M.ensure_io_view()
   local test_cases = cache.get_test_cases(platform, contest_id, problem_id)
   if test_cases and #test_cases > 0 then
     local input_lines = {}
-    for _, tc in ipairs(test_cases) do
-      for _, line in ipairs(vim.split(tc.input, '\n')) do
-        table.insert(input_lines, line)
+    local is_multi_test = contest_data.problems[contest_data.index_map[problem_id]].multi_test
+
+    if is_multi_test and #test_cases > 1 then
+      table.insert(input_lines, tostring(#test_cases))
+      for _, tc in ipairs(test_cases) do
+        local stripped = tc.input:gsub('^1\n', '')
+        for _, line in ipairs(vim.split(stripped, '\n')) do
+          table.insert(input_lines, line)
+        end
+      end
+    else
+      for _, tc in ipairs(test_cases) do
+        for _, line in ipairs(vim.split(tc.input, '\n')) do
+          table.insert(input_lines, line)
+        end
       end
     end
     utils.update_buffer_content(input_buf, input_lines, nil, nil)
@@ -437,6 +449,12 @@ function M.run_io_view(test_index, debug)
     )
   end
 
+  local is_multi_test = contest_data.problems[contest_data.index_map[problem_id]].multi_test
+
+  if is_multi_test and #test_indices > 1 then
+    table.insert(input_lines, tostring(#test_indices))
+  end
+
   for _, idx in ipairs(test_indices) do
     local tc = test_state.test_cases[idx]
 
@@ -479,7 +497,11 @@ function M.run_io_view(test_index, debug)
       end
     end
 
-    for _, line in ipairs(vim.split(tc.input, '\n')) do
+    local test_input = tc.input
+    if is_multi_test and #test_indices > 1 then
+      test_input = test_input:gsub('^1\n', '')
+    end
+    for _, line in ipairs(vim.split(test_input, '\n')) do
       table.insert(input_lines, line)
     end
   end