wafer-lsp 0.1.0__py3-none-any.whl

wafer_lsp/handlers/workspace_symbol.py

@@ -0,0 +1,87 @@
+
+from lsprotocol.types import Location, Position, Range, SymbolKind, WorkspaceSymbol
+
+from ..languages.registry import get_language_registry
+
+
+def _matches_query(name: str, query: str) -> bool:
+    if not query:
+        return True
+
+    name_lower = name.lower()
+    query_lower = query.lower()
+
+    query_idx = 0
+    for char in name_lower:
+        if query_idx < len(query_lower) and char == query_lower[query_idx]:
+            query_idx += 1
+
+    return query_idx == len(query_lower)
+
+
+def handle_workspace_symbol(query: str) -> list[WorkspaceSymbol]:
+    registry = get_language_registry()
+    symbols: list[WorkspaceSymbol] = []
+
+    return symbols
+
+
+def handle_workspace_symbol_with_documents(
+    query: str,
+    document_contents: dict[str, str]
+) -> list[WorkspaceSymbol]:
+    registry = get_language_registry()
+    symbols: list[WorkspaceSymbol] = []
+
+    for uri, content in document_contents.items():
+        language_info = registry.parse_file(uri, content)
+
+        if not language_info:
+            continue
+
+        for kernel in language_info.kernels:
+            if _matches_query(kernel.name, query):
+                symbols.append(WorkspaceSymbol(
+                    name=kernel.name,
+                    kind=SymbolKind.Function,
+                    location=Location(
+                        uri=uri,
+                        range=Range(
+                            start=Position(line=kernel.line, character=0),
+                            end=Position(line=kernel.line, character=0)
+                        )
+                    ),
+                    container_name=f"GPU Kernel ({registry.get_language_name(kernel.language)})"
+                ))
+
+        for layout in language_info.layouts:
+            if _matches_query(layout.name, query):
+                symbols.append(WorkspaceSymbol(
+                    name=layout.name,
+                    kind=SymbolKind.Variable,
+                    location=Location(
+                        uri=uri,
+                        range=Range(
+                            start=Position(line=layout.line, character=0),
+                            end=Position(line=layout.line, character=0)
+                        )
+                    ),
+                    container_name="Layout"
+                ))
+
+        for struct in language_info.structs:
+            if _matches_query(struct.name, query):
+                symbols.append(WorkspaceSymbol(
+                    name=struct.name,
+                    kind=SymbolKind.Struct,
+                    location=Location(
+                        uri=uri,
+                        range=Range(
+                            start=Position(line=struct.line, character=0),
+                            end=Position(line=struct.line, character=0)
+                        )
+                    ),
+                    container_name=f"Struct ({registry.get_language_name(struct.language)})"
+                ))
+
+    return symbols

wafer_lsp/languages/README.md

