Optimize some SemanticModel methods (#13091)

This commit is contained in:
Alex Waygood 2024-09-02 10:03:52 +01:00 committed by GitHub
parent 227fa4e035
commit 58c641c92f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -1087,7 +1087,7 @@ impl<'a> SemanticModel<'a> {
let id = self.node_id.expect("No current node");
self.nodes
.ancestor_ids(id)
.filter_map(move |id| self.nodes[id].as_expression())
.map_while(move |id| self.nodes[id].as_expression())
}
/// Return the current [`Expr`].
@ -1199,18 +1199,16 @@ impl<'a> SemanticModel<'a> {
/// Given a [`NodeId`], return its parent, if any.
#[inline]
pub fn parent_expression(&self, node_id: NodeId) -> Option<&'a Expr> {
self.nodes
.ancestor_ids(node_id)
.filter_map(|id| self.nodes[id].as_expression())
.nth(1)
let parent_node_id = self.nodes.ancestor_ids(node_id).nth(1)?;
self.nodes[parent_node_id].as_expression()
}
/// Given a [`NodeId`], return the [`NodeId`] of the parent expression, if any.
pub fn parent_expression_id(&self, node_id: NodeId) -> Option<NodeId> {
self.nodes
.ancestor_ids(node_id)
.filter(|id| self.nodes[*id].is_expression())
.nth(1)
let parent_node_id = self.nodes.ancestor_ids(node_id).nth(1)?;
self.nodes[parent_node_id]
.is_expression()
.then_some(parent_node_id)
}
/// Return the [`Stmt`] corresponding to the given [`NodeId`].
@ -1250,9 +1248,7 @@ impl<'a> SemanticModel<'a> {
/// Return the [`Expr`] corresponding to the given [`NodeId`].
#[inline]
pub fn expression(&self, node_id: NodeId) -> Option<&'a Expr> {
self.nodes
.ancestor_ids(node_id)
.find_map(|id| self.nodes[id].as_expression())
self.nodes[node_id].as_expression()
}
/// Returns an [`Iterator`] over the expressions, starting from the given [`NodeId`].
@ -1260,7 +1256,7 @@ impl<'a> SemanticModel<'a> {
pub fn expressions(&self, node_id: NodeId) -> impl Iterator<Item = &'a Expr> + '_ {
self.nodes
.ancestor_ids(node_id)
.filter_map(move |id| self.nodes[id].as_expression())
.map_while(move |id| self.nodes[id].as_expression())
}
/// Mark a Python module as "seen" by the semantic model. Future callers can quickly discount