gen_tasks.py

#!/usr/bin/env python3
"""ComposeBench task generator.

ComposeBench is a clean-room synthetic benchmark for isolating executable
composition. The domain is a deterministic sensor network with seven pure
functions. Chain tasks vary dependent-call depth; fan-in tasks vary the number
of intermediate values that must be carried forward. Ground truth is computed by
the same tool functions exposed to every substrate.
"""
import json
import pathlib

ROOT = pathlib.Path(__file__).resolve().parent


def _h(*xs):
    import hashlib
    s = "|".join(str(x) for x in xs).encode()
    return int.from_bytes(hashlib.sha256(s).digest()[:4], "big")


SENSOR_COUNT = 64
SENSOR_VALUES = {
    i: 50 + (_h("sensor", i) % 300) for i in range(SENSOR_COUNT)
}


def read_sensor(sensor_id: int) -> int:
    """Return the raw integer reading for one sensor."""
    return SENSOR_VALUES[int(sensor_id) % SENSOR_COUNT]


def calibrate(value: int, offset: int) -> int:
    """Add an integer calibration offset to a reading."""
    return int(value) + int(offset)


def scale(value: int, factor: int) -> int:
    """Multiply an integer value by an integer factor."""
    return int(value) * int(factor)


def combine(a: int, b: int) -> int:
    """Add two integer values."""
    return int(a) + int(b)


def difference(a: int, b: int) -> int:
    """Subtract b from a."""
    return int(a) - int(b)


def aggregate(values: list) -> int:
    """Sum a list of integer values."""
    return sum(int(v) for v in values)


def threshold_count(values: list, threshold: int) -> int:
    """Count values strictly above an integer threshold."""
    return sum(1 for v in values if int(v) > int(threshold))


TOOLS = {
    "read_sensor": ("read_sensor(sensor_id: int) -> int",
                    "Return the raw integer reading for one sensor.",
                    read_sensor),
    "calibrate": ("calibrate(value: int, offset: int) -> int",
                  "Add an integer calibration offset.", calibrate),
    "scale": ("scale(value: int, factor: int) -> int",
              "Multiply an integer value by an integer factor.", scale),
    "combine": ("combine(a: int, b: int) -> int",
                "Add two integer values.", combine),
    "difference": ("difference(a: int, b: int) -> int",
                   "Subtract b from a.", difference),
    "aggregate": ("aggregate(values: list[int]) -> int",
                  "Sum a list of integers.", aggregate),
    "threshold_count": ("threshold_count(values: list[int], threshold: int) -> int",
                        "Count values strictly above a threshold.",
                        threshold_count),
}


def _resolve_arg(arg, values):
    if isinstance(arg, str) and arg.startswith("$"):
        return values[arg[1:]]
    if isinstance(arg, list):
        return [_resolve_arg(x, values) for x in arg]
    return arg


def eval_dag(steps, return_ref):
    values = {}
    for step in steps:
        args = [_resolve_arg(a, values) for a in step.get("args", [])]
        values[step["id"]] = TOOLS[step["tool"]][2](*args)
    if isinstance(return_ref, str) and return_ref.startswith("$"):
        return values[return_ref[1:]]
    return values[return_ref]


def _chain_params(seed):
    return {
        "sensor": _h("chain-sensor", seed) % SENSOR_COUNT,
        "offset": -12 + (_h("chain-offset", seed) % 25),
        "factor": 2 + (_h("chain-factor", seed) % 5),
        "baseline": 120 + (_h("chain-base", seed) % 180),
        "anchor": 10 + (_h("chain-anchor", seed) % 90),
    }


