zerolang_editing/task_builders.py

"""Task construction helpers for synthetic Zerolang editing rows."""

from __future__ import annotations

from typing import Any


def _source(text: str) -> str:
    return text.strip() + "\n"


def _write_program(message: str, *, raises: bool = True) -> str:
    raises_suffix = " raises" if raises else ""
    return _source(
        f"""
pub fn main(world: World) -> Void{raises_suffix} {{
    check world.out.write("{message}\\n")
}}
"""
    )


def _literal_task(
    task_id: str, old: str, new: str, goal: str | None = None, *, split: str = "eval"
) -> dict[str, Any]:
    return {
        "id": task_id,
        "split": split,
        "category": "graph_patch_literal",
        "goal": goal or f'Replace the string literal "{old}\\n" with "{new}\\n".',
        "source": _write_program(old),
        "target_source": _write_program(new),
    }


def _branch_literal_task(
    task_id: str, helper: str, old: str, new: str, *, split: str = "eval"
) -> dict[str, Any]:
    return {
        "id": task_id,
        "split": split,
        "category": "graph_patch_literal",
        "goal": (
            "Keep the helper-controlled branch intact and update only the string "
            f'literal from "{old}\\n" to "{new}\\n".'
        ),
        "source": _source(
            f"""
fn {helper}() -> i32 {{
    return 1
}}

pub fn main(world: World) -> Void raises {{
    if {helper}() == 1 {{
        check world.out.write("{old}\\n")
    }}
}}
"""
        ),
        "target_source": _source(
            f"""
fn {helper}() -> i32 {{
    return 1
}}

pub fn main(world: World) -> Void raises {{
    if {helper}() == 1 {{
        check world.out.write("{new}\\n")
    }}
}}
"""
        ),
    }


def _helper_task(
    task_id: str,
    helper: str,
    source_expr: str,
    target_expr: str,
    expected: int,
    output: str,
    *,
    split: str = "eval",
) -> dict[str, Any]:
    return {
        "id": task_id,
        "split": split,
        "category": "semantic_update",
        "goal": (
            f"Update {helper}() so it returns {expected} and the existing main "
            f"branch prints {output}."
        ),
        "source": _source(
            f"""
fn {helper}() -> i32 {{
    return {source_expr}
}}

pub fn main(world: World) -> Void raises {{
    if {helper}() == {expected} {{
        check world.out.write("{output}\\n")
    }}
}}
"""
        ),
        "target_source": _source(
            f"""
fn {helper}() -> i32 {{
    return {target_expr}
}}

pub fn main(world: World) -> Void raises {{
    if {helper}() == {expected} {{
        check world.out.write("{output}\\n")
    }}
}}
"""
        ),
    }


def _two_helper_task(
    task_id: str,
    helper: str,
    other: str,
    source_expr: str,
    target_expr: str,
    other_expr: str,
    expected: int,
    *,
    split: str = "eval",
) -> dict[str, Any]:
    return {
        "id": task_id,
        "split": split,
        "category": "semantic_update",
        "goal": (
            f"Update only {helper}() so main writes ok when the comparison succeeds; "
            f"leave {other}() unchanged."
        ),
        "source": _source(
            f"""
fn {helper}() -> i32 {{
    return {source_expr}
}}

fn {other}() -> i32 {{
    return {other_expr}
}}

pub fn main(world: World) -> Void raises {{
    if {helper}() == {expected} {{
        check world.out.write("ok\\n")
    }}
}}
"""
        ),
        "target_source": _source(
            f"""
fn {helper}() -> i32 {{
    return {target_expr}
}}

fn {other}() -> i32 {{
    return {other_expr}
}}

pub fn main(world: World) -> Void raises {{
    if {helper}() == {expected} {{
        check world.out.write("ok\\n")
    }}
}}
"""
        ),
    }


def _call_task(
    task_id: str, source_args: str, target_args: str, expected: int, *, split: str = "eval"
) -> dict[str, Any]:
    return {
        "id": task_id,
        "split": split,
        "category": "call_update",
        "goal": "Keep add unchanged, but edit one call argument so the comparison is true.",
        "source": _source(
            f"""
fn add(a: i32, b: i32) -> i32 {{
    return a + b
}}

pub fn main(world: World) -> Void raises {{
    if add({source_args}) == {expected} {{
        check world.out.write("ok\\n")
    }}
}}
"""
        ),
        "target_source": _source(
            f"""
fn add(a: i32, b: i32) -> i32 {{
    return a + b
}}

pub fn main(world: World) -> Void raises {{
    if add({target_args}) == {expected} {{
        check world.out.write("ok\\n")
    }}
}}
"""
        ),
    }


def _condition_task(
    task_id: str,
    helper: str,
    returned: int,
    source_compare: int,
    output: str,
    *,
    split: str = "eval",
) -> dict[str, Any]:
    return {
        "id": task_id,
        "split": split,
        "category": "condition_update",
        "goal": (
            "Edit the comparison literal so the branch is true without changing "
            f"{helper}() or the output string."
        ),
        "source": _source(
            f"""
fn {helper}() -> i32 {{
    return {returned}
}}

pub fn main(world: World) -> Void raises {{
    if {helper}() == {source_compare} {{
        check world.out.write("{output}\\n")
    }}
}}
"""
        ),
        "target_source": _source(
            f"""
fn {helper}() -> i32 {{
    return {returned}
}}

pub fn main(world: World) -> Void raises {{
    if {helper}() == {returned} {{
        check world.out.write("{output}\\n")
    }}
}}
"""
        ),
    }


def _diagnostic_task(task_id: str, message: str, *, split: str = "eval") -> dict[str, Any]:
    return {
        "id": task_id,
        "split": split,
        "category": "diagnostic_repair",
        "goal": "Repair the main signature so the existing world.out.write check is valid.",
        "source": _write_program(message, raises=False),
        "target_source": _write_program(message, raises=True),
    }