1# BSD 3-Clause License
  2#
  3# Copyright (c) 2022-Present, nxtlo
  4# All rights reserved.
  5#
  6# Redistribution and use in source and binary forms, with or without
  7# modification, are permitted provided that the following conditions are met:
  8#
  9# * Redistributions of source code must retain the above copyright notice, this
 10#   list of conditions and the following disclaimer.
 11#
 12# * Redistributions in binary form must reproduce the above copyright notice,
 13#   this list of conditions and the following disclaimer in the documentation
 14#   and/or other materials provided with the distribution.
 15#
 16# * Neither the name of the copyright holder nor the names of its
 17#   contributors may be used to endorse or promote products derived from
 18#   this software without specific prior written permission.
 19#
 20# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 21# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 22# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 23# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 24# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 25# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 26# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 27# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 28# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 29# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 30"""A wrapper type to construct uninitialized instances of T."""
 31
 32from __future__ import annotations
 33
 34import typing
 35
 36from . import macros
 37from .macros import rustc_diagnostic_item
 38
 39if typing.TYPE_CHECKING:
 40    import collections.abc as collections
 41
 42    from typing_extensions import Self
 43
 44T = typing.TypeVar("T")
 45
 46
 47@rustc_diagnostic_item("MaybeUninit")
 48@typing.final
 49class MaybeUninit(typing.Generic[T]):
 50    """A wrapper type to construct uninitialized instances of `T`.
 51
 52    This is kinda equivalent to Rust's `MaybeUninit<T>` wrapper.
 53
 54    The difference is what's happening under the hood is when initializing an instance of this object,
 55    nothing really is being initialized, until you call `.write()` to initialize the inner value.
 56
 57    What??
 58    -----
 59    Ok, so `MaybeUninit<T>` is in a sense the unsafe equivalent of `Option<T>`, it either contains a value of type `T`,
 60    or it contains uninitialized memory (`the attribute doesn't exist`).
 61
 62    However, `MaybeUninit` is unable to tell you whether the value it contains is `Some(T)` or `None`, So you
 63    as a programmer is responsible for initializing it.
 64
 65    And by default, `Option<T>` always contain a default value, which is `None`.
 66
 67    Examples
 68    --------
 69    ```py
 70    # Let's write a simple wrapper around `MaybeUninit<T>` to create a dynamic array.
 71    class SizedArray[T]:
 72        def __init__(self, size: int) -> None:
 73            # Creates count `size` of `MaybeUninit<T>` in an uninitialized state.
 74            self.buf = MaybeUninit[T].uninit_array(size)
 75            # Track how many initialized elements are in the array.
 76            self.len = 0
 77
 78        def push(self, value: T) -> None:
 79            # if the length of the array is equal to the size of the array,
 80            # then we've reached the max size of the array.
 81            if self.len == len(self.buf):
 82                raise RuntimeError("max size reached")
 83
 84            # get the next uninitialized element and write into it.
 85            self.buf[self.len].write(value)
 86            self.len += 1
 87
 88        def copy_from(self, other: Sequence[T]) -> None:
 89            # get the remaining space in the array.
 90            spared = self.buf[self.len:]
 91            other_len = len(other)
 92            # Check if there's enough capacity to copy the other list into the array.
 93            if len(spared) < other_len:
 94                raise RuntimeError("not enough space")
 95
 96            # Copy the other list into the array.
 97            index = 0
 98            while index < other_len:
 99                self.buf[self.len + index].write(other[index])
100                index += 1
101
102            self.len += other_len
103
104        def read(self) -> Sequence[T]:
105            return MaybeUninit.array_assume_init(self.buf[:self.len])
106
107    array = SizedArray[int](3)
108    array.push(1)
109    array.copy_from([2, 3])
110    assert array.read() == (1, 2, 3)
111    ```
112    """
113
114    __slots__ = ("__value",)
115
116    @typing.overload
117    def __init__(self) -> None: ...
118
119    @typing.overload
120    def __init__(self, value: T) -> None:
121        """Creates a new MaybeUninit<T> initialized with the given value.
122        It is safe to call assume_init on the return value of this function.
123
124        Example
125        -------
126        ```py
127        v = MaybeUninit(1)
128        assert v.assume_init() == 1
129        ```
130        """
131
132    def __init__(self, value: T | None = None) -> None:
133        if value is None:
134            return None
135        else:
136            # we simply pre-initialize the value if it was passed
137            # when constructing the instance.
138            self.__value = value
139
140    @classmethod
141    def uninit(cls) -> Self:
142        """Creates a new `MaybeUninit<T>` in an uninitialized state.
143
144        Example
145        -------
146        ```py
147        v: MaybeUninit[str] = MaybeUninit.uninit() # or just MaybeUninit()
148        ```
149        """
150        return cls()
151
152    @classmethod
153    def uninit_array(cls, n: int) -> collections.Sequence[Self]:
154        """Creates an immutable sequence of `MaybeUninit<T>` in an uninitialized state.
155
156        Example
157        -------
158        ```py
159        v = MaybeUninit[str].uninit_array(3)
160        assert len(v) == 3
161        for uninit in v:
162            uninit.write('content')
163
164            initialized = uninit.assume_init()
165        ```
166        """
167        return tuple(cls() for _ in range(n))
168
169    @macros.unsafe
170    def assume_init(self) -> T:
171        """Get the inner value, assuming that it was initialized by the caller.
172
173        It is unsafe and undefined behavior to call this method on an uninitialized state,
174
175        Example
176        -------
177        ```py
178        uninit: MaybeUninit[int] = MaybeUninit.uninit()
179        val = uninit.assume_init() # This is UNSAFE ⚠️
180
181        # Initialize it first.
182        uninit.write(0)
183        val = uninit.assume_init() # This is safe to access.
184        ```
185        """
186        # SAFETY: the caller must guarantee that `self` is initialized.
187        return self.__value
188
189    @staticmethod
190    @macros.unsafe
191    def array_assume_init(
192        array: collections.Sequence[MaybeUninit[T]],
193    ) -> collections.Sequence[T]:
194        """Extracts a sequence of `MaybeUninit[T]` containers.
195
196        You must guarantee that all elements in the array are initialized.
197
198        Example
199        -------
200        ```py
201        array: list[MaybeUninit[int]] = [MaybeUninit(), MaybeUninit(), MaybeUninit()]
202        array[0].write(0)
203        array[1].write(1)
204        array[2].write(2)
205        # transposed into a tuple.
206        assert MaybeUninit.array_assume_init(array) == (0, 1, 2)
207        ```
208        """
209        # SAFETY: The caller guarantees that all elements of the array are initialized
210        return tuple(uninit.assume_init() for uninit in array)
211
212    @staticmethod
213    @macros.unsafe
214    def array_assume_init_mut(
215        array: collections.Sequence[MaybeUninit[T]],
216    ) -> collections.MutableSequence[T]:
217        """Extracts a sequence of `MaybeUninit[T]` to a list of `T`s.
218
219        You must guarantee that all elements in the array are initialized.
220
221        Example
222        -------
223        ```py
224        array: list[MaybeUninit[int]] = [MaybeUninit(), MaybeUninit(), MaybeUninit()]
225        array[0].write(0)
226        array[1].write(1)
227        array[2].write(2)
228        # transposed into a list.
229        assert MaybeUninit.array_assume_init_mut(array) == [0, 1, 2]
230        ```
231        """
232        # SAFETY: The caller guarantees that all elements of the array are initialized
233        return [uninit.assume_init() for uninit in array]
234
235    def write(self, value: T) -> T:
236        """Sets the value of the `MaybeUninit[T]`, Then return a reference to it.
237
238        This will overwrite any previous values, if was initialized.
239
240        Example
241        -------
242        ```py
243        buffer: MaybeUninit[bytes] = MaybeUninit()
244        buffer.write(b"Hello, World!")
245        print(buffer.assume_init())
246        ```
247        """
248        self.__value = value
249        # SAFETY: we just initialized the value, so it's safe to access.
250        return self.__value
251
252    def __repr__(self) -> str:
253        if self:
254            return f"MaybeUninit(value: {self.__value!r})"
255        return "<uninit>"
256
257    __str__ = __repr__
258
259    def __bool__(self) -> bool:
260        """Perform a boolean check on whether the value is initialized or not.
261
262        Example
263        -------
264        ```py
265        v = MaybeUninit[bool]()
266        assert not v
267        ```
268        """
269        return hasattr(self, "_MaybeUninit__value")
270
271    def __hash__(self) -> int:
272        return self.__value.__hash__()