SPB Git

spb/forge Public MIT

Forge — LLM training from scratch in pure C++20 + Metal on Apple Silicon.

C++ 61.2% C 23% Python 7.6% TeX 7.2% CMake 1.1%
8.7 KB · 325 lines cpp
Raw Blame History
1//-------------------------------------------------------------------------------------------------------------------------------------------------------------2//3// Foundation/NSSharedPtr.hpp4//5// Copyright 2020-2024 Apple Inc.6//7// Licensed under the Apache License, Version 2.0 (the "License");8// you may not use this file except in compliance with the License.9// You may obtain a copy of the License at10//11//     http://www.apache.org/licenses/LICENSE-2.012//13// Unless required by applicable law or agreed to in writing, software14// distributed under the License is distributed on an "AS IS" BASIS,15// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.16// See the License for the specific language governing permissions and17// limitations under the License.18//19//-------------------------------------------------------------------------------------------------------------------------------------------------------------2021#pragma once2223#include <cstddef>24#include "NSDefines.hpp"2526namespace NS27{28template <class _Class>29class SharedPtr30{31public:32    /**33     * Create a new null pointer.34     */35    SharedPtr();3637    /**38     * Destroy this SharedPtr, decreasing the reference count.39     */40    ~SharedPtr();4142    /**43     * Create a new null pointer.44     */45    SharedPtr(std::nullptr_t) noexcept;4647    /**48     * SharedPtr copy constructor.49     */50    SharedPtr(const SharedPtr<_Class>& other) noexcept;5152    /**53     * Construction from another pointee type.54     */55    template <class _OtherClass>56    SharedPtr(const SharedPtr<_OtherClass>& other, typename std::enable_if_t<std::is_convertible_v<_OtherClass *, _Class *>> * = nullptr) noexcept;5758    /**59     * SharedPtr move constructor.60     */61    SharedPtr(SharedPtr<_Class>&& other) noexcept;6263    /**64     * Move from another pointee type.65     */66    template <class _OtherClass>67    SharedPtr(SharedPtr<_OtherClass>&& other, typename std::enable_if_t<std::is_convertible_v<_OtherClass *, _Class *>> * = nullptr) noexcept;6869    /**70     * Copy assignment operator.71     * Copying increases reference count. Only releases previous pointee if objects are different.72     */73    SharedPtr& operator=(const SharedPtr<_Class>& other);7475    /**76     * Copy-assignment from different pointee.77     * Copying increases reference count. Only releases previous pointee if objects are different.78     */79    template <class _OtherClass>80    typename std::enable_if_t<std::is_convertible_v<_OtherClass *, _Class *>, SharedPtr &>81    operator=(const SharedPtr<_OtherClass>& other);8283    /**84     * Move assignment operator.85     * Move without affecting reference counts, unless pointees are equal. Moved-from object is reset to nullptr.86     */87    SharedPtr& operator=(SharedPtr<_Class>&& other);8889    /**90     * Move-asignment from different pointee.91     * Move without affecting reference counts, unless pointees are equal. Moved-from object is reset to nullptr.92     */93    template <class _OtherClass>94    typename std::enable_if_t<std::is_convertible_v<_OtherClass *, _Class *>, SharedPtr &>95    operator=(SharedPtr<_OtherClass>&& other);9697    /**98     * Access raw pointee.99     * @warning Avoid wrapping the returned value again, as it may lead double frees unless this object becomes detached.100     */101    _Class* get() const;102103    /**104     * Call operations directly on the pointee.105     */106    _Class* operator->() const;107108    /**109     * Implicit cast to bool.110     */111    explicit operator bool() const;112113    /**114     * Reset this SharedPtr to null, decreasing the reference count.115     */116    void reset();117118    /**119     * Detach the SharedPtr from the pointee, without decreasing the reference count.120     */121    void detach();122123    template <class _OtherClass>124    friend SharedPtr<_OtherClass> RetainPtr(_OtherClass* ptr);125126    template <class _OtherClass>127    friend SharedPtr<_OtherClass> TransferPtr(_OtherClass* ptr);128129private:130    _Class* m_pObject;131};132133/**134 * Create a SharedPtr by retaining an existing raw pointer.135 * Increases the reference count of the passed-in object.136 * If the passed-in object was in an AutoreleasePool, it will be removed from it.137 */138template <class _Class>139_NS_INLINE NS::SharedPtr<_Class> RetainPtr(_Class* pObject)140{141    NS::SharedPtr<_Class> ret;142    ret.m_pObject = pObject->retain();143    return ret;144}145146/*147 * Create a SharedPtr by transfering the ownership of an existing raw pointer to SharedPtr.148 * Does not increase the reference count of the passed-in pointer, it is assumed to be >= 1.149 * This method does not remove objects from an AutoreleasePool.150*/151template <class _Class>152_NS_INLINE NS::SharedPtr<_Class> TransferPtr(_Class* pObject)153{154    NS::SharedPtr<_Class> ret;155    ret.m_pObject = pObject;156    return ret;157}158159}160161template <class _Class>162_NS_INLINE NS::SharedPtr<_Class>::SharedPtr()163    : m_pObject(nullptr)164{165}166167template <class _Class>168_NS_INLINE NS::SharedPtr<_Class>::~SharedPtr<_Class>() __attribute__((no_sanitize("undefined")))169{170    m_pObject->release();171}172173template <class _Class>174_NS_INLINE NS::SharedPtr<_Class>::SharedPtr(std::nullptr_t) noexcept175    : m_pObject(nullptr)176{177}178179template <class _Class>180_NS_INLINE NS::SharedPtr<_Class>::SharedPtr(const SharedPtr<_Class>& other) noexcept181    : m_pObject(other.m_pObject->retain())182{183}184185template <class _Class>186template <class _OtherClass>187_NS_INLINE NS::SharedPtr<_Class>::SharedPtr(const SharedPtr<_OtherClass>& other, typename std::enable_if_t<std::is_convertible_v<_OtherClass *, _Class *>> *) noexcept188    : m_pObject(reinterpret_cast<_Class*>(other.get()->retain()))189{190}191192template <class _Class>193_NS_INLINE NS::SharedPtr<_Class>::SharedPtr(SharedPtr<_Class>&& other) noexcept194    : m_pObject(other.m_pObject)195{196    other.m_pObject = nullptr;197}198199template <class _Class>200template <class _OtherClass>201_NS_INLINE NS::SharedPtr<_Class>::SharedPtr(SharedPtr<_OtherClass>&& other, typename std::enable_if_t<std::is_convertible_v<_OtherClass *, _Class *>> *) noexcept202    : m_pObject(reinterpret_cast<_Class*>(other.get()))203{204    other.detach();205}206207template <class _Class>208_NS_INLINE _Class* NS::SharedPtr<_Class>::get() const209{210    return m_pObject;211}212213template <class _Class>214_NS_INLINE _Class* NS::SharedPtr<_Class>::operator->() const215{216    return m_pObject;217}218219template <class _Class>220_NS_INLINE NS::SharedPtr<_Class>::operator bool() const221{222    return nullptr != m_pObject;223}224225template <class _Class>226_NS_INLINE void NS::SharedPtr<_Class>::reset() __attribute__((no_sanitize("undefined")))227{228    m_pObject->release();229    m_pObject = nullptr;230}231232template <class _Class>233_NS_INLINE void NS::SharedPtr<_Class>::detach()234{235    m_pObject = nullptr;236}237238template <class _Class>239_NS_INLINE NS::SharedPtr<_Class>& NS::SharedPtr<_Class>::operator=(const SharedPtr<_Class>& other) __attribute__((no_sanitize("undefined")))240{241    _Class* pOldObject = m_pObject;242243    m_pObject = other.m_pObject->retain();244245    pOldObject->release();246247    return *this;248}249250template <class _Class>251template <class _OtherClass>252typename std::enable_if_t<std::is_convertible_v<_OtherClass *, _Class *>, NS::SharedPtr<_Class> &>253_NS_INLINE NS::SharedPtr<_Class>::operator=(const SharedPtr<_OtherClass>& other) __attribute__((no_sanitize("undefined")))254{255    _Class* pOldObject = m_pObject;256257    m_pObject = reinterpret_cast<_Class*>(other.get()->retain());258259    pOldObject->release();260261    return *this;262}263264template <class _Class>265_NS_INLINE NS::SharedPtr<_Class>& NS::SharedPtr<_Class>::operator=(SharedPtr<_Class>&& other) __attribute__((no_sanitize("undefined")))266{267    if (m_pObject != other.m_pObject)268    {269        m_pObject->release();270        m_pObject = other.m_pObject;271    }272    else273    {274        m_pObject = other.m_pObject;275        other.m_pObject->release();276    }277    other.m_pObject = nullptr;278    return *this;279}280281template <class _Class>282template <class _OtherClass>283typename std::enable_if_t<std::is_convertible_v<_OtherClass *, _Class *>, NS::SharedPtr<_Class> &>284_NS_INLINE NS::SharedPtr<_Class>::operator=(SharedPtr<_OtherClass>&& other) __attribute__((no_sanitize("undefined")))285{286    if (m_pObject != other.get())287    {288        m_pObject->release();289        m_pObject = reinterpret_cast<_Class*>(other.get());290        other.detach();291    }292    else293    {294        m_pObject = other.get();295        other.reset();296    }297    return *this;298}299300template <class _ClassLhs, class _ClassRhs>301_NS_INLINE bool operator==(const NS::SharedPtr<_ClassLhs>& lhs, const NS::SharedPtr<_ClassRhs>& rhs)302{303    return lhs.get() == rhs.get();304}305306template <class _ClassLhs, class _ClassRhs>307_NS_INLINE bool operator!=(const NS::SharedPtr<_ClassLhs>& lhs, const NS::SharedPtr<_ClassRhs>& rhs)308{309    return lhs.get() != rhs.get();310}311312namespace std313{314315template <class T>316struct hash<NS::SharedPtr<T>>317{318    size_t operator()(const NS::SharedPtr<T>& p) const319    {320        return std::hash<T*>{}(p.get());321    }322};323324} // namespace std325