@@ -0,0 +1,195 @@
+# Adding a New Language to Wafer LSP
+
+This guide explains how to add support for a new GPU programming language to the LSP server.
+
+## Architecture Overview
+
+The LSP server uses a modular, component-based architecture:
+
+- **`LanguageDetector`**: Detects language from file extensions/URIs
+- **`ParserManager`**: Manages parser instances for each language
+- **`ParserResultConverter`**: Converts parser-specific types to common types
+- **`LanguageRegistry`**: Coordinates all components (facade pattern)
+
+Each component is small, focused, and independently testable.
+
+## Overview
+
+To add a new language, you need to:
+
+1. Create a parser that implements `BaseParser`
+2. Register the parser with the language registry
+3. (Optional) Add language-specific handlers
+
+## Step 1: Create a Parser
+
+Create a new parser file in `wafer_lsp/parsers/`:
+
+```python
+# wafer_lsp/parsers/my_language_parser.py
+
+from typing import List, Dict, Any, Optional
+from dataclasses import dataclass
+from .base_parser import BaseParser
+
+@dataclass
+class MyLanguageKernel:
+    name: str
+    line: int
+    parameters: List[str]
+    docstring: Optional[str] = None
+
+class MyLanguageParser(BaseParser):
+    """Parser for MyLanguage GPU code."""
+    
+    def parse_file(self, content: str) -> Dict[str, Any]:
+        """Parse file and extract kernels, layouts, etc."""
+        kernels: List[MyLanguageKernel] = []
+        
+        # Your parsing logic here
+        # Example: find kernel definitions
+        for match in re.finditer(r'kernel\s+(\w+)', content):
+            kernels.append(MyLanguageKernel(
+                name=match.group(1),
+                line=content[:match.start()].count('\n'),
+                parameters=[],
+            ))
+        
+        return {
+            "kernels": kernels,
+            "layouts": [],  # If your language has layouts
+            "structs": [],  # If your language has structs
+        }
+```
+
+## Step 2: Register the Parser
+
+Update `wafer_lsp/languages/registry.py` to register your parser:
+
+```python
+from ..parsers.my_language_parser import MyLanguageParser
+
+class LanguageRegistry:
+    def _register_defaults(self):
+        # ... existing registrations ...
+        
+        # Register your new language
+        self.register_language(
+            language_id="mylang",
+            display_name="My Language",
+            parser=MyLanguageParser(),
+            extensions=[".mylang", ".ml"],
+            file_patterns=["*.mylang", "*.ml"]
+        )
+```
+
+## Step 3: Update VS Code Extension (if needed)
+
+If your language needs special handling in the VS Code extension, update `lspClient.ts`:
+
+```typescript
+const clientOptions: LanguageClientOptions = {
+    documentSelector: [
+        // ... existing selectors ...
+        { scheme: 'file', language: 'mylang' },  // Your language
+    ],
+    // ...
+};
+```
+
+## Parser Interface
+
+Your parser must implement `BaseParser`:
+
+```python
+class BaseParser(ABC):
+    @abstractmethod
+    def parse_file(self, content: str) -> Dict[str, Any]:
+        """Parse file content and extract language-specific constructs.
+        
+        Returns:
+            Dictionary with keys:
+            - "kernels": List of kernel objects (must have .name, .line, .parameters)
+            - "layouts": List of layout objects (must have .name, .line)
+            - "structs": List of struct objects (must have .name, .line)
+        """
+        pass
+```
+
+## Example: Adding OpenCL Support
+
+Here's a complete example for adding OpenCL support:
+
+```python
+# wafer_lsp/parsers/opencl_parser.py
+
+import re
+from typing import List, Dict, Any
+from dataclasses import dataclass
+from .base_parser import BaseParser
+
+@dataclass
+class OpenCLKernel:
+    name: str
+    line: int
+    parameters: List[str]
+
+class OpenCLParser(BaseParser):
+    def parse_file(self, content: str) -> Dict[str, Any]:
+        kernels: List[OpenCLKernel] = []
+        
+        # Pattern: __kernel void kernel_name(...)
+        pattern = r'__kernel\s+(?:__global\s+)?(?:void|.*?)\s+(\w+)\s*\('
+        
+        for match in re.finditer(pattern, content):
+            line = content[:match.start()].count('\n')
+            kernel_name = match.group(1)
+            params = self._extract_parameters(content, match.end())
+            
+            kernels.append(OpenCLKernel(
+                name=kernel_name,
+                line=line,
+                parameters=params
+            ))
+        
+        return {"kernels": kernels, "layouts": [], "structs": []}
+    
+    def _extract_parameters(self, content: str, start: int) -> List[str]:
+        # Extract parameter list logic
+        return []
+```
+
+Then register it:
+
+```python
+# In registry.py _register_defaults()
+self.register_language(
+    language_id="opencl",
+    display_name="OpenCL",
+    parser=OpenCLParser(),
+    extensions=[".cl"],
+    file_patterns=["*.cl"]
+)
+```
+
+## Benefits of This Architecture
+
+1. **Modular**: Each language parser is independent
+2. **Extensible**: Easy to add new languages without modifying existing code
+3. **Type-safe**: Common types (`KernelInfo`, `LayoutInfo`) ensure consistency
+4. **Language-agnostic handlers**: Handlers work with any registered language
+5. **Automatic detection**: Language is detected from file extension automatically
+
+## Testing
+
+Add tests for your parser:
+
+```python
+# tests/test_my_language_parser.py
+
+def test_parse_kernel():
+    parser = MyLanguageParser()
+    result = parser.parse_file("kernel my_kernel() { }")
+    assert len(result["kernels"]) == 1
+    assert result["kernels"][0].name == "my_kernel"
+```

wafer_lsp/languages/__init__.py

@@ -0,0 +1,17 @@
+from .converter import ParserResultConverter
+from .detector import LanguageDetector
+from .parser_manager import ParserManager
+from .registry import LanguageRegistry, get_language_registry
+from .types import KernelInfo, LanguageInfo, LayoutInfo, StructInfo
+
+__all__ = [
+    "KernelInfo",
+    "LanguageDetector",
+    "LanguageInfo",
+    "LanguageRegistry",
+    "LayoutInfo",
+    "ParserManager",
+    "ParserResultConverter",
+    "StructInfo",
+    "get_language_registry",
+]

