from typing import List


def build_branches(program: List[dict],
                   branches_end_nodes: List[int]) -> List[List[int]]:
    """
    Build branches (currently only 2 branches are possible) by iterating through
    the program. Stop once all branches_end_nodes are reached.

    Parameters
    ---
    program  (List[dict])
        Functional program
    branches_end_nodes (List[int])
        Indices of the last nodes in branches before the merge into a single node

    Result
    ---
    List[List[int]]
    List of branches (only 2) containing indices of nodes. 

    """

    # not really important since we know it's only 2, but in case
    # this changes in the future
    num_branches = len(branches_end_nodes)
    branches = [[] for i in range(num_branches)]

    for branch_idx, end_node_idx in enumerate(branches_end_nodes):
        branches[branch_idx].append(end_node_idx)
        inputs = program[end_node_idx]["inputs"]

        # stop when we reach empty inputs (i.e. scene program)
        while inputs:
            # there shouldn't be anymore branches
            assert len(inputs) == 1
            prev_node = inputs[0]
            # append current branch with previous node
            branches[branch_idx].append(prev_node)
            inputs = program[prev_node]["inputs"]

    return branches