@atomic-ehr/fhirpath 0.0.2 → 0.0.3

package/src/completion-provider.ts

@@ -2,10 +2,55 @@ import { parse } from './parser';
 import { Analyzer } from './analyzer';
 import type { TypeInfo, ModelProvider } from './types';
 
-import { CursorContext, isCursorNode } from './cursor-nodes';
-import type { AnyCursorNode } from './cursor-nodes';
+import { CursorContext, isCursorNode } from './parser/cursor-nodes';
+import type { AnyCursorNode } from './parser/cursor-nodes';
 import { registry } from './registry';
 
+/**
+ * Type guard for ErrorNode
+ */
+function isErrorNode(node: any): node is import('./types').ErrorNode {
+  return node.type === 'Error';
+}
+
+/**
+ * Find the cursor node in the AST
+ */
+function findCursorNode(node: import('./types').ASTNode): AnyCursorNode | undefined {
+  if (isCursorNode(node)) {
+    return node as AnyCursorNode;
+  }
+
+  switch (node.type) {
+    case 'Binary':
+      return findCursorNode(node.left) || findCursorNode(node.right);
+    case 'Unary':
+      return findCursorNode(node.operand);
+    case 'Function':
+      // It's more likely for the cursor to be in the arguments, so check them first.
+      for (const arg of node.arguments) {
+        const found = findCursorNode(arg);
+        if (found) return found;
+      }
+      return findCursorNode(node.name);
+    case 'Index':
+      return findCursorNode(node.expression) || findCursorNode(node.index);
+    case 'Collection':
+      for (const element of node.elements) {
+        const found = findCursorNode(element);
+        if (found) return found;
+      }
+      break;
+    case 'TypeCast':
+      return findCursorNode((node as any).expression);
+    case 'MembershipTest':
+      return findCursorNode((node as any).expression);
+  }
+
+  return undefined;
+}
+
+
 /**
  * Kind of completion item
  */
@@ -33,7 +78,7 @@ export interface CompletionItem {
   documentation?: string;
   /** Text to insert (if different from label) */
   insertText?: string;
-  /** Sort order */
+  /** Sort text for ordering (if different from label) */
   sortText?: string;
 }
 
