change from prefold to fold_listcomp, and simply the fold_listcomp

This commit is contained in:
CrescentonC 2021-07-16 18:13:38 +08:00
parent b961128367
commit 94ffe4dac2
1 changed files with 49 additions and 41 deletions

View File

@ -36,7 +36,7 @@ impl<'a> ast::fold::Fold<Option<Type>> for InferenceContext<'a> {
// pre-fold // pre-fold
let mut expr = node; let mut expr = node;
expr = match &expr.node { expr = match &expr.node {
ast::ExprKind::ListComp { .. } => self.prefold_list_comprehension(expr)?, ast::ExprKind::ListComp { .. } => self.fold_listcomp(expr)?,
_ => rustpython_parser::ast::fold::fold_expr(self, expr)? _ => rustpython_parser::ast::fold::fold_expr(self, expr)?
}; };
@ -55,7 +55,7 @@ impl<'a> ast::fold::Fold<Option<Type>> for InferenceContext<'a> {
ast::ExprKind::Call {func, args, keywords} => self.infer_call(func, args, keywords), ast::ExprKind::Call {func, args, keywords} => self.infer_call(func, args, keywords),
ast::ExprKind::Subscript {value, slice, ctx: _} => self.infer_subscript(value, slice), ast::ExprKind::Subscript {value, slice, ctx: _} => self.infer_subscript(value, slice),
ast::ExprKind::IfExp {test, body, orelse} => self.infer_if_expr(test, body, orelse), ast::ExprKind::IfExp {test, body, orelse} => self.infer_if_expr(test, body, orelse),
ast::ExprKind::ListComp {elt, generators} => self.infer_list_comprehesion(elt, generators), ast::ExprKind::ListComp {elt: _, generators: _} => Ok(expr.custom), // already folded
ast::ExprKind::Slice { .. } => Ok(None), // special handling for slice, which is supported ast::ExprKind::Slice { .. } => Ok(None), // special handling for slice, which is supported
_ => Err("not supported yet".into()) _ => Err("not supported yet".into())
}?, }?,
@ -318,7 +318,7 @@ impl<'a> InferenceContext<'a> {
} }
} }
fn infer_list_comprehesion(&self, elt: &Box<ast::Expr<Option<Type>>>, generators: &Vec<ast::Comprehension<Option<Type>>>) -> Result<Option<Type>, String> { fn _infer_list_comprehesion(&self, elt: &Box<ast::Expr<Option<Type>>>, generators: &Vec<ast::Comprehension<Option<Type>>>) -> Result<Option<Type>, String> {
if generators[0] if generators[0]
.ifs .ifs
.iter() .iter()
@ -328,39 +328,37 @@ impl<'a> InferenceContext<'a> {
vec![elt.custom.clone().ok_or_else(|| "elements should have value".to_string())?]).into())) vec![elt.custom.clone().ok_or_else(|| "elements should have value".to_string())?]).into()))
} else { } else {
Err("test must be bool".into()) Err("test must be bool".into())
} }
} }
// some pre-folds need special handling // some pre-folds need special handling
fn prefold_list_comprehension(&mut self, expr: ast::Expr<Option<Type>>) -> Result<ast::Expr<Option<Type>>, String> { fn fold_listcomp(&mut self, expr: ast::Expr<Option<Type>>) -> Result<ast::Expr<Option<Type>>, String> {
if let ast::Expr { if let ast::Expr {
location, location,
custom, custom: _,
node: ast::ExprKind::ListComp { node: ast::ExprKind::ListComp {
elt, elt,
generators}} = expr { mut generators}} = expr {
// if is list comprehension, need special pre-fold // if is list comprehension, need special pre-fold
if generators.len() != 1 { if generators.len() != 1 {
return Err("only 1 generator statement is supported".into()); return Err("only 1 generator statement is supported".into());
} }
if generators[0].is_async { let gen = generators.remove(0);
if gen.is_async {
return Err("async is not supported".into()); return Err("async is not supported".into());
} }
// fold iter first since it does not contain new identifiers // fold iter first since it does not contain new identifiers
let generators_first_folded = generators let gen_first_folded = ast::Comprehension {
.into_iter() target: gen.target,
.map(|x| -> Result<ast::Comprehension<Option<Type>>, String> {Ok(ast::Comprehension { iter: Box::new(self.fold_expr(*gen.iter)?), // fold here
target: x.target, ifs: gen.ifs,
iter: Box::new(self.fold_expr(*x.iter)?), // fold here is_async: gen.is_async
ifs: x.ifs, };
is_async: x.is_async
})})
.collect::<Result<Vec<_>, _>>()?;
if let TypeEnum::ParametricType( if let TypeEnum::ParametricType(
primitives::LIST_TYPE, primitives::LIST_TYPE,
ls) = generators_first_folded[0] ls) = gen_first_folded
.iter .iter
.custom .custom
.as_ref() .as_ref()
@ -368,29 +366,39 @@ impl<'a> InferenceContext<'a> {
.as_ref() .as_ref()
.clone() { .clone() {
self.with_scope(|ctx| -> Result<ast::Expr<Option<Type>>, String> { self.with_scope(|ctx| -> Result<ast::Expr<Option<Type>>, String> {
ctx.infer_simple_binding(&generators_first_folded[0].target, ls[0].clone())?; ctx.infer_simple_binding(&gen_first_folded.target, ls[0].clone())?;
Ok(ast::Expr { let elt_folded = Box::new(ctx.fold_expr(*elt)?);
location, let target_folded = Box::new(ctx.fold_expr(*gen_first_folded.target)?);
custom, let ifs_folded = gen_first_folded
node: ast::ExprKind::ListComp { // now fold things with new name .ifs
elt: .into_iter()
Box::new(ctx.fold_expr(*elt)?), .map(|x| ctx.fold_expr(x))
generators: .collect::<Result<Vec<ast::Expr<Option<Type>>>, _>>()?;
generators_first_folded
.into_iter() if ifs_folded
.map(|x| -> Result<ast::Comprehension<Option<Type>>, String> {Ok(ast::Comprehension { .iter()
target: Box::new(ctx.fold_expr(*x.target)?), .all(|x| x.custom == Some(ctx.get_primitive(primitives::BOOL_TYPE))) {
iter: x.iter, Ok(ast::Expr {
ifs: x location,
.ifs custom: Some(TypeEnum::ParametricType(
.into_iter() primitives::LIST_TYPE,
.map(|x| ctx.fold_expr(x)) vec![elt_folded
.collect::<Result<Vec<_>, _>>()?, .custom
is_async: x.is_async .clone()
})}) .ok_or_else(|| "elements cannot be typped".to_string())?]).into()),
.collect::<Result<Vec<_>, _>>()? node: ast::ExprKind::ListComp {
} elt: elt_folded,
}) generators: vec![ast::Comprehension {
target: target_folded,
ifs: ifs_folded,
iter: gen_first_folded.iter,
is_async: gen_first_folded.is_async
}]
}
})
} else {
Err("test must be bool".into())
}
}).1 }).1
} else { } else {
Err("iteration is supported for list only".into()) Err("iteration is supported for list only".into())