from __future__ import annotations

from .._utils.protocols import is_device_array
from ._base import IteratorBase

CUDA_PREAMBLE = """#include <cuda/std/cstdint>
#include <cuda_fp16.h>
#include <cuda/std/cstring>
using namespace cuda::std;
"""


def ensure_iterator(obj):
    """Wrap array in PointerIterator if needed."""
    from ._pointer import PointerIterator

    if isinstance(obj, IteratorBase):
        return obj
    if is_device_array(obj):
        return PointerIterator(obj)
    raise TypeError("Expected an iterator or a device array")
