@atomic-ehr/fhirpath 0.0.1 → 0.0.2-canary.2fee5d9.20250808110507

package/src/analyzer.ts

@@ -20,26 +20,15 @@ import type {
 } from './types';
 import { NodeType, DiagnosticSeverity } from './types';
 import { registry } from './registry';
+import { Errors, toDiagnostic } from './errors';
 
-// 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 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();
 
   constructor(modelProvider?: ModelProvider) {
     this.modelProvider = modelProvider;
@@ -101,10 +90,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,6 +109,32 @@ 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;
@@ -130,11 +145,8 @@ export class Analyzer {
           if (!this.isTypeCompatible(node.left.typeInfo, func.signature.input)) {
             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
+            this.diagnostics.push(
+              toDiagnostic(Errors.typeNotAssignable(inputTypeStr, expectedTypeStr, funcNode.range))
             );
           }
         }
@@ -150,7 +162,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,10 +181,73 @@ 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;
@@ -178,9 +255,13 @@ export class Analyzer {
         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,30 +274,215 @@ 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 {
     // Handle error nodes
@@ -341,11 +607,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))
         );
       }
     }
@@ -381,6 +644,81 @@ 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 = this.inferType(node.arguments[1]!, inputType);
+        if (node.arguments.length >= 3) {
+          const falseBranchType = 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 = 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 = 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') {
       // Functions like where() return the same type as input but always as collection
@@ -443,11 +781,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
@@ -626,11 +981,8 @@ 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))
       );
     }
   }
@@ -647,11 +999,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))
           );
         }
       }
@@ -676,12 +1025,13 @@ export class Analyzer {
       // 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;
     }
     
@@ -696,7 +1046,36 @@ export class Analyzer {
         
         // 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
+            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
+            this.annotateAST(binaryNode.right, binaryNode.left.typeInfo);
+          }
         } else {
           this.annotateAST(binaryNode.right, inputType);
         }
@@ -710,7 +1089,100 @@ export class Analyzer {
       case NodeType.Function:
         const funcNode = node as FunctionNode;
         this.annotateAST(funcNode.name, inputType);
-        funcNode.arguments.forEach(arg => this.annotateAST(arg, 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
+              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
+              this.annotateAST(funcNode.arguments[0]!, inputType);
+              
+              // Process remaining arguments
+              funcNode.arguments.slice(2).forEach(arg => this.annotateAST(arg, inputType));
+            } else {
+              // No init - first pass to infer aggregator type
+              this.systemVariableTypes.set('$total', { type: 'Any', singleton: false });
+              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
+                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
+            funcNode.arguments.forEach(arg => this.annotateAST(arg, inputType));
+            
+            // 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
+            funcNode.arguments.forEach(arg => this.annotateAST(arg, inputType));
+          }
+        }
         break;
         
       case NodeType.Collection: