forked from M-Labs/nac3
change from prefold to fold_listcomp, and simply the fold_listcomp
This commit is contained in:
parent
b961128367
commit
94ffe4dac2
@ -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())
|
||||||
|
Loading…
Reference in New Issue
Block a user