Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[stdlib] Remove attribute self._reserved from Dict #2905

Closed
Closed
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion stdlib/src/builtin/reversed.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -193,5 +193,5 @@ fn reversed[
"""
var src = Reference(value)[].src
return _DictEntryIter[K, V, dict_mutability, dict_lifetime, False](
src[]._reserved - 1, 0, src
src[]._reserved() - 1, 0, src
)
63 changes: 33 additions & 30 deletions stdlib/src/collections/dict.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -178,7 +178,7 @@ struct _DictValueIter[
var src = self.iter.src
return _DictValueIter(
_DictEntryIter[K, V, dict_mutability, dict_lifetime, False](
src[]._reserved - 1, 0, src
src[]._reserved() - 1, 0, src
)
)

Expand Down Expand Up @@ -429,10 +429,11 @@ struct Dict[K: KeyElement, V: CollectionElement](
"""The number of elements currently stored in the dict."""
var _n_entries: Int
"""The number of entries currently allocated."""
var _reserved: Int
"""The current reserved size of the dictionary."""

var _index: _DictIndex

# We use everything available in the list. Which means that
# len(self._entries) == self._entries.capacity == self._reserved()
var _entries: List[Optional[DictEntry[K, V]]]

# ===-------------------------------------------------------------------===#
Expand All @@ -444,9 +445,14 @@ struct Dict[K: KeyElement, V: CollectionElement](
"""Initialize an empty dictiontary."""
self.size = 0
self._n_entries = 0
self._reserved = Self._initial_reservation
self._index = _DictIndex(self._reserved)
self._entries = Self._new_entries(self._reserved)
self._entries = Self._new_entries(Self._initial_reservation)
self._index = _DictIndex(len(self._entries))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: I guess it makes more sense to use Self._initial_reservation here to make clear that it's compile time constant?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

len(self._entries) can be greater than or equal to Self._initial_reservation, but they're not equal in all cases. See the section "Using all the capacity available in the List" of the pull request description

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My point exactly: should we leave that to List at all? Your call anyways, shouldn't be blocking.


# TODO: add @property when Mojo supports it to make
# it possible to do `self._reserved`.
@always_inline
fn _reserved(self) -> Int:
return len(self._entries)

@always_inline
fn __init__(inout self, existing: Self):
Expand All @@ -457,8 +463,7 @@ struct Dict[K: KeyElement, V: CollectionElement](
"""
self.size = existing.size
self._n_entries = existing._n_entries
self._reserved = existing._reserved
self._index = existing._index.copy(existing._reserved)
self._index = existing._index.copy(existing._reserved())
self._entries = existing._entries

@staticmethod
Expand Down Expand Up @@ -503,8 +508,7 @@ struct Dict[K: KeyElement, V: CollectionElement](
"""
self.size = existing.size
self._n_entries = existing._n_entries
self._reserved = existing._reserved
self._index = existing._index.copy(existing._reserved)
self._index = existing._index.copy(existing._reserved())
self._entries = existing._entries

fn __moveinit__(inout self, owned existing: Self):
Expand All @@ -515,7 +519,6 @@ struct Dict[K: KeyElement, V: CollectionElement](
"""
self.size = existing.size
self._n_entries = existing._n_entries
self._reserved = existing._reserved
self._index = existing._index^
self._entries = existing._entries^

Expand Down Expand Up @@ -593,7 +596,7 @@ struct Dict[K: KeyElement, V: CollectionElement](
A reversed iterator of immutable references to the dict keys.
"""
return _DictKeyIter(
_DictEntryIter[forward=False](self[]._reserved - 1, 0, self)
_DictEntryIter[forward=False](self[]._reserved() - 1, 0, self)
)

fn __or__(self, other: Self) -> Self:
Expand Down Expand Up @@ -842,15 +845,15 @@ struct Dict[K: KeyElement, V: CollectionElement](
"""Remove all elements from the dictionary."""
self.size = 0
self._n_entries = 0
self._reserved = Self._initial_reservation
self._index = _DictIndex(self._reserved)
self._entries = Self._new_entries(self._reserved)
self._entries = Self._new_entries(Self._initial_reservation)
self._index = _DictIndex(self._reserved())

@staticmethod
@always_inline
fn _new_entries(reserved: Int) -> List[Optional[DictEntry[K, V]]]:
var entries = List[Optional[DictEntry[K, V]]](capacity=reserved)
for i in range(reserved):
fn _new_entries(reserve_at_least: Int) -> List[Optional[DictEntry[K, V]]]:
var entries = List[Optional[DictEntry[K, V]]](capacity=reserve_at_least)
# We have memory available, we'll use everything.
for i in range(entries.capacity):
entries.append(None)
return entries

Expand All @@ -871,18 +874,18 @@ struct Dict[K: KeyElement, V: CollectionElement](
self._n_entries += 1

fn _get_index(self, slot: Int) -> Int:
return self._index.get_index(self._reserved, slot)
return self._index.get_index(self._reserved(), slot)

fn _set_index(inout self, slot: Int, index: Int):
return self._index.set_index(self._reserved, slot, index)
return self._index.set_index(self._reserved(), slot, index)

fn _next_index_slot(self, inout slot: Int, inout perturb: UInt64):
alias PERTURB_SHIFT = 5
perturb >>= PERTURB_SHIFT
slot = ((5 * slot) + int(perturb + 1)) % self._reserved
slot = ((5 * slot) + int(perturb + 1)) % self._reserved()

fn _find_empty_index(self, hash: Int) -> Int:
var slot = hash % self._reserved
var slot = hash % self._reserved()
var perturb = bitcast[DType.uint64](Int64(hash))
while True:
var index = self._get_index(slot)
Expand All @@ -892,7 +895,7 @@ struct Dict[K: KeyElement, V: CollectionElement](

fn _find_index(self, hash: Int, key: K) -> (Bool, Int, Int):
# Return (found, slot, index)
var slot = hash % self._reserved
var slot = hash % self._reserved()
var perturb = bitcast[DType.uint64](Int64(hash))
while True:
var index = self._get_index(slot)
Expand All @@ -911,35 +914,35 @@ struct Dict[K: KeyElement, V: CollectionElement](
self._next_index_slot(slot, perturb)

fn _over_load_factor(self) -> Bool:
return 3 * self.size > 2 * self._reserved
return 3 * self.size > 2 * self._reserved()

fn _over_compact_factor(self) -> Bool:
return 4 * self._n_entries > 3 * self._reserved
return 4 * self._n_entries > 3 * self._reserved()

fn _maybe_resize(inout self):
if not self._over_load_factor():
if self._over_compact_factor():
self._compact()
return
self._reserved *= 2
var _reserved = self._reserved() * 2
self.size = 0
self._n_entries = 0
self._index = _DictIndex(self._reserved)
var old_entries = self._entries^
self._entries = self._new_entries(self._reserved)
self._entries = self._new_entries(_reserved)
self._index = _DictIndex(self._reserved())

for i in range(len(old_entries)):
var entry = old_entries.__get_ref(i)
if entry[]:
self._insert(entry[].unsafe_take())

fn _compact(inout self):
self._index = _DictIndex(self._reserved)
self._index = _DictIndex(self._reserved())
var right = 0
for left in range(self.size):
while not self._entries.__get_ref(right)[]:
right += 1
debug_assert(right < self._reserved, "Invalid dict state")
debug_assert(right < self._reserved(), "Invalid dict state")
var entry = self._entries.__get_ref(right)
debug_assert(entry[].__bool__(), "Logic error")
var slot = self._find_empty_index(entry[].value()[].hash)
Expand Down
Loading