typechecking

Reviewed By: stepancheg

Differential Revision: D60601185

fbshipit-source-id: 9126587fa7b1ad5825ee64ae6807f6e5984dae4e

starlark/src/eval/compiler/constants.rs

@@ -47,6 +47,7 @@ pub(crate) struct Constants {
     pub(crate) fn_dict: BuiltinFn,
     pub(crate) fn_tuple: BuiltinFn,
     pub(crate) fn_isinstance: BuiltinFn,
+    pub(crate) fn_set: BuiltinFn,
     // Technically, this is not a function.
     pub(crate) typing_callable: BuiltinFn,
 }
@@ -62,6 +63,7 @@ impl Constants {
                 fn_dict: BuiltinFn(g.get_frozen("dict").unwrap()),
                 fn_tuple: BuiltinFn(g.get_frozen("tuple").unwrap()),
                 fn_isinstance: BuiltinFn(g.get_frozen("isinstance").unwrap()),
+                fn_set: BuiltinFn(g.get_frozen("set").unwrap()),
                 typing_callable: {
                     let typing = g.get_frozen("typing").unwrap();
                     let typing = FrozenStructRef::from_value(typing).unwrap();

starlark/src/eval/compiler/types.rs

@@ -168,7 +168,9 @@ impl<'v> Compiler<'v, '_, '_, '_> {
             TypeExprUnpackP::Path(path) => self.eval_path(path),
             TypeExprUnpackP::Index(a, i) => {
                 let a = self.eval_ident_in_type_expr(a)?;
-                if !a.ptr_eq(Constants::get().fn_list.0.to_value()) {
+                if !a.ptr_eq(Constants::get().fn_list.0.to_value())
+                    && !a.ptr_eq(Constants::get().fn_set.0.to_value())
+                {
                     return Err(EvalException::new_anyhow(
                         TypesError::TypeIndexOnNonList.into(),
                         expr.span,

starlark/src/typing/basic.rs

@@ -56,6 +56,8 @@ pub enum TyBasic {
     Dict(ArcTy, ArcTy),
     /// Custom type.
     Custom(TyCustom),
+    /// A set.
+    Set(ArcTy),
 }
 
 impl TyBasic {
@@ -94,9 +96,9 @@ impl TyBasic {
         Self::dict(Ty::any(), Ty::any())
     }
 
-    // pub(crate) fn any_set() -> Self {
-    //     Self::set(Ty::any())
-    // }
+    pub(crate) fn any_set() -> Self {
+        TyBasic::Set(ArcTy::any())
+    }
 
     /// Create a iterable type.
     pub(crate) fn iter(item: Ty) -> Self {
@@ -108,6 +110,11 @@ impl TyBasic {
         TyBasic::Dict(ArcTy::new(key), ArcTy::new(value))
     }
 
+    /// Create a set type.
+    pub(crate) fn set(item: Ty) -> Self {
+        TyBasic::Set(ArcTy::new(item))
+    }
+
     pub(crate) fn custom(custom: impl TyCustomImpl) -> Self {
         TyBasic::Custom(TyCustom::new(custom))
     }
@@ -124,6 +131,7 @@ impl TyBasic {
             TyBasic::Type => Some("type"),
             TyBasic::Custom(c) => c.as_name(),
             TyBasic::Any | TyBasic::Iter(_) | TyBasic::Callable(_) => None,
+            TyBasic::Set(_) => Some("set"),
         }
     }
 
@@ -176,6 +184,7 @@ impl Display for TyBasic {
             }
             TyBasic::Type => write!(f, "type"),
             TyBasic::Custom(c) => Display::fmt(c, f),
+            TyBasic::Set(x) => write!(f, "set[{}]", x),
         }
     }
 }

starlark/src/typing/oracle/ctx.rs

@@ -43,6 +43,7 @@ use crate::typing::TypingBinOp;
 use crate::typing::TypingUnOp;
 use crate::values::dict::value::MutableDict;
 use crate::values::list::value::List;
+use crate::values::set::value::MutableSet;
 use crate::values::tuple::value::Tuple;
 
 #[derive(Debug, thiserror::Error)]
@@ -272,7 +273,7 @@ impl<'a> TypingOracleCtx<'a> {
         match fun {
             TyBasic::Any => Ok(Ty::any()),
             TyBasic::StarlarkValue(t) => Ok(t.validate_call(span, *self)?),
-            TyBasic::List(_) | TyBasic::Dict(..) | TyBasic::Tuple(_) => Err(self
+            TyBasic::List(_) | TyBasic::Dict(..) | TyBasic::Tuple(_) | TyBasic::Set(_) => Err(self
                 .mk_error_as_maybe_internal(
                     span,
                     TypingOracleCtxError::CallToNonCallable {
@@ -335,6 +336,7 @@ impl<'a> TypingOracleCtx<'a> {
             TyBasic::Type => Ok(Ty::any()),
             TyBasic::Iter(ty) => Ok(ty.to_ty()),
             TyBasic::Custom(ty) => ty.0.iter_item_dyn(),
+            TyBasic::Set(item) => Ok((**item).dupe()),
         }
     }
 
@@ -376,6 +378,12 @@ impl<'a> TypingOracleCtx<'a> {
                 }
                 Ok((**v).dupe())
             }
+            TyBasic::Set(item) => {
+                if !self.intersects(&Ty::basic(index.node.dupe()), item)? {
+                    return Err(TypingNoContextOrInternalError::Typing);
+                }
+                Ok((**item).dupe())
+            }
             TyBasic::StarlarkValue(array) => Ok(array.index(index.node)?),
             TyBasic::Custom(c) => Ok(c.0.index_dyn(index.node, self)?),
         }
@@ -501,6 +509,8 @@ impl<'a> TypingOracleCtx<'a> {
                 }
             }
             TyBasic::Custom(custom) => custom.0.attribute_dyn(attr),
+            //TODO(romanp) add match on attr similar to Dict
+            TyBasic::Set(_) => TyStarlarkValue::new::<MutableSet>().attr(attr),
         }
     }
 
@@ -614,6 +624,29 @@ impl<'a> TypingOracleCtx<'a> {
                 bin_op => Ok(TyStarlarkValue::new::<MutableDict>().bin_op(bin_op, rhs.node)?),
             },
             TyBasic::Custom(lhs) => Ok(lhs.0.bin_op_dyn(bin_op, rhs.node, self)?),
+            TyBasic::Set(elem) => match bin_op {
+                TypingBinOp::In => {
+                    if self.intersects(&Ty::basic(rhs.node.dupe()), elem)? {
+                        Ok(Ty::bool())
+                    } else {
+                        Err(TypingNoContextOrInternalError::Typing)
+                    }
+                }
+                TypingBinOp::BitXor
+                | TypingBinOp::BitAnd
+                | TypingBinOp::Sub
+                | TypingBinOp::BitOr => {
+                    if self.intersects_basic(rhs.node, &TyBasic::any_set())? {
+                        Ok(Ty::union2(
+                            Ty::set(elem.to_ty()),
+                            Ty::basic(rhs.node.dupe()),
+                        ))
+                    } else {
+                        Err(TypingNoContextOrInternalError::Typing)
+                    }
+                }
+                bin_op => Ok(TyStarlarkValue::new::<MutableSet>().bin_op(bin_op, rhs.node)?),
+            },
         }
     }
 
@@ -908,6 +941,9 @@ impl<'a> TypingOracleCtx<'a> {
             (TyBasic::List(x), TyBasic::List(y)) => self.intersects(x, y),
             (TyBasic::List(_), TyBasic::StarlarkValue(y)) => Ok(y.is_list()),
             (TyBasic::List(_), _) => Ok(false),
+            (TyBasic::Set(x), TyBasic::Set(y)) => self.intersects(x, y),
+            (TyBasic::Set(_), TyBasic::StarlarkValue(y)) => Ok(y.is_set()),
+            (TyBasic::Set(_), _) => Ok(false),
             (TyBasic::Dict(x_k, x_v), TyBasic::Dict(y_k, y_v)) => {
                 Ok(self.intersects(x_k, y_k)? && self.intersects(x_v, y_v)?)
             }

starlark/src/typing/ty.rs

@@ -214,6 +214,15 @@ impl Ty {
         Self::dict(Ty::any(), Ty::any())
     }
 
+    /// Create a set type.
+    pub fn set(item: Ty) -> Self {
+        Ty::basic(TyBasic::set(item))
+    }
+
+    pub(crate) fn any_set() -> Self {
+        Self::set(Ty::any())
+    }
+
     /// Create a tuple of two elements
     pub fn tuple2(a: Ty, b: Ty) -> Self {
         Ty::tuple(vec![a, b])

starlark/src/values/type_repr.rs

@@ -78,7 +78,7 @@ impl<T: StarlarkTypeRepr> StarlarkTypeRepr for SetType<T> {
     type Canonical = SetType<T::Canonical>;
 
     fn starlark_type_repr() -> Ty {
-        Ty::todo()
+        Ty::set(T::starlark_type_repr())
     }
 }
 

starlark/src/values/types/function.rs

@@ -86,6 +86,7 @@ pub enum SpecialBuiltinFunction {
     List,
     Dict,
     Tuple,
+    Set,
 }
 
 /// A native function that can be evaluated.
@@ -209,6 +210,10 @@ impl<'v> StarlarkValue<'v> for NativeFunction {
                 let index = TypeCompiled::new(index, heap)?;
                 Ok(TypeCompiled::type_list_of(index, heap).to_inner())
             }
+            Some(SpecialBuiltinFunction::Set) => {
+                let index = TypeCompiled::new(index, heap)?;
+                Ok(TypeCompiled::type_set_of(index, heap).to_inner())
+            }
             _ => ValueError::unsupported(self, "[]"),
         }
     }

starlark/src/values/types/set/set.rs

@@ -19,6 +19,7 @@ use starlark_derive::starlark_module;
 
 use crate as starlark;
 use crate::environment::GlobalsBuilder;
+use crate::values::function::SpecialBuiltinFunction;
 use crate::values::set::refs::SetRef;
 use crate::values::set::value::SetData;
 use crate::values::typing::StarlarkIter;
@@ -28,7 +29,10 @@ use crate::values::ValueOfUnchecked;
 
 #[starlark_module]
 pub(crate) fn register_set(globals: &mut GlobalsBuilder) {
-    #[starlark(speculative_exec_safe)]
+    #[starlark(
+        speculative_exec_safe,
+        special_builtin_function = SpecialBuiltinFunction::Set,
+    )]
     fn set<'v>(
         #[starlark(require = pos)] arg: Option<ValueOfUnchecked<'v, StarlarkIter<Value<'v>>>>,
         heap: &'v Heap,
@@ -51,3 +55,52 @@ pub(crate) fn register_set(globals: &mut GlobalsBuilder) {
         Ok(set)
     }
 }
+#[cfg(test)]
+mod tests {
+    use crate::assert;
+
+    #[test]
+    fn test_set_type_as_type_compile_time() {
+        assert::fail(
+            r"
+def f_fail_ct(x: set[int]):
+    return x
+
+s = set(['not_int'])
+
+f_fail_ct(s)
+",
+            //Is it actually runtime or compile time error?
+            r#"Value `set(["not_int"])` of type `set` does not match the type annotation `set[int]` for argument `x`"#,
+        );
+    }
+
+    #[test]
+    fn test_return_set_type_as_type_compile_time() {
+        assert::fail(
+            r"
+def f_fail_ct(x: str) -> set[int]:
+    return set([x])
+
+f_fail_ct('not_int')
+",
+            //Is it actually runtime or compile time error?
+            r#"Value `set(["not_int"])` of type `set` does not match the type annotation `set[int]` for return type"#,
+        );
+    }
+
+    #[test]
+    fn test_set_type_as_type_run_time() {
+        assert::fail(
+            r"
+def f_fail_rt(x: set[int]):
+    return x
+
+s = set(['not_int'])
+
+noop(f_fail_rt)(s)
+",
+            r#"Value `set(["not_int"])` of type `set` does not match the type annotation `set[int]` for argument `x`"#,
+        );
+    }
+}

starlark/src/values/types/set/value.rs

@@ -345,6 +345,14 @@ where
         }
         Ok(heap.alloc(data))
     }
