From 793063a68edfaab71947bdcf06a6f95498360370 Mon Sep 17 00:00:00 2001
From: Barrett Ruth <br.barrettruth@gmail.com>
Date: Fri, 19 Sep 2025 20:41:19 -0400
Subject: [PATCH] feat(test_panel): integrate scraped data

---
 doc/cp.txt             |  3 --
 lua/cp/cache.lua       | 34 ++++++++++++++++++++-
 lua/cp/config.lua      |  2 --
 lua/cp/execute.lua     |  4 +--
 lua/cp/scrape.lua      | 24 ++++++++++++++-
 lua/cp/test.lua        | 43 +++++++++++++++++++++++++--
 lua/cp/test_render.lua | 67 +++++++++++++++++++++++++++++++++---------
 scrapers/atcoder.py    |  5 ++--
 scrapers/codeforces.py |  4 +--
 scrapers/cses.py       |  4 +--
 scrapers/models.py     | 10 +------
 11 files changed, 160 insertions(+), 40 deletions(-)

diff --git a/doc/cp.txt b/doc/cp.txt
index 32bb734..9ac961d 100644
--- a/doc/cp.txt
+++ b/doc/cp.txt
@@ -100,7 +100,6 @@ Here's an example configuration with lazy.nvim: >
                 extension = "py",
               },
               default_language = "cpp",
