diff --git a/starlark-rust/starlark_map/src/sorted_set.rs b/starlark-rust/starlark_map/src/sorted_set.rs index d29197a1a2495..68627ea516d6f 100644 --- a/starlark-rust/starlark_map/src/sorted_set.rs +++ b/starlark-rust/starlark_map/src/sorted_set.rs @@ -46,8 +46,15 @@ where } /// Construct without checking that the elements are sorted. + /// + /// # Safety + /// Caller must guarantee that `inner` is already sorted. #[inline] pub fn new_unchecked(inner: OrderedSet) -> Self { + debug_assert!( + inner.is_sorted(), + "SortedSet::new_unchecked called with unsorted OrderedSet" + ); Self { inner } } @@ -93,11 +100,19 @@ where self.inner.get_index(index) } - /// Iterate over the union of two sets. + /// Return a sorted union of two `SortedSet`s. #[inline] - pub fn union<'a>(&'a self, other: &'a Self) -> impl Iterator { - // TODO(nga): return sorted. - self.inner.union(&other.inner) + pub fn union(&self, other: &Self) -> SortedSet + where + T: Clone, + { + let mut inner = OrderedSet::new(); + + inner.extend(self.inner.iter().cloned()); + inner.extend(other.inner.iter().cloned()); + + inner.sort(); + SortedSet { inner } } }