+
+    fn typechecker_ty(&self) -> Option<Ty> {
+        Some(Ty::any_set())
+    }
+
+    fn get_type_starlark_repr() -> Ty {
+        Ty::any_set()
+    }
 }
 
 impl<'v, T: SetLike<'v>> Serialize for SetGen<T> {

starlark/src/values/typing/type_compiled/alloc.rs

@@ -36,6 +36,8 @@ use crate::values::typing::type_compiled::matchers::IsList;
 use crate::values::typing::type_compiled::matchers::IsListOf;
 use crate::values::typing::type_compiled::matchers::IsNever;
 use crate::values::typing::type_compiled::matchers::IsNone;
+use crate::values::typing::type_compiled::matchers::IsSet;
+use crate::values::typing::type_compiled::matchers::IsSetOf;
 use crate::values::typing::type_compiled::matchers::IsStr;
 use crate::values::typing::type_compiled::matchers::IsType;
 use crate::values::typing::type_compiled::matchers::StarlarkTypeIdMatcher;
@@ -143,6 +145,7 @@ pub trait TypeMatcherAlloc: Sized {
             TyBasic::Callable(_c) => self.alloc(IsCallable),
             TyBasic::Type => self.alloc(IsType),
             TyBasic::Custom(custom) => self.custom(custom),
+            TyBasic::Set(item) => self.set_of(item),
         }
     }
 
