import pathlib
from typing import Dict
from typing import FrozenSet
import itertools


def _source_files_in_same_targets(path: pathlib.Path, related_source_files: Dict[pathlib.Path, FrozenSet[pathlib.Path]]) -> FrozenSet[pathlib.Path]:
    """Returns `.rs` file paths relating to `path`.

    :param path: Path to a `.rs` file
    :param related_source_files: Output of `_related_source_files`
    :returns: Relating `.rs` file paths
    """
    # If `p` is `src_path` of a target, it does not belong to any other target unless it's weirdly symlinked,
    if path in related_source_files:
        return frozenset({path, *related_source_files[path]})

    # Otherwise, it may be used by multiple targets with `#[path = ".."] mod foo;` or something.
    return frozenset(itertools.chain.from_iterable({k, *v} for (k, v) in related_source_files.items() if path in v)) or frozenset({path})