@atomic-ehr/fhirpath 0.0.1 → 0.0.2-canary.261bb9b.20250811212042

package/src/analyzer.ts

@@ -20,34 +20,53 @@ import type {
 } from './types';
 import { NodeType, DiagnosticSeverity } from './types';
 import { registry } from './registry';
+import { Errors, toDiagnostic } from './errors';
+import { isCursorNode, CursorContext } from './cursor-nodes';
+import type { AnyCursorNode } from './cursor-nodes';
 
-// Diagnostic codes as constants to avoid magic strings
-const DiagnosticCode = {
-  UnknownOperator: 'UNKNOWN_OPERATOR',
-  UnknownVariable: 'UNKNOWN_VARIABLE',
-  UnknownUserVariable: 'UNKNOWN_USER_VARIABLE',
-  UnknownFunction: 'UNKNOWN_FUNCTION',
-  TooFewArgs: 'TOO_FEW_ARGS',
-  TooManyArgs: 'TOO_MANY_ARGS',
-  TypeMismatch: 'TYPE_MISMATCH',
-  InputTypeMismatch: 'INPUT_TYPE_MISMATCH',
-  ArgumentTypeMismatch: 'ARGUMENT_TYPE_MISMATCH',
-  ParseError: 'PARSE_ERROR',
-} as const;
+
+export interface AnalyzerOptions {
+  cursorMode?: boolean;
+}
+
+export interface AnalysisResultWithCursor extends AnalysisResult {
+  stoppedAtCursor?: boolean;
+  cursorContext?: {
+    typeBeforeCursor?: TypeInfo;
+    expectedType?: TypeInfo;
+    cursorNode?: AnyCursorNode;
+  };
+}
 
 export class Analyzer {
   private diagnostics: Diagnostic[] = [];
-  private variables: Set<string> = new Set(['$this', '$index', '$total']);
+  private variables: Set<string> = new Set(['$this', '$index', '$total', 'context', 'resource', 'rootResource']);
   private modelProvider?: ModelProvider;
   private userVariableTypes: Map<string, TypeInfo> = new Map();
+  private systemVariableTypes: Map<string, TypeInfo> = new Map();
+  private cursorMode: boolean = false;
+  private stoppedAtCursor: boolean = false;
+  private cursorContext?: {
+    typeBeforeCursor?: TypeInfo;
+    expectedType?: TypeInfo;
+    cursorNode?: AnyCursorNode;
+  };
 
   constructor(modelProvider?: ModelProvider) {
     this.modelProvider = modelProvider;
   }
 
-  analyze(ast: ASTNode, userVariables?: Record<string, any>, inputType?: TypeInfo): AnalysisResult {
+  async analyze(
+    ast: ASTNode, 
+    userVariables?: Record<string, any>, 
+    inputType?: TypeInfo,
+    options?: AnalyzerOptions
+  ): Promise<AnalysisResultWithCursor> {
     this.diagnostics = [];
     this.userVariableTypes.clear();
+    this.cursorMode = options?.cursorMode ?? false;
+    this.stoppedAtCursor = false;
+    this.cursorContext = undefined;
     
     if (userVariables) {
       Object.keys(userVariables).forEach(name => {
@@ -61,24 +80,43 @@ export class Analyzer {
     }
     
     // Annotate AST with type information
-    this.annotateAST(ast, inputType);
+    await this.annotateAST(ast, inputType);
     
-    // Perform validation with type checking
-    this.visitNode(ast);
+    // Perform validation with type checking (if not stopped at cursor)
+    if (!this.stoppedAtCursor) {
+      this.visitNode(ast);
+    }
     
     return {
       diagnostics: this.diagnostics,
-      ast
+      ast,
+      stoppedAtCursor: this.cursorMode ? this.stoppedAtCursor : undefined,
+      cursorContext: this.cursorMode ? this.cursorContext : undefined
     };
   }
 
   private visitNode(node: ASTNode): void {
+    // Check for cursor node in cursor mode
+    if (this.cursorMode && isCursorNode(node)) {
+      this.stoppedAtCursor = true;
+      this.cursorContext = {
+        cursorNode: node as AnyCursorNode,
+        typeBeforeCursor: (node as any).typeInfo
+      };
+      return; // Short-circuit
+    }
+    
     // Handle error nodes - process them for diagnostics but don't traverse
     if (node.type === 'Error') {
       // Diagnostics already added in annotateAST
       return;
     }
     
+    // If we've already stopped at cursor, don't continue
+    if (this.stoppedAtCursor) {
+      return;
+    }
+    
     switch (node.type) {
       case NodeType.Binary:
         this.visitBinaryOperator(node as BinaryNode);
@@ -101,10 +139,10 @@ export class Analyzer {
         this.visitNode((node as UnaryNode).operand);
         break;
       case NodeType.MembershipTest:
-        this.visitNode((node as MembershipTestNode).expression);
+        this.visitMembershipTest(node as MembershipTestNode);
         break;
       case NodeType.TypeCast:
-        this.visitNode((node as TypeCastNode).expression);
+        this.visitTypeCast(node as TypeCastNode);
         break;
       case NodeType.Variable:
         this.validateVariable((node as VariableNode).name, node);
@@ -120,22 +158,84 @@ export class Analyzer {
   private visitBinaryOperator(node: BinaryNode): void {
     this.visitNode(node.left);
     
+    // Track defineVariable for validation - collect all variables defined in the chain
+    if (node.operator === '.') {
+      const definedVars = this.collectDefinedVariables(node.left);
+      if (definedVars.size > 0) {
+        // Track which variables were already known
+        const previouslyKnown = new Set<string>();
+        definedVars.forEach(varName => {
+          if (this.variables.has(varName)) {
+            previouslyKnown.add(varName);
+          }
+          this.variables.add(varName);
+        });
+        
+        // Visit right side with new variables in scope
+        this.visitNode(node.right);
+        
+        // Restore previous state
+        definedVars.forEach(varName => {
+          if (!previouslyKnown.has(varName)) {
+            this.variables.delete(varName);
+          }
+        });
+        return;
+      }
+    }
+    
     // Special handling for dot operator with function on right side
     if (node.operator === '.' && node.right.type === NodeType.Function) {
       const funcNode = node.right as FunctionNode;
       if (funcNode.name.type === NodeType.Identifier) {
         const funcName = (funcNode.name as IdentifierNode).name;
         const func = registry.getFunction(funcName);
-        if (func && func.signature.input && node.left.typeInfo) {
-          if (!this.isTypeCompatible(node.left.typeInfo, func.signature.input)) {
+        if (func && func.signatures && func.signatures.length > 0 && node.left.typeInfo) {
+          // Check if any signature matches the input type
+          let matchFound = false;
+          let expectedTypes: string[] = [];
+          
+          for (const signature of func.signatures) {
+            if (signature.input) {
+              if (this.isTypeCompatible(node.left.typeInfo, signature.input)) {
+                matchFound = true;
+                break;
+              }
+              expectedTypes.push(this.typeToString(signature.input));
+            } else {
+              // If any signature has no input constraint, it matches
+              matchFound = true;
+              break;
+            }
+          }
+          
+          if (!matchFound) {
             const inputTypeStr = this.typeToString(node.left.typeInfo);
-            const expectedTypeStr = this.typeToString(func.signature.input);
-            this.addDiagnostic(
-              DiagnosticSeverity.Error,
-              `Type mismatch: function '${func.name}' expects input type ${expectedTypeStr} but got ${inputTypeStr}`,
-              funcNode,
-              DiagnosticCode.InputTypeMismatch
+            const firstSignature = func.signatures[0];
+            
+            if (!firstSignature) return;
+            
+            // Check if this is specifically a singleton/collection mismatch
+            const inputIsCollection = !node.left.typeInfo.singleton;
+            const expectedIsSingleton = firstSignature.input?.singleton;
+            
+            // Check if the base types are compatible (same type or subtype)
+            const typesCompatible = firstSignature.input && (
+              node.left.typeInfo.type === firstSignature.input.type ||
+              this.isSubtypeOf(node.left.typeInfo.type, firstSignature.input.type)
             );
+            
+            if (inputIsCollection && expectedIsSingleton && typesCompatible) {
+              // Compatible base types but collection vs singleton mismatch
+              this.diagnostics.push(
+                toDiagnostic(Errors.singletonTypeRequired(funcName, inputTypeStr, funcNode.range))
+              );
+            } else {
+              // Function received invalid operand type - report as runtime error
+              this.diagnostics.push(
+                toDiagnostic(Errors.invalidOperandType(funcName + '()', inputTypeStr, funcNode.range))
+              );
+            }
           }
         }
       }
@@ -150,7 +250,9 @@ export class Analyzer {
     
     const op = registry.getOperatorDefinition(node.operator);
     if (!op) {
-      this.addDiagnostic(DiagnosticSeverity.Error, `Unknown operator: ${node.operator}`, node, DiagnosticCode.UnknownOperator);
+      this.diagnostics.push(
+        toDiagnostic(Errors.unknownOperator(node.operator, node.range))
+      );
       return;
     }
     
@@ -167,20 +269,87 @@ export class Analyzer {
   private visitFunctionCall(node: FunctionNode): void {
     if (node.name.type === NodeType.Identifier) {
       const funcName = (node.name as IdentifierNode).name;
+      
+      // Check if this is a type operation that requires ModelProvider
+      if (funcName === 'ofType' && !this.modelProvider) {
+        // Check if the type argument is a primitive type
+        const primitiveTypes = ['String', 'Integer', 'Decimal', 'Boolean', 'Date', 'DateTime', 'Time', 'Quantity'];
+        let isPrimitive = false;
+        
+        if (node.arguments.length > 0) {
+          const typeArg = node.arguments[0]!;
+          if (typeArg.type === NodeType.Identifier) {
+            isPrimitive = primitiveTypes.includes((typeArg as IdentifierNode).name);
+          } else if ((typeArg as any).type === NodeType.TypeOrIdentifier || (typeArg as any).type === NodeType.TypeReference) {
+            isPrimitive = primitiveTypes.includes((typeArg as any).name);
+          }
+        }
+        
+        if (!isPrimitive) {
+          this.diagnostics.push(
+            toDiagnostic(Errors.modelProviderRequired('ofType', node.range))
+          );
+        }
+      }
+      
+      // Check ofType with union types
+      if (funcName === 'ofType' && node.typeInfo) {
+        const inputType = node.typeInfo;
+        if (node.arguments.length > 0 && inputType.modelContext && 
+            typeof inputType.modelContext === 'object' &&
+            'isUnion' in inputType.modelContext && 
+            inputType.modelContext.isUnion &&
+            'choices' in inputType.modelContext &&
+            Array.isArray(inputType.modelContext.choices)) {
+          
+          // Extract target type from argument
+          let targetType: string | undefined;
+          const typeArg = node.arguments[0]!;
+          if (typeArg.type === NodeType.Identifier) {
+            targetType = (typeArg as IdentifierNode).name;
+          } else if ((typeArg as any).type === NodeType.TypeOrIdentifier || (typeArg as any).type === NodeType.TypeReference) {
+            targetType = (typeArg as any).name;
+          }
+          
+          if (targetType) {
+            const validChoice = inputType.modelContext.choices.find((choice: any) => 
+              choice.type === targetType || choice.code === targetType
+            );
+            
+            if (!validChoice) {
+              this.diagnostics.push({
+                severity: DiagnosticSeverity.Warning,
+                code: 'invalid-type-filter',
+                message: `Type '${targetType}' is not present in the union type. Available types: ${
+                  inputType.modelContext.choices.map((c: any) => c.type || c.code).join(', ')
+                }`,
+                range: node.range
+              });
+            }
+          }
+        }
+      }
+      
       const func = registry.getFunction(funcName);
       
       if (!func) {
-        this.addDiagnostic(DiagnosticSeverity.Error, `Unknown function: ${funcName}`, node, DiagnosticCode.UnknownFunction);
+        this.diagnostics.push(
+          toDiagnostic(Errors.unknownFunction(funcName, node.range))
+        );
       } else {
         // Check argument count based on signature
-        const params = func.signature.parameters;
+        const params = func.signatures?.[0]?.parameters || [];
         const requiredParams = params.filter(p => !p.optional).length;
         const maxParams = params.length;
         
         if (node.arguments.length < requiredParams) {
-          this.addDiagnostic(DiagnosticSeverity.Error, `Function '${funcName}' requires at least ${requiredParams} arguments, got ${node.arguments.length}`, node, DiagnosticCode.TooFewArgs);
+          this.diagnostics.push(
+            toDiagnostic(Errors.wrongArgumentCount(funcName, requiredParams, node.arguments.length, node.range))
+          );
         } else if (node.arguments.length > maxParams) {
-          this.addDiagnostic(DiagnosticSeverity.Error, `Function '${funcName}' accepts at most ${maxParams} arguments, got ${node.arguments.length}`, node, DiagnosticCode.TooManyArgs);
+          this.diagnostics.push(
+            toDiagnostic(Errors.wrongArgumentCount(funcName, maxParams, node.arguments.length, node.range))
+          );
         }
         
         // Type check arguments if we have type information
@@ -193,32 +362,217 @@ export class Analyzer {
     node.arguments.forEach(arg => this.visitNode(arg));
   }
 
+  private visitMembershipTest(node: MembershipTestNode): void {
+    // Check if ModelProvider is required
+    // Basic primitive types can be checked without ModelProvider
+    const primitiveTypes = ['String', 'Integer', 'Decimal', 'Boolean', 'Date', 'DateTime', 'Time', 'Quantity'];
+    if (!this.modelProvider && !primitiveTypes.includes(node.targetType)) {
+      this.diagnostics.push(
+        toDiagnostic(Errors.modelProviderRequired('is', node.range))
+      );
+    }
+    
+    // Check 'is' with union types
+    if (node.expression.typeInfo) {
+      const leftType = node.expression.typeInfo;
+      if (leftType.modelContext && 
+          typeof leftType.modelContext === 'object' &&
+          'isUnion' in leftType.modelContext && 
+          leftType.modelContext.isUnion &&
+          'choices' in leftType.modelContext &&
+          Array.isArray(leftType.modelContext.choices)) {
+        
+        const targetTypeName = node.targetType;
+        const validChoice = leftType.modelContext.choices.find((choice: any) =>
+          choice.type === targetTypeName || choice.code === targetTypeName
+        );
+        
+        if (!validChoice) {
+          this.diagnostics.push({
+            severity: DiagnosticSeverity.Warning,
+            code: 'invalid-type-test',
+            message: `Type test 'is ${targetTypeName}' will always be false. Type '${targetTypeName}' is not in the union. Available types: ${
+              leftType.modelContext.choices.map((c: any) => c.type || c.code).join(', ')
+            }`,
+            range: node.range
+          });
+        }
+      }
+    }
+    
+    this.visitNode(node.expression);
+  }
+
+  private visitTypeCast(node: TypeCastNode): void {
+    // Check if ModelProvider is required
+    // Basic primitive types can be checked without ModelProvider
+    const primitiveTypes = ['String', 'Integer', 'Decimal', 'Boolean', 'Date', 'DateTime', 'Time', 'Quantity'];
+    if (!this.modelProvider && !primitiveTypes.includes(node.targetType)) {
+      this.diagnostics.push(
+        toDiagnostic(Errors.modelProviderRequired('as', node.range))
+      );
+    }
+    
+    // Check 'as' with union types
+    if (node.expression.typeInfo) {
+      const leftType = node.expression.typeInfo;
+      if (leftType.modelContext && 
+          typeof leftType.modelContext === 'object' &&
+          'isUnion' in leftType.modelContext && 
+          leftType.modelContext.isUnion &&
+          'choices' in leftType.modelContext &&
+          Array.isArray(leftType.modelContext.choices)) {
+        
+        const targetTypeName = node.targetType;
+        const validChoice = leftType.modelContext.choices.find((choice: any) =>
+          choice.type === targetTypeName || choice.code === targetTypeName
+        );
+        
+        if (!validChoice) {
+          this.diagnostics.push({
+            severity: DiagnosticSeverity.Warning,
+            code: 'invalid-type-cast',
+            message: `Type cast 'as ${targetTypeName}' may fail. Type '${targetTypeName}' is not guaranteed in the union. Available types: ${
+              leftType.modelContext.choices.map((c: any) => c.type || c.code).join(', ')
+            }`,
+            range: node.range
+          });
+        }
+      }
+    }
+    
+    this.visitNode(node.expression);
+  }
+
   // Unified variable validation to eliminate duplication
   private validateVariable(name: string, node: ASTNode): void {
     if (name.startsWith('$')) {
       if (!this.variables.has(name)) {
-        this.addDiagnostic(DiagnosticSeverity.Error, `Unknown variable: ${name}`, node, DiagnosticCode.UnknownVariable);
+        this.diagnostics.push(
+          toDiagnostic(Errors.unknownVariable(name, node.range))
+        );
       }
     } else if (name.startsWith('%')) {
       const varName = name.substring(1);
       if (!this.variables.has(varName)) {
-        this.addDiagnostic(DiagnosticSeverity.Error, `Unknown user variable: ${name}`, node, DiagnosticCode.UnknownUserVariable);
+        this.diagnostics.push(
+          toDiagnostic(Errors.unknownUserVariable(name, node.range))
+        );
       }
     }
   }
-
-  private addDiagnostic(severity: DiagnosticSeverity, message: string, node: ASTNode, code: string): void {
-    this.diagnostics.push({
-      range: node.range,
-      severity,
-      message,
-      code,
-      source: 'fhirpath-analyzer'
-    });
+  
+  private collectDefinedVariables(node: ASTNode): Set<string> {
+    const vars = new Set<string>();
+    
+    // If this is a defineVariable call, extract the variable name
+    if (node.type === NodeType.Function) {
+      const funcNode = node as FunctionNode;
+      if (funcNode.name.type === NodeType.Identifier && 
+          (funcNode.name as IdentifierNode).name === 'defineVariable' &&
+          funcNode.arguments.length >= 1) {
+        const nameArg = funcNode.arguments[0];
+        if (nameArg && nameArg.type === NodeType.Literal && nameArg.valueType === 'string') {
+          vars.add(nameArg.value as string);
+        }
+      }
+    }
+    
+    // If this is a binary dot operator, collect from left side recursively
+    if (node.type === NodeType.Binary) {
+      const binaryNode = node as BinaryNode;
+      if (binaryNode.operator === '.') {
+        // Collect from left side
+        const leftVars = this.collectDefinedVariables(binaryNode.left);
+        leftVars.forEach(v => vars.add(v));
+        
+        // Check if right side is also defineVariable
+        if (binaryNode.right.type === NodeType.Function) {
+          const rightFunc = binaryNode.right as FunctionNode;
+          if (rightFunc.name.type === NodeType.Identifier && 
+              (rightFunc.name as IdentifierNode).name === 'defineVariable' &&
+              rightFunc.arguments.length >= 1) {
+            const nameArg = rightFunc.arguments[0];
+            if (nameArg && nameArg.type === NodeType.Literal && nameArg.valueType === 'string') {
+              vars.add(nameArg.value as string);
+            }
+          }
+        }
+      }
+    }
+    
+    return vars;
+  }
+  
+  private collectDefinedVariablesWithTypes(node: ASTNode): Map<string, TypeInfo> {
+    const varsWithTypes = new Map<string, TypeInfo>();
+    
+    // If this is a defineVariable call, extract the variable name and type
+    if (node.type === NodeType.Function) {
+      const funcNode = node as FunctionNode;
+      if (funcNode.name.type === NodeType.Identifier && 
+          (funcNode.name as IdentifierNode).name === 'defineVariable' &&
+          funcNode.arguments.length >= 1) {
+        const nameArg = funcNode.arguments[0];
+        if (nameArg && nameArg.type === NodeType.Literal && nameArg.valueType === 'string') {
+          const varName = nameArg.value as string;
+          let varType: TypeInfo;
+          
+          if (funcNode.arguments.length >= 2 && funcNode.arguments[1]!.typeInfo) {
+            // Has value expression - use its type
+            varType = funcNode.arguments[1]!.typeInfo;
+          } else if (node.typeInfo) {
+            // No value expression - uses input as value (defineVariable returns input)
+            varType = node.typeInfo;
+          } else {
+            varType = { type: 'Any', singleton: false };
+          }
+          
+          varsWithTypes.set(varName, varType);
+        }
+      }
+    }
+    
+    // If this is a binary dot operator, collect from entire chain
+    if (node.type === NodeType.Binary) {
+      const binaryNode = node as BinaryNode;
+      if (binaryNode.operator === '.') {
+        // Collect from left side recursively
+        const leftVars = this.collectDefinedVariablesWithTypes(binaryNode.left);
+        leftVars.forEach((type, name) => varsWithTypes.set(name, type));
+        
+        // Check if right side is also defineVariable
+        if (binaryNode.right.type === NodeType.Function) {
+          const rightFunc = binaryNode.right as FunctionNode;
+          if (rightFunc.name.type === NodeType.Identifier && 
+              (rightFunc.name as IdentifierNode).name === 'defineVariable' &&
+              rightFunc.arguments.length >= 1) {
+            const nameArg = rightFunc.arguments[0];
+            if (nameArg && nameArg.type === NodeType.Literal && nameArg.valueType === 'string') {
+              const varName = nameArg.value as string;
+              let varType: TypeInfo;
+              
+              if (rightFunc.arguments.length >= 2 && rightFunc.arguments[1]!.typeInfo) {
+                varType = rightFunc.arguments[1]!.typeInfo;
+              } else if (binaryNode.typeInfo) {
+                varType = binaryNode.typeInfo;
+              } else {
+                varType = { type: 'Any', singleton: false };
+              }
+              
+              varsWithTypes.set(varName, varType);
+            }
+          }
+        }
+      }
+    }
+    
+    return varsWithTypes;
   }
 
+
   // Type inference methods
-  private inferType(node: ASTNode, inputType?: TypeInfo): TypeInfo {
+  private async inferType(node: ASTNode, inputType?: TypeInfo): Promise<TypeInfo> {
     // Handle error nodes
     if (node.type === 'Error') {
       return this.inferErrorNodeType(node as ErrorNode, inputType);
@@ -229,31 +583,31 @@ export class Analyzer {
         return this.inferLiteralType(node as LiteralNode);
         
       case NodeType.Binary:
-        return this.inferBinaryType(node as BinaryNode, inputType);
+        return await this.inferBinaryType(node as BinaryNode, inputType);
         
       case NodeType.Unary:
         return this.inferUnaryType(node as UnaryNode);
         
       case NodeType.Function:
-        return this.inferFunctionType(node as FunctionNode, inputType);
+        return await this.inferFunctionType(node as FunctionNode, inputType);
         
       case NodeType.Identifier:
-        return this.inferIdentifierType(node as IdentifierNode, inputType);
+        return await this.inferIdentifierType(node as IdentifierNode, inputType);
         
       case NodeType.Variable:
         return this.inferVariableType(node as VariableNode);
         
       case NodeType.Collection:
-        return this.inferCollectionType(node as CollectionNode);
+        return await this.inferCollectionType(node as CollectionNode);
         
       case NodeType.TypeCast:
-        return this.inferTypeCastType(node as TypeCastNode);
+        return await this.inferTypeCastType(node as TypeCastNode);
         
       case NodeType.MembershipTest:
         return { type: 'Boolean', singleton: true };
         
       case NodeType.TypeOrIdentifier:
-        return this.inferTypeOrIdentifierType(node as TypeOrIdentifierNode, inputType);
+        return await this.inferTypeOrIdentifierType(node as TypeOrIdentifierNode, inputType);
         
       default:
         return { type: 'Any', singleton: false };
@@ -291,7 +645,7 @@ export class Analyzer {
     }
   }
   
-  private inferBinaryType(node: BinaryNode, inputType?: TypeInfo): TypeInfo {
+  private async inferBinaryType(node: BinaryNode, inputType?: TypeInfo): Promise<TypeInfo> {
     const operator = registry.getOperatorDefinition(node.operator);
     if (!operator) {
       return { type: 'Any', singleton: false };
@@ -299,12 +653,12 @@ export class Analyzer {
     
     // For navigation (dot operator), we need special handling
     if (node.operator === '.') {
-      return this.inferNavigationType(node, inputType);
+      return await this.inferNavigationType(node, inputType);
     }
     
     // Infer types of operands
-    const leftType = this.inferType(node.left, inputType);
-    const rightType = this.inferType(node.right, inputType);
+    const leftType = await this.inferType(node.left, inputType);
+    const rightType = await this.inferType(node.right, inputType);
     
     // Find matching signature
     for (const sig of operator.signatures) {
@@ -319,12 +673,12 @@ export class Analyzer {
     return this.resolveResultType(defaultResult, inputType, leftType, rightType);
   }
   
-  private inferNavigationType(node: BinaryNode, inputType?: TypeInfo): TypeInfo {
-    const leftType = this.inferType(node.left, inputType);
+  private async inferNavigationType(node: BinaryNode, inputType?: TypeInfo): Promise<TypeInfo> {
+    const leftType = await this.inferType(node.left, inputType);
     
     // If the right side is a function, return the function's type
     if (node.right.type === NodeType.Function) {
-      return this.inferType(node.right, leftType);
+      return await this.inferType(node.right, leftType);
     }
     
     // If we have a model provider and the right side is an identifier
@@ -332,7 +686,7 @@ export class Analyzer {
       const propertyName = (node.right as IdentifierNode).name;
       
       // Use getElementType to navigate the property
-      const resultType = this.modelProvider.getElementType(leftType, propertyName);
+      const resultType = await this.modelProvider.getElementType(leftType, propertyName);
       if (resultType) {
         return resultType;
       }
@@ -341,11 +695,8 @@ export class Analyzer {
       // Skip diagnostics for union types - they may have dynamic properties
       if (leftType.namespace && leftType.name && leftType.modelContext && 
           !(leftType.modelContext as any).isUnion) {
-        this.addDiagnostic(
-          DiagnosticSeverity.Error,
-          `Unknown property '${propertyName}' on type ${leftType.namespace}.${leftType.name}`,
-          node.right,
-          'UNKNOWN_PROPERTY'
+        this.diagnostics.push(
+          toDiagnostic(Errors.unknownProperty(propertyName, `${leftType.namespace}.${leftType.name}`, node.right.range))
         );
       }
     }
@@ -369,7 +720,7 @@ export class Analyzer {
     return { type: 'Any', singleton: false };
   }
   
-  private inferFunctionType(node: FunctionNode, inputType?: TypeInfo): TypeInfo {
+  private async inferFunctionType(node: FunctionNode, inputType?: TypeInfo): Promise<TypeInfo> {
     if (node.name.type !== NodeType.Identifier) {
       return { type: 'Any', singleton: false };
     }
@@ -381,28 +732,112 @@ export class Analyzer {
       return { type: 'Any', singleton: false };
     }
     
+    // Special handling for iif function
+    if (funcName === 'iif') {
+      // iif returns the common type of the true and false branches
+      if (node.arguments.length >= 2) {
+        const trueBranchType = await this.inferType(node.arguments[1]!, inputType);
+        if (node.arguments.length >= 3) {
+          const falseBranchType = await this.inferType(node.arguments[2]!, inputType);
+          // If both branches have the same type, use that
+          if (trueBranchType.type === falseBranchType.type && 
+              trueBranchType.singleton === falseBranchType.singleton) {
+            return trueBranchType;
+          }
+          // If types are the same but singleton differs, return as collection
+          if (trueBranchType.type === falseBranchType.type) {
+            // One is singleton, one is collection - result must be collection
+            return { type: trueBranchType.type, singleton: false };
+          }
+          // Otherwise, check if one is a subtype of the other
+          if (this.isTypeCompatible(trueBranchType, falseBranchType)) {
+            return falseBranchType;
+          }
+          if (this.isTypeCompatible(falseBranchType, trueBranchType)) {
+            return trueBranchType;
+          }
+        } else {
+          // Only true branch, result can be that type or empty
+          return { ...trueBranchType, singleton: false };
+        }
+      }
+      return { type: 'Any', singleton: false };
+    }
+    
+    // Special handling for defineVariable function
+    if (funcName === 'defineVariable') {
+      // defineVariable returns its input type unchanged
+      return inputType || { type: 'Any', singleton: false };
+    }
+    
+    // Special handling for aggregate function
+    if (funcName === 'aggregate') {
+      // If init parameter is provided, use its type to infer result type
+      if (node.arguments.length >= 2) {
+        const initType = await this.inferType(node.arguments[1]!, inputType);
+        // The result type is the same as init type
+        return initType;
+      }
+      // Without init, we can't fully infer the type without running annotation
+      // This is a limitation - the actual type will be set during annotateAST
+      if (node.arguments.length >= 1) {
+        // We could try to infer, but it would require setting up system variables
+        // For now, return Any and let annotateAST handle proper typing
+        return { type: 'Any', singleton: false };
+      }
+      // No arguments at all
+      return { type: 'Any', singleton: false };
+    }
+    
+    // Special handling for children function
+    if (funcName === 'children') {
+      if (inputType && this.modelProvider && 'getChildrenType' in this.modelProvider) {
+        const childrenType = await this.modelProvider.getChildrenType(inputType);
+        if (childrenType) {
+          return childrenType;
+        }
+      }
+      // Fallback to Any collection
+      return { type: 'Any', singleton: false };
+    }
+    
+    // Special handling for descendants function
+    // Returns Any type due to combinatorial explosion of possible types
+    if (funcName === 'descendants') {
+      return { type: 'Any', singleton: false };
+    }
+    
     // Special handling for functions with dynamic result types
-    if (func.signature.result === 'inputType') {
+    // Use first matching signature's result type
+    const matchingSignature = func.signatures?.find(sig => 
+      !sig.input || !inputType || this.isTypeCompatible(inputType, sig.input)
+    ) || func.signatures?.[0];
+    
+    if (!matchingSignature) {
+      return { type: 'Any', singleton: false };
+    }
+    
+    if (matchingSignature.result === 'inputType') {
       // Functions like where() return the same type as input but always as collection
       return inputType ? { ...inputType, singleton: false } : { type: 'Any', singleton: false };
-    } else if (func.signature.result === 'inputTypeSingleton') {
+    } else if (matchingSignature.result === 'inputTypeSingleton') {
       // Functions like first(), last() return the same type as input but as singleton
       return inputType ? { ...inputType, singleton: true } : { type: 'Any', singleton: true };
-    } else if (func.signature.result === 'parameterType' && node.arguments.length > 0) {
+    } else if (matchingSignature.result === 'parameterType' && node.arguments.length > 0) {
       // Functions like select() return the type of the first parameter expression as collection
-      const paramType = this.inferType(node.arguments[0]!, inputType);
+      const paramType = await this.inferType(node.arguments[0]!, inputType);
       return { ...paramType, singleton: false };
-    } else if (typeof func.signature.result === 'object') {
-      return func.signature.result;
+    } else if (typeof matchingSignature.result === 'object') {
+      return matchingSignature.result;
     }
     
     return { type: 'Any', singleton: false };
   }
   
-  private inferIdentifierType(node: IdentifierNode, inputType?: TypeInfo): TypeInfo {
+  private async inferIdentifierType(node: IdentifierNode, inputType?: TypeInfo): Promise<TypeInfo> {
     // First, try to navigate from input type (most common case)
     if (inputType && this.modelProvider) {
-      const elementType = this.modelProvider.getElementType(inputType, node.name);
+      const elementType = await this.modelProvider.getElementType(inputType, node.name);
       if (elementType) {
         return elementType;
       }
@@ -410,8 +845,11 @@ export class Analyzer {
     
     // Only check if it's a type name if it starts with uppercase (FHIR convention)
     // or if there's no input type context
-    if (this.modelProvider && (!inputType || /^[A-Z]/.test(node.name))) {
-      const typeInfo = this.modelProvider.getType(node.name);
+    // Skip common FHIRPath keywords and function names that aren't types
+    const fhirPathKeywords = ['Boolean', 'String', 'Integer', 'Decimal', 'Date', 'DateTime', 'Time', 'Quantity', 'ofType'];
+    if (this.modelProvider && (!inputType || /^[A-Z]/.test(node.name)) && !fhirPathKeywords.includes(node.name)) {
+      // Try to get type from model provider
+      const typeInfo = await this.modelProvider.getType(node.name);
       if (typeInfo) {
         return typeInfo;
       }
@@ -420,20 +858,23 @@ export class Analyzer {
     return { type: 'Any', singleton: false };
   }
   
-  private inferTypeOrIdentifierType(node: TypeOrIdentifierNode, inputType?: TypeInfo): TypeInfo {
+  private async inferTypeOrIdentifierType(node: TypeOrIdentifierNode, inputType?: TypeInfo): Promise<TypeInfo> {
     // TypeOrIdentifier can be either a type name or a property navigation
     
     // First, try navigation from input type (most common case)
     if (inputType && this.modelProvider) {
-      const elementType = this.modelProvider.getElementType(inputType, node.name);
+      const elementType = await this.modelProvider.getElementType(inputType, node.name);
       if (elementType) {
         return elementType;
       }
     }
     
     // Then check if it's a type name (only for uppercase names or no input context)
-    if (this.modelProvider && (!inputType || /^[A-Z]/.test(node.name))) {
-      const typeInfo = this.modelProvider.getType(node.name);
+    // Skip common FHIRPath keywords and function names that aren't types
+    const fhirPathKeywords = ['Boolean', 'String', 'Integer', 'Decimal', 'Date', 'DateTime', 'Time', 'Quantity', 'ofType'];
+    if (this.modelProvider && (!inputType || /^[A-Z]/.test(node.name)) && !fhirPathKeywords.includes(node.name)) {
+      // Try to get type from model provider
+      const typeInfo = await this.modelProvider.getType(node.name);
       if (typeInfo) {
         return typeInfo;
       }
@@ -443,11 +884,28 @@ export class Analyzer {
   }
   
   private inferVariableType(node: VariableNode): TypeInfo {
-    // Built-in variables
-    if (node.name === '$this') {
-      return { type: 'Any', singleton: false }; // $this is usually a collection
-    } else if (node.name === '$index' || node.name === '$total') {
-      return { type: 'Integer', singleton: true };
+    // System variables - check temporary context
+    if (node.name.startsWith('$')) {
+      const systemType = this.systemVariableTypes.get(node.name);
+      if (systemType) {
+        return systemType;
+      }
+      // Fallback defaults for system variables
+      switch (node.name) {
+        case '$this':
+          return { type: 'Any', singleton: false };
+        case '$index':
+          return { type: 'Integer', singleton: true };
+        case '$total':
+          return { type: 'Any', singleton: false };
+        default:
+          return { type: 'Any', singleton: false };
+      }
+    }
+    
+    // Special FHIRPath environment variables
+    if (node.name === '%context' || node.name === '%resource' || node.name === '%rootResource') {
+      return { type: 'Any', singleton: false }; // These return the original input
     }
     
     // User-defined variables - check with or without % prefix
@@ -464,13 +922,13 @@ export class Analyzer {
     return { type: 'Any', singleton: true };
   }
   
-  private inferCollectionType(node: CollectionNode): TypeInfo {
+  private async inferCollectionType(node: CollectionNode): Promise<TypeInfo> {
     if (node.elements.length === 0) {
       return { type: 'Any', singleton: false };
     }
     
     // Infer types of all elements
-    const elementTypes = node.elements.map(el => this.inferType(el));
+    const elementTypes = await Promise.all(node.elements.map(el => this.inferType(el)));
     
     // If all elements have the same type, use that
     const firstType = elementTypes[0];
@@ -490,12 +948,12 @@ export class Analyzer {
     return { type: 'Any', singleton: false };
   }
   
-  private inferTypeCastType(node: TypeCastNode): TypeInfo {
+  private async inferTypeCastType(node: TypeCastNode): Promise<TypeInfo> {
     const targetType = node.targetType;
     
-    // If we have a model provider, try to get the type info
+    // If we have a model provider, try to get the type
     if (this.modelProvider) {
-      const typeInfo = this.modelProvider.getType(targetType);
+      const typeInfo = await this.modelProvider.getType(targetType);
       if (typeInfo) {
         return typeInfo;
       }
@@ -626,17 +1084,14 @@ export class Analyzer {
     if (!foundMatch) {
       const leftTypeStr = this.typeToString(leftType);
       const rightTypeStr = this.typeToString(rightType);
-      this.addDiagnostic(
-        DiagnosticSeverity.Error,
-        `Type mismatch: operator '${node.operator}' cannot be applied to types ${leftTypeStr} and ${rightTypeStr}`,
-        node,
-        DiagnosticCode.TypeMismatch
+      this.diagnostics.push(
+        toDiagnostic(Errors.operatorTypeMismatch(node.operator, leftTypeStr, rightTypeStr, node.range))
       );
     }
   }
   
   private checkFunctionArgumentTypes(node: FunctionNode, func: import('./types').FunctionDefinition): void {
-    const params = func.signature.parameters;
+    const params = func.signatures?.[0]?.parameters || [];
     
     for (let i = 0; i < Math.min(node.arguments.length, params.length); i++) {
       const arg = node.arguments[i]!;
@@ -647,11 +1102,8 @@ export class Analyzer {
         if (!this.isTypeCompatible(arg.typeInfo, param.type)) {
           const argTypeStr = this.typeToString(arg.typeInfo);
           const paramTypeStr = this.typeToString(param.type);
-          this.addDiagnostic(
-            DiagnosticSeverity.Error,
-            `Type mismatch: argument ${i + 1} of function '${func.name}' expects ${paramTypeStr} but got ${argTypeStr}`,
-            arg,
-            DiagnosticCode.ArgumentTypeMismatch
+          this.diagnostics.push(
+            toDiagnostic(Errors.argumentTypeMismatch(i + 1, func.name, paramTypeStr, argTypeStr, arg.range))
           );
         }
       }
@@ -666,72 +1118,279 @@ export class Analyzer {
     return `${type.type}${singletonStr}`;
   }
   
+  /**
+   * Infer the expected type at a cursor position based on context
+   */
+  private inferExpectedTypeForCursor(cursorNode: AnyCursorNode, inputType?: TypeInfo): TypeInfo | undefined {
+    const context = cursorNode.context;
+    
+    switch (context) {
+      case CursorContext.Identifier:
+        // After dot, expecting a member of the input type
+        return inputType;
+        
+      case CursorContext.Type:
+        // After is/as/ofType, expecting a type name
+        return { type: 'System.String' as TypeName, singleton: true };
+        
+      case CursorContext.Argument:
+        // In function argument, would need to look up function signature
+        // For now, return Any
+        return { type: 'Any' as TypeName, singleton: false };
+        
+      case CursorContext.Index:
+        // In indexer, expecting Integer
+        return { type: 'Integer' as TypeName, singleton: true };
+        
+      case CursorContext.Operator:
+        // Between expressions, could be any operator
+        // Return the input type as context
+        return inputType;
+        
+      default:
+        return undefined;
+    }
+  }
+
   /**
    * Annotate AST with type information
    */
-  private annotateAST(node: ASTNode, inputType?: TypeInfo): void {
+  private async annotateAST(node: ASTNode, inputType?: TypeInfo): Promise<void> {
+    // Check for cursor node in cursor mode
+    if (this.cursorMode && isCursorNode(node)) {
+      this.stoppedAtCursor = true;
+      this.cursorContext = {
+        cursorNode: node as AnyCursorNode,
+        typeBeforeCursor: inputType,
+        expectedType: this.inferExpectedTypeForCursor(node as AnyCursorNode, inputType)
+      };
+      // Still attach a type to the cursor node for consistency
+      (node as any).typeInfo = inputType || { type: 'Any', singleton: false };
+      return; // Short-circuit
+    }
+    
+    // If we've already stopped at cursor, don't continue
+    if (this.stoppedAtCursor) {
+      return;
+    }
+    
     // Handle error nodes
     if (node.type === 'Error') {
       const errorNode = node as ErrorNode;
       // Infer a reasonable type for error nodes
       node.typeInfo = this.inferErrorNodeType(errorNode, inputType);
       // Add diagnostic for the error
-      this.addDiagnostic(
-        errorNode.severity || DiagnosticSeverity.Error,
-        errorNode.message,
-        errorNode,
-        errorNode.code?.toString() || DiagnosticCode.ParseError
-      );
+      this.diagnostics.push({
+        severity: errorNode.severity || DiagnosticSeverity.Error,
+        message: errorNode.message,
+        range: errorNode.range,
+        code: errorNode.code?.toString() || 'FP5003',
+        source: 'fhirpath'
+      });
       return;
     }
     
     // Infer and attach type info
-    node.typeInfo = this.inferType(node, inputType);
+    node.typeInfo = await this.inferType(node, inputType);
 
     // Recursively annotate children
     switch (node.type) {
       case NodeType.Binary:
         const binaryNode = node as BinaryNode;
-        this.annotateAST(binaryNode.left, inputType);
+        await this.annotateAST(binaryNode.left, inputType);
+        
+        // If we stopped at cursor, don't continue
+        if (this.stoppedAtCursor) {
+          break;
+        }
+        
+        // Check if right side is a cursor node - if so, set type from left
+        if (this.cursorMode && isCursorNode(binaryNode.right)) {
+          this.stoppedAtCursor = true;
+          this.cursorContext = {
+            cursorNode: binaryNode.right as AnyCursorNode,
+            typeBeforeCursor: binaryNode.left.typeInfo,
+            expectedType: this.inferExpectedTypeForCursor(binaryNode.right as AnyCursorNode, binaryNode.left.typeInfo)
+          };
+          // Still attach type to cursor node
+          (binaryNode.right as any).typeInfo = binaryNode.left.typeInfo || { type: 'Any', singleton: false };
+          break;
+        }
         
         // For navigation, pass the left's type as input to the right
         if (binaryNode.operator === '.') {
-          this.annotateAST(binaryNode.right, binaryNode.left.typeInfo);
+          // Collect all variables defined in the left side chain
+          const definedVarsWithTypes = this.collectDefinedVariablesWithTypes(binaryNode.left);
+          
+          if (definedVarsWithTypes.size > 0) {
+            // Save current variable types
+            const savedTypes = new Map<string, TypeInfo>();
+            definedVarsWithTypes.forEach((type, varName) => {
+              const currentType = this.userVariableTypes.get(varName);
+              if (currentType) {
+                savedTypes.set(varName, currentType);
+              }
+              this.userVariableTypes.set(varName, type);
+            });
+            
+            // Annotate right side with new variables in scope
+            await this.annotateAST(binaryNode.right, binaryNode.left.typeInfo);
+            
+            // Restore previous types
+            definedVarsWithTypes.forEach((_, varName) => {
+              const savedType = savedTypes.get(varName);
+              if (savedType) {
+                this.userVariableTypes.set(varName, savedType);
+              } else {
+                this.userVariableTypes.delete(varName);
+              }
+            });
+          } else {
+            // No defineVariable in chain, proceed normally
+            await this.annotateAST(binaryNode.right, binaryNode.left.typeInfo);
+          }
         } else {
-          this.annotateAST(binaryNode.right, inputType);
+          await this.annotateAST(binaryNode.right, inputType);
         }
         break;
         
       case NodeType.Unary:
         const unaryNode = node as UnaryNode;
-        this.annotateAST(unaryNode.operand, inputType);
+        await this.annotateAST(unaryNode.operand, inputType);
         break;
         
       case NodeType.Function:
         const funcNode = node as FunctionNode;
-        this.annotateAST(funcNode.name, inputType);
-        funcNode.arguments.forEach(arg => this.annotateAST(arg, inputType));
+        await this.annotateAST(funcNode.name, inputType);
+        
+        // Special handling for aggregate function arguments
+        if (funcNode.name.type === NodeType.Identifier && 
+            (funcNode.name as IdentifierNode).name === 'aggregate') {
+          // Aggregate establishes both $this and $total
+          if (funcNode.arguments.length >= 1) {
+            const itemType = inputType ? { ...inputType, singleton: true } : { type: 'Any' as TypeName, singleton: true };
+            
+            // Save current system variable context
+            const savedThis = this.systemVariableTypes.get('$this');
+            const savedTotal = this.systemVariableTypes.get('$total');
+            
+            // Set $this for iteration
+            this.systemVariableTypes.set('$this', itemType);
+            
+            if (funcNode.arguments.length >= 2) {
+              // Has init parameter - evaluate it first
+              await this.annotateAST(funcNode.arguments[1]!, inputType);
+              const initType = funcNode.arguments[1]!.typeInfo;
+              
+              // Set $total to init type
+              if (initType) {
+                this.systemVariableTypes.set('$total', initType);
+              } else {
+                this.systemVariableTypes.set('$total', { type: 'Any', singleton: false });
+              }
+              
+              // Process aggregator with both variables set
+              await this.annotateAST(funcNode.arguments[0]!, inputType);
+              
+              // Process remaining arguments
+              for (const arg of funcNode.arguments.slice(2)) {
+                await this.annotateAST(arg, inputType);
+                if (this.stoppedAtCursor) break;
+              }
+            } else {
+              // No init - first pass to infer aggregator type
+              this.systemVariableTypes.set('$total', { type: 'Any', singleton: false });
+              await this.annotateAST(funcNode.arguments[0]!, inputType);
+              
+              // Second pass with inferred type
+              const aggregatorType = funcNode.arguments[0]!.typeInfo;
+              if (aggregatorType) {
+                this.systemVariableTypes.set('$total', aggregatorType);
+                // Re-annotate with proper $total type
+                await this.annotateAST(funcNode.arguments[0]!, inputType);
+              }
+            }
+            
+            // Restore previous context
+            if (savedThis) {
+              this.systemVariableTypes.set('$this', savedThis);
+            } else {
+              this.systemVariableTypes.delete('$this');
+            }
+            if (savedTotal) {
+              this.systemVariableTypes.set('$total', savedTotal);
+            } else {
+              this.systemVariableTypes.delete('$total');
+            }
+          }
+        } else {
+          // Special handling for functions that pass their input as context to arguments
+          const funcName = funcNode.name.type === NodeType.Identifier ? 
+            (funcNode.name as IdentifierNode).name : null;
+          
+          if (funcName && ['where', 'select', 'all', 'exists'].includes(funcName)) {
+            // These functions establish $this as each element of the input collection
+            const elementType = inputType ? { ...inputType, singleton: true } : { type: 'Any' as TypeName, singleton: true };
+            
+            // Save current system variable context
+            const savedThis = this.systemVariableTypes.get('$this');
+            const savedIndex = this.systemVariableTypes.get('$index');
+            
+            // Set system variables for expression evaluation
+            this.systemVariableTypes.set('$this', elementType);
+            this.systemVariableTypes.set('$index', { type: 'Integer', singleton: true });
+            
+            // Process arguments with system variables in scope
+            for (const arg of funcNode.arguments) {
+              await this.annotateAST(arg, inputType);
+              if (this.stoppedAtCursor) break;
+            }
+            
+            // Restore previous context
+            if (savedThis) {
+              this.systemVariableTypes.set('$this', savedThis);
+            } else {
+              this.systemVariableTypes.delete('$this');
+            }
+            if (savedIndex) {
+              this.systemVariableTypes.set('$index', savedIndex);
+            } else {
+              this.systemVariableTypes.delete('$index');
+            }
+          } else {
+            // Regular function argument annotation
+            for (const arg of funcNode.arguments) {
+              await this.annotateAST(arg, inputType);
+              if (this.stoppedAtCursor) break;
+            }
+          }
+        }
         break;
         
       case NodeType.Collection:
         const collNode = node as CollectionNode;
-        collNode.elements.forEach(el => this.annotateAST(el, inputType));
+        for (const el of collNode.elements) {
+          await this.annotateAST(el, inputType);
+          if (this.stoppedAtCursor) break;
+        }
         break;
         
       case NodeType.TypeCast:
         const castNode = node as TypeCastNode;
-        this.annotateAST(castNode.expression, inputType);
+        await this.annotateAST(castNode.expression, inputType);
         break;
         
       case NodeType.MembershipTest:
         const memberNode = node as MembershipTestNode;
-        this.annotateAST(memberNode.expression, inputType);
+        await this.annotateAST(memberNode.expression, inputType);
         break;
         
       case NodeType.Index:
         const indexNode = node as IndexNode;
-        this.annotateAST(indexNode.expression, inputType);
-        this.annotateAST(indexNode.index, inputType);
+        await this.annotateAST(indexNode.expression, inputType);
+        if (!this.stoppedAtCursor) {
+          await this.annotateAST(indexNode.index, inputType);
+        }
         break;
         
       case NodeType.TypeOrIdentifier: