Keeping SyntaxNode parent's data id for optimizations

Author: eytan-starkware

crates/cairo-lang-defs/src/ids.rs

@@ -774,8 +774,7 @@ define_language_element_id_basic!(GenericParamId, GenericParamLongId, ast::Gener
 impl<'db> GenericParamLongId<'db> {
     pub fn name(&self, db: &'db dyn Database) -> Option<SmolStrId<'db>> {
         let node = self.1.0.0;
-        // TODO(eytan-starkware): have a non-query accessor to parent and use it here.
-        assert!(node.parent(db).is_some());
+        assert!(!node.is_root());
         let key_fields = node.key_fields(db);
         let kind = node.kind(db);
         require(!matches!(
@@ -792,7 +791,7 @@ impl<'db> GenericParamLongId<'db> {
     }
     pub fn kind(&self, db: &dyn Database) -> GenericKind {
         let node = self.1.0.0;
-        assert!(node.parent(db).is_some());
+        assert!(!node.is_root());
         let kind = node.kind(db);
         match kind {
             SyntaxKind::GenericParamType => GenericKind::Type,
@@ -842,7 +841,7 @@ impl<'db> GenericParamId<'db> {
     pub fn format(&self, db: &'db dyn Database) -> SmolStrId<'db> {
         let long_ids = self.long(db);
         let node = long_ids.1.0.0;
-        assert!(node.parent(db).is_some());
+        assert!(!node.is_root());
         let key_fields = node.key_fields(db);
         let kind = node.kind(db);
 

crates/cairo-lang-semantic/src/resolve/mod.rs

@@ -571,9 +571,7 @@ impl<'db> Resolver<'db> {
             && last.identifier(self.db).long(self.db) == SELF_PARAM_KW
         {
             // If the `self` keyword is used in a non-multi-use path, report an error.
-            if use_path.as_syntax_node().parent(self.db).unwrap().kind(self.db)
-                != SyntaxKind::UsePathList
-            {
+            if use_path.as_syntax_node().parent_kind(self.db).unwrap() != SyntaxKind::UsePathList {
                 diagnostics.report(use_path.stable_ptr(self.db), UseSelfNonMulti);
             }
             segments.segments.pop();

crates/cairo-lang-starknet/src/plugin/starknet_module/mod.rs

@@ -246,7 +246,7 @@ fn grand_grand_parent_starknet_module<'db>(
 ) -> Option<(ast::ItemModule<'db>, StarknetModuleKind, ast::Attribute<'db>)> {
     // Get the containing module node. The parent is the item list, the grand parent is the module
     // body, and the grand grand parent is the module.
-    let module_node = item_node.parent(db)?.parent(db)?.parent(db)?;
+    let module_node = item_node.grandparent(db)?.parent(db)?;
     let module_ast = ast::ItemModule::cast(db, module_node)?;
     let (module_kind, attr) = StarknetModuleKind::from_module(db, &module_ast)?;
     Some((module_ast, module_kind, attr))

crates/cairo-lang-syntax/src/node/mod.rs

@@ -65,6 +65,16 @@ struct SyntaxNodeData<'a> {
     id: SyntaxNodeId<'a>,
 }
 
+impl<'db> SyntaxNodeData<'db> {
+    /// Gets the kind of the given node's parent if it exists.
+    pub fn parent(&self, db: &'db dyn Database) -> Option<SyntaxNode<'db>> {
+        match self.id(db) {
+            SyntaxNodeId::Root(_) => None,
+            SyntaxNodeId::Child { parent, .. } => Some(*parent),
+        }
+    }
+}
+
 impl<'db> cairo_lang_debug::DebugWithDb<'db> for SyntaxNodeData<'db> {
     type Db = dyn Database;
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>, db: &'db Self::Db) -> std::fmt::Result {
@@ -82,57 +92,81 @@ impl<'db> cairo_lang_debug::DebugWithDb<'db> for SyntaxNodeData<'db> {
 /// tracked functions to ensure uniqueness of SyntaxNodes.
 /// Use `SyntaxNode::new_root` or `SyntaxNode::new_root_with_offset` to create root nodes.
 #[derive(Clone, Copy, PartialEq, Eq, Hash, salsa::Update)]
-pub struct SyntaxNode<'a>(SyntaxNodeData<'a>);
+pub struct SyntaxNode<'a> {
+    data: SyntaxNodeData<'a>,
+    /// Cached parent data to avoid database lookups. None for root nodes.
+    parent: Option<SyntaxNodeData<'a>>,
+}
 
 impl<'db> std::fmt::Debug for SyntaxNode<'db> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "SyntaxNode({:x})", self.0.as_id().index())
+        write!(f, "SyntaxNode({:x})", self.data.as_id().index())
     }
 }
 
 impl<'db> cairo_lang_debug::DebugWithDb<'db> for SyntaxNode<'db> {
     type Db = dyn Database;
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>, db: &'db Self::Db) -> std::fmt::Result {
-        self.0.fmt(f, db)
+        self.data.fmt(f, db)
     }
 }
 
 impl<'db> SyntaxNode<'db> {
     /// Get the offset of this syntax node from the beginning of the file.
     pub fn offset(self, db: &'db dyn Database) -> TextOffset {
-        self.0.offset(db)
+        self.data.offset(db)
     }
 
     /// Get the parent syntax node, if any.
     pub fn parent(self, db: &'db dyn Database) -> Option<SyntaxNode<'db>> {
-        match self.0.id(db) {
+        match self.data.id(db) {
             SyntaxNodeId::Child { parent, .. } => Some(*parent),
             SyntaxNodeId::Root(_) => None,
         }
     }
 
+    /// Get the grandparent syntax node, if any.
+    /// This uses a cached parent when available, avoiding some database lookups.
+    pub fn grandparent(self, db: &'db dyn Database) -> Option<SyntaxNode<'db>> {
+        self.parent?.parent(db)
+    }
+
+    /// Check if this syntax node is the root of the syntax tree.
+    pub fn is_root(self) -> bool {
+        self.parent.is_none()
+    }
+
     /// Get the stable pointer for this syntax node.
     pub fn stable_ptr(self, _db: &'db dyn Database) -> SyntaxStablePtrId<'db> {
         SyntaxStablePtrId(self)
     }
 
     pub fn raw_id(&self, db: &'db dyn Database) -> &'db SyntaxNodeId<'db> {
-        self.0.id(db)
+        self.data.id(db)
     }
 
     /// Get the key fields of this syntax node. These define the unique identifier of the node.
     pub fn key_fields(self, db: &'db dyn Database) -> &'db [GreenId<'db>] {
-        match self.0.id(db) {
+        match self.data.id(db) {
             SyntaxNodeId::Child { key_fields, .. } => key_fields,
             SyntaxNodeId::Root(_) => &[],
         }
     }
 
     /// Returns the file id of the file containing this node.
     pub fn file_id(&self, db: &'db dyn Database) -> FileId<'db> {
-        match self.0.id(db) {
-            SyntaxNodeId::Child { parent, .. } => parent.file_id(db),
+        // Use cached parent if available to avoid database lookup
+        if let Some(parent_data) = self.parent {
+            return match parent_data.id(db) {
+                SyntaxNodeId::Child { parent, .. } => parent.file_id(db),
+                SyntaxNodeId::Root(file_id) => *file_id,
+            };
+        }
+        match self.data.id(db) {
             SyntaxNodeId::Root(file_id) => *file_id,
+            SyntaxNodeId::Child { .. } => {
+                unreachable!("Parent already checked and found to not exist.")
+            }
         }
     }
 
@@ -142,18 +176,26 @@ impl<'db> SyntaxNode<'db> {
     /// n = 2: return the grand parent.
     /// And so on...
     /// Assumes that the `n`th parent exists. Panics otherwise.
-    pub fn nth_parent<'r: 'db>(&self, db: &'r dyn Database, n: usize) -> SyntaxNode<'db> {
-        let mut ptr = *self;
-        for _ in 0..n {
-            ptr = ptr.parent(db).unwrap_or_else(|| {
-                panic!(
-                    "N'th parent did not exist. File {} Offset {:?}",
-                    self.file_id(db).long(db).full_path(db),
-                    self.offset(db)
-                )
-            });
+    pub fn nth_parent<'r: 'db>(&self, db: &'r dyn Database, mut n: usize) -> SyntaxNode<'db> {
+        let mut result = Some(*self);
+        while let Some(p) = result
+            && n > 1
+        {
+            result = p.grandparent(db);
+            n -= 2;
         }
-        ptr
+        if let Some(p) = result
+            && n == 1
+        {
+            result = p.parent(db);
+        }
+        result.unwrap_or_else(|| {
+            panic!(
+                "N'th parent did not exist. File {} Offset {:?}",
+                self.file_id(db).long(db).full_path(db),
+                self.offset(db)
+            )
+        })
     }
 }
 
@@ -165,7 +207,12 @@ pub fn new_syntax_node<'db>(
     offset: TextOffset,
     id: SyntaxNodeId<'db>,
 ) -> SyntaxNode<'db> {
-    SyntaxNode(SyntaxNodeData::new(db, green, offset, id))
+    let parent = match &id {
+        SyntaxNodeId::Child { parent, .. } => Some(parent.data),
+        SyntaxNodeId::Root(_) => None,
+    };
+    let data = SyntaxNodeData::new(db, green, offset, id);
+    SyntaxNode { data, parent }
 }
 
 // Construction methods
@@ -212,7 +259,7 @@ impl<'a> SyntaxNode<'a> {
 
     /// Returns the green node of the syntax node.
     pub fn green_node(&self, db: &'a dyn Database) -> &'a GreenNode<'a> {
-        self.0.green(db).long(db)
+        self.data.green(db).long(db)
     }
 
     /// Returns the span of the syntax node without trivia.
@@ -507,17 +554,17 @@ impl<'a> SyntaxNode<'a> {
 
     /// Gets the kind of the given node's parent if it exists.
     pub fn parent_kind(&self, db: &dyn Database) -> Option<SyntaxKind> {
-        Some(self.parent(db)?.kind(db))
+        Some(self.parent?.green(db).long(db).kind)
     }
 
     /// Gets the kind of the given node's grandparent if it exists.
     pub fn grandparent_kind(&self, db: &dyn Database) -> Option<SyntaxKind> {
-        Some(self.parent(db)?.parent(db)?.kind(db))
+        self.parent(db)?.parent_kind(db)
     }
 
     /// Gets the kind of the given node's grandgrandparent if it exists.
     pub fn grandgrandparent_kind(&self, db: &dyn Database) -> Option<SyntaxKind> {
-        Some(self.parent(db)?.parent(db)?.parent(db)?.kind(db))
+        self.grandparent(db)?.parent_kind(db)
     }
 }