@@ -63,10 +108,19 @@ export async function provideCompletions(
   
   try {
     // Parse with cursor
-    const parseResult = parse(expression, { cursorPosition });
+    const parseResult = parse(expression, {
+      mode: 'lsp',
+      cursorPosition,
+    });
     if (!parseResult.ast) {
       return [];
     }
+
+    // Find the cursor node
+    const cursorNode = findCursorNode(parseResult.ast);
+    if (!cursorNode || isErrorNode(cursorNode)) {
+      return [];
+    }
     
     // Analyze with cursor mode
     const analyzer = new Analyzer(modelProvider);
@@ -78,89 +132,22 @@ export async function provideCompletions(
     );
     
     // Extract cursor context
-    let cursorNode = analysis.cursorContext?.cursorNode;
-    let typeBeforeCursor = analysis.cursorContext?.typeBeforeCursor;
+    const { typeBeforeCursor, functionCall } = analysis.cursorContext || {};
     
-    // Check if cursor is in a binary expression as the right operand
-    // This handles cases where analyzer found the cursor but we need to adjust context
-    if (parseResult.ast) {
-      const ast = parseResult.ast as any;
-      if (ast.type === 'Binary' && ast.right && isCursorNode(ast.right)) {
-        cursorNode = ast.right;
-        // Get type from left side
-        typeBeforeCursor = ast.left?.typeInfo;
-        
-        // Special case for navigation: cursor right after identifier
-        if (ast.left?.type === 'Binary' && ast.left.operator === '.') {
-          if (ast.left.right?.type === 'Identifier') {
-            const identifierEnd = ast.left.right.range?.end?.offset;
-            
-            // If cursor is right after identifier with no space
-            if (identifierEnd === cursorPosition) {
-              // Check if this could be a partial identifier by seeing if there are
-              // any properties that would match if we treated it as a prefix
-              const partialText = extractPartialText(expression, cursorPosition);
-              if (partialText && modelProvider) {
-                // Try to get elements from the type before the identifier
-                const baseType = ast.left.left?.typeInfo;
-                if (baseType) {
-                  const typeName = baseType.name || baseType.type;
-                  // Skip if type is 'Any' as it's not a real FHIR type
-                  if (typeName && typeName !== 'Any') {
-                    const elements = await modelProvider.getElements(typeName);
-                    if (elements.length > 0) {
-                      // Check if any elements start with the partial text
-                      const hasMatches = elements.some(p => 
-                        p.name.toLowerCase().startsWith(partialText.toLowerCase()) &&
-                        p.name.toLowerCase() !== partialText.toLowerCase()
-                      );
-                      
-                      if (hasMatches) {
-                        // There are potential completions - treat as identifier context
-                        cursorNode = {
-                          ...cursorNode,
-                          context: CursorContext.Identifier
-                        } as any;
-                        typeBeforeCursor = baseType;
-                      }
-                    }
-                  }
-                }
-              }
-            }
-          }
-        }
-      }
-    }
-    
-    // Check if cursor is in a function argument and fix the function name
-    if (cursorNode && cursorNode.context === CursorContext.Argument) {
-      const functionName = findFunctionName(parseResult.ast as any, cursorNode);
-      if (functionName) {
-        (cursorNode as any).functionName = functionName;
-        
-        // Special case: ofType, as, is should treat their arguments as type context
-        if (functionName === 'ofType' || functionName === 'as' || functionName === 'is') {
-          // Convert to type cursor node
-          cursorNode = {
-            ...cursorNode,
-            context: CursorContext.Type,
-            typeOperator: functionName
-          } as any;
-        }
+    // Get partial text for filtering
+    // Determine partial text for filtering (prefer AST-provided, else infer from source)
+    let partialText = (cursorNode as any).partialText || '';
+    if (!partialText && cursorNode.context === CursorContext.Type) {
+      // Heuristic: scan left from cursor to preceding '(' or whitespace and take the identifier fragment
+      const cp = cursorPosition;
+      const left = expression.slice(0, cp);
+      // Take the trailing run of identifier characters
+      const m = left.match(/[A-Za-z][A-Za-z0-9.]*$/);
+      if (m) {
+        partialText = m[0];
       }
     }
     
-    if (!cursorNode) {
-      // Fallback: provide general completions if no cursor node found
-      return getGeneralCompletions();
-    }
-    
-    const expectedType = analysis.cursorContext?.expectedType;
-    
-    // Get partial text for filtering
-    const partialText = extractPartialText(expression, cursorPosition);
-    
     // Generate completions based on cursor context
     let completions: CompletionItem[] = [];
     
@@ -170,7 +157,13 @@ export async function provideCompletions(
         break;
       
       case CursorContext.Operator:
-        completions = getOperatorCompletions(typeBeforeCursor);
+        // Fallback to provided inputType if analyzer did not provide typeBeforeCursor
+        completions = getOperatorCompletions(typeBeforeCursor || (inputType as any));
+        // Post-filter: when left side is a collection, exclude operators that require singleton left (e.g., 'in')
+        const leftType = (typeBeforeCursor || (inputType as any));
+        if (leftType && leftType.singleton === false) {
+          completions = completions.filter(c => c.label !== 'in');
+        }
         break;
       
       case CursorContext.Type:
@@ -178,7 +171,7 @@ export async function provideCompletions(
         break;
       
       case CursorContext.Argument:
-        completions = await getArgumentCompletions(cursorNode, typeBeforeCursor, modelProvider, variables);
+        completions = await getArgumentCompletions(cursorNode, typeBeforeCursor, modelProvider, variables, functionCall);
         break;
       
       case CursorContext.Index:
@@ -245,7 +238,16 @@ function isFunctionApplicable(funcDef: any, typeInfo: TypeInfo): boolean {
  */
 function isOperatorApplicable(opDef: any, typeInfo: TypeInfo): boolean {
   if (!typeInfo || !typeInfo.type) return true;
-  return registry.isOperatorApplicableToType(opDef.symbol, typeInfo.type);
+  const isCollection = typeInfo.singleton === false;
+  // If left operand must be singleton for all signatures, do not suggest for collections
+  if (isCollection && opDef.signatures && opDef.signatures.length > 0) {
+    const allowsCollection = opDef.signatures.some((s: any) => s.left && s.left.singleton === false);
+    if (!allowsCollection) {
+      return false;
+    }
+  }
+  const typeForRegistry = isCollection ? `${typeInfo.type}[]` : typeInfo.type;
+  return registry.isOperatorApplicableToType(opDef.symbol, typeForRegistry);
 }
 
 /**
@@ -261,69 +263,6 @@ function getSortTextForOperator(opDef: any): string {
   return `1_${opDef.symbol}`;
 }
 
-/**
- * Find function name for a cursor node in an argument position
- */
-function findFunctionName(ast: any, cursorNode: any): string | null {
-  // Check if AST is a function with cursor in arguments
-  if (ast.type === 'Function' && ast.arguments?.includes(cursorNode)) {
-    return ast.name?.name || null;
-  }
-  
-  // Check if AST is binary with function on right
-  if (ast.type === 'Binary' && ast.right?.type === 'Function') {
-    if (ast.right.arguments?.includes(cursorNode)) {
-      return ast.right.name?.name || null;
-    }
-  }
-  
-  // Recursively check children
-  if (ast.left) {
-    const leftResult = findFunctionName(ast.left, cursorNode);
-    if (leftResult) return leftResult;
-  }
-  
-  if (ast.right) {
-    const rightResult = findFunctionName(ast.right, cursorNode);
-    if (rightResult) return rightResult;
-  }
-  
-  if (ast.arguments) {
-    for (const arg of ast.arguments) {
-      if (arg !== cursorNode) {
-        const argResult = findFunctionName(arg, cursorNode);
-        if (argResult) return argResult;
-      }
-    }
-  }
-  
-  return null;
-}
-
-/**
- * Extract partial text before cursor for filtering
- */
-function extractPartialText(expression: string, cursorPosition: number): string {
-  // Handle case where cursor is right after a dot
-  if (cursorPosition > 0 && expression[cursorPosition - 1] === '.') {
-    return '';
-  }
-  
-  let start = cursorPosition;
-  
-  // Move back to find start of identifier
-  while (start > 0) {
-    const char = expression[start - 1];
-    if (char && /[a-zA-Z0-9_$%]/.test(char)) {
-      start--;
-    } else {
-      break;
-    }
-  }
-  
-  return expression.substring(start, cursorPosition);
-}
-
 /**
  * Get completions for identifier context (after dot)
  */
@@ -362,15 +301,15 @@ async function getIdentifierCompletions(
       const isApplicable = !typeBeforeCursor || isFunctionApplicable(funcDef, typeBeforeCursor);
       
       if (isApplicable) {
-        // Determine if function takes parameters
-        const hasParams = funcDef.signatures?.[0]?.parameters && funcDef.signatures[0].parameters.length > 0;
+        // Determine if any signature takes parameters
+        const hasParams = funcDef.signatures?.some(sig => (sig.parameters?.length ?? 0) > 0) ?? false;
         const funcDescription = funcDef.description || `FHIRPath ${name} function`;
         
         completions.push({
           label: name,
           kind: CompletionKind.Function,
           detail: funcDescription,
-          insertText: name + (hasParams ? '()' : '()')
+          insertText: name + (hasParams ? '()' : '')
         });
       }
     }
@@ -386,12 +325,12 @@ async function getIdentifierCompletions(
     for (const func of typeFunctions) {
       // Check if function is already added from general functions
       if (!completions.some(c => c.label === func.name)) {
-        const hasParams = func.signatures?.[0]?.parameters && func.signatures[0].parameters.length > 0;
+        const hasParams = func.signatures?.some(sig => (sig.parameters?.length ?? 0) > 0) ?? false;
         completions.push({
           label: func.name,
           kind: CompletionKind.Function,
           detail: func.description || `FHIRPath ${func.name} function`,
-          insertText: func.name + (hasParams ? '()' : '()')
+          insertText: func.name + (hasParams ? '()' : '')
         });
       }
     }
@@ -413,6 +352,13 @@ function getOperatorCompletions(typeBeforeCursor?: TypeInfo): CompletionItem[] {
   for (const opName of operatorNames) {
     const opDef = registry.getOperatorDefinition(opName);
     if (opDef) {
+      // If we clearly have a collection to the left and no signature allows a collection left operand, skip early
+      if (typeBeforeCursor && typeBeforeCursor.singleton === false && opDef.signatures && opDef.signatures.length > 0) {
+        const allowsCollection = opDef.signatures.some((s: any) => s.left && s.left.singleton === false);
+        if (!allowsCollection) {
+          continue;
+        }
+      }
       // Check if operator is applicable to the current type
       const isApplicable = !typeBeforeCursor || isOperatorApplicable(opDef, typeBeforeCursor);
       
@@ -447,7 +393,8 @@ async function getTypeCompletions(
       completions.push({
         label: type,
         kind: CompletionKind.Type,
-        detail: 'FHIRPath primitive type'
+        detail: 'FHIRPath primitive type',
+        sortText: `0_${type}` // Priority 0 for primitives
       });
     }
     
@@ -457,20 +404,22 @@ async function getTypeCompletions(
       completions.push({
         label: type,
         kind: CompletionKind.Type,
-        detail: 'FHIR complex type'
+        detail: 'FHIR complex type',
+        sortText: `1_${type}` // Priority 1 for complex types
       });
     }
   }
   
-  // For ofType, add resource types
-  const typeOperator = (cursorNode as any).typeOperator;
-  if (typeOperator === 'ofType' && modelProvider) {
+  // Always add resource types when we have a model provider
+  // They're valid type names in contexts like ofType()
+  if (modelProvider) {
     const resourceTypes = await modelProvider.getResourceTypes();
     for (const type of resourceTypes) {
       completions.push({
         label: type,
         kind: CompletionKind.Type,
-        detail: 'FHIR resource type'
+        detail: 'FHIR resource type',
+        sortText: `2_${type}` // Priority 2 for resource types
       });
     }
   }
@@ -485,12 +434,91 @@ async function getArgumentCompletions(
   cursorNode: AnyCursorNode,
   typeBeforeCursor?: TypeInfo,
   modelProvider?: ModelProvider,
-  variables?: Record<string, any>
+  variables?: Record<string, any>,
+  functionCall?: {
+    definition: import('./types').FunctionDefinition;
+    argumentIndex: number;
+  }
 ): Promise<CompletionItem[]> {
+  // Inspect all overloads to determine parameter kinds for this argument position
+  if (functionCall?.definition && functionCall.argumentIndex >= 0) {
+    const argIndex = functionCall.argumentIndex;
+    let expectsTypeReference = false;
+    let expectsLambda = false;
+
+    for (const sig of functionCall.definition.signatures) {
+      const param = sig.parameters[argIndex];
+      if (!param) {
+        continue;
+      }
+      if (param.typeReference) {
+        expectsTypeReference = true;
+      }
+      if (param.expression) {
+        expectsLambda = true;
+      }
+    }
+
+    if (expectsTypeReference) {
+      // This parameter expects a type name, provide type completions
+      return getTypeCompletions(cursorNode, modelProvider);
+    }
+
+    // If lambda is expected, suggest element properties from the item type below
+    // (continue to build completions; variables will also be added)
+    const completions: CompletionItem[] = [];
+
+    // Add user variables if available
+    if (variables) {
+      for (const varName of Object.keys(variables)) {
+        if (varName === '$this') {
+          completions.push({
+            label: '$this',
+            kind: CompletionKind.Variable,
+            detail: 'Current item in iteration'
+          });
+        } else if (varName === '$index') {
+          completions.push({
+            label: '$index',
+            kind: CompletionKind.Variable,
+            detail: 'Current index in iteration'
+          });
+        } else {
+          completions.push({
+            label: varName.startsWith('%') ? varName : '%' + varName,
+            kind: CompletionKind.Variable,
+            detail: 'User-defined variable'
+          });
+        }
+      }
+    }
+
+    if (expectsLambda && typeBeforeCursor && modelProvider) {
+      const itemType = { ...typeBeforeCursor, singleton: true };
+      const typeName = itemType.name || itemType.type;
+      if (typeName && typeName !== 'Any') {
+        const elements = await modelProvider.getElements(typeName);
+        if (elements.length > 0) {
+          for (const element of elements) {
+            completions.push({
+              label: element.name,
+              kind: CompletionKind.Property,
+              detail: element.type
+            });
+          }
+        }
+      }
+    }
+
+    // If overloads do not indicate typeRef or lambda, fall through to variables-only completions below
+    if (completions.length > 0) {
+      return completions;
+    }
+  }
   const completions: CompletionItem[] = [];
   const argNode = cursorNode as any;
-  
-  // Add user variables if available
+
+  // Default: suggest variables in scope
   if (variables) {
     for (const varName of Object.keys(variables)) {
       if (varName === '$this') {
@@ -514,34 +542,8 @@ async function getArgumentCompletions(
       }
     }
   }
-  
-  // Add elements if in lambda context
-  const functionName = argNode.functionName;
-  // Check if the function accepts lambda expressions (could be determined from function signature in registry)
-  const lambdaFunctions = ['where', 'select', 'all', 'exists', 'any', 'repeat'];
-  if (functionName && lambdaFunctions.includes(functionName)) {
-    // In lambda function, provide elements of item type
-    if (typeBeforeCursor && modelProvider) {
-      const itemType = { ...typeBeforeCursor, singleton: true };
-      const typeName = itemType.name || itemType.type;
-      // Skip if type is 'Any' as it's not a real FHIR type
-      if (typeName && typeName !== 'Any') {
-        const elements = await modelProvider.getElements(typeName);
-        if (elements.length > 0) {
-          for (const element of elements) {
-            completions.push({
-              label: element.name,
-              kind: CompletionKind.Property,
-              detail: element.type
-            });
-          }
-        }
-      }
-    }
-  }
-  
+
   // No hardcoded constants - these should come from context or be typed by user
-  
   return completions;
 }
 
@@ -556,6 +558,15 @@ function getIndexCompletions(
   
   // Add user variables if available
   if (variables) {
+    // Add $this if it's in scope (for consistency with argument context)
+    if ('$this' in variables) {
+      completions.push({
+        label: '$this',
+        kind: CompletionKind.Variable,
+        detail: 'Current item'
+      });
+    }
+
     // Add $index if it's in scope (should be provided by context)
     if ('$index' in variables) {
       completions.push({
@@ -612,6 +623,13 @@ function filterCompletions(completions: CompletionItem[], partialText: string):
  */
 function rankCompletions(completions: CompletionItem[]): CompletionItem[] {
   return completions.sort((a, b) => {
+    // First check if items have sortText
+    if (a.sortText || b.sortText) {
+      const aSortText = a.sortText || a.label;
+      const bSortText = b.sortText || b.label;
+      return aSortText.localeCompare(bSortText);
+    }
+    
     // Sort by kind priority
     const kindPriority: Record<CompletionKind, number> = {
       [CompletionKind.Property]: 1,
@@ -633,4 +651,4 @@ function rankCompletions(completions: CompletionItem[]): CompletionItem[] {
     // Then alphabetically
     return a.label.localeCompare(b.label);
   });
-}
+}