Bump tree-sitter-swift to 0.7.3 and overhaul Swift query stack

ast/Cargo.toml

@@ -46,7 +46,7 @@ streaming-iterator = "0.1.9"
 git-url-parse = "0.4.5"
 tree-sitter-kotlin-sg = "0.*"
 # tree-sitter-swift = "0.*"
-tree-sitter-swift = "0.7.1"
+tree-sitter-swift = "0.7.3"
 tree-sitter-java = "0.23.5"
 tree-sitter-svelte-ng = "1.*"
 # gitoxide-core = { version = "0.42.0", features = ["blocking-client"] }

ast/src/lang/parse/collect.rs

@@ -1,12 +1,14 @@
-use super::utils::trim_quotes;
+use super::super::queries::consts::{CLASS_DEFINITION, CLASS_NAME};
+use super::utils::{clean_class_name, trim_quotes};
 use crate::builder::utils::log_stage_timing;
 use crate::lang::{graphs::Graph, *};
 use lsp::{Cmd as LspCmd, Position, Res as LspRes};
 use shared::error::{Error, Result};
+use std::collections::HashSet;
 use std::time::Instant;
 use streaming_iterator::StreamingIterator;
 use tracing::warn;
-use tree_sitter::Node as TreeNode;
+use tree_sitter::{Node as TreeNode, QueryMatch};
 impl Lang {
     pub fn collect<G: Graph>(
         &self,
@@ -40,6 +42,25 @@ impl Lang {
         Ok(res)
     }
 
+    fn class_match_name_and_kind(
+        &self,
+        q: &Query,
+        m: &QueryMatch,
+        code: &str,
+    ) -> Result<(Option<String>, Option<String>)> {
+        let mut name: Option<String> = None;
+        let mut kind: Option<String> = None;
+        Self::loop_captures(q, m, code, |body, node, o| {
+            if o == CLASS_NAME {
+                name = Some(clean_class_name(&body));
+            } else if o == CLASS_DEFINITION {
+                kind = self.lang.class_declaration_kind(node, code);
+            }
+            Ok(())
+        })?;
+        Ok((name, kind))
+    }
+
     pub fn collect_classes<G: Graph>(
         &self,
         q: &Query,
@@ -48,10 +69,35 @@ impl Lang {
         graph: &G,
     ) -> Result<Vec<(NodeData, Vec<Edge>)>> {
         let tree = self.parse(code, &NodeType::Class)?;
+
+        // Avoid a redundant node for e.g. Swift's `extension Foo { ... }` when `Foo` is
+        // already declared elsewhere in the same file.
+        let mut names_with_primary_decl: HashSet<String> = HashSet::new();
+        {
+            let mut cursor = QueryCursor::new();
+            let mut matches = cursor.matches(q, tree.root_node(), code.as_bytes());
+            while let Some(m) = matches.next() {
+                let (name, kind) = self.class_match_name_and_kind(q, m, code)?;
+                if let Some(name) = name {
+                    if kind.as_deref() != Some("extension") {
+                        names_with_primary_decl.insert(name);
+                    }
+                }
+            }
+        }
+
         let mut cursor = QueryCursor::new();
         let mut matches = cursor.matches(q, tree.root_node(), code.as_bytes());
         let mut res = Vec::new();
         while let Some(m) = matches.next() {
+            let (name, kind) = Self::class_match_name_and_kind(self, q, m, code)?;
+            if kind.as_deref() == Some("extension") {
+                if let Some(name) = &name {
+                    if names_with_primary_decl.contains(name) {
+                        continue;
+                    }
+                }
+            }
             if let Some((cls, edges)) =
                 self.format_class_with_associations(m, code, file, q, graph)?
             {

ast/src/lang/queries/mod.rs

@@ -82,6 +82,11 @@ pub trait Stack {
         None
     }
     fn class_definition_query(&self) -> String;
+    // Discriminator (e.g. "extension") for a CLASS_DEFINITION node, for grammars
+    // that reuse one node kind for several declaration forms.
+    fn class_declaration_kind(&self, _node: TreeNode, _code: &str) -> Option<String> {
+        None
+    }
     fn instance_definition_query(&self) -> Option<String> {
         None
     }

ast/src/lang/queries/swift.rs

@@ -48,12 +48,22 @@ impl Stack for Swift {
         Some(format!(
             r#"
             (source_file
-                (property_declaration 
+                (property_declaration
                     (pattern
                         (simple_identifier) @{VARIABLE_NAME}
                     )
                 )@{VARIABLE_DECLARATION}
             )
+
+            (class_declaration
+                (class_body
+                    (property_declaration
+                        (pattern
+                            (simple_identifier) @{VARIABLE_NAME}
+                        )
+                    )@{VARIABLE_DECLARATION}
+                )
+            )
             "#
         ))
     }
@@ -71,6 +81,40 @@ impl Stack for Swift {
         )
     }
 
+    fn class_declaration_kind(&self, node: TreeNode, code: &str) -> Option<String> {
+        node.child_by_field_name("declaration_kind")
+            .and_then(|n| n.utf8_text(code.as_bytes()).ok())
+            .map(|s| s.to_string())
+    }
+
+    fn trait_query(&self) -> Option<String> {
+        Some(format!(
+            r#"
+            (protocol_declaration
+                name: (type_identifier) @{TRAIT_NAME}
+            ) @{TRAIT}
+            "#
+        ))
+    }
+
+    fn implements_query(&self) -> Option<String> {
+        // Swift's grammar doesn't distinguish a superclass from a protocol
+        // conformance (both are `inheritance_specifier`), so treat every entry
+        // as a conformance/inheritance relationship.
+        Some(format!(
+            r#"
+            (class_declaration
+                name: [(type_identifier)(user_type)] @{CLASS_NAME}
+                (inheritance_specifier
+                    (user_type
+                        (type_identifier) @{TRAIT_NAME}
+                    )
+                )
+            ) @{IMPLEMENTS}
+            "#
+        ))
+    }
+
     fn function_definition_query(&self) -> String {
         format!(
             r#"
@@ -118,23 +162,30 @@ impl Stack for Swift {
         code: &str,
         file: &str,
         func_name: &str,
-        _callback: &dyn Fn(&str) -> Option<(NodeData, NodeType)>,
+        find_class: &dyn Fn(&str) -> Option<(NodeData, NodeType)>,
         _parent_type: Option<&str>,
     ) -> Result<Option<Operand>> {
         let mut parent = node.parent();
         while let Some(current) = parent {
-            if current.kind() == "class_declaration" {
+            if current.kind() == "class_declaration" || current.kind() == "protocol_declaration" {
                 break;
             }
             parent = current.parent();
         }
         let parent_of = match parent {
             Some(p) => {
-                let query = self.q("name: (type_identifier) @class-name", &NodeType::Class);
-                query_to_ident(query, p, code)?.map(|parent_name| Operand {
-                    source: NodeKeys::new(&parent_name, file, p.start_position().row),
-                    target: NodeKeys::new(func_name, file, node.start_position().row),
-                    source_type: NodeType::Class,
+                let query =
+                    self.q("name: [(type_identifier)(user_type)] @class-name", &NodeType::Class);
+                let parent_name = query_to_ident(query, p, code)?;
+                // Resolve by name (not this block's own position) so an
+                // `extension Foo { ... }` block's methods land on the same
+                // node as `Foo`'s own declaration.
+                parent_name.and_then(|name| {
+                    find_class(&name).map(|(class, source_type)| Operand {
+                        source: NodeKeys::new(&class.name, &class.file, class.start),
+                        target: NodeKeys::new(func_name, file, node.start_position().row),
+                        source_type,
+                    })
                 })
             }
             None => None,

ast/src/testing/coverage/swift.rs

@@ -71,7 +71,7 @@ async fn test_btreemap_graph_structure() -> Result<()> {
     assert_eq!(endpoints.len(), 0);
 
     let functions = graph.find_nodes_by_type(NodeType::Function);
-    assert_eq!(functions.len(), 36);
+    assert_eq!(functions.len(), 37);
 
     let unit_tests = graph.find_nodes_by_type(NodeType::UnitTest);
     assert_eq!(unit_tests.len(), 5);
@@ -89,7 +89,7 @@ async fn test_btreemap_graph_structure() -> Result<()> {
     assert_eq!(data_models.len(), 1);
 
     let traits = graph.find_nodes_by_type(NodeType::Trait);
-    assert_eq!(traits.len(), 0);
+    assert_eq!(traits.len(), 1);
 
     Ok(())
 }
@@ -108,16 +108,16 @@ async fn test_btreemap_test_to_function_edges() -> Result<()> {
     let graph = repo.build_graph_inner::<BTreeMapGraph>().await?;
 
     let calls_edges = graph.count_edges_of_type(EdgeType::Calls);
-    assert_eq!(calls_edges, 24);
+    assert_eq!(calls_edges, 25);
 
     let contains_edges = graph.count_edges_of_type(EdgeType::Contains);
-    assert_eq!(contains_edges, 168);
+    assert_eq!(contains_edges, 210);
 
     let handler_edges = graph.count_edges_of_type(EdgeType::Handler);
     assert_eq!(handler_edges, 0);
 
     let implements_edges = graph.count_edges_of_type(EdgeType::Implements);
-    assert_eq!(implements_edges, 0);
+    assert_eq!(implements_edges, 1);
 
     let unit_test_to_function_edges =
         graph.find_nodes_with_edge_type(NodeType::UnitTest, NodeType::Function, EdgeType::Calls);

ast/src/testing/swift/LegacyApp/IntegrationTests/APIIntegrationTests.swift

@@ -12,6 +12,7 @@ import XCTest
 // @ast edge: Calls -> Function "updatePeopleProfileWith" "API.swift"
 // @ast edge: Calls -> Function "createRequest" "API.swift"
 // @ast node: Import "import-imports-srctestingswiftlegacyappintegrationtestsapiintegrationtestsswift-0"
+// @ast node: Var "mockAPI"
 
 final class APIIntegrationTests: XCTestCase {
 

ast/src/testing/swift/LegacyApp/SphinxTestApp/API.swift

@@ -22,6 +22,9 @@ import SwiftyJSON
 // @ast node: Request "/people"
 // @ast node: Request "/person"
 // @ast node: Import "import-imports-srctestingswiftlegacyappsphinxtestappapiswift-7"
+// @ast node: Var "sharedInstance"
+// @ast node: Var "instance"
+// @ast node: Var "host"
 
 class API {
     class var sharedInstance : API {

ast/src/testing/swift/LegacyApp/SphinxTestApp/AppDelegate.swift

@@ -18,6 +18,7 @@ import CoreData
 // @ast node: Function "saveContext"
 // @ast node: Var "Name"
 // @ast node: Var "Version"
+// @ast node: Var "persistentContainer"
 // @ast node: Import "import-imports-srctestingswiftlegacyappsphinxtestappappdelegateswift-7"
 
 let Name = "StakGraph"

ast/src/testing/swift/LegacyApp/SphinxTestApp/CoreData/Person+CoreDataProperties.swift

@@ -11,6 +11,10 @@ import CoreData
 // @ast node: Class "Person"
 // @ast node: Function "fetchRequest"
 // @ast node: Import "import-imports-srctestingswiftlegacyappsphinxtestappcoredatapersoncoredatapropertiesswift-8"
+// @ast node: Var "alias"
+// @ast node: Var "imageUrl"
+// @ast node: Var "publicKey"
+// @ast node: Var "routeHint"
 
 
 extension Person {

ast/src/testing/swift/LegacyApp/SphinxTestApp/PersonTableViewCell.swift

@@ -11,6 +11,7 @@ import UIKit
 // @ast edge: Operand -> Function "setSelected" "PersonTableViewCell.swift"
 // @ast node: Function "awakeFromNib"
 // @ast node: Function "setSelected"
+// @ast node: Var "nameLabel"
 // @ast node: Import "import-imports-srctestingswiftlegacyappsphinxtestapppersontableviewcellswift-7"
 
 class PersonTableViewCell: UITableViewCell {

ast/src/testing/swift/LegacyApp/SphinxTestApp/SceneDelegate.swift

@@ -19,6 +19,7 @@ import UIKit
 // @ast node: Function "sceneWillResignActive"
 // @ast node: Function "sceneWillEnterForeground"
 // @ast node: Function "sceneDidEnterBackground"
+// @ast node: Var "window"
 // @ast node: Import "import-imports-srctestingswiftlegacyappsphinxtestappscenedelegateswift-7"
 
 class SceneDelegate: UIResponder, UIWindowSceneDelegate {

ast/src/testing/swift/LegacyApp/SphinxTestApp/ViewController.swift

@@ -8,7 +8,6 @@
 import UIKit
 import SwiftyJSON
 // @ast node: Class "ViewController"
-// @ast node: Class "ViewController"
 // @ast edge: Operand -> Function "viewDidLoad" "ViewController.swift"
 // @ast edge: Operand -> Function "configureTableView" "ViewController.swift"
 // @ast edge: Operand -> Function "getPeopleAndSave" "ViewController.swift"
@@ -28,6 +27,8 @@ import SwiftyJSON
 // @ast node: Function "tableView"
 // @ast node: Function "tableView"
 // @ast edge: Calls -> Function "updateProfile" "ViewController.swift"
+// @ast node: Var "peopleTableView"
+// @ast node: Var "persons"
 // @ast node: Import "import-imports-srctestingswiftlegacyappsphinxtestappviewcontrollerswift-7"
 
 class ViewController: UIViewController {

ast/src/testing/swift/ModernApp/Sources/Core/Extensions/String+Extensions.swift

@@ -1,6 +1,7 @@
 import Foundation
 // @ast node: Class "String"
 // @ast node: Function "isValidEmail"
+// @ast node: Var "localized"
 // @ast node: Import "import-imports-srctestingswiftmodernappsourcescoreextensionsstringextensionsswift-0"
 
 extension String {

ast/src/testing/swift/ModernApp/Sources/Core/Network/APIClient.swift

@@ -5,6 +5,7 @@ import Foundation
 // @ast edge: Operand -> Function "post" "APIClient.swift"
 // @ast node: Function "fetch"
 // @ast node: Function "post"
+// @ast node: Var "session"
 // @ast node: Import "import-imports-srctestingswiftmodernappsourcescorenetworkapiclientswift-0"
 
 enum APIError: Error {

ast/src/testing/swift/ModernApp/Sources/Features/Chat/ChatManager.swift

@@ -2,6 +2,13 @@ import Foundation
 // @ast node: Class "ChatManager"
 // @ast edge: Operand -> Function "send" "ChatManager.swift"
 // @ast node: Function "send"
+// @ast node: Trait "ChatManagerDelegate"
+// @ast node: Class "ConsoleChatObserver"
+// @ast edge: Implements -> Trait "ChatManagerDelegate" "ChatManager.swift"
+// @ast edge: Operand -> Function "didReceiveMessage" "ChatManager.swift"
+// @ast node: Function "didReceiveMessage"
+// @ast node: Var "delegate"
+// @ast node: Var "onStatusChange"
 // @ast node: Import "import-imports-srctestingswiftmodernappsourcesfeatureschatchatmanagerswift-0"
 
 protocol ChatManagerDelegate: AnyObject {
@@ -11,7 +18,7 @@ protocol ChatManagerDelegate: AnyObject {
 class ChatManager {
     weak var delegate: ChatManagerDelegate?
     var onStatusChange: ((Bool) -> Void)?
-    
+
     func send(message: String) {
         // Simulate network delay
         DispatchQueue.global().asyncAfter(deadline: .now() + 0.5) { [weak self] in
@@ -20,3 +27,9 @@ class ChatManager {
         }
     }
 }
+
+class ConsoleChatObserver: ChatManagerDelegate {
+    func didReceiveMessage(_ message: String) {
+        print(message)
+    }
+}

ast/src/testing/swift/ModernApp/Sources/Features/Chat/ChatView.swift

@@ -1,6 +1,10 @@
 import SwiftUI
 // @ast node: Class "ChatView"
 // @ast node: Class "ChatViewModel"
+// @ast node: Var "viewModel"
+// @ast node: Var "body"
+// @ast node: Var "messages"
+// @ast node: Var "searchText"
 // @ast node: Import "import-imports-srctestingswiftmodernappsourcesfeatureschatchatviewswift-0"
 
 struct ChatView: View {

ast/src/testing/swift/ModernApp/Sources/Features/Profile/ProfileModel.swift

@@ -1,6 +1,11 @@
 import Foundation
 // @ast node: Class "Profile"
 // @ast node: Class "Status"
+// @ast node: Var "id"
+// @ast node: Var "username"
+// @ast node: Var "bio"
+// @ast node: Var "avatarURL"
+// @ast node: Var "status"
 // @ast node: Import "import-imports-srctestingswiftmodernappsourcesfeaturesprofileprofilemodelswift-0"
 
 struct Profile: Codable, Identifiable {