Ty::as_function

Summary: Expose some `Ty` internals for the following diff D48894467.

Reviewed By: ianlevesque

Differential Revision: D48894468

fbshipit-source-id: 83a6380610a2a4cd7cdff910001900b59b3697e0

Author: stepancheg

starlark/src/typing/basic.rs

@@ -27,6 +27,7 @@ use crate::typing::custom::TyCustomImpl;
 use crate::typing::starlark_value::TyStarlarkValue;
 use crate::typing::tuple::TyTuple;
 use crate::typing::Ty;
+use crate::typing::TyFunction;
 use crate::typing::TyName;
 use crate::values::none::NoneType;
 use crate::values::string::StarlarkStr;
@@ -128,6 +129,14 @@ impl TyBasic {
         self.as_name() == Some("function")
     }
 
+    /// If this type is function, return the function type.
+    pub(crate) fn as_function(&self) -> Option<&TyFunction> {
+        match self {
+            TyBasic::Custom(c) => c.0.as_function_dyn(),
+            _ => None,
+        }
+    }
+
     /// Type is a tuple, with specified or unspecified member types.
     pub(crate) fn is_tuple(&self) -> bool {
         matches!(self, TyBasic::Tuple(_))

starlark/src/typing/custom.rs

@@ -37,6 +37,7 @@ use crate::typing::error::TypingOrInternalError;
 use crate::typing::Arg;
 use crate::typing::Ty;
 use crate::typing::TyBasic;
+use crate::typing::TyFunction;
 use crate::typing::TypingBinOp;
 use crate::typing::TypingOracleCtx;
 use crate::values::typing::type_compiled::alloc::TypeMatcherAlloc;
@@ -61,6 +62,9 @@ pub trait TyCustomImpl: Debug + Display + Hash + Ord + Allocative + Send + Sync
     fn is_callable(&self) -> bool {
         false
     }
+    fn as_function(&self) -> Option<&TyFunction> {
+        None
+    }
     fn bin_op(&self, bin_op: TypingBinOp, rhs: &TyBasic, ctx: &TypingOracleCtx) -> Result<Ty, ()> {
         let _unused = (bin_op, rhs, ctx);
         Err(())
@@ -100,6 +104,7 @@ pub(crate) trait TyCustomDyn: Debug + Display + Allocative + Send + Sync + 'stat
         oracle: TypingOracleCtx,
     ) -> Result<Ty, TypingOrInternalError>;
     fn is_callable_dyn(&self) -> bool;
+    fn as_function_dyn(&self) -> Option<&TyFunction>;
     fn iter_item_dyn(&self) -> Result<Ty, ()>;
     fn index_dyn(&self, index: &TyBasic, ctx: &TypingOracleCtx) -> Result<Ty, ()>;
     fn attribute_dyn(&self, attr: &str) -> Result<Ty, ()>;
@@ -163,6 +168,10 @@ impl<T: TyCustomImpl> TyCustomDyn for T {
         self.is_callable()
     }
 
+    fn as_function_dyn(&self) -> Option<&TyFunction> {
+        self.as_function()
+    }
+
     fn attribute_dyn(&self, attr: &str) -> Result<Ty, ()> {
         self.attribute(attr)
     }

starlark/src/typing/function.rs

@@ -173,6 +173,10 @@ pub trait TyCustomFunctionImpl:
         args: &[Spanned<Arg>],
         oracle: TypingOracleCtx,
     ) -> Result<Ty, TypingOrInternalError>;
+
+    fn as_function(&self) -> Option<&TyFunction> {
+        None
+    }
 }
 
 #[derive(
@@ -206,6 +210,10 @@ impl<F: TyCustomFunctionImpl> TyCustomImpl for TyCustomFunction<F> {
         true
     }
 
+    fn as_function(&self) -> Option<&TyFunction> {
+        self.0.as_function()
+    }
+
     fn bin_op(
         &self,
         bin_op: TypingBinOp,
@@ -332,4 +340,8 @@ impl TyCustomFunctionImpl for TyFunction {
     ) -> Result<Ty, TypingOrInternalError> {
         oracle.validate_fn_call(span, self, args)
     }
+
+    fn as_function(&self) -> Option<&TyFunction> {
+        Some(self)
+    }
 }

starlark/src/typing/ty.rs

@@ -362,6 +362,17 @@ impl Ty {
         self.as_name() == Some("function")
     }
 
+    /// If this type is function, return the function type.
+    ///
+    /// This is exposed for buck2 providers implementation,
+    /// probably it does not do what you think.
+    pub fn as_function(&self) -> Option<&TyFunction> {
+        match self.iter_union() {
+            [x] => x.as_function(),
+            _ => None,
+        }
+    }
+
     /// Create a unions type, which will be normalised before being created.
     pub fn unions(xs: Vec<Self>) -> Self {
         // Handle common cases first.