From 0e052ab8970777e8f418c4ccf495845804aeae90 Mon Sep 17 00:00:00 2001
From: Vytautas Astrauskas <astrauv@amazon.com>
Date: Wed, 29 Apr 2020 15:12:09 -0700
Subject: [PATCH] Use Entry API in set_dtors_running.

---
 src/shims/tls.rs | 25 ++++++++++++-------------
 1 file changed, 12 insertions(+), 13 deletions(-)

diff --git a/src/shims/tls.rs b/src/shims/tls.rs
index 8a5bb7b42c5d..57b041e685f7 100644
--- a/src/shims/tls.rs
+++ b/src/shims/tls.rs
@@ -1,7 +1,8 @@
 //! Implement thread-local storage.
 
 use std::collections::BTreeMap;
-use std::collections::btree_map::Entry;
+use std::collections::btree_map::Entry as BTreeEntry;
+use std::collections::hash_map::Entry as HashMapEntry;
 
 use log::trace;
 
@@ -186,7 +187,7 @@ impl<'tcx> TlsData<'tcx> {
             thread_local.range_mut((start, Unbounded))
         {
             match data.entry(thread_id) {
-                Entry::Occupied(entry) => {
+                BTreeEntry::Occupied(entry) => {
                     if let Some(dtor) = dtor {
                         // Set TLS data to NULL, and call dtor with old value.
                         let data_scalar = entry.remove();
@@ -194,7 +195,7 @@ impl<'tcx> TlsData<'tcx> {
                         return ret;
                     }
                 }
-                Entry::Vacant(_) => {}
+                BTreeEntry::Vacant(_) => {}
             }
         }
         None
@@ -204,16 +205,14 @@ impl<'tcx> TlsData<'tcx> {
     /// the existing values stored in `dtors_running` for this thread. Returns
     /// `true` if dtors for `thread` are already running.
     fn set_dtors_running_for_thread(&mut self, thread: ThreadId) -> bool {
-        if self.dtors_running.contains_key(&thread) {
-            true
-        } else {
-            // We need to guard this `insert` with a check because otherwise we
-            // would risk to overwrite `last_dtor_key` with `None`.
-            self.dtors_running.insert(
-                thread,
-                RunningDtorsState { last_dtor_key: None }
-            );
-            false
+        match self.dtors_running.entry(thread) {
+            HashMapEntry::Occupied(_) => true,
+            HashMapEntry::Vacant(entry) => {
+                // We cannot just do `self.dtors_running.insert` because that
+                // would overwrite `last_dtor_key` with `None`.
+                entry.insert(RunningDtorsState { last_dtor_key: None });
+                false
+            }
         }
     }