[stdlib] Implement Weak[T] refs for Arc[T] #3450

szbergeron · 2024-09-04T23:17:17Z

This implements a Weak[T] non-owning ref for Arc[T] that should allow for more ergonomic implementation of cyclic data structures without having to fall back to UnsafePointer[T] or Optional[Arc[T]] + complex directed drop impls.

This increases the footprint of _ArcInner by (minimally) 8 bytes. If we wish to cut this down, specialization and conditional inclusion of the field (once implemented) may allow for Arc[T, Kind=StrongOnly] constructs.

Signed-off-by: Sawyer Bergeron <sawyerbergeron@gmail.com>

szbergeron · 2024-09-23T13:11:48Z

Needs some fixes after parameterizing Arc for Weak, but recent compiler changes have fixed the parameter propagation/inference issues so it's just a matter of fixing the broken ref count.

ConnorGray

Hey @szbergeron, thanks for tackling this! 🙂 This will be a nice QoL thing for folks doing self-referential data structures in Mojo.

I've left a few initial comments on the implementation strategy, particularly about how we might be able to make stronger guarantees around the enable_weak: Bool parameter using the type system.

Let me know if my suggestions are workable! 🙂

Thanks again for your contribution, awesome stuff 🔥

ConnorGray · 2024-10-15T21:30:46Z

stdlib/src/memory/arc.mojo

+        T: The type of the value that this Weak may be upgraded into an Arc[T] of.
+    """
+
+    alias _inner_type = _ArcInner[T]


Question Should this be:

Suggested change

alias _inner_type = _ArcInner[T]

alias _inner_type = _ArcInner[T, True]

since this Weak must have been derived from an Arc[T, True]?

ConnorGray · 2024-10-15T21:30:51Z

stdlib/src/memory/arc.mojo

-        return self._inner[].payload
+        return self._inner[].payload.assume_initialized()
+
+    fn downgrade(self) -> Weak[T, enable_weak]:


Suggested change

fn downgrade(self) -> Weak[T, enable_weak]:

fn downgrade(self: Arc[T, True]) -> Weak[T]:

ConnorGray · 2024-10-15T21:31:54Z

stdlib/src/memory/arc.mojo

+        else:
+            constrained[False, "A typical Arc[T] does not support downgrading." + 
+                "If you want to enable downgrading of this Arc[T], convert it to an Arc[T, enable_weak = True]"]()
+            return abort[Weak[T, enable_weak]]("unreachable")


Suggestion By constraining self: Arc[T, True (shown in my comment above), I think this whole statement could be simplified to just:

Suggested change

return abort[Weak[T, enable_weak]]("unreachable")

return Weak(self)

I like the simplification, but is there a good way with that change of presenting specific, actionable instructions to the user that what they are trying to do can only be done if they make their Arc enable_weak = True?

ConnorGray · 2024-10-15T21:35:25Z

stdlib/src/memory/arc.mojo

@@ -110,6 +309,16 @@ struct Arc[T: Movable](CollectionElement, CollectionElementNew):
        other._inner[].add_ref()
        self._inner = other._inner

+    fn __init__(inout self, ptr: UnsafePointer[Self._inner_type], private: ()):


Suggestion We could annotate this with @doc_private instead of using a private parameter.

ConnorGray · 2024-10-15T21:35:57Z

stdlib/src/os/atomic.mojo

@@ -24,7 +24,7 @@ from memory import UnsafePointer, bitcast
 from sys.info import triple_is_nvidia_cuda


-struct Atomic[type: DType]:
+struct Atomic[type: DType](ExplicitlyCopyable):


Question Why is this needed?

This was to allow me to construct the [0|1] length array for weak. I wanted to revisit the specifics, since having an array of atomics can be useful in a lot of contexts but there isn't currently a good way of constructing one without arcane MaybeUninit incantations. I do have some questions about the atomic type that can be had out-of-thread too, that (I think) make Arc technically unsound in a single threaded context and relatively impossible to use in a multi-threaded one

ConnorGray · 2024-10-15T22:29:56Z

stdlib/src/memory/arc.mojo

+    alias _inner_type = _ArcInner[T]
+    var _inner: UnsafePointer[Self._inner_type]
+
+    fn __init__(inout self, strong: Arc[T, enable_weak]):


Suggestion Likewise, we can change this initializer so that Weak[T] can only be constructed frfom an Arc[T, True]:

fn __init__(inout self, strong: Arc[T, True]):

This would let you remove the constrained[enable_weak, "..."]() below as well 🙂

ConnorGray · 2024-10-15T22:31:24Z

stdlib/src/memory/arc.mojo

+        other._inner[].add_weak()
+        self._inner = other._inner
+
+    fn upgrade(inout self) -> Optional[Arc[T]]:


Suggestion I think this needs to encode that this is an Arc containing a weak count:

Suggested change

fn upgrade(inout self) -> Optional[Arc[T]]:

fn upgrade(inout self) -> Optional[Arc[T, True]]:

ConnorGray · 2024-10-15T22:39:44Z

stdlib/src/memory/arc.mojo

            (self._inner).destroy_pointee()
            self._inner.free()
+        self._inner = UnsafePointer[Self._inner_type]()


Request This "nulling out" of fields is not required within __del__() methods. After the __del__() completes, this object instance is no longer alive, so there's no code that could observe self._inner one way or the other 🙂

ConnorGray · 2024-10-15T22:41:27Z

stdlib/src/memory/arc.mojo

+
+    alias _weak_refcount_type = __mlir_type[
+        `!pop.array<`, Int(enable_weak).value, `, `, Atomic[DType.int64], `>`
+    ]


Comment This is clever using the boolean to make this a zero-sized field, neat trick! 🙂

ConnorGray · 2024-10-15T22:46:00Z

stdlib/src/memory/arc.mojo

+        """
+        self.payload.assume_initialized_destroy()
+
+    fn try_upgrade_weak(inout self) -> Bool:


