def infer_or_check_handler_id(
    tag: str,
    subtree: Any,
    explicit_id: str | None = None,
    min_hash_length: int = 5,
) -> HandlerId:
  """Tries to generate a unique handler ID from the structure of a subtree.

  Each handler in Penzai's effect system needs to have a unique ID that
  identifies which effects it is responsible for handling. Since effect requests
  and effect handlers are both PyTree nodes that appear inside the PyTree
  structure, we can generate handler IDs by hashing the structure of the
  subtree, which should almost always result in a unique ID for each handler.
  (Failing this would mean there exists a hash such that hashing that hash
  inside its structure yields the structure again.)

  Args:
    tag: User-readable tag for this handler ID. Mostly to help identify what
      kind of handler this is for debugging purposes.
    subtree: Subtree with effects that is being wrapped by the effect handler.
    explicit_id: An explicit ID requested by the user. We will use this if
      provided and if it does not already appear in the structure.
    min_hash_length: Minimum number of hash characters to append to try to make
      IDs unique. Note that we may append more characters as needed to ensure
      uniqueness.

  Returns:
    A unique handler ID which does not already appear in this structure.

  Raises:
    RuntimeError: If there is a hash collision, and increasing the length of
      the hash did not resolve the collision.
    ValueError: If there was an explicit ID provided, but it conflicts with an
      existing ID in the model.
  """
  existing = all_handler_ids(subtree)
  if explicit_id is not None:
    if explicit_id in existing:
      raise ValueError(
          f"Explicit handler ID {explicit_id} already exists in the subtree!"
      )
    return explicit_id
  # Extract and hash the structure. We just hash the string representation of
  # the pytreedef, which isn't necessarily unique but should be close enough
  # for our purposes. Note that any already-bound effects will have their
  # handler IDs actually shown inside this structure, since handler IDs are not
  # PyTree nodes.
  structure = jax.tree_util.tree_structure(subtree)
  digest = hashlib.sha256(repr(structure).encode("utf-8")).hexdigest()
  long_id = f"{tag}_{digest}"
  # Avoid conflicts.
  potential_conflicts = {hid for hid in existing if long_id.startswith(hid)}
  for length in range(min_hash_length, len(digest)):
    handler_id = f"{tag}_{digest[:length]}"
    if handler_id not in potential_conflicts:
      return handler_id
  # If we get here, we couldn't find a unique ID.
  raise RuntimeError(
      f"Hash collision while generating a unique ID for tag {tag} with subtree"
      f" hash {digest}! This usually means you have a very large number of"
      " handlers and are also extremely unlucky. Try providing an explicit"
      " name for the handler or making a small change to your model's"
      " structure."
  )
Tries to generate a unique handler ID from the structure of a subtree.

  Each handler in Penzai's effect system needs to have a unique ID that
  identifies which effects it is responsible for handling. Since effect requests
  and effect handlers are both PyTree nodes that appear inside the PyTree
  structure, we can generate handler IDs by hashing the structure of the
  subtree, which should almost always result in a unique ID for each handler.
  (Failing this would mean there exists a hash such that hashing that hash
  inside its structure yields the structure again.)

  Args:
    tag: User-readable tag for this handler ID. Mostly to help identify what
      kind of handler this is for debugging purposes.
    subtree: Subtree with effects that is being wrapped by the effect handler.
    explicit_id: An explicit ID requested by the user. We will use this if
      provided and if it does not already appear in the structure.
    min_hash_length: Minimum number of hash characters to append to try to make
      IDs unique. Note that we may append more characters as needed to ensure
      uniqueness.

  Returns:
    A unique handler ID which does not already appear in this structure.

  Raises:
    RuntimeError: If there is a hash collision, and increasing the length of
      the hash did not resolve the collision.
    ValueError: If there was an explicit ID provided, but it conflicts with an
      existing ID in the model.
  
penzai/deprecated/v1/data_effects/effect_base.py