wafer_lsp/languages/converter.py

@@ -0,0 +1,88 @@
+from typing import Any
+
+from .types import KernelInfo, LanguageInfo, LayoutInfo, StructInfo
+
+
+class ParserResultConverter:
+
+    def convert(
+        self,
+        parsed_data: dict[str, Any],
+        language_id: str
+    ) -> LanguageInfo:
+        kernels = self._convert_kernels(
+            parsed_data.get("kernels", []),
+            language_id
+        )
+        layouts = self._convert_layouts(
+            parsed_data.get("layouts", []),
+            language_id
+        )
+        structs = self._convert_structs(
+            parsed_data.get("structs", []),
+            language_id
+        )
+
+        return LanguageInfo(
+            kernels=kernels,
+            layouts=layouts,
+            structs=structs,
+            language=language_id,
+            raw_data=parsed_data
+        )
+
+    def _convert_kernels(
+        self,
+        kernels: list[Any],
+        language_id: str
+    ) -> list[KernelInfo]:
+        result: list[KernelInfo] = []
+
+        for kernel in kernels:
+            if hasattr(kernel, "name") and hasattr(kernel, "line"):
+                result.append(KernelInfo(
+                    name=kernel.name,
+                    line=kernel.line,
+                    parameters=getattr(kernel, "parameters", []),
+                    docstring=getattr(kernel, "docstring", None),
+                    language=language_id
+                ))
+
+        return result
+
+    def _convert_layouts(
+        self,
+        layouts: list[Any],
+        language_id: str
+    ) -> list[LayoutInfo]:
+        result: list[LayoutInfo] = []
+
+        for layout in layouts:
+            if hasattr(layout, "name") and hasattr(layout, "line"):
+                result.append(LayoutInfo(
+                    name=layout.name,
+                    line=layout.line,
+                    shape=getattr(layout, "shape", None),
+                    stride=getattr(layout, "stride", None),
+                    language=language_id
+                ))
+
+        return result
+
+    def _convert_structs(
+        self,
+        structs: list[Any],
+        language_id: str
+    ) -> list[StructInfo]:
+        result: list[StructInfo] = []
+
+        for struct in structs:
+            if hasattr(struct, "name") and hasattr(struct, "line"):
+                result.append(StructInfo(
+                    name=struct.name,
+                    line=struct.line,
+                    docstring=getattr(struct, "docstring", None),
+                    language=language_id
+                ))
+
+        return result

wafer_lsp/languages/detector.py

@@ -0,0 +1,34 @@
+from pathlib import Path
+
+
+class LanguageDetector:
+
+    def __init__(self):
+        self._extensions: dict[str, str] = {}
+
+    def register_extension(self, extension: str, language_id: str):
+        normalized_ext = extension if extension.startswith(".") else f".{ext}"
+        self._extensions[normalized_ext] = language_id
+
+    def detect_from_uri(self, uri: str) -> str | None:
+        if uri.startswith("file://"):
+            file_path = uri[7:]
+        else:
+            file_path = uri
+
+        return self.detect_from_path(file_path)
+
+    def detect_from_path(self, file_path: str) -> str | None:
+        path = Path(file_path)
+        ext = path.suffix.lower()
+        return self._extensions.get(ext)
+
+    def detect_from_extension(self, extension: str) -> str | None:
+        normalized_ext = extension if extension.startswith(".") else f".{ext}"
+        return self._extensions.get(normalized_ext)
+
+    def get_supported_extensions(self) -> list[str]:
+        return list(self._extensions.keys())
+
+    def is_supported(self, uri: str) -> bool:
+        return self.detect_from_uri(uri) is not None

wafer_lsp/languages/parser_manager.py

@@ -0,0 +1,33 @@
+
+from ..parsers.base_parser import BaseParser
+
+
+class ParserManager:
+
+    def __init__(self):
+        self._parsers: dict[str, BaseParser] = {}
+        self._language_names: dict[str, str] = {}
+
+    def register_parser(
+        self,
+        language_id: str,
+        display_name: str,
+        parser: BaseParser
+    ):
+        assert language_id not in self._parsers, \
+            f"Language {language_id} already registered"
+
+        self._parsers[language_id] = parser
+        self._language_names[language_id] = display_name
+
+    def get_parser(self, language_id: str) -> BaseParser | None:
+        return self._parsers.get(language_id)
+
+    def get_language_name(self, language_id: str) -> str | None:
+        return self._language_names.get(language_id)
+
+    def has_parser(self, language_id: str) -> bool:
+        return language_id in self._parsers
+
+    def list_languages(self) -> list[str]:
+        return list(self._parsers.keys())

