From e51ac614c5e5be7670d1c1103ba37f5137e7bc22 Mon Sep 17 00:00:00 2001
From: Karol Herbst <kherbst@redhat.com>
Date: Mon, 8 Jul 2024 10:33:25 +0200
Subject: [PATCH] rusticl/ptr: add a few APIs to TrackedPointers

Cc: mesa-stable
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/30082>
(cherry picked from commit 00180933ad69daa1f526bdd23dc96af408a2c3cb)
---
 .pick_status.json                             |  2 +-
 src/gallium/frontends/rusticl/core/context.rs |  2 +-
 src/gallium/frontends/rusticl/util/ptr.rs     | 18 +++++++++++++++---
 3 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/.pick_status.json b/.pick_status.json
index a8e1990f528..92b1dedc93c 100644
--- a/.pick_status.json
+++ b/.pick_status.json
@@ -1104,7 +1104,7 @@
         "description": "rusticl/ptr: add a few APIs to TrackedPointers",
         "nominated": true,
         "nomination_type": 0,
-        "resolution": 0,
+        "resolution": 1,
         "main_sha": null,
         "because_sha": null,
         "notes": null
diff --git a/src/gallium/frontends/rusticl/core/context.rs b/src/gallium/frontends/rusticl/core/context.rs
index fb5575fc0f8..544d5011a7c 100644
--- a/src/gallium/frontends/rusticl/core/context.rs
+++ b/src/gallium/frontends/rusticl/core/context.rs
@@ -200,7 +200,7 @@ impl Context {
     }
 
     pub fn remove_svm_ptr(&self, ptr: usize) -> Option<Layout> {
-        self.svm_ptrs.lock().unwrap().remove(&ptr)
+        self.svm_ptrs.lock().unwrap().remove(ptr)
     }
 
     pub fn import_gl_buffer(
diff --git a/src/gallium/frontends/rusticl/util/ptr.rs b/src/gallium/frontends/rusticl/util/ptr.rs
index 32b43883b05..7e6b31d624b 100644
--- a/src/gallium/frontends/rusticl/util/ptr.rs
+++ b/src/gallium/frontends/rusticl/util/ptr.rs
@@ -1,6 +1,6 @@
 use std::{
     alloc::Layout,
-    collections::BTreeMap,
+    collections::{btree_map::Entry, BTreeMap},
     hash::{Hash, Hasher},
     mem,
     ops::{Add, Deref},
@@ -162,6 +162,14 @@ impl<P, T: AllocSize<P>> TrackedPointers<P, T>
 where
     P: Ord + Add<Output = P> + Copy,
 {
+    pub fn contains_key(&self, ptr: P) -> bool {
+        self.ptrs.contains_key(&ptr)
+    }
+
+    pub fn entry(&mut self, ptr: P) -> Entry<P, T> {
+        self.ptrs.entry(ptr)
+    }
+
     pub fn find_alloc(&self, ptr: P) -> Option<(P, &T)> {
         if let Some((&base, val)) = self.ptrs.range(..=ptr).next_back() {
             let size = val.size();
@@ -174,12 +182,16 @@ where
         None
     }
 
+    pub fn find_alloc_precise(&self, ptr: P) -> Option<&T> {
+        self.ptrs.get(&ptr)
+    }
+
     pub fn insert(&mut self, ptr: P, val: T) -> Option<T> {
         self.ptrs.insert(ptr, val)
     }
 
-    pub fn remove(&mut self, ptr: &P) -> Option<T> {
-        self.ptrs.remove(ptr)
+    pub fn remove(&mut self, ptr: P) -> Option<T> {
+        self.ptrs.remove(&ptr)
     }
 }