rbx.cp 0.5.28__py3-none-any.whl → 0.5.30__py3-none-any.whl

rbx/box/generators.py

@@ -1,10 +1,12 @@
+import abc
 import pathlib
 import shlex
 import shutil
 from pathlib import PosixPath
-from typing import Callable, Dict, List, Optional, Set
+from typing import Dict, List, Optional, Set
 
 import typer
+from pydantic import BaseModel
 
 from rbx import console
 from rbx.box import checkers, package, testcases, validators
@@ -15,14 +17,12 @@ from rbx.box.environment import (
 )
 from rbx.box.schema import (
     CodeItem,
-    Generator,
     GeneratorCall,
     Testcase,
     TestcaseSubgroup,
 )
 from rbx.box.stressing import generator_parser
-from rbx.box.testcases import find_built_testcases
-from rbx.grading.judge.cacher import FileCacher
+from rbx.box.testcases import TestcaseEntry, find_built_testcases
 from rbx.grading.steps import (
     DigestHolder,
     DigestOrDest,
@@ -47,51 +47,38 @@ def _get_group_output(
     return group_path / f'{subgroup_prefix}{i:03d}.out'
 
 
+def _fill_output_for_defined_testcase(testcase: Testcase) -> Testcase:
+    res = testcase.model_copy()
+    if res.outputPath is not None:
+        return res
+    output_path = res.inputPath.with_suffix('.ans')
+    if output_path.is_file():
+        res.outputPath = output_path
+    return res
+
+
 def _copy_testcase_over(
-    testcase: Testcase, group_path: pathlib.Path, subgroup_prefix: str, i: int
+    testcase: Testcase,
+    dest: Testcase,
 ):
+    testcase = _fill_output_for_defined_testcase(testcase)
+    dest.inputPath.parent.mkdir(parents=True, exist_ok=True)
     shutil.copy(
         str(testcase.inputPath),
-        _get_group_input(group_path, subgroup_prefix, i),
+        str(dest.inputPath),
     )
-    if testcase.outputPath is not None and testcase.outputPath.is_file():
+    if (
+        testcase.outputPath is not None
+        and testcase.outputPath.is_file()
+        and dest.outputPath is not None
+    ):
+        dest.outputPath.parent.mkdir(parents=True, exist_ok=True)
         shutil.copy(
             str(testcase.outputPath),
-            _get_group_output(group_path, subgroup_prefix, i),
+            str(dest.outputPath),
         )
 
 
-def _run_generator(
-    generator: Generator,
-    args: Optional[str],
-    compiled_digest: str,
-    group_path: pathlib.Path,
-    subgroup_prefix: str,
-    i: int = 0,
-):
-    generation_stderr = DigestHolder()
-    run_log = run_item(
-        generator,
-        DigestOrSource.create(compiled_digest),
-        stdout=DigestOrDest.create(_get_group_input(group_path, subgroup_prefix, i)),
-        stderr=DigestOrDest.create(generation_stderr),
-        extra_args=args or None,
-    )
-
-    if not run_log or run_log.exitcode != 0:
-        console.console.print(
-            f'[error]Failed generating test {i} from group path {group_path}[/error]',
-        )
-        if run_log is not None:
-            console.console.print(f'[error]Summary:[/error] {run_log.get_summary()}')
-        if generation_stderr.value is not None:
-            console.console.print('[error]Stderr:[/error]')
-            console.console.print(
-                package.get_digest_as_string(generation_stderr.value) or ''
-            )
-        raise typer.Exit(1)
-
-
 def get_all_built_testcases() -> Dict[str, List[Testcase]]:
     pkg = package.find_problem_package_or_die()
     res = {group.name: find_built_testcases(group) for group in pkg.testcases}
@@ -103,108 +90,11 @@ def get_call_from_string(call_str: str) -> GeneratorCall:
     return GeneratorCall(name=name, args=args)
 
 
-def generate_output_for_testcase(
-    main_solution_digest: str,
-    testcase: Testcase,
-    stderr_path: Optional[pathlib.Path] = None,
-):
-    assert testcase.outputPath is not None
-    pkg = package.find_problem_package_or_die()
-    main_solution = package.get_main_solution()
-    if main_solution is None:
-        return
-
-    # Obey no limits when generating testcases.
-    sandbox = EnvironmentSandbox()
-    sandbox.fileSizeLimit = pkg.outputLimit
-    extra_config = ExecutionConfig(sandbox=sandbox)
-
-    try:
-        run_log = run_item(
-            main_solution,
-            DigestOrSource.create(main_solution_digest),
-            stdin=DigestOrSource.create(testcase.inputPath),
-            stdout=DigestOrDest.create(testcase.outputPath),
-            stderr=DigestOrDest.create(stderr_path)
-            if stderr_path is not None
-            else None,
-            extra_config=extra_config,
-        )
-    except:
-        console.console.print(
-            '[error]Failed running main solution to generate testcase.[/error]'
-        )
-        raise
-
-    if run_log is None or run_log.exitcode != 0:
-        console.console.print(
-            f'[error]Failed generating output for [item]{testcase.inputPath}[/item][/error]',
-        )
-        if run_log is not None:
-            console.console.print(f'[error]Summary:[/error] {run_log.get_summary()}')
-            checker_result = checkers.check_with_no_output(run_log)
-            console.console.print(
-                f'[warning]Verdict: [item]{checker_result.outcome.value}[/item][/warning]',
-            )
-            console.console.print(
-                f'[warning]Message: [info]{checker_result.message}[/info][/warning]',
-            )
-            console.console.print(f'Input written at [item]{testcase.inputPath}[/item]')
-            console.console.print(
-                f'Output written at [item]{testcase.outputPath}[/item]'
-            )
-            console.console.print(f'Stderr written at [item]{stderr_path}[/item]')
-        raise typer.Exit(1)
-
-
-def generate_outputs_for_testcases(
-    progress: Optional[StatusProgress] = None, groups: Optional[Set[str]] = None
-):
-    def step():
-        if progress is not None:
-            progress.step()
-
-    pkg = package.find_problem_package_or_die()
-
-    built_testcases = get_all_built_testcases()
-    main_solution = package.get_main_solution()
-    solution_digest: Optional[str] = None
-
-    if main_solution is not None:
-        if progress:
-            progress.update('Compiling main solution...')
-        try:
-            solution_digest = compile_item(main_solution)
-        except:
-            console.console.print('[error]Failed compiling main solution.[/error]')
-            raise
-
-    gen_runs_dir = package.get_problem_runs_dir() / '.gen'
-    shutil.rmtree(str(gen_runs_dir), ignore_errors=True)
-    gen_runs_dir.mkdir(parents=True, exist_ok=True)
-
-    for group in pkg.testcases:
-        if groups is not None and group.name not in groups:
-            continue
-        group_testcases = built_testcases[group.name]
-
-        for testcase in group_testcases:
-            stderr_path = gen_runs_dir / 'main.stderr'
-
-            assert testcase.outputPath is not None
-            if main_solution is None or solution_digest is None:
-                console.console.print(
-                    '[error]No main solution found to generate outputs for testcases.[/error]',
-                )
-                raise typer.Exit(1)
-
-            generate_output_for_testcase(solution_digest, testcase, stderr_path)
-            step()
-
-
-def _run_generator_script(testcase: TestcaseSubgroup, cacher: FileCacher) -> str:
+def _run_generator_script(testcase: TestcaseSubgroup) -> str:
     assert testcase.generatorScript is not None
 
+    cacher = package.get_file_cacher()
+
     if not testcase.generatorScript.path.is_file():
         console.console.print(
             f'[error]Generator script not found: [item]{testcase.generatorScript.path}[/item][/error]'
@@ -262,22 +152,171 @@ def _extract_script_lines(script: str):
         yield shlex.split(line)[0], shlex.join(shlex.split(line)[1:])
 
 
-def _get_necessary_generators(groups: Set[str], cacher: FileCacher) -> Set[str]:
+class GenerationMetadata(BaseModel):
+    copied_to: Testcase
+
+    copied_from: Optional[Testcase] = None
+    generator_call: Optional[GeneratorCall] = None
+
+
+class GenerationTestcaseEntry(BaseModel):
+    group_entry: TestcaseEntry
+    subgroup_entry: TestcaseEntry
+
+    metadata: GenerationMetadata
+
+
+class TestcaseVisitor(abc.ABC):
+    @abc.abstractmethod
+    def visit(self, entry: GenerationTestcaseEntry):
+        pass
+
+    def should_visit_group(self, group_name: str) -> bool:
+        return True
+
+    def should_visit_subgroup(self, subgroup_path: str) -> bool:
+        return True
+
+    def should_visit_generator_scripts(
+        self, group_name: str, subgroup_path: str
+    ) -> bool:
+        return True
+
+
+class TestcaseGroupVisitor(TestcaseVisitor):
+    def __init__(self, groups: Optional[Set[str]] = None):
+        self.groups = groups
+
+    def should_visit_group(self, group_name: str) -> bool:
+        return self.groups is None or group_name in self.groups
+
+
+def run_testcase_visitor(visitor: TestcaseVisitor):
     pkg = package.find_problem_package_or_die()
-    existing_generators = set(generator.name for generator in pkg.generators)
 
-    necessary_generators = set()
+    def _explore_subgroup(
+        subgroup: TestcaseSubgroup, subgroup_index: Optional[int], prefix: List[str]
+    ):
+        assert prefix and len(prefix) >= 1 and len(prefix) <= 2
+        group_path = prefix[0]
+        subgroup_path = '/'.join(prefix)
+        if not visitor.should_visit_subgroup(subgroup_path):
+            return
+
+        def _entry(i: int) -> TestcaseEntry:
+            return TestcaseEntry(group=group_path, index=i)
+
+        def _sub_entry(i: int) -> TestcaseEntry:
+            return TestcaseEntry(group=subgroup_path, index=i)
+
+        def _copied_to(i: int) -> Testcase:
+            group_fs_path = package.get_build_testgroup_path(group_path)
+            group_prefix = ''
+            if subgroup_index is not None:
+                group_prefix = f'{subgroup_index}-'
+            if len(prefix) == 2:
+                group_prefix += f'{prefix[1]}-'
+            return Testcase(
+                inputPath=_get_group_input(group_fs_path, group_prefix, i),
+                outputPath=_get_group_output(group_fs_path, group_prefix, i),
+            )
+
+        # Go through testcases.
+        i = 0
+        # Individual testcases.
+        for tc in subgroup.testcases or []:
+            visitor.visit(
+                GenerationTestcaseEntry(
+                    group_entry=_entry(i),
+                    subgroup_entry=_sub_entry(i),
+                    metadata=GenerationMetadata(
+                        copied_from=_fill_output_for_defined_testcase(tc),
+                        copied_to=_copied_to(i),
+                    ),
+                )
+            )
+            i += 1
+
+        # Glob testcases.
+        if subgroup.testcaseGlob:
+            matched_inputs = sorted(PosixPath().glob(subgroup.testcaseGlob))
+
+            for input_path in matched_inputs:
+                if not input_path.is_file() or input_path.suffix != '.in':
+                    continue
+
+                tc = Testcase(inputPath=input_path)
+                visitor.visit(
+                    GenerationTestcaseEntry(
+                        group_entry=_entry(i),
+                        subgroup_entry=_sub_entry(i),
+                        metadata=GenerationMetadata(
+                            copied_from=_fill_output_for_defined_testcase(tc),
+                            copied_to=_copied_to(i),
+                        ),
+                    )
+                )
+                i += 1
+
+        # Single generators.
+        for generator_call in subgroup.generators:
+            visitor.visit(
+                GenerationTestcaseEntry(
+                    group_entry=_entry(i),
+                    subgroup_entry=_sub_entry(i),
+                    metadata=GenerationMetadata(
+                        generator_call=generator_call,
+                        copied_to=_copied_to(i),
+                    ),
+                )
+            )
+            i += 1
+
+        if not visitor.should_visit_generator_scripts(group_path, subgroup_path):
+            return
+
+        # Run generator script.
+        if subgroup.generatorScript is not None:
+            script = _run_generator_script(subgroup)
+
+            # Run each line from generator script.
+            for generator_name, args in _extract_script_lines(script):
+                call = GeneratorCall(name=generator_name, args=args)
+                visitor.visit(
+                    GenerationTestcaseEntry(
+                        group_entry=_entry(i),
+                        subgroup_entry=_sub_entry(i),
+                        metadata=GenerationMetadata(
+                            generator_call=call,
+                            copied_to=_copied_to(i),
+                        ),
+                    )
+                )
+                i += 1
+
     for group in pkg.testcases:
-        if groups is not None and group.name not in groups:
+        if not visitor.should_visit_group(group.name):
             continue
 
-        for generator_call in group.generators:
-            necessary_generators.add(generator_call.name)
+        _explore_subgroup(group, 0 if group.subgroups else None, [group.name])
+
+        for i, subgroup in enumerate(group.subgroups):
+            _explore_subgroup(subgroup, i + 1, [group.name, subgroup.name])
 
-        if group.generatorScript is not None:
-            script = _run_generator_script(group, cacher)
-            for generator_name, _ in _extract_script_lines(script):
-                necessary_generators.add(generator_name)
+
+def _get_necessary_generators_for_groups(
+    groups: Optional[Set[str]] = None,
+) -> Set[str]:
+    pkg = package.find_problem_package_or_die()
+    existing_generators = set(generator.name for generator in pkg.generators)
+    necessary_generators = set()
+
+    class NecessaryGeneratorsVisitor(TestcaseGroupVisitor):
+        def visit(self, entry: GenerationTestcaseEntry):
+            if entry.metadata.generator_call is not None:
+                necessary_generators.add(entry.metadata.generator_call.name)
+
+    run_testcase_visitor(NecessaryGeneratorsVisitor(groups))
 
     return existing_generators.intersection(necessary_generators)
 
@@ -309,175 +348,265 @@ def compile_generators(
     return generator_to_compiled_digest
 
 
+def expand_generator_call(call: GeneratorCall) -> GeneratorCall:
+    vars = package.find_problem_package_or_die().expanded_vars
+    generator_for_args = generator_parser.Generator(vars)
+    parsed_args = generator_parser.parse(call.args or '')
+    return call.model_copy(update={'args': generator_for_args.generate(parsed_args)})
+
+
 def generate_standalone(
-    call: GeneratorCall,
-    output: pathlib.Path,
+    spec: GenerationMetadata,
     validate: bool = True,
+    group_entry: Optional[TestcaseEntry] = None,
     generator_digest: Optional[str] = None,
     validator_digest: Optional[str] = None,
-) -> GeneratorCall:
-    # Generator args parser
-    parsed_args = generator_parser.parse(call.args or '')
-    vars = package.find_problem_package_or_die().expanded_vars
-    generator_for_args = generator_parser.Generator(vars)
-    expanded_args_str = generator_for_args.generate(parsed_args)
+    progress: Optional[StatusProgress] = None,
+):
+    def _print_error_header(text: Optional[str] = None):
+        prefix = 'Failed generating test'
+        if group_entry is not None:
+            prefix += (
+                f' [item]{group_entry.group}[/item]/[item]{group_entry.index}[/item]'
+            )
+        suffix = '.'
+        if text:
+            suffix = f': {text}'
+        if spec.generator_call is not None:
+            console.console.print(
+                f'[error]{prefix} using generator call [info]{spec.generator_call.name} {spec.generator_call.args}[/info]{suffix}[/error]'
+            )
+        else:
+            console.console.print(f'[error]{prefix}{suffix}[/error]')
 
-    generation_stderr = DigestHolder()
+    if spec.generator_call is not None:
+        call = spec.generator_call
 
-    # Get generator item
-    generator = package.get_generator(call.name)
-    if generator_digest is None:
-        generator_digest = _compile_generator(generator)
+        generation_stderr = DigestHolder()
 
-    generation_log = run_item(
-        generator,
-        DigestOrSource.create(generator_digest),
-        stdout=DigestOrDest.create(output),
-        stderr=DigestOrDest.create(generation_stderr),
-        extra_args=expanded_args_str or None,
-    )
-    if not generation_log or generation_log.exitcode != 0:
-        console.console.print(
-            f'[error]Failed generating test using generator call [info]{call.name} {expanded_args_str}[/info][/error]',
-        )
-        if generation_log is not None:
-            console.console.print(
-                f'[error]Summary:[/error] {generation_log.get_summary()}'
-            )
-        if generation_stderr.value is not None:
-            console.console.print('[error]Stderr:[/error]')
-            console.console.print(
-                package.get_digest_as_string(generation_stderr.value) or ''
+        # Get generator item
+        generator = package.get_generator(call.name)
+        if generator_digest is None:
+            if progress:
+                progress.update(f'Compiling generator {generator.name}...')
+            generator_digest = _compile_generator(generator)
+
+        if progress:
+            progress.update(
+                f'Generating testcase [status]{generator.name} {call.args}[/status]...'
             )
+        generation_log = run_item(
+            generator,
+            DigestOrSource.create(generator_digest),
+            stdout=DigestOrDest.create(spec.copied_to.inputPath),
+            stderr=DigestOrDest.create(generation_stderr),
+            extra_args=call.args or None,
+        )
+        if not generation_log or generation_log.exitcode != 0:
+            _print_error_header()
+            if generation_log is not None:
+                console.console.print(
+                    f'[error]Summary:[/error] {generation_log.get_summary()}'
+                )
+            if generation_stderr.value is not None:
+                console.console.print('[error]Stderr:[/error]')
+                console.console.print(
+                    package.get_digest_as_string(generation_stderr.value) or ''
+                )
 
-        raise typer.Exit(1)
+            raise typer.Exit(1)
+    elif spec.copied_from is not None:
+        _copy_testcase_over(spec.copied_from, spec.copied_to)
 
     validator = package.get_validator_or_nil()
     # Run validator, if it is available.
     if validator is not None and validate:
         if validator_digest is None:
+            if progress:
+                progress.update('Compiling validator...')
             validator_tp = validators.compile_main_validator()
             assert validator_tp is not None
             _, validator_digest = validator_tp
-        ok, message, *_ = validators.validate_test(output, validator, validator_digest)
+        if progress:
+            progress.update('Validating test...')
+        ok, message, *_ = validators.validate_test(
+            spec.copied_to.inputPath,
+            validator,
+            validator_digest,
+        )
         if not ok:
+            _print_error_header('Failed validating testcase.')
+            console.console.print(f'[error]Message:[/error] {message}')
             console.console.print(
-                f'[error]Failed validating testcase generated by call [info]{call.name} {expanded_args_str}[/info][/error]'
+                f'Testcase written at [item]{spec.copied_to.inputPath}[/item]'
             )
-            console.console.print(f'[error]Message:[/error] {message}')
-            console.console.print(f'Testcase written at [item]{output}[/item]')
             raise typer.Exit(1)
 
-    return call.model_copy(update={'args': expanded_args_str})
 
-
-def _generate_testcases_for_subgroup(
-    subgroup: TestcaseSubgroup,
-    group_path: pathlib.Path,
-    subgroup_prefix: str,
-    compiled_generators: Dict[str, str],
-    step: Callable,
+def generate_testcases(
+    progress: Optional[StatusProgress] = None, groups: Optional[Set[str]] = None
 ):
-    cacher = package.get_file_cacher()
+    def step():
+        if progress is not None:
+            progress.step()
 
-    group_path.mkdir(parents=True, exist_ok=True)
-
-    i = 0
-    # Individual testcases.
-    for tc in subgroup.testcases or []:
-        _copy_testcase_over(tc, group_path, subgroup_prefix, i)
-        i += 1
-        step()
-
-    # Glob testcases.
-    if subgroup.testcaseGlob:
-        matched_inputs = sorted(PosixPath().glob(subgroup.testcaseGlob))
-
-        for input_path in matched_inputs:
-            if not input_path.is_file() or input_path.suffix != '.in':
-                continue
-            output_path = input_path.parent / f'{input_path.stem}.out'
-            tc = Testcase(inputPath=input_path, outputPath=output_path)
-            _copy_testcase_over(tc, group_path, subgroup_prefix, i)
-            i += 1
+    compiled_generators = compile_generators(
+        progress=progress,
+        tracked_generators=_get_necessary_generators_for_groups(groups)
+        if groups is not None
+        else None,
+    )
+
+    testcases.clear_built_testcases()
+
+    class BuildTestcaseVisitor(TestcaseGroupVisitor):
+        def visit(self, entry: GenerationTestcaseEntry):
+            if entry.metadata.copied_from is not None:
+                _copy_testcase_over(
+                    entry.metadata.copied_from,
+                    entry.metadata.copied_to,
+                )
+
+            if entry.metadata.generator_call is not None:
+                generate_standalone(
+                    entry.metadata,
+                    group_entry=entry.group_entry,
+                    validate=False,
+                    generator_digest=compiled_generators[
+                        entry.metadata.generator_call.name
+                    ],
+                )
             step()
 
-    # Run single generators.
-    for generator_call in subgroup.generators:
-        generator = package.get_generator(generator_call.name)
-        if generator.name not in compiled_generators:
-            console.console.print(f'Generator {generator.name} not compiled')
-            raise typer.Exit(1)
+    run_testcase_visitor(BuildTestcaseVisitor(groups))
 
-        _run_generator(
-            generator,
-            generator_call.args,
-            compiled_generators[generator.name],
-            group_path,
-            subgroup_prefix,
-            i,
+
+def generate_output_for_testcase(
+    main_solution_digest: str,
+    testcase: Testcase,
+    stderr_path: Optional[pathlib.Path] = None,
+):
+    assert testcase.outputPath is not None
+
+    if testcase.outputPath.is_file():
+        # Output file was already copied over from manual tests.
+        return
+
+    pkg = package.find_problem_package_or_die()
+    main_solution = package.get_main_solution()
+    if main_solution is None:
+        return
+
+    # Obey no limits when generating testcases.
+    sandbox = EnvironmentSandbox()
+    sandbox.fileSizeLimit = pkg.outputLimit
+    extra_config = ExecutionConfig(sandbox=sandbox)
+
+    try:
+        run_log = run_item(
+            main_solution,
+            DigestOrSource.create(main_solution_digest),
+            stdin=DigestOrSource.create(testcase.inputPath),
+            stdout=DigestOrDest.create(testcase.outputPath),
+            stderr=DigestOrDest.create(stderr_path)
+            if stderr_path is not None
+            else None,
+            extra_config=extra_config,
         )
-        i += 1
-        step()
-
-    # Run generator script.
-    if subgroup.generatorScript is not None:
-        script = _run_generator_script(subgroup, cacher)
-
-        # Run each line from generator script.
-        for generator_name, args in _extract_script_lines(script):
-            generator = package.get_generator(generator_name)
-            if generator.name not in compiled_generators:
-                console.console.print(f'Generator {generator.name} not compiled')
-                raise typer.Exit(1)
+    except:
+        console.console.print(
+            '[error]Failed running main solution to generate testcase.[/error]'
+        )
+        raise
 
-            _run_generator(
-                generator,
-                args,
-                compiled_generators[generator.name],
-                group_path,
-                subgroup_prefix,
-                i,
+    if run_log is None or run_log.exitcode != 0:
+        console.console.print(
+            f'[error]Failed generating output for [item]{testcase.inputPath}[/item][/error]',
+        )
+        if run_log is not None:
+            console.console.print(f'[error]Summary:[/error] {run_log.get_summary()}')
+            checker_result = checkers.check_with_no_output(run_log)
+            console.console.print(
+                f'[warning]Verdict: [item]{checker_result.outcome.value}[/item][/warning]',
             )
-            i += 1
-            step()
+            console.console.print(
+                f'[warning]Message: [info]{checker_result.message}[/info][/warning]',
+            )
+            console.console.print(f'Input written at [item]{testcase.inputPath}[/item]')
+            console.console.print(
+                f'Output written at [item]{testcase.outputPath}[/item]'
+            )
+            console.console.print(f'Stderr written at [item]{stderr_path}[/item]')
+        raise typer.Exit(1)
 
 
-def generate_testcases(
+def generate_outputs_for_testcases(
     progress: Optional[StatusProgress] = None, groups: Optional[Set[str]] = None
 ):
     def step():
         if progress is not None:
             progress.step()
 
-    pkg = package.find_problem_package_or_die()
-    cacher = package.get_file_cacher()
+    main_solution = package.get_main_solution()
+    solution_digest: Optional[str] = None
 
-    compiled_generators = compile_generators(
-        progress=progress,
-        tracked_generators=_get_necessary_generators(groups, cacher)
-        if groups is not None
-        else None,
-    )
+    if main_solution is not None:
+        if progress:
+            progress.update('Compiling main solution...')
+        try:
+            solution_digest = compile_item(main_solution)
+        except:
+            console.console.print('[error]Failed compiling main solution.[/error]')
+            raise
 
-    testcases.clear_built_testcases()
+    gen_runs_dir = package.get_problem_runs_dir() / '.gen'
+    shutil.rmtree(str(gen_runs_dir), ignore_errors=True)
+    gen_runs_dir.mkdir(parents=True, exist_ok=True)
 
-    for testcase in pkg.testcases:
-        if groups is not None and testcase.name not in groups:
-            continue
-        group_path = package.get_build_testgroup_path(testcase.name)
+    class GenerateOutputsVisitor(TestcaseGroupVisitor):
+        def visit(self, entry: GenerationTestcaseEntry):
+            tc = entry.metadata.copied_to
+            if not tc.inputPath.is_file():
+                return
+            assert tc.outputPath is not None
 
-        if not testcase.subgroups:
-            # Testcase group is itself a test subgroup.
-            _generate_testcases_for_subgroup(
-                testcase, group_path, '', compiled_generators, step
-            )
-            continue
+            if (
+                main_solution is None or solution_digest is None
+            ) and not tc.outputPath.is_file():
+                console.console.print(
+                    '[error]No main solution found to generate outputs for testcases.[/error]',
+                )
+                raise typer.Exit(1)
 
-        renamed_testcase = testcase.model_copy(update={'name': 'main'})
-        subgroups = [renamed_testcase] + testcase.subgroups
-        for i, subgroup in enumerate(subgroups):
-            # Test subgroups were specified, use them.
-            _generate_testcases_for_subgroup(
-                subgroup, group_path, f'{i}-{subgroup.name}-', compiled_generators, step
+            assert solution_digest is not None
+            generate_output_for_testcase(
+                solution_digest,
+                tc,
+                gen_runs_dir / 'main.stderr',
             )
+            step()
+
+    run_testcase_visitor(GenerateOutputsVisitor(groups))
+
+
+def extract_generation_testcases(
+    entries: List[TestcaseEntry],
+) -> List[GenerationTestcaseEntry]:
+    # TODO: support subgroups.
+    groups = set(entry.group for entry in entries)
+    entry_keys = set(entry.key() for entry in entries)
+
+    res: List[GenerationTestcaseEntry] = []
+
+    class ExtractGenerationTestcasesVisitor(TestcaseVisitor):
+        def should_visit_group(self, group_name: str) -> bool:
+            return group_name in groups
+
+        def visit(self, entry: GenerationTestcaseEntry):
+            # TODO: support subgroups.
+            if entry.group_entry.key() not in entry_keys:
+                return
+            res.append(entry)
+
+    run_testcase_visitor(ExtractGenerationTestcasesVisitor())
+    return res