def chain_task(depth, index):
    p = _chain_params(depth * 100 + index)
    steps = [
        {"id": "a", "tool": "read_sensor", "args": [p["sensor"]]},
    ]
    if depth >= 2:
        steps.append({"id": "b", "tool": "calibrate",
                      "args": ["$a", p["offset"]]})
    if depth >= 3:
        steps.append({"id": "c", "tool": "scale",
                      "args": ["$b", p["factor"]]})
    if depth >= 4:
        steps.append({"id": "d", "tool": "difference",
                      "args": ["$c", p["baseline"]]})
    if depth >= 5:
        steps.append({"id": "e", "tool": "aggregate",
                      "args": [["$d", p["anchor"]]]})
    return_ref = "$" + steps[-1]["id"]
    answer = eval_dag(steps, return_ref)
    if depth == 1:
        instruction = (
            f"Read sensor {p['sensor']} and report the exact integer reading."
        )
    elif depth == 2:
        instruction = (
            f"Read sensor {p['sensor']}, then calibrate that reading with "
            f"offset {p['offset']}. Report the exact integer calibrated value."
        )
    elif depth == 3:
        instruction = (
            f"Read sensor {p['sensor']}, calibrate that reading with offset "
            f"{p['offset']}, then scale the calibrated value by factor "
            f"{p['factor']}. Report the exact integer scaled value."
        )
    elif depth == 4:
        instruction = (
            f"Read sensor {p['sensor']}, calibrate that reading with offset "
            f"{p['offset']}, scale the calibrated value by factor "
            f"{p['factor']}, then subtract baseline {p['baseline']}. Report "
            f"the exact integer final value."
        )
    else:
        instruction = (
            f"Read sensor {p['sensor']}, calibrate that reading with offset "
            f"{p['offset']}, scale the calibrated value by factor "
            f"{p['factor']}, subtract baseline {p['baseline']}, then aggregate "
            f"the result with anchor {p['anchor']}. Report the exact integer "
            f"final value."
        )
    return {
        "task_id": f"chain_d{depth}_{index:02d}",
        "kind": "chain",
        "depth": depth,
        "width": None,
        "instruction": instruction,
        "answer": answer,
        "answer_format": "exact integer",
        "oracle_dag": steps,
        "oracle_return": return_ref,
    }


def _fanin_params(width, index):
    start = _h("fanin-start", width, index) % SENSOR_COUNT
    sensor_ids = [(start + 3 * j + index) % SENSOR_COUNT for j in range(width)]
    offsets = [-5 + (_h("fanin-offset", width, index, j) % 11)
               for j in range(width)]
    return {
        "sensor_ids": sensor_ids,
        "offsets": offsets,
        "factor": 2 + (_h("fanin-factor", width, index) % 4),
    }


def fanin_task(width, index):
    p = _fanin_params(width, index)
    steps = []
    calibrated_ids = []
    for j, (sid, off) in enumerate(zip(p["sensor_ids"], p["offsets"])):
        rid = f"r{j}"
        cid = f"c{j}"
        steps.append({"id": rid, "tool": "read_sensor", "args": [sid]})
        steps.append({"id": cid, "tool": "calibrate", "args": [f"${rid}", off]})
        calibrated_ids.append(f"${cid}")
    steps.append({"id": "agg", "tool": "aggregate", "args": [calibrated_ids]})
    steps.append({"id": "out", "tool": "scale", "args": ["$agg", p["factor"]]})
    answer = eval_dag(steps, "$out")
    pairs = ", ".join(f"{sid}:{off}" for sid, off in
                      zip(p["sensor_ids"], p["offsets"]))
    instruction = (
        f"For these sensor_id:offset pairs [{pairs}], read each sensor, "
        f"calibrate each reading with its offset, aggregate all calibrated "
        f"values by summing them, then scale the aggregate by factor "
        f"{p['factor']}. Report the exact integer final value."
    )
    return {
        "task_id": f"fanin_w{width}_{index:02d}",
        "kind": "fanin",
        "depth": None,
        "width": width,
        "instruction": instruction,
        "answer": answer,
        "answer_format": "exact integer",
        "oracle_dag": steps,
        "oracle_return": "$out",
    }


def build():
    tasks = []
    for depth in range(1, 6):
        for index in range(12):
            tasks.append(chain_task(depth, index))
    for width in (3, 5, 7, 9):
        for index in range(12):
            tasks.append(fanin_task(width, index))
    return {"tasks": tasks}


def main():
    out = ROOT / "tasks.json"
    data = build()
    for task in data["tasks"]:
        got = eval_dag(task["oracle_dag"], task["oracle_return"])
        if got != task["answer"]:
            raise SystemExit(f"gold mismatch for {task['task_id']}: {got} != {task['answer']}")
    out.write_text(json.dumps(data, indent=2) + "\n")
    print(f"wrote {len(data['tasks'])} tasks -> {out}")


if __name__ == "__main__":
    main()