feat(parser): add support for parsing and replacing custom functions

Author: sgulseth

src/markProcessor.ts

@@ -20,6 +20,7 @@ export type MarkName =
   | 'float'
   | 'func_args_end'
   | 'func_call'
+  | 'func_decl'
   | 'ident'
   | 'inc_range'
   | 'integer'

src/nodeTypes.ts

@@ -1,7 +1,12 @@
 import type {GroqFunction, GroqPipeFunction} from './evaluator/functions'
 
 /** Any sort of node which appears as syntax */
-export type SyntaxNode = ExprNode | ArrayElementNode | ObjectAttributeNode | SelectAlternativeNode
+export type SyntaxNode =
+  | ExprNode
+  | ArrayElementNode
+  | ObjectAttributeNode
+  | SelectAlternativeNode
+  | FunctionDeclarationNode
 
 export type ObjectAttributeNode =
   | ObjectAttributeValueNode
@@ -280,6 +285,14 @@ export interface ValueNode<P = any> {
   value: P
 }
 
+export interface FunctionDeclarationNode extends BaseNode {
+  type: 'FuncDeclaration'
+  namespace: string
+  name: string
+  params: ParameterNode[]
+  body: ExprNode
+}
+
 export interface FlatMapNode extends BaseNode {
   type: 'FlatMap'
   base: ExprNode

src/parser.ts

@@ -3,14 +3,21 @@ import {tryConstantEvaluate} from './evaluator'
 import {type GroqFunctionArity, namespaces, pipeFunctions} from './evaluator/functions'
 import {MarkProcessor, type MarkVisitor} from './markProcessor'
 import {
+  type ArrayCoerceNode,
   type ArrayElementNode,
+  type DerefNode,
   type ExprNode,
   type FuncCallNode,
+  type FunctionDeclarationNode,
+  type InlineFuncCallNode,
   isSelectorNested,
+  type MapNode,
   type ObjectAttributeNode,
   type ObjectSplatNode,
   type OpCall,
+  type ParameterNode,
   type ParentNode,
+  type ProjectionNode,
   type SelectNode,
   type SelectorNode,
   walk,
@@ -50,9 +57,14 @@ class GroqQueryError extends Error {
   public override name = 'GroqQueryError'
 }
 
+type FunctionId = `${string}::${string}`
+type CustomFunctions = Record<FunctionId, FunctionDeclarationNode>
+
 function createExpressionBuilder(): {
   exprBuilder: MarkVisitor<ExprNode>
+  customFunctions: CustomFunctions
 } {
+  const customFunctions: CustomFunctions = {}
   const exprBuilder: MarkVisitor<ExprNode> = {
     group(p) {
       const inner = p.process(exprBuilder)
@@ -522,6 +534,37 @@ function createExpressionBuilder(): {
         name,
       }
     },
+
+    func_decl(p) {
+      const namespace = p.processString()
+      const name = p.processString()
+      const params: ParameterNode[] = []
+      while (p.getMark().name !== 'func_params_end') {
+        const param = p.process(exprBuilder)
+        if (param.type !== 'Parameter') throw new Error('expected parameter')
+        params.push(param)
+      }
+
+      if (params.length !== 1) {
+        throw new GroqQueryError('Custom functions can only have one parameter')
+      }
+
+      p.shift() // func_params_end
+
+      const body = p.process(exprBuilder)
+
+      const decl = {
+        type: 'FuncDeclaration',
+        namespace,
+        name,
+        params,
+        body,
+      } satisfies FunctionDeclarationNode
+
+      customFunctions[`${namespace}::${name}`] = decl
+
+      return p.process(exprBuilder)
+    },
   }
 
   const OBJECT_BUILDER: MarkVisitor<ObjectAttributeNode> = {
@@ -860,7 +903,7 @@ function createExpressionBuilder(): {
     },
   }
 
-  return {exprBuilder}
+  return {exprBuilder, customFunctions}
 }
 
 function extractPropertyKey(node: ExprNode): string {
@@ -899,6 +942,71 @@ function validateArity(name: string, arity: GroqFunctionArity, count: number) {
   }
 }
 
+/**
+ * The function body is one of the forms:
+ * - $param{…}
+ * - $param->{…}
+ * - $param[]{…}
+ * - $param[]->{…}
+ *
+ * https://github.com/sanity-io/go-groq/blob/b7fb57f5aefe080becff9e3522c0b7b52a79ffd0/parser/internal/parserv2/parser.go#L975-L981
+ */
+function resolveFunctionParameter(
+  parameter: ParameterNode,
+  funcDeclaration: FunctionDeclarationNode,
+  funcCall: InlineFuncCallNode,
+) {
+  const index = funcDeclaration.params.findIndex((p) => p.name === parameter.name)
+  if (index === -1) {
+    throw new GroqQueryError(`Missing argument for parameter ${parameter.name} in function call`)
+  }
+  return funcCall.args[index]
+}
+function replaceCustomFunctionBody(
+  funcDeclaration: FunctionDeclarationNode,
+  funcCall: InlineFuncCallNode,
+): ExprNode {
+  const {body} = funcDeclaration
+
+  if (body.type === 'Projection') {
+    if (body.base.type === 'Parameter') {
+      return {
+        type: 'Projection',
+        base: resolveFunctionParameter(body.base, funcDeclaration, funcCall),
+        expr: body.expr,
+      }
+    }
+
+    if (body.base.type === 'Deref') {
+      if (body.base.base.type === 'Parameter') {
+        return {
+          type: 'Projection',
+          base: {
+            type: 'Deref',
+            base: resolveFunctionParameter(body.base.base, funcDeclaration, funcCall),
+          },
+          expr: body.expr,
+        }
+      }
+    }
+  }
+
+  if (body.type === 'Map' && body.base.type === 'ArrayCoerce') {
+    if (body.base.base.type === 'Parameter') {
+      return {
+        type: 'Map',
+        base: {
+          type: 'ArrayCoerce',
+          base: resolveFunctionParameter(body.base.base, funcDeclaration, funcCall),
+        },
+        expr: body.expr,
+      }
+    }
+  }
+
+  throw new GroqQueryError(`Unexpected function body, must be a projection. Got "${body.type}"`)
+}
+
 function argumentShouldBeSelector(namespace: string, functionName: string, argCount: number) {
   const functionsRequiringSelectors = ['changedAny', 'changedOnly']
 
@@ -924,20 +1032,41 @@ export function parse(input: string, options: ParseOptions = {}): ExprNode {
     throw new GroqSyntaxError(result.position, result.message)
   }
   const processor = new MarkProcessor(input, result.marks, options)
-  const {exprBuilder} = createExpressionBuilder()
+  const {exprBuilder, customFunctions} = createExpressionBuilder()
   const procssed = processor.process(exprBuilder)
-  const replaceInlineFuncCalls = createReplaceInlineFuncCalls(options)
-  return walk(procssed, replaceInlineFuncCalls)
+  const replaceInlineFuncCalls = createReplaceInlineFuncCalls(options, customFunctions)
+  return walk(procssed, (node) => replaceInlineFuncCalls(node))
 }
 
-function createReplaceInlineFuncCalls(options: ParseOptions) {
-  const replacer = (node: ExprNode): ExprNode => {
+function createReplaceInlineFuncCalls(
+  options: ParseOptions,
+  customFunctions: Record<string, FunctionDeclarationNode>,
+) {
+  const replacer = (
+    node: ExprNode,
+    recurssion: Set<FunctionId> = new Set<FunctionId>(),
+  ): ExprNode => {
     if (node.type !== 'InlineFuncCall') {
       return node
     }
 
     const {namespace, name} = node
 
+    const functionId: FunctionId = `${namespace}::${name}`
+    if (recurssion.has(functionId)) {
+      throw new GroqQueryError(`Recursion detected in function ${name}`)
+    }
+
+    // Check for custom function first
+    const customFunction = customFunctions[functionId]
+    if (customFunction) {
+      validateArity(name, customFunction.params.length, node.args.length)
+
+      return walk(replaceCustomFunctionBody(customFunction, node), (node) =>
+        replacer(node, new Set([...recurssion, functionId])),
+      )
+    }
+
     const funcs = namespaces[namespace]
     if (!funcs) {
       throw new GroqQueryError(`Undefined namespace: ${namespace}`)

src/rawParser.js

@@ -25,9 +25,22 @@ const PREC_NEG = 8
 function parse(str) {
   let pos = 0
   pos = skipWS(str, pos)
+
+  let fnMarks = []
+
+  // Parse function declarations first
+  while (pos < str.length && str.substring(pos, pos + 2) === 'fn') {
+    let funcResult = parseFunctionDeclaration(str, pos)
+    if (funcResult.type === 'error') return funcResult
+    fnMarks = fnMarks.concat(funcResult.marks)
+    pos = skipWS(str, funcResult.position)
+  }
+
+  // Parse the main query expression
   let result = parseExpr(str, pos, 0)
   if (result.type === 'error') return result
   pos = skipWS(str, result.position)
+
   if (pos !== str.length) {
     if (result.failPosition) {
       pos = result.failPosition - 1
@@ -36,6 +49,7 @@ function parse(str) {
   }
   delete result.position
   delete result.failPosition
+  result.marks = fnMarks.concat(result.marks)
   return result
 }
 
@@ -831,4 +845,109 @@ function parseRegexStr(str, pos, re) {
   return m ? m[0] : null
 }
 
+/**
+ * Parses a function declaration: fn namespace::name(params) = body;
+ */
+function parseFunctionDeclaration(str, startPos) {
+  let pos = startPos
+  let marks = []
+
+  // Parse "fn"
+  if (str.substring(pos, pos + 2) !== 'fn') {
+    return {
+      type: 'success',
+      position: pos,
+      marks: marks,
+    }
+  }
+  pos = skipWS(str, pos + 2)
+
+  marks.push({name: 'func_decl', position: startPos})
+
+  let identStart = pos
+  let identResult = parseRegexStr(str, pos, IDENT)
+  if (!identResult) {
+    return {type: 'error', message: 'Expected function name', position: pos}
+  }
+  pos += identResult.length
+  pos = skipWS(str, pos)
+
+  marks.push({name: 'ident', position: identStart}, {name: 'ident_end', position: pos})
+
+  // Check for "::"
+  if (str.substring(pos, pos + 2) !== '::') {
+    return {type: 'error', message: 'Expected "::" after namespace', position: pos}
+  }
+
+  pos = skipWS(str, pos + 2)
+  let nameLen = parseRegex(str, pos, IDENT)
+  if (!nameLen) return {type: 'error', message: 'Expected function name', position: pos}
+  marks.push({name: 'ident', position: pos}, {name: 'ident_end', position: pos + nameLen})
+  pos = skipWS(str, pos + nameLen)
+
+  if (str[pos] !== '(') {
+    return {type: 'error', message: 'Expected "("', position: pos}
+  }
+  pos = skipWS(str, pos + 1)
+
+  // Parse parameters
+  while (pos < str.length && str[pos] !== ')') {
+    // Parse parameter (should start with $)
+    if (str[pos] !== '$') {
+      return {type: 'error', message: 'Parameter should start with "$"', position: pos}
+    }
+    const startPos = pos
+    pos++
+
+    const paramName = parseRegexStr(str, pos, IDENT)
+    if (!paramName) {
+      return {type: 'error', message: 'Expected function name', position: pos}
+    }
+    pos += paramName.length
+    marks.push(
+      {name: 'param', position: startPos},
+      {name: 'ident', position: startPos + 1},
+      {name: 'ident_end', position: pos},
+    )
+    pos = skipWS(str, pos)
+
+    // Check for comma
+    if (str[pos] === ',') {
+      pos = skipWS(str, pos + 1)
+    } else if (str[pos] !== ')') {
+      return {type: 'error', message: 'Expected "," or ")"', position: pos}
+    }
+  }
+
+  if (str[pos] !== ')') {
+    return {type: 'error', message: 'Expected ")"', position: pos}
+  }
+  marks.push({name: 'func_params_end', position: pos})
+  pos = skipWS(str, pos + 1)
+
+  if (str[pos] !== '=') {
+    return {type: 'error', message: 'Expected "="', position: pos}
+  }
+  pos = skipWS(str, pos + 1)
+
+  // Parse function body (expression)
+  // marks.push({name: 'func_body', position: pos})
+  let bodyResult = parseExpr(str, pos, 0)
+  if (bodyResult.type === 'error') return bodyResult
+  marks = marks.concat(bodyResult.marks)
+  pos = skipWS(str, bodyResult.position)
+
+  // Parse ";"
+  if (str[pos] !== ';') {
+    return {type: 'error', message: 'Expected ";" after function declaration', position: pos}
+  }
+  pos++
+
+  return {
+    type: 'success',
+    position: pos,
+    marks: marks,
+  }
+}
+
 export {parse}