add context level functionality

Author: chrisknapp98

Core/Sources/HostApp/Benchmark/Data/Manager/MultiFileContextBenchmarkManager.swift

@@ -29,6 +29,17 @@ class MultiFileContextBenchmarkManager: BenchmarkManager {
     var selectedGenAIModel: AnyPublisher<GenAILanguageModel, Never> {
         selectedGenAIModelSubject.eraseToAnyPublisher()
     }
+    
+    private let contextLevelLimitSubject = CurrentValueSubject<ContextLevelLimit, Never>(.firstLevel)
+    var contextLevelLimit: AnyPublisher<ContextLevelLimit, Never> {
+        contextLevelLimitSubject.eraseToAnyPublisher()
+    }
+    
+    private var contextFileAmountLimitSubject = CurrentValueSubject<Int?, Never>(nil)
+    var contextFileAmountLimit: AnyPublisher<Int?, Never> {
+        contextFileAmountLimitSubject.eraseToAnyPublisher()
+    }
+    
     private var cancellables = Set<AnyCancellable>()
 
     init(benchmarkSettingsRepository: BenchmarkSettingsRepository) {
@@ -239,25 +250,23 @@ class MultiFileContextBenchmarkManager: BenchmarkManager {
         }
     }
 
-    private func retrieveRelevantSymbolsForFileContent(file: FileContent, workspace: Workspace) async -> RelevantSymbolsSummary? {
-        let multiFileContextManager = MultiFileContextManager(
+    private func retrieveRelevantSymbolsForFileContent(
+        file: FileContent,
+        workspace: Workspace
+    ) async -> RelevantSymbolsSummary? {
+        guard isMultiFileEnabledSubject.value else { return nil }
+
+        let manager = MultiFileContextManager(
             workspaceProvider: ManualWorkspaceProvider(workspace: workspace),
             parser: SwiftProgrammingLanguageSyntaxParser()
         )
-        if isMultiFileEnabledSubject.value {
-            let start = Date()
-            let symbols = await multiFileContextManager.retrieveRelevantSymbolsForFileContent(file: file, ignoreWithinPaths: ["/Benchmark/"])
-            let end = Date()
-            let timeTakenInSeconds = end.timeIntervalSince(start)
-            return RelevantSymbolsSummary(
-                symbolsAtLevel: [
-                    Array(symbols.values)
-                ],
-                durationInSeconds: timeTakenInSeconds
-            )
-        } else {
-            return nil
-        }
+
+        return await manager.expandRelevantSymbols(
+            from: file,
+            ignoreWithinPaths: ["/Benchmark/"],
+            maxDepth: contextLevelLimitSubject.value.indexLimit,   // e.g. First=0, Second=1, …, nil=no limit
+            maxFiles: contextFileAmountLimitSubject.value             // Int?, nil=no cap
+        )
     }
 
     private func applyCodeSuggestion(suggestion: SuggestionBasic.CodeSuggestion, at fileURL: URL) async {
@@ -470,6 +479,16 @@ class MultiFileContextBenchmarkManager: BenchmarkManager {
         currentStates[directory] = currentStatesInDirectory
         taskStatesSubject.send(currentStates)
     }
+    
+    @MainActor
+    func saveContextLevelLimit(_ limit: ContextLevelLimit) {
+        contextLevelLimitSubject.send(limit)
+    }
+    
+    @MainActor
+    func saveContextFileAmountLimit(_ limit: Int?) {
+        contextFileAmountLimitSubject.send(limit)
+    }
 }
 
 extension MetadataDTO {

Core/Sources/HostApp/Benchmark/Domain/Entity/ContextLevelLimit.swift

@@ -0,0 +1,21 @@
+enum ContextLevelLimit: CaseIterable {
+    case firstLevel
+    case secondLevel
+    case thirdLevel
+    
+    var name: String {
+        switch self {
+        case .firstLevel: return "First Level"
+        case .secondLevel: return "Second Level"
+        case .thirdLevel: return "Third Level"
+        }
+    }
+    
+    var indexLimit: Int? {
+        switch self {
+        case .firstLevel: return 0
+        case .secondLevel: return 1
+        case .thirdLevel: return 2
+        }
+    }
+}

Core/Sources/HostApp/Benchmark/Domain/Entity/RelevantSymbolsSummary.swift

@@ -10,4 +10,12 @@ protocol BenchmarkManager {
     func changeGenAIModel(to newModel: GenAILanguageModel)
     @MainActor
     func updateMultiFileContextState(_ newValue: Bool)
+    
+    var contextLevelLimit: AnyPublisher<ContextLevelLimit, Never> { get }
+    @MainActor
+    func saveContextLevelLimit(_ limit: ContextLevelLimit)
+    
+    var contextFileAmountLimit: AnyPublisher<Int?, Never> { get }
+    @MainActor
+    func saveContextFileAmountLimit(_ limit: Int?)
 }

Core/Sources/HostApp/Benchmark/Domain/Manager/BenchmarkManager.swift

@@ -28,6 +28,28 @@ struct BenchmarkView: View {
                         ConfigurationButtonView(module: module)
                     }
                     .padding(.vertical)
+                    
+                    HStack(spacing: 20) {
+                        Spacer()
+                        Picker("Context Level Limit: ", selection: $viewModel.contextLevelLimit) {
+                            ForEach(ContextLevelLimit.allCases, id: \.self) { limit in
+                                Text(limit.name).tag(limit)
+                            }
+                        }
+                        .fixedSize(horizontal: true, vertical: false)
+                        
+                        Text("Context File Amount Limit:")
+                        let numberFormatter: NumberFormatter = {
+                            let f = NumberFormatter()
+                            f.numberStyle = .none
+                            return f
+                        }()
+                        TextField("Optional", value: $viewModel.contextFileAmountLimit, formatter: numberFormatter)
+                            .textFieldStyle(RoundedBorderTextFieldStyle())
+                            .frame(width: 100)
+                    }
+                    .padding(.bottom)
+                    
                     ForEach(viewModel.benchmarkDirectories, id: \.self) { directory in
                         BenchmarkDirectoryEntryView(
                             directory: directory,

Core/Sources/HostApp/Benchmark/Presentation/View/BenchmarkView.swift

@@ -16,6 +16,18 @@ class BenchmarkViewModel: ObservableObject {
             Task { await benchmarkManager.changeGenAIModel(to: selectedLanguageModel) }
         }
     }
+    @Published var contextLevelLimit: ContextLevelLimit = .firstLevel {
+        didSet {
+            guard oldValue != contextLevelLimit else { return }
+            Task { await benchmarkManager.saveContextLevelLimit(contextLevelLimit) }
+        }
+    }
+    @Published var contextFileAmountLimit: Int? = nil {
+        didSet {
+            guard oldValue != contextFileAmountLimit else { return }
+            Task { await benchmarkManager.saveContextFileAmountLimit(contextFileAmountLimit) }
+        }
+    }
 
     private let benchmarkSettingsRepository: BenchmarkSettingsRepository
     private let benchmarkManager: BenchmarkManager
@@ -35,6 +47,12 @@ class BenchmarkViewModel: ObservableObject {
         benchmarkManager.selectedGenAIModel
             .assign(to: \.selectedLanguageModel, on: self)
             .store(in: &cancellables)
+        benchmarkManager.contextLevelLimit
+            .assign(to: \.contextLevelLimit, on: self)
+            .store(in: &cancellables)
+        benchmarkManager.contextFileAmountLimit
+            .assign(to: \.contextFileAmountLimit, on: self)
+            .store(in: &cancellables)
     }
 
     func loadBenchmarkDirectories() {

Core/Sources/HostApp/Benchmark/Presentation/ViewModel/BenchmarkViewModel.swift

@@ -0,0 +1,4 @@
+public struct RelevantSymbolsSummary {
+    public let symbolsAtLevel: [[SymbolContent]]
+    public let durationInSeconds: Double
+}

Core/Sources/Service/MultiFileContext/Entity/RelevantSymbolsSummary.swift

@@ -62,6 +62,87 @@ public class MultiFileContextManager {
         return result
     }
 
+    private func readFile(_ urlString: String) -> FileContent? {
+        guard let url = URL(string: urlString) else { return nil }
+        guard let content = try? String(contentsOf: url, encoding: .utf8) else { return nil }
+        return FileContent(fileURL: urlString, content: content)
+    }
+
+    public func expandRelevantSymbols(
+        from root: FileContent,
+        ignoreWithinPaths: [String] = [],
+        maxDepth: Int? = nil,
+        maxFiles: Int? = nil
+    ) async -> RelevantSymbolsSummary {
+        @inline(__always)
+        func isIgnored(_ url: String) -> Bool {
+            ignoreWithinPaths.contains { url.contains($0) }
+        }
+
+        let start = Date()
+
+        let level0 = await retrieveRelevantSymbolsForFileContent(file: root, ignoreWithinPaths: ignoreWithinPaths)
+        var level0Unique: [SymbolContent] = []
+        var includedFiles = Set<String>()
+        for s in level0.values {
+            if includedFiles.insert(s.fileURL).inserted {
+                level0Unique.append(s)
+            }
+        }
+        var levels: [[SymbolContent]] = [level0Unique]
+
+        var visited = Set<String>([root.fileURL])
+        var depth = 0
+        var remainingFiles = maxFiles
+
+        while maxDepth == nil || depth < maxDepth! {
+            let prev = levels.last ?? []
+            var nextURLs = Set(prev.map(\.fileURL))
+            nextURLs.subtract(visited)
+            nextURLs = nextURLs.filter { !isIgnored($0) }
+
+            if nextURLs.isEmpty { break }
+
+            if let remaining = remainingFiles, nextURLs.count > remaining {
+                nextURLs = Set(nextURLs.prefix(remaining))
+            }
+
+            var rawNextLevel: [SymbolContent] = []
+            for url in nextURLs {
+                visited.insert(url)
+                guard let fc = readFile(url) else { continue }
+                let dict = await retrieveRelevantSymbolsForFileContent(file: fc, ignoreWithinPaths: ignoreWithinPaths)
+                rawNextLevel.append(contentsOf: dict.values)
+            }
+
+            var nextLevel: [SymbolContent] = []
+            var seenThisLevel = Set<String>()
+            for s in rawNextLevel {
+                guard !includedFiles.contains(s.fileURL) else { continue }
+                if seenThisLevel.insert(s.fileURL).inserted {
+                    nextLevel.append(s)
+                }
+            }
+
+            if nextLevel.isEmpty { break }
+
+            includedFiles.formUnion(nextLevel.map(\.fileURL))
+
+            levels.append(nextLevel)
+            depth += 1
+
+            if remainingFiles != nil {
+                remainingFiles! -= nextURLs.count
+                if remainingFiles! <= 0 { break }
+            }
+        }
+
+        return RelevantSymbolsSummary(
+            symbolsAtLevel: levels,
+            durationInSeconds: Date().timeIntervalSince(start)
+        )
+    }
+    
     public func retrieveRelevantSymbolsForFileContent(
         file: FileContent,
         ignoreWithinPaths: [String] = []
@@ -171,7 +252,7 @@ extension String {
             let before = (range.lowerBound == startIndex) ? nil : self[index(before: range.lowerBound)]
             let after  = (range.upperBound == endIndex)   ? nil : self[range.upperBound]
 
-            let boundaryBefore = before.map { !isIdentChar($0) } ?? true
+            let boundaryBefore = before.map { !isIdentChar($0) && $0 != "." } ?? true
             let boundaryAfter  = after.map  { !isIdentChar($0) } ?? true
 
             if boundaryBefore && boundaryAfter {

Core/Sources/Service/MultiFileContext/MultiFileContextManager.swift

@@ -10,43 +10,45 @@ class SwiftDeclarationCollector: SyntaxVisitor {
         self.sourceText = sourceText
         super.init(viewMode: .all)
     }
-    
-//    override func visit(_ node: ImportDeclSyntax) -> SyntaxVisitorContinueKind {
-//        recordSymbol(name: node.name.text, kind: "import", node: node)
-//        return .skipChildren
-//    }
 
     override func visit(_ node: ClassDeclSyntax) -> SyntaxVisitorContinueKind {
+        guard !isPrivateOrFilePrivate(node.modifiers) else { return .skipChildren }
         recordSymbol(name: node.name.text, kind: ClassificationKeywords.classWord, node: node)
         return .skipChildren
     }
 
     override func visit(_ node: StructDeclSyntax) -> SyntaxVisitorContinueKind {
+        guard !isPrivateOrFilePrivate(node.modifiers) else { return .skipChildren }
         recordSymbol(name: node.name.text, kind: ClassificationKeywords.structWord, node: node)
         return .skipChildren
     }
 
     override func visit(_ node: EnumDeclSyntax) -> SyntaxVisitorContinueKind {
+        guard !isPrivateOrFilePrivate(node.modifiers) else { return .skipChildren }
         recordSymbol(name: node.name.text, kind: ClassificationKeywords.enumWord, node: node)
         return .skipChildren
     }
 
     override func visit(_ node: ProtocolDeclSyntax) -> SyntaxVisitorContinueKind {
+        guard !isPrivateOrFilePrivate(node.modifiers) else { return .skipChildren }
         recordSymbol(name: node.name.text, kind: ClassificationKeywords.protocolWord, node: node)
         return .skipChildren
     }
 
     override func visit(_ node: ActorDeclSyntax) -> SyntaxVisitorContinueKind {
+        guard !isPrivateOrFilePrivate(node.modifiers) else { return .skipChildren }
         recordSymbol(name: node.name.text, kind: ClassificationKeywords.actorWord, node: node)
         return .skipChildren
     }
 
     override func visit(_ node: FunctionDeclSyntax) -> SyntaxVisitorContinueKind {
+        guard !isPrivateOrFilePrivate(node.modifiers) else { return .skipChildren }
         recordSymbol(name: node.name.text, kind: ClassificationKeywords.funcWord, node: node)
         return .skipChildren
     }
 
     override func visit(_ node: VariableDeclSyntax) -> SyntaxVisitorContinueKind {
+        guard !isPrivateOrFilePrivate(node.modifiers) else { return .skipChildren }
         guard let binding = node.bindings.first,
               let pattern = binding.pattern.as(IdentifierPatternSyntax.self),
               let keyword: ClassificationKeywords = ClassificationKeywords(rawValue: node.bindingSpecifier.text) else {
@@ -58,11 +60,14 @@ class SwiftDeclarationCollector: SyntaxVisitor {
 
     override func visit(_ node: ExtensionDeclSyntax) -> SyntaxVisitorContinueKind {
         let name = node.extendedType.trimmedDescription
+        guard !isPrivateOrFilePrivate(node.modifiers) && name != "View" else { return .skipChildren }
+        
         recordSymbol(name: name, kind: ClassificationKeywords.extensionWord, node: node)
         return .skipChildren
     }
 
     override func visit(_ node: TypeAliasDeclSyntax) -> SyntaxVisitorContinueKind {
+        guard !isPrivateOrFilePrivate(node.modifiers) else { return .skipChildren }
         recordSymbol(name: node.name.text, kind: ClassificationKeywords.typealiasWord, node: node)
         return .skipChildren
     }
@@ -87,4 +92,12 @@ class SwiftDeclarationCollector: SyntaxVisitor {
         )
         symbols.append(symbol)
     }
+    
+    private func isPrivateOrFilePrivate(_ modifiers: DeclModifierListSyntax?) -> Bool {
+        guard let modifiers else { return false }
+        return modifiers.contains { modifier in
+            let text = modifier.name.text
+            return text == "private" || text == "fileprivate"
+        }
+    }
 }