Suggestion I think this could use the "custom self types" / "conditional conformance" trick to be statically constrained to:

Suggested change

fn try_upgrade_weak(inout self) -> Bool:

fn try_upgrade_weak(inout self: Arc[T, True]) -> Bool

and then the @parameter if enable_weak statement below could be simplified.

szbergeron added 4 commits September 4, 2024 17:38

Initial impl of Weak[T] for Arc[T]

133c5c4

Signed-off-by: Sawyer Bergeron <sawyerbergeron@gmail.com>

Fix premature free with strong failing to imply weak

0d17fd3

Signed-off-by: Sawyer Bergeron <sawyerbergeron@gmail.com>

Fix 'todo' docstr

d69b04c

Signed-off-by: Sawyer Bergeron <sawyerbergeron@gmail.com>

Fixing docstrs

043c1b6

Signed-off-by: Sawyer Bergeron <sawyerbergeron@gmail.com>

szbergeron requested a review from a team as a code owner September 4, 2024 23:17

szbergeron changed the title ~~Implement Weak[T] refs for Arc[T]~~ [stdlib] Implement Weak[T] refs for Arc[T] Sep 4, 2024

szbergeron added 3 commits September 5, 2024 00:43

Fix nit

21c69e2

Signed-off-by: Sawyer Bergeron <sawyerbergeron@gmail.com>

Forgot to put back test_arc::main()

757e97a

Signed-off-by: Sawyer Bergeron <sawyerbergeron@gmail.com>

Merge branch 'nightly' of github.com:modularml/mojo into arc_weak_01

b19789b

szbergeron force-pushed the arc_weak_01 branch from 5a66a4e to b19789b Compare September 5, 2024 14:32

szbergeron added 6 commits September 5, 2024 15:14

Fix fmt

c25a3f4

Signed-off-by: Sawyer Bergeron <sawyerbergeron@gmail.com>

Work on adding type param to remove overhead when Weak not used

c49cf9f

Signed-off-by: Sawyer Bergeron <sawyerbergeron@gmail.com>

Merge branch 'nightly' into arc_weak_01

da7b4ab

Try parameterizing weak on enable

e77e8b0

Signed-off-by: Sawyer Bergeron <sawyerbergeron@gmail.com>

Merge branch 'nightly' of github.com:modularml/mojo into arc_weak_01

ddb0bb5

Merge branch 'nightly' of github.com:modularml/mojo into arc_weak_01

219bc69

ConnorGray reviewed Oct 15, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[stdlib] Implement Weak[T] refs for Arc[T] #3450

[stdlib] Implement Weak[T] refs for Arc[T] #3450

szbergeron commented Sep 4, 2024

szbergeron commented Sep 23, 2024

ConnorGray left a comment

ConnorGray Oct 15, 2024

ConnorGray Oct 15, 2024

ConnorGray Oct 15, 2024

szbergeron Oct 16, 2024

ConnorGray Oct 15, 2024

ConnorGray Oct 15, 2024

szbergeron Oct 16, 2024

ConnorGray Oct 15, 2024

ConnorGray Oct 15, 2024

ConnorGray Oct 15, 2024

ConnorGray Oct 15, 2024

ConnorGray Oct 15, 2024

	alias _inner_type = _ArcInner[T]
	alias _inner_type = _ArcInner[T, True]

	fn downgrade(self) -> Weak[T, enable_weak]:
	fn downgrade(self: Arc[T, True]) -> Weak[T]:

	return abort[Weak[T, enable_weak]]("unreachable")
	return Weak(self)

	fn upgrade(inout self) -> Optional[Arc[T]]:
	fn upgrade(inout self) -> Optional[Arc[T, True]]:

	fn try_upgrade_weak(inout self) -> Bool:
	fn try_upgrade_weak(inout self: Arc[T, True]) -> Bool

[stdlib] Implement Weak[T] refs for Arc[T] #3450

Are you sure you want to change the base?

[stdlib] Implement Weak[T] refs for Arc[T] #3450

Conversation

szbergeron commented Sep 4, 2024

szbergeron commented Sep 23, 2024

ConnorGray left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment