Refactor the ExprDict node (#11267)

Co-authored-by: Micha Reiser <micha@reiser.io>
This commit is contained in:
Alex Waygood 2024-05-07 12:46:10 +01:00 committed by GitHub
parent de270154a1
commit 6774f27f4b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
52 changed files with 2425 additions and 2240 deletions

View file

@ -687,10 +687,24 @@ pub struct ExprIf<'a> {
orelse: Box<ComparableExpr<'a>>,
}
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct ComparableDictItem<'a> {
key: Option<ComparableExpr<'a>>,
value: ComparableExpr<'a>,
}
impl<'a> From<&'a ast::DictItem> for ComparableDictItem<'a> {
fn from(ast::DictItem { key, value }: &'a ast::DictItem) -> Self {
Self {
key: key.as_ref().map(ComparableExpr::from),
value: value.into(),
}
}
}
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct ExprDict<'a> {
keys: Vec<Option<ComparableExpr<'a>>>,
values: Vec<ComparableExpr<'a>>,
items: Vec<ComparableDictItem<'a>>,
}
#[derive(Debug, PartialEq, Eq, Hash)]
@ -933,16 +947,8 @@ impl<'a> From<&'a ast::Expr> for ComparableExpr<'a> {
body: body.into(),
orelse: orelse.into(),
}),
ast::Expr::Dict(ast::ExprDict {
keys,
values,
range: _,
}) => Self::Dict(ExprDict {
keys: keys
.iter()
.map(|expr| expr.as_ref().map(Into::into))
.collect(),
values: values.iter().map(Into::into).collect(),
ast::Expr::Dict(ast::ExprDict { items, range: _ }) => Self::Dict(ExprDict {
items: items.iter().map(ComparableDictItem::from).collect(),
}),
ast::Expr::Set(ast::ExprSet { elts, range: _ }) => Self::Set(ExprSet {
elts: elts.iter().map(Into::into).collect(),

View file

@ -155,14 +155,12 @@ pub fn any_over_expr(expr: &Expr, func: &dyn Fn(&Expr) -> bool) -> bool {
orelse,
range: _,
}) => any_over_expr(test, func) || any_over_expr(body, func) || any_over_expr(orelse, func),
Expr::Dict(ast::ExprDict {
keys,
values,
range: _,
}) => values
.iter()
.chain(keys.iter().flatten())
.any(|expr| any_over_expr(expr, func)),
Expr::Dict(ast::ExprDict { items, range: _ }) => {
items.iter().any(|ast::DictItem { key, value }| {
any_over_expr(value, func)
|| key.as_ref().is_some_and(|key| any_over_expr(key, func))
})
}
Expr::Set(ast::ExprSet { elts, range: _ })
| Expr::List(ast::ExprList { elts, range: _, .. })
| Expr::Tuple(ast::ExprTuple { elts, range: _, .. }) => {
@ -1188,8 +1186,8 @@ impl Truthiness {
Self::Truthy
}
}
Expr::Dict(ast::ExprDict { keys, .. }) => {
if keys.is_empty() {
Expr::Dict(ast::ExprDict { items, .. }) => {
if items.is_empty() {
Self::Falsey
} else {
Self::Truthy

View file

@ -2301,13 +2301,9 @@ impl AstNode for ast::ExprDict {
where
V: PreorderVisitor<'a> + ?Sized,
{
let ast::ExprDict {
keys,
values,
range: _,
} = self;
let ast::ExprDict { items, range: _ } = self;
for (key, value) in keys.iter().zip(values) {
for ast::DictItem { key, value } in items {
if let Some(key) = key {
visitor.visit_expr(key);
}

View file

@ -767,12 +767,89 @@ impl From<ExprIf> for Expr {
}
}
/// Represents an item in a [dictionary literal display][1].
///
/// Consider the following Python dictionary literal:
/// ```python
/// {key1: value1, **other_dictionary}
/// ```
///
/// In our AST, this would be represented using an `ExprDict` node containing
/// two `DictItem` nodes inside it:
/// ```ignore
/// [
/// DictItem {
/// key: Some(Expr::Name(ExprName { id: "key1" })),
/// value: Expr::Name(ExprName { id: "value1" }),
/// },
/// DictItem {
/// key: None,
/// value: Expr::Name(ExprName { id: "other_dictionary" }),
/// }
/// ]
/// ```
///
/// [1]: https://docs.python.org/3/reference/expressions.html#displays-for-lists-sets-and-dictionaries
#[derive(Debug, Clone, PartialEq)]
pub struct DictItem {
pub key: Option<Expr>,
pub value: Expr,
}
impl DictItem {
fn key(&self) -> Option<&Expr> {
self.key.as_ref()
}
fn value(&self) -> &Expr {
&self.value
}
}
impl Ranged for DictItem {
fn range(&self) -> TextRange {
TextRange::new(
self.key.as_ref().map_or(self.value.start(), Ranged::start),
self.value.end(),
)
}
}
/// See also [Dict](https://docs.python.org/3/library/ast.html#ast.Dict)
#[derive(Clone, Debug, PartialEq)]
pub struct ExprDict {
pub range: TextRange,
pub keys: Vec<Option<Expr>>,
pub values: Vec<Expr>,
pub items: Vec<DictItem>,
}
impl ExprDict {
/// Returns an `Iterator` over the AST nodes representing the
/// dictionary's keys.
pub fn iter_keys(&self) -> DictKeyIterator {
DictKeyIterator::new(&self.items)
}
/// Returns an `Iterator` over the AST nodes representing the
/// dictionary's values.
pub fn iter_values(&self) -> DictValueIterator {
DictValueIterator::new(&self.items)
}
/// Returns the AST node representing the *n*th key of this
/// dictionary.
///
/// Panics: If the index `n` is out of bounds.
pub fn key(&self, n: usize) -> Option<&Expr> {
self.items[n].key()
}
/// Returns the AST node representing the *n*th value of this
/// dictionary.
///
/// Panics: If the index `n` is out of bounds.
pub fn value(&self, n: usize) -> &Expr {
self.items[n].value()
}
}
impl From<ExprDict> for Expr {
@ -781,6 +858,90 @@ impl From<ExprDict> for Expr {
}
}
#[derive(Debug, Clone)]
pub struct DictKeyIterator<'a> {
items: Iter<'a, DictItem>,
}
impl<'a> DictKeyIterator<'a> {
fn new(items: &'a [DictItem]) -> Self {
Self {
items: items.iter(),
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
impl<'a> Iterator for DictKeyIterator<'a> {
type Item = Option<&'a Expr>;
fn next(&mut self) -> Option<Self::Item> {
self.items.next().map(DictItem::key)
}
fn last(mut self) -> Option<Self::Item> {
self.next_back()
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.items.size_hint()
}
}
impl<'a> DoubleEndedIterator for DictKeyIterator<'a> {
fn next_back(&mut self) -> Option<Self::Item> {
self.items.next_back().map(DictItem::key)
}
}
impl<'a> FusedIterator for DictKeyIterator<'a> {}
impl<'a> ExactSizeIterator for DictKeyIterator<'a> {}
#[derive(Debug, Clone)]
pub struct DictValueIterator<'a> {
items: Iter<'a, DictItem>,
}
impl<'a> DictValueIterator<'a> {
fn new(items: &'a [DictItem]) -> Self {
Self {
items: items.iter(),
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
}
impl<'a> Iterator for DictValueIterator<'a> {
type Item = &'a Expr;
fn next(&mut self) -> Option<Self::Item> {
self.items.next().map(DictItem::value)
}
fn last(mut self) -> Option<Self::Item> {
self.next_back()
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.items.size_hint()
}
}
impl<'a> DoubleEndedIterator for DictValueIterator<'a> {
fn next_back(&mut self) -> Option<Self::Item> {
self.items.next_back().map(DictItem::value)
}
}
impl<'a> FusedIterator for DictValueIterator<'a> {}
impl<'a> ExactSizeIterator for DictValueIterator<'a> {}
/// See also [Set](https://docs.python.org/3/library/ast.html#ast.Set)
#[derive(Clone, Debug, PartialEq)]
pub struct ExprSet {
@ -4358,7 +4519,7 @@ mod tests {
assert_eq!(std::mem::size_of::<ExprBytesLiteral>(), 40);
assert_eq!(std::mem::size_of::<ExprCall>(), 56);
assert_eq!(std::mem::size_of::<ExprCompare>(), 48);
assert_eq!(std::mem::size_of::<ExprDict>(), 56);
assert_eq!(std::mem::size_of::<ExprDict>(), 32);
assert_eq!(std::mem::size_of::<ExprDictComp>(), 48);
assert_eq!(std::mem::size_of::<ExprEllipsisLiteral>(), 8);
// 56 for Rustc < 1.76

View file

@ -389,16 +389,12 @@ pub fn walk_expr<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, expr: &'a Expr) {
visitor.visit_expr(body);
visitor.visit_expr(orelse);
}
Expr::Dict(ast::ExprDict {
keys,
values,
range: _,
}) => {
for expr in keys.iter().flatten() {
visitor.visit_expr(expr);
}
for expr in values {
visitor.visit_expr(expr);
Expr::Dict(ast::ExprDict { items, range: _ }) => {
for ast::DictItem { key, value } in items {
if let Some(key) = key {
visitor.visit_expr(key);
}
visitor.visit_expr(value);
}
}
Expr::Set(ast::ExprSet { elts, range: _ }) => {

View file

@ -376,16 +376,12 @@ pub fn walk_expr<V: Transformer + ?Sized>(visitor: &V, expr: &mut Expr) {
visitor.visit_expr(body);
visitor.visit_expr(orelse);
}
Expr::Dict(ast::ExprDict {
keys,
values,
range: _,
}) => {
for expr in keys.iter_mut().flatten() {
visitor.visit_expr(expr);
}
for expr in values {
visitor.visit_expr(expr);
Expr::Dict(ast::ExprDict { items, range: _ }) => {
for ast::DictItem { key, value } in items {
if let Some(key) = key {
visitor.visit_expr(key);
}
visitor.visit_expr(value);
}
}
Expr::Set(ast::ExprSet { elts, range: _ }) => {