use Entry

compiler/mono/src/layout.rs

@@ -8,6 +8,7 @@ use roc_types::subs::{
     Content, FlatType, RecordFields, Subs, UnionTags, Variable, VariableSubsSlice,
 };
 use roc_types::types::{gather_fields_unsorted_iter, RecordField};
+use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use ven_pretty::{DocAllocator, DocBuilder};
 
@@ -2534,6 +2535,64 @@ struct IdsByLayout<'a> {
     next_id: u32,
 }
 
+impl<'a> IdsByLayout<'a> {
+    #[inline(always)]
+    fn insert_layout(&mut self, layout: Layout<'a>) -> LayoutId {
+        match self.by_id.entry(layout) {
+            Entry::Vacant(vacant) => {
+                let answer = self.next_id;
+                vacant.insert(answer);
+                self.next_id += 1;
+
+                LayoutId(answer)
+            }
+            Entry::Occupied(occupied) => LayoutId(*occupied.get()),
+        }
+    }
+
+    #[inline(always)]
+    fn singleton_layout(layout: Layout<'a>) -> (Self, LayoutId) {
+        let mut by_id = HashMap::with_capacity_and_hasher(1, default_hasher());
+        by_id.insert(layout, 1);
+
+        let ids_by_layout = IdsByLayout {
+            by_id,
+            toplevels_by_id: Default::default(),
+            next_id: 2,
+        };
+
+        (ids_by_layout, LayoutId(1))
+    }
+
+    #[inline(always)]
+    fn insert_toplevel(&mut self, layout: crate::ir::ProcLayout<'a>) -> LayoutId {
+        match self.toplevels_by_id.entry(layout) {
+            Entry::Vacant(vacant) => {
+                let answer = self.next_id;
+                vacant.insert(answer);
+                self.next_id += 1;
+
+                LayoutId(answer)
+            }
+            Entry::Occupied(occupied) => LayoutId(*occupied.get()),
+        }
+    }
+
+    #[inline(always)]
+    fn singleton_toplevel(layout: crate::ir::ProcLayout<'a>) -> (Self, LayoutId) {
+        let mut toplevels_by_id = HashMap::with_capacity_and_hasher(1, default_hasher());
+        toplevels_by_id.insert(layout, 1);
+
+        let ids_by_layout = IdsByLayout {
+            by_id: Default::default(),
+            toplevels_by_id,
+            next_id: 2,
+        };
+
+        (ids_by_layout, LayoutId(1))
+    }
+}
+
 #[derive(Default)]
 pub struct LayoutIds<'a> {
     by_symbol: MutMap<Symbol, IdsByLayout<'a>>,
@@ -2542,60 +2601,38 @@ pub struct LayoutIds<'a> {
 impl<'a> LayoutIds<'a> {
     /// Returns a LayoutId which is unique for the given symbol and layout.
     /// If given the same symbol and same layout, returns the same LayoutId.
+    #[inline(always)]
     pub fn get<'b>(&mut self, symbol: Symbol, layout: &'b Layout<'a>) -> LayoutId {
-        // Note: this function does some weird stuff to satisfy the borrow checker.
-        // There's probably a nicer way to write it that still works.
-        let ids = self.by_symbol.entry(symbol).or_insert_with(|| IdsByLayout {
-            by_id: HashMap::with_capacity_and_hasher(1, default_hasher()),
-            toplevels_by_id: Default::default(),
-            next_id: 1,
-        });
+        match self.by_symbol.entry(symbol) {
+            Entry::Vacant(vacant) => {
+                let (ids_by_layout, layout_id) = IdsByLayout::singleton_layout(*layout);
 
-        // Get the id associated with this layout, or default to next_id.
-        let answer = ids.by_id.get(layout).copied().unwrap_or(ids.next_id);
+                vacant.insert(ids_by_layout);
 
-        // If we had to default to next_id, it must not have been found;
-        // store the ID we're going to return and increment next_id.
-        if answer == ids.next_id {
-            ids.by_id.insert(*layout, ids.next_id);
-
-            ids.next_id += 1;
+                layout_id
+            }
+            Entry::Occupied(mut occupied_ids) => occupied_ids.get_mut().insert_layout(*layout),
         }
-
-        LayoutId(answer)
     }
 
     /// Returns a LayoutId which is unique for the given symbol and layout.
     /// If given the same symbol and same layout, returns the same LayoutId.
+    #[inline(always)]
     pub fn get_toplevel<'b>(
         &mut self,
         symbol: Symbol,
         layout: &'b crate::ir::ProcLayout<'a>,
     ) -> LayoutId {
-        // Note: this function does some weird stuff to satisfy the borrow checker.
-        // There's probably a nicer way to write it that still works.
-        let ids = self.by_symbol.entry(symbol).or_insert_with(|| IdsByLayout {
-            by_id: Default::default(),
-            toplevels_by_id: HashMap::with_capacity_and_hasher(1, default_hasher()),
-            next_id: 1,
-        });
+        match self.by_symbol.entry(symbol) {
+            Entry::Vacant(vacant) => {
+                let (ids_by_layout, layout_id) = IdsByLayout::singleton_toplevel(*layout);
 
-        // Get the id associated with this layout, or default to next_id.
-        let answer = ids
-            .toplevels_by_id
-            .get(layout)
-            .copied()
-            .unwrap_or(ids.next_id);
+                vacant.insert(ids_by_layout);
 
-        // If we had to default to next_id, it must not have been found;
-        // store the ID we're going to return and increment next_id.
-        if answer == ids.next_id {
-            ids.toplevels_by_id.insert(*layout, ids.next_id);
-
-            ids.next_id += 1;
+                layout_id
+            }
+            Entry::Occupied(mut occupied_ids) => occupied_ids.get_mut().insert_toplevel(*layout),
         }
-
-        LayoutId(answer)
     }
 }