forked from M-Labs/nac3
1
0
Fork 0

core/typecheck: Support tuple arg type in len()

This commit is contained in:
lyken 2024-08-13 10:27:27 +08:00 committed by sb10q
parent d0da688aa7
commit 973dc5041a
2 changed files with 58 additions and 21 deletions

View File

@ -1464,36 +1464,21 @@ impl<'a> BuiltinBuilder<'a> {
fn build_len_function(&mut self) -> TopLevelDef {
let prim = PrimDef::FunLen;
let PrimitiveStore { uint64, int32, .. } = *self.primitives;
// Type handled in [`Inferencer::try_fold_special_call`]
let arg_tvar = self.unifier.get_dummy_var();
let tvar = self.unifier.get_fresh_var(Some("L".into()), None);
let list = self
.unifier
.subst(
self.primitives.list,
&into_var_map([TypeVar { id: self.list_tvar.id, ty: tvar.ty }]),
)
.unwrap();
let ndims = self.unifier.get_fresh_const_generic_var(uint64, Some("N".into()), None);
let ndarray = make_ndarray_ty(self.unifier, self.primitives, Some(tvar.ty), Some(ndims.ty));
let arg_ty = self.unifier.get_fresh_var_with_range(
&[list, ndarray, self.primitives.range],
Some("I".into()),
None,
);
TopLevelDef::Function {
name: prim.name().into(),
simple_name: prim.simple_name().into(),
signature: self.unifier.add_ty(TypeEnum::TFunc(FunSignature {
args: vec![FuncArg {
name: "ls".into(),
ty: arg_ty.ty,
name: "obj".into(),
ty: arg_tvar.ty,
default_value: None,
is_vararg: false,
}],
ret: int32,
vars: into_var_map([tvar, arg_ty]),
ret: self.primitives.int32,
vars: into_var_map([arg_tvar]),
})),
var_id: Vec::default(),
instance_to_symbol: HashMap::default(),

View File

@ -1070,6 +1070,58 @@ impl<'a> Inferencer<'a> {
}));
}
if id == &"len".into() && args.len() == 1 {
let obj = self.fold_expr(args.remove(0))?;
let obj_ty = obj.custom.unwrap();
match &*self.unifier.get_ty(obj_ty) {
TypeEnum::TObj { obj_id, .. }
if *obj_id == self.primitives.range.obj_id(self.unifier).unwrap() => {}
TypeEnum::TObj { obj_id, .. }
if *obj_id == self.primitives.list.obj_id(self.unifier).unwrap() => {}
TypeEnum::TObj { obj_id, .. }
if *obj_id == self.primitives.ndarray.obj_id(self.unifier).unwrap() => {}
TypeEnum::TTuple { .. } => {}
_ => {
return report_error(
format!(
"len() only accepts range, list, ndarray, or tuple. Got {}",
self.unifier.stringify(obj_ty)
)
.as_str(),
obj.location,
)
}
}
let ret_ty = self.primitives.int32;
let func_ty = self.unifier.add_ty(TypeEnum::TFunc(FunSignature {
args: vec![FuncArg {
name: "obj".into(),
ty: obj_ty,
default_value: None,
is_vararg: false,
}],
ret: ret_ty,
vars: VarMap::new(),
}));
return Ok(Some(Located {
location,
custom: Some(ret_ty),
node: ExprKind::Call {
func: Box::new(Located {
custom: Some(func_ty),
location: func.location,
node: ExprKind::Name { id: *id, ctx: *ctx },
}),
args: vec![obj],
keywords: vec![],
},
}));
}
if ["int32", "float", "bool", "round", "round64", "np_isnan", "np_isinf"]
.iter()
.any(|fun_id| id == &(*fun_id).into())