Add function argument count validation to language service

index.ts

@@ -43,8 +43,10 @@ export type {
   HoverV2,
   GetCompletionsParams,
   GetHoverParams,
+  GetDiagnosticsParams,
   HighlightToken,
-  LanguageServiceOptions
+  LanguageServiceOptions,
+  ArityInfo
 } from './src/language-service/index.js';
 
 export { createLanguageService, Expression, Parser };

src/language-service/language-service.models.ts

@@ -1,5 +1,6 @@
 import { Parser } from '../parsing/parser';
 import { BUILTIN_FUNCTION_DOCS, FunctionDoc } from './language-service.documentation';
+import type { ArityInfo } from './language-service.types';
 
 export class FunctionDetails {
   private readonly builtInFunctionDoc : FunctionDoc | undefined;
@@ -17,6 +18,41 @@ export class FunctionDetails {
     return typeof f === 'function' ? f.length : undefined;
   }
 
+  /**
+   * Returns the arity information for this function:
+   * - min: minimum number of required arguments
+   * - max: maximum number of arguments, or undefined if variadic
+   */
+  public arityInfo(): ArityInfo | undefined {
+    if (this.builtInFunctionDoc) {
+      const params = this.builtInFunctionDoc.params || [];
+      if (params.length === 0) {
+        return { min: 0, max: 0 };
+      }
+
+      // Check if any parameter is variadic
+      const hasVariadic = params.some(p => p.isVariadic);
+      // Count required (non-optional, non-variadic) parameters
+      const requiredParams = params.filter(p => !p.optional && !p.isVariadic);
+      const optionalParams = params.filter(p => p.optional && !p.isVariadic);
+
+      const min = requiredParams.length;
+      // If variadic, max is undefined (unlimited); otherwise, it's all non-variadic params
+      const max = hasVariadic ? undefined : (requiredParams.length + optionalParams.length);
+
+      return { min, max };
+    }
+
+    // For functions without documentation, use the JavaScript function's .length property
+    const f: unknown = (this.parser.functions && this.parser.functions[this.name]) || (this.parser.unaryOps && this.parser.unaryOps[this.name]);
+    if (typeof f === 'function') {
+      // JavaScript's .length gives number of expected arguments (doesn't account for variadic)
+      return { min: f.length, max: f.length };
+    }
+
+    return undefined;
+  }
+
   public docs() {
     if (this.builtInFunctionDoc) {
       const description = this.builtInFunctionDoc.description || '';

src/language-service/language-service.ts

@@ -21,11 +21,12 @@ import type {
   LanguageServiceOptions,
   GetCompletionsParams,
   GetHoverParams,
+  GetDiagnosticsParams,
   LanguageServiceApi,
   HoverV2
 } from './language-service.types';
-import type { CompletionItem, Range } from 'vscode-languageserver-types';
-import { CompletionItemKind, MarkupKind, InsertTextFormat } from 'vscode-languageserver-types';
+import type { CompletionItem, Range, Diagnostic } from 'vscode-languageserver-types';
+import { CompletionItemKind, MarkupKind, InsertTextFormat, DiagnosticSeverity } from 'vscode-languageserver-types';
 import { TextDocument } from 'vscode-languageserver-textdocument';
 import { BUILTIN_KEYWORD_DOCS, DEFAULT_CONSTANT_DOCS } from './language-service.documentation';
 import { FunctionDetails } from './language-service.models';
@@ -295,10 +296,153 @@ export function createLanguageService(options: LanguageServiceOptions | undefine
     }));
   }
 
