refactor: centralizing shared-allocation accounting for Arc DFHeapSize impls

datafusion/common/src/heap_size.rs

@@ -76,6 +76,12 @@ pub struct DFHeapSizeCtx {
     seen: HashSet<usize>,
 }
 
+impl DFHeapSizeCtx {
+    fn count_allocation_once(&mut self, ptr: usize) -> bool {
+        self.seen.insert(ptr)
+    }
+}
+
 impl DFHeapSize for Statistics {
     fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
         self.num_rows.heap_size(ctx)
@@ -281,11 +287,21 @@ impl<K: DFHeapSize, V: DFHeapSize> DFHeapSize for HashMap<K, V> {
     }
 }
 
+fn arc_ptr<T>(arc: &Arc<T>) -> usize {
+    Arc::as_ptr(arc) as usize
+}
+
+/// For unsized types, `Arc::as_ptr` returns the data address + metadata - we only need the thin address
+/// Casting through `*const i32` gets us the thin pointer
+fn arc_unsized_ptr<T: ?Sized>(arc: &Arc<T>) -> usize {
+    Arc::as_ptr(arc) as *const i32 as usize
+}
+
 impl<T: DFHeapSize> DFHeapSize for Arc<T> {
     fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
-        let ptr = Arc::as_ptr(self) as usize;
+        let ptr = arc_ptr(self);
 
-        if !ctx.seen.insert(ptr) {
+        if !ctx.count_allocation_once(ptr) {
             return 0;
         }
 
@@ -296,9 +312,9 @@ impl<T: DFHeapSize> DFHeapSize for Arc<T> {
 
 impl DFHeapSize for Arc<str> {
     fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
-        let ptr = Arc::as_ptr(self) as *const i32 as usize;
+        let ptr = arc_unsized_ptr(self);
 
-        if !ctx.seen.insert(ptr) {
+        if !ctx.count_allocation_once(ptr) {
             return 0;
         }
 
@@ -309,9 +325,9 @@ impl DFHeapSize for Arc<str> {
 
 impl DFHeapSize for Arc<dyn DFHeapSize> {
     fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
-        let ptr = Arc::as_ptr(self) as *const i32 as usize;
+        let ptr = arc_unsized_ptr(self);
 
-        if !ctx.seen.insert(ptr) {
+        if !ctx.count_allocation_once(ptr) {
             return 0;
         }