wafer_lsp/languages/registry.py

@@ -0,0 +1,99 @@
+
+from ..parsers.cuda_parser import CUDAParser
+from ..parsers.cutedsl_parser import CuTeDSLParser
+from .converter import ParserResultConverter
+from .detector import LanguageDetector
+from .parser_manager import ParserManager
+from .types import LanguageInfo
+
+
+class LanguageRegistry:
+
+    def __init__(self):
+        self._detector = LanguageDetector()
+        self._parser_manager = ParserManager()
+        self._converter = ParserResultConverter()
+
+        self._register_defaults()
+
+    def _register_defaults(self):
+        self.register_language(
+            language_id="cutedsl",
+            display_name="CuTeDSL",
+            parser=CuTeDSLParser(),
+            extensions=[".py"],
+            file_patterns=["*.py"]
+        )
+
+        self.register_language(
+            language_id="cuda",
+            display_name="CUDA",
+            parser=CUDAParser(),
+            extensions=[".cu", ".cuh"],
+            file_patterns=["*.cu", "*.cuh"]
+        )
+
+        self.register_language(
+            language_id="cpp",
+            display_name="C++",
+            parser=CUDAParser(),
+            extensions=[".cpp", ".hpp", ".cc", ".cxx"],
+            file_patterns=["*.cpp", "*.hpp", "*.cc", "*.cxx"]
+        )
+
+    def register_language(
+        self,
+        language_id: str,
+        display_name: str,
+        parser,
+        extensions: list[str],
+        file_patterns: list[str] | None = None
+    ):
+        self._parser_manager.register_parser(language_id, display_name, parser)
+
+        for ext in extensions:
+            self._detector.register_extension(ext, language_id)
+
+    def detect_language(self, uri: str) -> str | None:
+        return self._detector.detect_from_uri(uri)
+
+    def get_parser(self, language_id: str):
+        return self._parser_manager.get_parser(language_id)
+
+    def parse_file(self, uri: str, content: str) -> LanguageInfo | None:
+        language_id = self.detect_language(uri)
+        if not language_id:
+            return None
+
+        parser = self.get_parser(language_id)
+        if not parser:
+            return None
+
+        try:
+            parsed_data = parser.parse_file(content)
+        except Exception:
+            return LanguageInfo(
+                kernels=[],
+                layouts=[],
+                structs=[],
+                language=language_id,
+                raw_data={}
+            )
+
+        return self._converter.convert(parsed_data, language_id)
+
+    def get_supported_extensions(self) -> list[str]:
+        return self._detector.get_supported_extensions()
+
+    def get_language_name(self, language_id: str) -> str | None:
+        return self._parser_manager.get_language_name(language_id)
+
+
+_registry: LanguageRegistry | None = None
+
+
+def get_language_registry() -> LanguageRegistry:
+    global _registry
+    if _registry is None:
+        _registry = LanguageRegistry()
+    return _registry

wafer_lsp/languages/types.py

@@ -0,0 +1,37 @@
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class KernelInfo:
+    name: str
+    line: int
+    parameters: list[str]
+    docstring: str | None = None
+    language: str = ""
+
+
+@dataclass
+class LayoutInfo:
+    name: str
+    line: int
+    shape: str | None = None
+    stride: str | None = None
+    language: str = ""
+
+
+@dataclass
+class StructInfo:
+    name: str
+    line: int
+    docstring: str | None = None
+    language: str = ""
+
+
+@dataclass
+class LanguageInfo:
+    kernels: list[KernelInfo]
+    layouts: list[LayoutInfo]
+    structs: list[StructInfo]
+    language: str
+    raw_data: dict[str, Any]

wafer_lsp/parsers/__init__.py

@@ -0,0 +1,18 @@
+from .base_parser import BaseParser
+from .cuda_parser import CUDAKernel, CUDAParser
+from .cutedsl_parser import (
+    CuTeDSLKernel,
+    CuTeDSLLayout,
+    CuTeDSLParser,
+    CuTeDSLStruct,
+)
+
+__all__ = [
+    "BaseParser",
+    "CUDAKernel",
+    "CUDAParser",
+    "CuTeDSLKernel",
+    "CuTeDSLLayout",
+    "CuTeDSLParser",
+    "CuTeDSLStruct",
+]

wafer_lsp/parsers/base_parser.py

@@ -0,0 +1,9 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class BaseParser(ABC):
+
+    @abstractmethod
+    def parse_file(self, content: str) -> dict[str, Any]:
+        pass