Use Any for plan downcasting

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

vortex-layout/src/plan/display.rs

@@ -120,7 +120,7 @@ impl PlanTreeExtractor for PlanExpressionExtractor {
         _context: &PlanTreeContext,
         formatter: &mut fmt::Formatter<'_>,
     ) -> fmt::Result {
-        let Some(expression_plan) = plan.as_any().downcast_ref::<ExpressionPlan>() else {
+        let Some(expression_plan) = plan.downcast_ref::<ExpressionPlan>() else {
             return Ok(());
         };
         write!(formatter, " expr={}", expression_plan.expression())

vortex-layout/src/plan/mod.rs

@@ -44,10 +44,7 @@ pub type PlanRef = Arc<dyn Plan>;
 /// Layout plans expose their optimizer-facing children in a stable logical order. Optional child
 /// slots count toward [`child_count`](Self::child_count) and are returned as `None` by
 /// [`child`](Self::child) when absent. Accessing a child may initialize and cache its plan.
-pub trait Plan: 'static + Send + Sync {
-    /// Returns this plan as [`Any`] for plan-specific optimization rules.
-    fn as_any(&self) -> &dyn Any;
-
+pub trait Plan: Any + Send + Sync {
     /// Returns the display name of this plan kind.
     ///
     /// Plans should override this when the fully qualified Rust type name is not appropriate.
@@ -109,6 +106,16 @@ pub fn new_plan(layout: &LayoutRef) -> VortexResult<PlanRef> {
 }
 
 impl dyn Plan + '_ {
+    /// Returns whether this plan has concrete type `T`.
+    pub fn is<T: Plan>(&self) -> bool {
+        (self as &dyn Any).is::<T>()
+    }
+
+    /// Downcasts this plan to concrete type `T`.
+    pub fn downcast_ref<T: Plan>(&self) -> Option<&T> {
+        (self as &dyn Any).downcast_ref::<T>()
+    }
+
     /// Displays this plan and its descendants with the default plan extractors.
     pub fn display_tree(&self) -> PlanTreeDisplay<'_> {
         PlanTreeDisplay::default_display(self)

vortex-layout/src/plan/plans/chunked.rs

@@ -46,10 +46,6 @@ impl ChunkedPlan {
 }
 
 impl Plan for ChunkedPlan {
-    fn as_any(&self) -> &dyn std::any::Any {
-        self
-    }
-
     fn name(&self) -> &'static str {
         "ChunkedPlan"
     }

vortex-layout/src/plan/plans/dict.rs

@@ -53,10 +53,6 @@ impl DictPlan {
 }
 
 impl Plan for DictPlan {
-    fn as_any(&self) -> &dyn std::any::Any {
-        self
-    }
-
     fn name(&self) -> &'static str {
         "DictPlan"
     }

vortex-layout/src/plan/plans/expression.rs

@@ -1,7 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use std::any::Any;
 use std::borrow::Cow;
 use std::sync::Arc;
 
@@ -46,10 +45,6 @@ impl ExpressionPlan {
 }
 
 impl Plan for ExpressionPlan {
-    fn as_any(&self) -> &dyn Any {
-        self
-    }
-
     fn name(&self) -> &'static str {
         "ExpressionPlan"
     }
@@ -60,7 +55,7 @@ impl Plan for ExpressionPlan {
         if is_root(&expression) {
             return Ok(child);
         }
-        if let Some(inner) = child.as_any().downcast_ref::<Self>() {
+        if let Some(inner) = child.downcast_ref::<Self>() {
             let expression = replace(expression, &root(), inner.expression.clone());
             return Ok(Arc::new(Self::try_new(
                 expression,

vortex-layout/src/plan/plans/flat.rs

@@ -23,10 +23,6 @@ impl FlatPlan {
 }
 
 impl Plan for FlatPlan {
-    fn as_any(&self) -> &dyn std::any::Any {
-        self
-    }
-
     fn name(&self) -> &'static str {
         "FlatPlan"
     }

vortex-layout/src/plan/plans/list.rs

@@ -70,10 +70,6 @@ impl ListPlan {
 }
 
 impl Plan for ListPlan {
-    fn as_any(&self) -> &dyn std::any::Any {
-        self
-    }
-
     fn name(&self) -> &'static str {
         "ListPlan"
     }

vortex-layout/src/plan/plans/row_idx.rs

@@ -1,7 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use std::any::Any;
 use std::borrow::Cow;
 use std::sync::Arc;
 
@@ -26,10 +25,6 @@ impl RowIdxPlan {
 }
 
 impl Plan for RowIdxPlan {
-    fn as_any(&self) -> &dyn Any {
-        self
-    }
-
     fn name(&self) -> &'static str {
         "RowIdxPlan"
     }

vortex-layout/src/plan/plans/struct_.rs

@@ -55,10 +55,6 @@ impl StructPlan {
 }
 
 impl Plan for StructPlan {
-    fn as_any(&self) -> &dyn std::any::Any {
-        self
-    }
-
     fn name(&self) -> &'static str {
         "StructPlan"
     }

vortex-layout/src/plan/tests.rs

@@ -69,7 +69,7 @@ fn unsupported_layout_has_no_plan() -> VortexResult<()> {
 fn flat_plan_has_no_children() -> VortexResult<()> {
     let plan = make_plan(flat(3, primitive(PType::I32, Nullability::NonNullable), 0))?;
 
-    assert!(plan.as_any().is::<FlatPlan>());
+    assert!(plan.is::<FlatPlan>());
     assert_eq!(plan.child_count(), 0);
     assert!(plan.child(0).is_err());
     Ok(())
@@ -86,7 +86,7 @@ fn chunked_plan_exposes_chunks() -> VortexResult<()> {
     .into_layout();
     let plan = make_plan(layout)?;
 
-    assert!(plan.as_any().is::<ChunkedPlan>());
+    assert!(plan.is::<ChunkedPlan>());
     assert_eq!(plan.child_count(), 2);
     assert_eq!(
         plan.child(0)?
@@ -120,7 +120,7 @@ fn chunked_plan_defers_unrequested_chunks_through_optimization() -> VortexResult
     let first = plan
         .child(0)?
         .ok_or_else(|| vortex_err!("missing first chunk"))?;
-    assert!(first.as_any().is::<FlatPlan>());
+    assert!(first.is::<FlatPlan>());
     let cached = plan
         .child(0)?
         .ok_or_else(|| vortex_err!("missing cached first chunk"))?;
@@ -149,7 +149,7 @@ fn dict_plan_orders_codes_before_values() -> VortexResult<()> {
     .into_layout();
     let plan = make_plan(layout)?;
 
-    assert!(plan.as_any().is::<DictPlan>());
+    assert!(plan.is::<DictPlan>());
     assert_eq!(plan.child_count(), 2);
     assert_eq!(
         plan.child(0)?
@@ -179,7 +179,7 @@ fn list_plan_has_stable_optional_validity_slot() -> VortexResult<()> {
     .into_layout();
     let plan = make_plan(non_nullable)?;
 
-    assert!(plan.as_any().is::<ListPlan>());
+    assert!(plan.is::<ListPlan>());
     assert_eq!(plan.child_count(), 3);
     assert_eq!(
         plan.child(0)?
@@ -228,7 +228,7 @@ fn struct_plan_orders_fields_before_optional_validity() -> VortexResult<()> {
     .into_layout();
     let plan = make_plan(non_nullable)?;
 
-    assert!(plan.as_any().is::<StructPlan>());
+    assert!(plan.is::<StructPlan>());
     assert_eq!(plan.child_count(), 3);
     assert_eq!(
         plan.child(0)?
@@ -282,7 +282,6 @@ fn struct_plan_defers_unrequested_fields_through_optimization() -> VortexResult<
     assert!(
         plan.child(0)?
             .ok_or_else(|| vortex_err!("missing field a"))?
-            .as_any()
             .is::<FlatPlan>()
     );
     let error = plan
@@ -445,12 +444,11 @@ fn row_idx_plan_preserves_row_index_expressions() -> VortexResult<()> {
     let plan = RowIdxPlan::new_ref(10, make_plan(layout)?);
     let plan = ExpressionPlan::try_new(row_idx(), plan)?.optimize()?;
     let expression = plan
-        .as_any()
         .downcast_ref::<ExpressionPlan>()
         .ok_or_else(|| vortex_err!("optimized plan is not an expression plan"))?;
 
     assert_eq!(expression.expression(), &row_idx());
-    assert!(expression.child_plan().as_any().is::<RowIdxPlan>());
+    assert!(expression.child_plan().is::<RowIdxPlan>());
     assert_eq!(expression.row_count(), 3);
     Ok(())
 }