```python:inference/core/workflows/execution_engine/v1/executor/execution_data_manager/step_input_assembler.py
from typing import Dict, List, Optional, Set

from inference.core.workflows.execution_engine.v1.executor.execution_data_manager.dynamic_batches_manager import \
    DynamicBatchIndex


def intersect_masks_per_dimension(
    batch_masks: List[Set[DynamicBatchIndex]],
    dimensions: Set[int],
) -> Dict[int, Set[DynamicBatchIndex]]:
    """Intersect masks at each dimensionality level.

    For each dimension d, returns the set of indices (with length d) that appear
    in every mask that has at least one index at that dimension. Masks with no
    indices at d are ignored for that dimension. Used for intra-dimensional
    intersection.
    """
    sorted_dims = sorted(dimensions)
    result: Dict[int, Set[DynamicBatchIndex]] = {}
    for dim in sorted_dims:
        acc: Optional[Set[DynamicBatchIndex]] = None
        # Iterate masks once per dimension, but avoid creating intermediate lists
        for mask in batch_masks:
            # Build the set of indices in this mask that have length == dim
            s = {idx for idx in mask if len(idx) == dim}
            if not s:
                # Masks with no indices at d are ignored
                continue
            if acc is None:
                # First non-empty set becomes the accumulator (copy to avoid mutating inputs)
                acc = set(s)
            else:
                # Intersect in place to avoid extra temporaries
                acc &= s
                if not acc:
                    # Early break: intersection is empty, further intersections cannot add elements
                    break
        result[dim] = acc if acc is not None else set()
    return result


def filter_to_valid_prefix_chains(
    per_dim_sets: Dict[int, Set[DynamicBatchIndex]],
    dimensions: Set[int],
) -> Dict[int, Set[DynamicBatchIndex]]:
    """Keep only indices that form a complete parent-child chain across dimensions.

    Given per-dimension sets (e.g. from intersect_masks_per_dimension), retains
    only indices that have a full lineage from the smallest to the largest
    dimension. Used for inter-level intersection.
    """
    sorted_dims = sorted(dimensions)
    by_dim: Dict[int, Set[DynamicBatchIndex]] = {
        dim: per_dim_sets.get(dim, set()) for dim in sorted_dims
    }

    if len(sorted_dims) <= 1:
        return dict(by_dim)

    prev_dim = {sorted_dims[i]: sorted_dims[i - 1] for i in range(1, len(sorted_dims))}

    # Bottom-up: mark indices that have at least one descendant
    has_child: Set[DynamicBatchIndex] = set()
    last_dim = sorted_dims[-1]
    # Iterate from largest to smallest dimension
    for dim in reversed(sorted_dims):
        # Localize for speed
        elems = by_dim[dim]
        if not elems:
            continue
        if dim == last_dim:
            # Any index at the last level contributes its parent
            for idx in elems:
                parent = idx[:-1]
                if parent:
                    has_child.add(parent)

    # Top-down: keep indices only if full prefix chain exists
        else:
            # Only indices that either are in has_child (i.e., have descendant) should propagate parents
            for idx in elems:
                if idx in has_child:
                    parent = idx[:-1]
                    if parent:
                        has_child.add(parent)

    # Top-down: keep indices only if full prefix chain exists
    valid: Dict[int, Set[DynamicBatchIndex]] = {dim: set() for dim in sorted_dims}
    first_dim = sorted_dims[0]
    last_dim = sorted_dims[-1]
    for dim in sorted_dims:
        elems = by_dim[dim]
        if not elems:
            continue
        if dim == first_dim:
            # Root-level indices are valid only if they have at least one descendant
            for idx in elems:
                if idx in has_child:
                    valid[dim].add(idx)
        else:
            parent_set = valid[prev_dim[dim]]
            if not parent_set:
                # No valid parents at this level, skip early
                continue
            if dim == last_dim:
                # For last dimension, any child of a valid parent is valid
                for idx in elems:
                    parent = idx[:-1]
                    if parent in parent_set:
                        valid[dim].add(idx)
            else:
                # For intermediate dims, require both a valid parent and that this index has a descendant
                for idx in elems:
                    parent = idx[:-1]
                    if parent in parent_set and idx in has_child:
                        valid[dim].add(idx)


    return valid


def get_masks_intersection_for_dimensions(
    batch_masks: List[Set[DynamicBatchIndex]],
    dimensions: Set[int],
) -> Dict[int, Optional[Set[DynamicBatchIndex]]]:
    """Intersect masks at each dimension and filter to valid prefix chains."""
    if not batch_masks:
        return {dim: None for dim in dimensions}

    sorted_dims = sorted(dimensions)

    if len(sorted_dims) <= 1:
        result = intersect_masks_per_dimension(batch_masks, dimensions)
        return {dim: result[dim] for dim in sorted_dims}

    per_dim = intersect_masks_per_dimension(batch_masks, dimensions)
    return filter_to_valid_prefix_chains(per_dim, dimensions)
```