-              timeout_ms = 2000,
             },
           },
           hooks = {
@@ -156,8 +155,6 @@ Here's an example configuration with lazy.nvim: >
       • {python}            (`LanguageConfig`) Python language configuration.
       • {default_language}  (`string`, default: `"cpp"`) Default language when
                             `--lang` not specified.
-      • {timeout_ms}        (`number`, default: `2000`) Execution timeout in
-                            milliseconds.
 
 *cp.LanguageConfig*
 
diff --git a/lua/cp/cache.lua b/lua/cp/cache.lua
index 86d26d8..8d5dd4a 100644
--- a/lua/cp/cache.lua
+++ b/lua/cp/cache.lua
@@ -7,6 +7,8 @@
 ---@field expires_at? number
 ---@field test_cases? CachedTestCase[]
 ---@field test_cases_cached_at? number
+---@field timeout_ms? number
+---@field memory_mb? number
 
 ---@class Problem
 ---@field id string
@@ -167,12 +169,16 @@ end
 ---@param contest_id string
 ---@param problem_id? string
 ---@param test_cases CachedTestCase[]
-function M.set_test_cases(platform, contest_id, problem_id, test_cases)
+---@param timeout_ms? number
+---@param memory_mb? number
+function M.set_test_cases(platform, contest_id, problem_id, test_cases, timeout_ms, memory_mb)
   vim.validate({
     platform = { platform, 'string' },
     contest_id = { contest_id, 'string' },
     problem_id = { problem_id, { 'string', 'nil' }, true },
     test_cases = { test_cases, 'table' },
+    timeout_ms = { timeout_ms, { 'number', 'nil' }, true },
+    memory_mb = { memory_mb, { 'number', 'nil' }, true },
   })
 
   local problem_key = problem_id and (contest_id .. '_' .. problem_id) or contest_id
@@ -185,7 +191,33 @@ function M.set_test_cases(platform, contest_id, problem_id, test_cases)
 
   cache_data[platform][problem_key].test_cases = test_cases
   cache_data[platform][problem_key].test_cases_cached_at = os.time()
+  if timeout_ms then
+    cache_data[platform][problem_key].timeout_ms = timeout_ms
+  end
+  if memory_mb then
+    cache_data[platform][problem_key].memory_mb = memory_mb
+  end
   M.save()
 end
 
+---@param platform string
+---@param contest_id string
+---@param problem_id? string
+---@return number?, number?
+function M.get_constraints(platform, contest_id, problem_id)
+  vim.validate({
+    platform = { platform, 'string' },
+    contest_id = { contest_id, 'string' },
+    problem_id = { problem_id, { 'string', 'nil' }, true },
+  })
+
+  local problem_key = problem_id and (contest_id .. '_' .. problem_id) or contest_id
+  if not cache_data[platform] or not cache_data[platform][problem_key] then
+    return nil, nil
+  end
+
+  local problem_data = cache_data[platform][problem_key]
+  return problem_data.timeout_ms, problem_data.memory_mb
+end
+
 return M
diff --git a/lua/cp/config.lua b/lua/cp/config.lua
index 0616e9f..caddd69 100644
--- a/lua/cp/config.lua
+++ b/lua/cp/config.lua
@@ -18,13 +18,11 @@
 ---@field cpp LanguageConfig
 ---@field python LanguageConfig
 ---@field default_language string
----@field timeout_ms number
 
 ---@class PartialContestConfig
 ---@field cpp? PartialLanguageConfig
 ---@field python? PartialLanguageConfig
 ---@field default_language? string
----@field timeout_ms? number
 
 ---@class Hooks
 ---@field before_run? fun(ctx: ProblemContext)
diff --git a/lua/cp/execute.lua b/lua/cp/execute.lua
index b59a954..ba6a4fe 100644
--- a/lua/cp/execute.lua
+++ b/lua/cp/execute.lua
@@ -103,7 +103,7 @@ end
 
 ---@param cmd string[]
 ---@param input_data string
----@param timeout_ms integer
+---@param timeout_ms number
 ---@return ExecuteResult
 local function execute_command(cmd, input_data, timeout_ms)
   vim.validate({
@@ -279,7 +279,7 @@ function M.run_problem(ctx, contest_config, is_debug)
   end
 
   local run_cmd = build_command(language_config.test, language_config.executable, substitutions)
-  local exec_result = execute_command(run_cmd, input_data, contest_config.timeout_ms)
+  local exec_result = execute_command(run_cmd, input_data, timeout_ms)
   local formatted_output = format_output(exec_result, ctx.expected_file, is_debug)
 
   local output_buf = vim.fn.bufnr(ctx.output_file)
diff --git a/lua/cp/scrape.lua b/lua/cp/scrape.lua
index 56a9227..c49eee4 100644
--- a/lua/cp/scrape.lua
+++ b/lua/cp/scrape.lua
@@ -7,6 +7,8 @@
 ---@field problem_id string
 ---@field url? string
 ---@field tests? ScraperTestCase[]
+---@field timeout_ms? number
+---@field memory_mb? number
 ---@field error? string
 
 local M = {}
@@ -152,7 +154,7 @@ function M.scrape_contest_metadata(platform, contest_id)
 end
 
 ---@param ctx ProblemContext
----@return {success: boolean, problem_id: string, test_count?: number, test_cases?: ScraperTestCase[], url?: string, error?: string}
+---@return {success: boolean, problem_id: string, test_count?: number, test_cases?: ScraperTestCase[], timeout_ms?: number, memory_mb?: number, url?: string, error?: string}
 function M.scrape_problem(ctx)
   vim.validate({
     ctx = { ctx, 'table' },
@@ -277,6 +279,24 @@ function M.scrape_problem(ctx)
       vim.fn.writefile(vim.split(input_content, '\n', true), input_file)
       vim.fn.writefile(vim.split(expected_content, '\n', true), expected_file)
     end
+
+    local cached_test_cases = {}
+    for i, test_case in ipairs(data.tests) do
+      table.insert(cached_test_cases, {
+        index = i,
+        input = test_case.input,
+        expected = test_case.expected,
+      })
+    end
+
+    cache.set_test_cases(
+      ctx.contest,
+      ctx.contest_id,
+      ctx.problem_id,
+      cached_test_cases,
+      data.timeout_ms,
+      data.memory_mb
+    )
   end
 
   return {
@@ -284,6 +304,8 @@ function M.scrape_problem(ctx)
     problem_id = ctx.problem_name,
     test_count = data.tests and #data.tests or 0,
     test_cases = data.tests,
+    timeout_ms = data.timeout_ms,
+    memory_mb = data.memory_mb,
     url = data.url,
   }
 end
diff --git a/lua/cp/test.lua b/lua/cp/test.lua
index 3e0e147..e894d4a 100644
--- a/lua/cp/test.lua
+++ b/lua/cp/test.lua
@@ -12,6 +12,10 @@
 ---@field signal string?
 ---@field timed_out boolean?
 
+---@class ProblemConstraints
+---@field timeout_ms number
+---@field memory_mb number
+
 ---@class RunPanelState
 ---@field test_cases TestCase[]
 ---@field current_index number
@@ -19,6 +23,7 @@
 ---@field namespace number?
 ---@field is_active boolean
 ---@field saved_layout table?
+---@field constraints ProblemConstraints?
 
 local M = {}
 local constants = require('cp.constants')
@@ -32,6 +37,7 @@ local run_panel_state = {
   namespace = nil,
   is_active = false,
   saved_layout = nil,
+  constraints = nil,
 }
 
 ---@param index number
@@ -114,6 +120,25 @@ local function parse_test_cases_from_files(input_file, expected_file)
   return test_cases
 end
 
+---@param platform string
+---@param contest_id string
+---@param problem_id string?
+---@return ProblemConstraints?
+local function load_constraints_from_cache(platform, contest_id, problem_id)
+  local cache = require('cp.cache')
+  cache.load()
+  local timeout_ms, memory_mb = cache.get_constraints(platform, contest_id, problem_id)
+
+  if timeout_ms and memory_mb then
+    return {
+      timeout_ms = timeout_ms,
+      memory_mb = memory_mb,
+    }
+  end
+
+  return nil
+end
+
 ---@param ctx ProblemContext
 ---@param contest_config ContestConfig
 ---@param test_case TestCase
@@ -177,10 +202,15 @@ local function run_single_test_case(ctx, contest_config, cp_config, test_case)
   local stdin_content = test_case.input .. '\n'
 
   local start_time = vim.uv.hrtime()
+  local timeout_ms = run_panel_state.constraints and run_panel_state.constraints.timeout_ms or 2000
+
+  if not run_panel_state.constraints then
+    logger.log('no problem constraints available, using default 2000ms timeout')
+  end
   local result = vim
     .system(run_cmd, {
       stdin = stdin_content,
-      timeout = contest_config.timeout_ms or 2000,
+      timeout = timeout_ms,
       text = true,
     })
     :wait()
@@ -241,8 +271,17 @@ function M.load_test_cases(ctx, state)
 
   run_panel_state.test_cases = test_cases
   run_panel_state.current_index = 1
+  run_panel_state.constraints =
+    load_constraints_from_cache(state.platform, state.contest_id, state.problem_id)
 
-  logger.log(('loaded %d test case(s)'):format(#test_cases))
+  local constraint_info = run_panel_state.constraints
+      and string.format(
+        ' with %dms/%dMB limits',
+        run_panel_state.constraints.timeout_ms,
+        run_panel_state.constraints.memory_mb
+      )
+    or ''
+  logger.log(('loaded %d test case(s)%s'):format(#test_cases, constraint_info))
   return #test_cases > 0
 end
 
diff --git a/lua/cp/test_render.lua b/lua/cp/test_render.lua
index 9bbe699..749e89e 100644
--- a/lua/cp/test_render.lua
+++ b/lua/cp/test_render.lua
@@ -55,24 +55,34 @@ end
 
 -- Compute column widths + aggregates
 local function compute_cols(test_state)
-  local w = { num = 3, status = 8, time = 6, exit = 11 }
+  local w = { num = 3, status = 8, time = 6, exit = 11, limits = 12 }
+
+  local limits_str = ''
+  if test_state.constraints then
+    limits_str =
+      string.format('%d/%.0f', test_state.constraints.timeout_ms, test_state.constraints.memory_mb)
+  else
+    limits_str = '—'
+  end
 
   for i, tc in ipairs(test_state.test_cases) do
     local prefix = (i == test_state.current_index) and '>' or ' '
     w.num = math.max(w.num, #(prefix .. i))
     w.status = math.max(w.status, #(' ' .. M.get_status_info(tc).text))
-    local time_str = tc.time_ms and (string.format('%.2f', tc.time_ms) .. 'ms') or '—'
+    local time_str = tc.time_ms and string.format('%.2f', tc.time_ms) or '—'
     w.time = math.max(w.time, #time_str)
     w.exit = math.max(w.exit, #(' ' .. format_exit_code(tc.code)))
+    w.limits = math.max(w.limits, #limits_str)
   end
 
   w.num = math.max(w.num, #' #')
   w.status = math.max(w.status, #' Status')
-  w.time = math.max(w.time, #' Time')
+  w.time = math.max(w.time, #' Runtime (ms)')
   w.exit = math.max(w.exit, #' Exit Code')
+  w.limits = math.max(w.limits, #' Time (ms)/Mem (MB)')
 
-  local sum = w.num + w.status + w.time + w.exit
-  local inner = sum + 3 -- three inner vertical dividers
+  local sum = w.num + w.status + w.time + w.exit + w.limits
+  local inner = sum + 4 -- four inner vertical dividers
   local total = inner + 2 -- two outer borders
   return { w = w, sum = sum, inner = inner, total = total }
 end
@@ -86,6 +96,14 @@ local function center(text, width)
   return string.rep(' ', left) .. text .. string.rep(' ', pad - left)
 end
 
+local function right_align(text, width)
+  local pad = width - #text
+  if pad <= 0 then
+    return text
+  end
+  return string.rep(' ', pad) .. text
+end
+
 local function top_border(c)
   local w = c.w
   return '┌'
@@ -96,6 +114,8 @@ local function top_border(c)
     .. string.rep('─', w.time)
     .. '┬'
     .. string.rep('─', w.exit)
+    .. '┬'
+    .. string.rep('─', w.limits)
     .. '┐'
 end
 
@@ -109,6 +129,8 @@ local function row_sep(c)
     .. string.rep('─', w.time)
     .. '┼'
     .. string.rep('─', w.exit)
+    .. '┼'
+    .. string.rep('─', w.limits)
     .. '┤'
 end
 
@@ -122,6 +144,8 @@ local function bottom_border(c)
     .. string.rep('─', w.time)
     .. '┴'
     .. string.rep('─', w.exit)
+    .. '┴'
+    .. string.rep('─', w.limits)
     .. '┘'
 end
 
@@ -135,6 +159,8 @@ local function flat_fence_above(c)
     .. string.rep('─', w.time)
     .. '┴'
     .. string.rep('─', w.exit)
+    .. '┴'
+    .. string.rep('─', w.limits)
     .. '┤'
 end
 
@@ -148,6 +174,8 @@ local function flat_fence_below(c)
     .. string.rep('─', w.time)
     .. '┬'
     .. string.rep('─', w.exit)
+    .. '┬'
+    .. string.rep('─', w.limits)
     .. '┤'
 end
 
@@ -162,34 +190,45 @@ local function header_line(c)
     .. '│'
     .. center('Status', w.status)
     .. '│'
-    .. center('Time', w.time)
+    .. center('Runtime (ms)', w.time)
     .. '│'
     .. center('Exit Code', w.exit)
     .. '│'
+    .. center('Time (ms)/Mem (MB)', w.limits)
+    .. '│'
 end
 
-local function data_row(c, idx, tc, is_current)
+local function data_row(c, idx, tc, is_current, test_state)
   local w = c.w
   local prefix = is_current and '>' or ' '
   local status = M.get_status_info(tc)
-  local time = tc.time_ms and (string.format('%.2f', tc.time_ms) .. 'ms') or '—'
+  local time = tc.time_ms and string.format('%.2f', tc.time_ms) or '—'
   local exit = format_exit_code(tc.code)
 
+  local limits = ''
+  if test_state.constraints then
+    limits =
+      string.format('%d/%.0f', test_state.constraints.timeout_ms, test_state.constraints.memory_mb)
+  else
+    limits = '—'
+  end
+
   local line = '│'
     .. center(prefix .. idx, w.num)
     .. '│'
-    .. center(status.text, w.status)
+    .. right_align(status.text, w.status)
     .. '│'
-    .. center(time, w.time)
+    .. right_align(time, w.time)
     .. '│'
-    .. center(exit, w.exit)
+    .. right_align(exit, w.exit)
+    .. '│'
+    .. right_align(limits, w.limits)
     .. '│'
 
   local hi
   if status.text ~= '' then
     local pad = w.status - #status.text
-    local left = math.floor(pad / 2)
-    local status_start_col = 1 + w.num + 1 + left
+    local status_start_col = 1 + w.num + 1 + pad
     local status_end_col = status_start_col + #status.text
     hi = {
       col_start = status_start_col,
@@ -213,7 +252,7 @@ function M.render_test_list(test_state)
 
   for i, tc in ipairs(test_state.test_cases) do
     local is_current = (i == test_state.current_index)
-    local row, hi = data_row(c, i, tc, is_current)
+    local row, hi = data_row(c, i, tc, is_current, test_state)
     table.insert(lines, row)
     if hi then
       hi.line = #lines - 1
diff --git a/scrapers/atcoder.py b/scrapers/atcoder.py
index d9bf3c5..e251c44 100644
--- a/scrapers/atcoder.py
+++ b/scrapers/atcoder.py
@@ -11,7 +11,7 @@ from bs4 import BeautifulSoup, Tag
 from .models import MetadataResult, ProblemSummary, TestCase, TestsResult
 
 
-def extract_problem_limits(soup: BeautifulSoup) -> tuple[int, int]:
+def extract_problem_limits(soup: BeautifulSoup) -> tuple[int, float]:
     timeout_ms = None
     memory_mb = None
 
@@ -26,7 +26,8 @@ def extract_problem_limits(soup: BeautifulSoup) -> tuple[int, int]:
 
             memory_match = re.search(r"Memory Limit:\s*(\d+)\s*MiB", text)
             if memory_match:
-                memory_mb = int(memory_match.group(1))
+                memory_mib = int(memory_match.group(1))
+                memory_mb = round(memory_mib * 1.048576, 2)
             break
 
     if timeout_ms is None:
diff --git a/scrapers/codeforces.py b/scrapers/codeforces.py
index 54a51a1..a66acbd 100644
--- a/scrapers/codeforces.py
+++ b/scrapers/codeforces.py
@@ -140,7 +140,7 @@ def parse_problem_url(contest_id: str, problem_letter: str) -> str:
     )
 
 
-def extract_problem_limits(soup: BeautifulSoup) -> tuple[int, int]:
+def extract_problem_limits(soup: BeautifulSoup) -> tuple[int, float]:
     import re
 
     timeout_ms = None
@@ -162,7 +162,7 @@ def extract_problem_limits(soup: BeautifulSoup) -> tuple[int, int]:
         text = memory_limit_div.get_text().strip()
         match = re.search(r"(\d+) megabytes", text)
         if match:
-            memory_mb = int(match.group(1))
+            memory_mb = float(match.group(1))
 
     if memory_mb is None:
         raise ValueError("Could not find valid memory limit in memory-limit section")
diff --git a/scrapers/cses.py b/scrapers/cses.py
index 760197f..edf3224 100755
--- a/scrapers/cses.py
+++ b/scrapers/cses.py
@@ -19,7 +19,7 @@ def parse_problem_url(problem_input: str) -> str | None:
     return None
 
 
-def extract_problem_limits(soup: BeautifulSoup) -> tuple[int, int]:
+def extract_problem_limits(soup: BeautifulSoup) -> tuple[int, float]:
     timeout_ms = None
     memory_mb = None
 
@@ -39,7 +39,7 @@ def extract_problem_limits(soup: BeautifulSoup) -> tuple[int, int]:
         if "Memory limit:" in text:
             match = re.search(r"Memory limit:\s*(\d+)\s*MB", text)
             if match:
-                memory_mb = int(match.group(1))
+                memory_mb = float(match.group(1))
 
     if timeout_ms is None:
         raise ValueError("Could not find valid timeout in task-constraints section")
diff --git a/scrapers/models.py b/scrapers/models.py
index a37d186..728e9bb 100644
--- a/scrapers/models.py
+++ b/scrapers/models.py
@@ -13,14 +13,6 @@ class ProblemSummary:
     name: str
 
 
-@dataclass
-class Problem:
-    id: str
-    name: str
-    timeout_ms: int
-    memory_mb: int
-
-
 @dataclass
 class ScrapingResult:
     success: bool
@@ -40,4 +32,4 @@ class TestsResult(ScrapingResult):
     url: str
     tests: list[TestCase]
     timeout_ms: int
-    memory_mb: int
+    memory_mb: float