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%
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