@@ -252,4 +255,47 @@ pub trait TypeMatcherAlloc: Sized {
             self.dict_of_matcher(k, v)
         }
     }
+
+    /// `set`.
+    fn set(self) -> Self::Result {
+        self.alloc(IsSet)
+    }
+
+    /// `set[Item]`.
+    fn set_of_matcher(self, item: impl TypeMatcher) -> Self::Result {
+        if item.is_wildcard() {
+            self.set()
+        } else {
+            self.alloc(IsSetOf(item))
+        }
+    }
+
+    /// `set[Item]`.
+    fn set_of_starlark_value(self, item: TyStarlarkValue) -> Self::Result {
+        if item.is_str() {
+            self.set_of_matcher(IsStr)
+        } else {
+            self.set_of_matcher(StarlarkTypeIdMatcher::new(item))
+        }
+    }
+
+    /// `set[Item]`.
+    fn set_of_basic(self, item: &TyBasic) -> Self::Result {
+        match item {
+            TyBasic::Any => self.set(),
+            TyBasic::StarlarkValue(ty) => self.set_of_starlark_value(*ty),
+            ty => self.set_of_matcher(TypeMatcherBoxAlloc.ty_basic(ty)),
+        }
+    }
+
+    fn set_of(self, item: &Ty) -> Self::Result {
+        if item.is_any() {
+            self.set()
+        } else if let [ty] = item.iter_union() {
+            self.set_of_basic(ty)
+        } else {
+            let matcher = TypeMatcherBoxAlloc.ty(item);
+            self.set_of_matcher(matcher)
+        }
+    }
 }

starlark/src/values/typing/type_compiled/compiled.rs

@@ -394,6 +394,13 @@ impl<'v> TypeCompiled<Value<'v>> {
         TypeCompiledFactory::alloc_ty(&Ty::list(t.as_ty().clone()), heap)
     }
 
+    pub(crate) fn type_set_of(
+        t: TypeCompiled<Value<'v>>,
+        heap: &'v Heap,
+    ) -> TypeCompiled<Value<'v>> {
+        TypeCompiledFactory::alloc_ty(&Ty::set(t.as_ty().clone()), heap)
+    }
+
     pub(crate) fn type_any_of_two(
         t0: TypeCompiled<Value<'v>>,
         t1: TypeCompiled<Value<'v>>,