+  /**
+   * Analyzes the document for function calls and checks if they have the correct number of arguments.
+   * Returns diagnostics for function calls with incorrect argument counts.
+   */
+  function getDiagnostics(params: GetDiagnosticsParams): Diagnostic[] {
+    const { textDocument } = params;
+    const text = textDocument.getText();
+    const diagnostics: Diagnostic[] = [];
+
+    const ts = makeTokenStream(parser, text);
+    const spans = iterateTokens(ts);
+
+    // Build a map from function name to FunctionDetails for quick lookup
+    const funcDetailsMap = new Map<string, FunctionDetails>();
+    for (const func of allFunctions()) {
+      funcDetailsMap.set(func.name, func);
+    }
+
+    // Find function calls: TNAME followed by TPAREN '('
+    for (let i = 0; i < spans.length; i++) {
+      const span = spans[i];
+      const token = span.token;
+
+      // Check if this is a function name followed by '('
+      if (token.type === TNAME && functionNamesSet().has(String(token.value))) {
+        const funcName = String(token.value);
+
+        // Look for the next token being '('
+        if (i + 1 < spans.length && spans[i + 1].token.type === TPAREN && spans[i + 1].token.value === '(') {
+          const openParenIndex = i + 1;
+          const openParenSpan = spans[openParenIndex];
+
+          // Count arguments by tracking parentheses/brackets depth and commas
+          let argCount = 0;
+          let parenDepth = 1;
+          let bracketDepth = 0;
+          let braceDepth = 0;
+          let foundClosingParen = false;
+          let closeParenSpan = openParenSpan;
+          let hasSeenArgumentToken = false;
+
+          for (let j = openParenIndex + 1; j < spans.length && parenDepth > 0; j++) {
+            const currentToken = spans[j].token;
+
+            if (currentToken.type === TPAREN) {
+              if (currentToken.value === '(') {
+                parenDepth++;
+                // Opening paren can start an argument (e.g., nested function call)
+                if (parenDepth === 2 && bracketDepth === 0 && braceDepth === 0 && !hasSeenArgumentToken) {
+                  hasSeenArgumentToken = true;
+                  argCount = 1;
+                }
+              } else if (currentToken.value === ')') {
+                parenDepth--;
+                if (parenDepth === 0) {
+                  foundClosingParen = true;
+                  closeParenSpan = spans[j];
+                }
+              }
+            } else if (currentToken.type === TBRACKET) {
+              if (currentToken.value === '[') {
+                bracketDepth++;
+                // Opening bracket starts an argument (array literal)
+                if (parenDepth === 1 && bracketDepth === 1 && braceDepth === 0 && !hasSeenArgumentToken) {
+                  hasSeenArgumentToken = true;
+                  argCount = 1;
+                }
+              } else if (currentToken.value === ']') {
+                bracketDepth--;
+              }
+            } else if (currentToken.type === TBRACE) {
+              if (currentToken.value === '{') {
+                braceDepth++;
+                // Opening brace starts an argument (object literal)
+                if (parenDepth === 1 && bracketDepth === 0 && braceDepth === 1 && !hasSeenArgumentToken) {
+                  hasSeenArgumentToken = true;
+                  argCount = 1;
+                }
+              } else if (currentToken.value === '}') {
+                braceDepth--;
+              }
+            } else if (currentToken.type === TCOMMA && parenDepth === 1 && bracketDepth === 0 && braceDepth === 0) {
+              // Only count commas at the top level of the function call
+              argCount++;
+              hasSeenArgumentToken = false; // Reset for next argument
+            } else if (parenDepth === 1 && bracketDepth === 0 && braceDepth === 0 && !hasSeenArgumentToken) {
+              // First non-comma, non-paren, non-bracket, non-brace token at depth 1 means we have at least one argument
+              hasSeenArgumentToken = true;
+              argCount = Math.max(argCount, 1);
+            }
+          }
+
+          // If we found a closing paren and there was content, argCount is commas + 1
+          // If there were no arguments (empty parens), argCount stays 0
+          if (foundClosingParen && argCount > 0) {
+            // argCount currently holds the count from counting commas
+            // When we saw first token at depth 1, we set argCount = 1
+            // Each comma adds 1 more, so argCount is correct
+          }
+
+          // Get the function's expected arity
+          const funcDetails = funcDetailsMap.get(funcName);
+          if (funcDetails) {
+            const arityInfo = funcDetails.arityInfo();
+            if (arityInfo) {
+              const { min, max } = arityInfo;
+
+              // Check if argument count is too few
+              if (argCount < min) {
+                const range: Range = {
+                  start: textDocument.positionAt(span.start),
+                  end: textDocument.positionAt(closeParenSpan.end)
+                };
+                diagnostics.push({
+                  range,
+                  severity: DiagnosticSeverity.Error,
+                  message: `Function '${funcName}' expects at least ${min} argument${min !== 1 ? 's' : ''}, but got ${argCount}.`,
+                  source: 'expr-eval'
+                });
+              }
+              // Check if argument count is too many (only if max is defined, i.e., not variadic)
+              else if (max !== undefined && argCount > max) {
+                const range: Range = {
+                  start: textDocument.positionAt(span.start),
+                  end: textDocument.positionAt(closeParenSpan.end)
+                };
+                diagnostics.push({
+                  range,
+                  severity: DiagnosticSeverity.Error,
+                  message: `Function '${funcName}' expects at most ${max} argument${max !== 1 ? 's' : ''}, but got ${argCount}.`,
+                  source: 'expr-eval'
+                });
+              }
+            }
+          }
+        }
+      }
+    }
+
+    return diagnostics;
+  }
+
   return {
     getCompletions,
     getHover,
-    getHighlighting
+    getHighlighting,
+    getDiagnostics
   };
 
 }

src/language-service/language-service.types.ts

@@ -1,5 +1,5 @@
 import type { Values } from '../types';
-import type { Position, Hover, CompletionItem, MarkupContent } from 'vscode-languageserver-types';
+import type { Position, Hover, CompletionItem, MarkupContent, Diagnostic } from 'vscode-languageserver-types';
 import type { TextDocument } from 'vscode-languageserver-textdocument';
 
 /**
@@ -23,6 +23,13 @@ export interface LanguageServiceApi {
      * @param textDocument - The text document to analyze
      */
     getHighlighting(textDocument: TextDocument): HighlightToken[];
+
+    /**
+     * Returns a list of diagnostics for the given text document.
+     * This includes errors like incorrect number of function arguments.
+     * @param params - Parameters for the diagnostics request
+     */
+    getDiagnostics(params: GetDiagnosticsParams): Diagnostic[];
 }
 
 export interface HighlightToken {
@@ -53,3 +60,17 @@ export interface GetHoverParams {
 export interface HoverV2 extends Hover {
     contents: MarkupContent; // Type narrowing since we know we are not going to return deprecated content
 }
+
+export interface GetDiagnosticsParams {
+    textDocument: TextDocument;
+}
+
+/**
+ * Describes the arity (expected number of arguments) for a function.
+ */
+export interface ArityInfo {
+    /** Minimum number of required arguments */
+    min: number;
+    /** Maximum number of arguments, or undefined if variadic (unlimited) */
+    max: number | undefined;
+}