From 69226fbf98da4bf926d98aee07a4704a9ab26728 Mon Sep 17 00:00:00 2001
From: Ivan Smirnov <i.s.smirnov@gmail.com>
Date: Sat, 9 Sep 2023 00:20:58 +0100
Subject: [PATCH] feat: add `array::{downcast_ref, downcast_mut}`

---
 Cargo.toml       |  3 ++
 src/array/mod.rs | 96 +++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 97 insertions(+), 2 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 1bb20a6955..df93db9616 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -100,6 +100,9 @@ odbc-api = { version = "0.36", optional = true }
 # Faster hashing
 ahash = "0.8"
 
+# workaround for pre-polonius borrow checker
+polonius-the-crab = "0.3"
+
 # Support conversion to/from arrow-rs
 arrow-buffer = { version = ">=40", optional = true }
 arrow-schema = { version = ">=40", optional = true }
diff --git a/src/array/mod.rs b/src/array/mod.rs
index 02735c3d0b..e96f341bb1 100644
--- a/src/array/mod.rs
+++ b/src/array/mod.rs
@@ -16,10 +16,10 @@
 //!
 //! Most arrays contain a [`MutableArray`] counterpart that is neither clonable nor slicable, but
 //! can be operated in-place.
-use std::any::Any;
+use std::any::{type_name, Any};
 use std::sync::Arc;
 
-use crate::error::Result;
+use crate::error::{Error, Result};
 use crate::{
     bitmap::{Bitmap, MutableBitmap},
     datatypes::DataType,
@@ -142,6 +142,12 @@ pub trait Array: Send + Sync + dyn_clone::DynClone + 'static {
 
     /// Clone a `&dyn Array` to an owned `Box<dyn Array>`.
     fn to_boxed(&self) -> Box<dyn Array>;
+
+    #[doc(hidden)]
+    #[inline]
+    fn as_mut_any(&mut self) -> &mut dyn Any {
+        self.as_any_mut()
+    }
 }
 
 dyn_clone::clone_trait_object!(Array);
@@ -703,6 +709,92 @@ impl<'a> AsRef<(dyn Array + 'a)> for dyn Array {
     }
 }
 
+mod downcast {
+    use super::{Array, DataType, MutableArray};
+    use std::any::Any;
+
+    /// Arrays that can be downcasted to a concrete type ([`Array`] and [`MutableArray`]).
+    pub trait ArrayAny {
+        /// The [`DataType`] of the array.
+        fn data_type(&self) -> &DataType;
+        /// Converts itself to a reference of [`Any`].
+        fn as_any(&self) -> &dyn Any;
+        /// Converts itself to a mutable reference of [`Any`].
+        fn as_mut_any(&mut self) -> &mut dyn Any;
+    }
+
+    macro_rules! impl_array_any {
+        ($ty:ident) => {
+            impl ArrayAny for dyn $ty {
+                #[inline]
+                fn data_type(&self) -> &DataType {
+                    $ty::data_type(self)
+                }
+
+                #[inline]
+                fn as_any(&self) -> &dyn Any {
+                    $ty::as_any(self)
+                }
+
+                #[inline]
+                fn as_mut_any(&mut self) -> &mut dyn Any {
+                    $ty::as_mut_any(self)
+                }
+            }
+
+            impl ArrayAny for Box<dyn $ty> {
+                #[inline]
+                fn data_type(&self) -> &DataType {
+                    $ty::data_type(self.as_ref())
+                }
+
+                #[inline]
+                fn as_any(&self) -> &dyn Any {
+                    $ty::as_any(self.as_ref())
+                }
+
+                #[inline]
+                fn as_mut_any(&mut self) -> &mut dyn Any {
+                    $ty::as_mut_any(self.as_mut())
+                }
+            }
+        };
+    }
+
+    impl_array_any!(Array);
+    impl_array_any!(MutableArray);
+}
+
+/// Downcast an array reference to a concrete type.
+#[inline]
+pub fn downcast_ref<M: 'static>(array: &(impl downcast::ArrayAny + ?Sized)) -> Result<&M> {
+    array.as_any().downcast_ref().ok_or_else(|| {
+        Error::oos(format!(
+            "Unable to downcast array of data type {:?} into {}",
+            array.data_type(),
+            type_name::<M>()
+        ))
+    })
+}
+
+/// Downcast a mutable array reference to a concrete type.
+#[inline]
+pub fn downcast_mut<M: 'static>(
+    mut array: &mut (impl downcast::ArrayAny + ?Sized),
+) -> Result<&mut M> {
+    use polonius_the_crab::{polonius, polonius_return};
+    polonius!(|array| -> Result<&'polonius mut M> {
+        if let Some(array) = array.as_mut_any().downcast_mut::<M>() {
+            polonius_return!(Ok(array));
+        }
+    });
+    Err(Error::oos(format!(
+        "Unable to downcast array of data type {:?} into {}",
+        array.data_type(),
+        type_name::<M>()
+    )))
+}
+
 mod binary;
 mod boolean;
 mod dictionary;