Task builder

def add_subtask(
    self,
    kernel: ir.Method[CallArgs, T],
    num_shots: int,
    metadata: dict | None = None,
    *args: CallArgs.args,
    **kernel_args: CallArgs.kwargs,
) -> int:
    """Append one kernel invocation and return its subtask index.

    Positional and keyword kernel arguments are bound against the kernel's
    Python signature and stored by parameter name. ``metadata`` is
    JSON-encoded as per-subtask user metadata. Reusing the same ``kernel``
    object reuses its program index while retaining independent arguments,
    metadata, and shot counts for this subtask.

    Args:
        kernel (ir.Method): Kernel to invoke.
        num_shots (int): Number of shots for this invocation. Must be at
            least one.
        metadata (dict | None): Optional JSON-serializable user metadata
            associated with this subtask. Defaults to None.
        *args: Positional arguments for ``kernel``.
        **kernel_args: Keyword arguments for ``kernel``.

    Returns:
        int: Zero-based index of the appended subtask.

    Raises:
        TypeError: If the kernel has no inspectable Python function or the
            supplied arguments do not match its signature.
        ValueError: If ``metadata`` is not a dictionary or ``None``.
        SubtaskValidationError: If the shot count or bound kernel arguments
            are invalid.
    """
    # TODO: define equality? we COULD implement hashing on kirin kernels. as a first pass maybe it's OK to
    # just do exact object check
    if metadata is not None and not isinstance(metadata, dict):
        raise ValueError("expecting `metadata` to be a dictionary or None. ")

    base_kernel_name = kernel.sym_name if kernel.sym_name is not None else "kernel"
    kernel_name_index = self._kernel_name_counter.get(base_kernel_name, 0)
    kernel_name = (
        base_kernel_name
        if kernel_name_index == 0
        else f"{base_kernel_name}_{kernel_name_index}"
    )

    if kernel in self._programs:
        program_idx = self._programs[kernel]
        is_new_program = False
    else:
        program_idx = self._max_program_idx
        is_new_program = True

    if metadata is not None:
        subtask_metadata = TaskMetadata(user_metadata=json.dumps(metadata))
    else:
        subtask_metadata = None

    if kernel.py_func is None:
        raise TypeError(
            f"Cannot determine the argument names for kernel {kernel.sym_name!r}"
        )

    signature = inspect.signature(kernel.py_func)
    bound = signature.bind(*args, **kernel_args)

    arguments = dict(bound.arguments)

    qlam_subtask = Subtask(
        program_index=program_idx,
        num_shots=num_shots,
        arguments=arguments,
        subtask_metadata=subtask_metadata,
    )

    new_subtask = _Subtask(kernel_name=kernel_name, qlam_subtask=qlam_subtask)
    try:
        self._validate_subtask(subtask=new_subtask)
    except ValueError as e:
        raise SubtaskValidationError(f"Error during subtask validation: {e}")

    # Commit builder state only after every operation above has succeeded.
    if is_new_program:
        self._programs[kernel] = program_idx
        self._max_program_idx += 1
    self._kernel_name_counter[base_kernel_name] = kernel_name_index + 1
    self._subtasks.append(new_subtask)
    return len(self._subtasks) - 1

def copy(self) -> "TaskBuilder":
    """Return a shallow, independently editable copy of this builder.

    The mutable program, subtask, and name-index containers are copied, so
    later calls to ``add_subtask()`` affect only the copy. Existing kernel
    objects and subtask records are shared rather than deep-copied.

    Returns:
        TaskBuilder: Builder containing the same programs and subtasks.
    """
    return TaskBuilder(
        _programs=dict(self._programs),
        _subtasks=list(self._subtasks),
        _max_program_idx=self._max_program_idx,
        _kernel_name_counter=dict(self._kernel_name_counter),
    )

def print_detailed(self) -> str:
    """Return a detailed summary containing program IR and subtask arguments."""
    ret_str = ""
    ret_str += "Task:\n"
    ret_str += "\tPrograms:\n"
    for program, program_idx in sorted(self._programs.items(), key=lambda x: x[1]):
        ret_str += f"\t\t{program_idx}. {program.print_str()}"

    ret_str += "\tSubtasks:\n"
    for subtask_idx, subtask in enumerate(self._subtasks):
        subtask_metadata = subtask.qlam_subtask.subtask_metadata
        if subtask_metadata is not None:
            user_metadata = subtask_metadata.user_metadata
            if user_metadata is not None:
                user_metadata = json.loads(user_metadata)
        else:
            user_metadata = None
        ret_str += f"\t\t{subtask_idx}. Subtask: {subtask.kernel_name}, Program {subtask.qlam_subtask.program_index}, Args {subtask.qlam_subtask.arguments}, Metadata {user_metadata} -> {subtask.qlam_subtask.num_shots} shots\n"

    return ret_str

def summary(self) -> str:
    """Return a human-readable summary printed on dry-run.

    Returns:
        str: Summary describing what would be submitted.
    """
    ret_str = ""
    ret_str += "Task:\n"
    ret_str += "\tPrograms:\n"
    for program, program_idx in sorted(self._programs.items(), key=lambda x: x[1]):
        ret_str += f"\t\t{program_idx}. {program.sym_name}\n"

    ret_str += "\tSubtasks:\n"
    for idx, subtask in enumerate(self._subtasks):
        subtask_metadata = subtask.qlam_subtask.subtask_metadata
        if subtask_metadata is not None:
            user_metadata = subtask_metadata.user_metadata
            if user_metadata is not None:
                user_metadata = json.loads(user_metadata)
        else:
            user_metadata = None
        subtask_str = f"\t\t{idx}. Subtask: {subtask.kernel_name}, Program {subtask.qlam_subtask.program_index}, Args {subtask.qlam_subtask.arguments}, Metadata {user_metadata} -> {subtask.qlam_subtask.num_shots} shots\n"
        ret_str += subtask_str

    return ret_str