feat(mongo-query-ast): add aggregation expression AST

packages/2-mongo-family/4-query/query-ast/src/exports/index.ts

@@ -1,3 +1,17 @@
+export type { MongoAggExpr, MongoAggSwitchBranch } from '../aggregation-expressions';
+export {
+  MongoAggAccumulator,
+  MongoAggArrayFilter,
+  MongoAggCond,
+  MongoAggFieldRef,
+  MongoAggLet,
+  MongoAggLiteral,
+  MongoAggMap,
+  MongoAggMergeObjects,
+  MongoAggOperator,
+  MongoAggReduce,
+  MongoAggSwitch,
+} from '../aggregation-expressions';
 export type { AggregatePipelineEntry, AnyMongoCommand } from '../commands';
 export {
   AggregateCommand,
@@ -14,6 +28,7 @@ export type { MongoFilterExpr } from '../filter-expressions';
 export {
   MongoAndExpr,
   MongoExistsExpr,
+  MongoExprFilter,
   MongoFieldFilter,
   MongoNotExpr,
   MongoOrExpr,
@@ -42,6 +57,8 @@ export {
   MongoUnwindStage,
 } from '../stages';
 export type {
+  MongoAggExprRewriter,
+  MongoAggExprVisitor,
   MongoFilterRewriter,
   MongoFilterVisitor,
   MongoStageVisitor,

packages/2-mongo-family/4-query/query-ast/src/filter-expressions.ts

@@ -1,4 +1,5 @@
 import type { MongoValue } from '@prisma-next/mongo-value';
+import type { MongoAggExpr } from './aggregation-expressions';
 import { MongoAstNode } from './ast-node';
 import type { MongoFilterRewriter, MongoFilterVisitor } from './visitors';
 
@@ -177,9 +178,33 @@ export class MongoExistsExpr extends MongoFilterExpression {
   }
 }
 
+export class MongoExprFilter extends MongoFilterExpression {
+  readonly kind = 'expr' as const;
+  readonly aggExpr: MongoAggExpr;
+
+  constructor(aggExpr: MongoAggExpr) {
+    super();
+    this.aggExpr = aggExpr;
+    this.freeze();
+  }
+
+  static of(aggExpr: MongoAggExpr): MongoExprFilter {
+    return new MongoExprFilter(aggExpr);
+  }
+
+  accept<R>(visitor: MongoFilterVisitor<R>): R {
+    return visitor.expr(this);
+  }
+
+  rewrite(rewriter: MongoFilterRewriter): MongoFilterExpr {
+    return rewriter.expr ? rewriter.expr(this) : this;
+  }
+}
+
 export type MongoFilterExpr =
   | MongoFieldFilter
   | MongoAndExpr
   | MongoOrExpr
   | MongoNotExpr
-  | MongoExistsExpr;
+  | MongoExistsExpr
+  | MongoExprFilter;

packages/2-mongo-family/4-query/query-ast/src/visitors.ts

@@ -1,6 +1,21 @@
+import type {
+  MongoAggAccumulator,
+  MongoAggArrayFilter,
+  MongoAggCond,
+  MongoAggExpr,
+  MongoAggFieldRef,
+  MongoAggLet,
+  MongoAggLiteral,
+  MongoAggMap,
+  MongoAggMergeObjects,
+  MongoAggOperator,
+  MongoAggReduce,
+  MongoAggSwitch,
+} from './aggregation-expressions';
 import type {
   MongoAndExpr,
   MongoExistsExpr,
+  MongoExprFilter,
   MongoFieldFilter,
   MongoFilterExpr,
   MongoNotExpr,
@@ -16,12 +31,41 @@ import type {
   MongoUnwindStage,
 } from './stages';
 
+export interface MongoAggExprVisitor<R> {
+  fieldRef(expr: MongoAggFieldRef): R;
+  literal(expr: MongoAggLiteral): R;
+  operator(expr: MongoAggOperator): R;
+  accumulator(expr: MongoAggAccumulator): R;
+  cond(expr: MongoAggCond): R;
+  switch_(expr: MongoAggSwitch): R;
+  filter(expr: MongoAggArrayFilter): R;
+  map(expr: MongoAggMap): R;
+  reduce(expr: MongoAggReduce): R;
+  let_(expr: MongoAggLet): R;
+  mergeObjects(expr: MongoAggMergeObjects): R;
+}
+
+export interface MongoAggExprRewriter {
+  fieldRef?(expr: MongoAggFieldRef): MongoAggExpr;
+  literal?(expr: MongoAggLiteral): MongoAggExpr;
+  operator?(expr: MongoAggOperator): MongoAggExpr;
+  accumulator?(expr: MongoAggAccumulator): MongoAggExpr;
+  cond?(expr: MongoAggCond): MongoAggExpr;
+  switch_?(expr: MongoAggSwitch): MongoAggExpr;
+  filter?(expr: MongoAggArrayFilter): MongoAggExpr;
+  map?(expr: MongoAggMap): MongoAggExpr;
+  reduce?(expr: MongoAggReduce): MongoAggExpr;
+  let_?(expr: MongoAggLet): MongoAggExpr;
+  mergeObjects?(expr: MongoAggMergeObjects): MongoAggExpr;
+}
+
 export interface MongoFilterVisitor<R> {
   field(expr: MongoFieldFilter): R;
   and(expr: MongoAndExpr): R;
   or(expr: MongoOrExpr): R;
   not(expr: MongoNotExpr): R;
   exists(expr: MongoExistsExpr): R;
+  expr(expr: MongoExprFilter): R;
 }
 
 export interface MongoFilterRewriter {
@@ -30,6 +74,7 @@ export interface MongoFilterRewriter {
   or?(expr: MongoOrExpr): MongoFilterExpr;
   not?(expr: MongoNotExpr): MongoFilterExpr;
   exists?(expr: MongoExistsExpr): MongoFilterExpr;
+  expr?(expr: MongoExprFilter): MongoFilterExpr;
 }
 
 export interface MongoStageVisitor<R> {

packages/2-mongo-family/4-query/query-ast/test/aggregation-expressions.test-d.ts

@@ -0,0 +1,168 @@
+import { assertType, expectTypeOf, test } from 'vitest';
+import type { MongoAggExpr } from '../src/aggregation-expressions';
+import {
+  type MongoAggAccumulator,
+  type MongoAggArrayFilter,
+  type MongoAggCond,
+  MongoAggFieldRef,
+  type MongoAggLet,
+  type MongoAggLiteral,
+  type MongoAggMap,
+  type MongoAggMergeObjects,
+  type MongoAggOperator,
+  type MongoAggReduce,
+  type MongoAggSwitch,
+} from '../src/aggregation-expressions';
+import type { MongoAggExprRewriter, MongoAggExprVisitor } from '../src/visitors';
+
+test('each concrete class is assignable to MongoAggExpr', () => {
+  expectTypeOf<MongoAggFieldRef>().toExtend<MongoAggExpr>();
+  expectTypeOf<MongoAggLiteral>().toExtend<MongoAggExpr>();
+  expectTypeOf<MongoAggOperator>().toExtend<MongoAggExpr>();
+  expectTypeOf<MongoAggAccumulator>().toExtend<MongoAggExpr>();
+  expectTypeOf<MongoAggCond>().toExtend<MongoAggExpr>();
+  expectTypeOf<MongoAggSwitch>().toExtend<MongoAggExpr>();
+  expectTypeOf<MongoAggArrayFilter>().toExtend<MongoAggExpr>();
+  expectTypeOf<MongoAggMap>().toExtend<MongoAggExpr>();
+  expectTypeOf<MongoAggReduce>().toExtend<MongoAggExpr>();
+  expectTypeOf<MongoAggLet>().toExtend<MongoAggExpr>();
+  expectTypeOf<MongoAggMergeObjects>().toExtend<MongoAggExpr>();
+});
+
+test('MongoAggExpr kind union covers all 11 kinds', () => {
+  expectTypeOf<MongoAggExpr['kind']>().toEqualTypeOf<
+    | 'fieldRef'
+    | 'literal'
+    | 'operator'
+    | 'accumulator'
+    | 'cond'
+    | 'switch'
+    | 'filter'
+    | 'map'
+    | 'reduce'
+    | 'let'
+    | 'mergeObjects'
+  >();
+});
+
+test('switching on kind is exhaustive', () => {
+  function exhaustiveSwitch(expr: MongoAggExpr): string {
+    switch (expr.kind) {
+      case 'fieldRef':
+        return 'fieldRef';
+      case 'literal':
+        return 'literal';
+      case 'operator':
+        return 'operator';
+      case 'accumulator':
+        return 'accumulator';
+      case 'cond':
+        return 'cond';
+      case 'switch':
+        return 'switch';
+      case 'filter':
+        return 'filter';
+      case 'map':
+        return 'map';
+      case 'reduce':
+        return 'reduce';
+      case 'let':
+        return 'let';
+      case 'mergeObjects':
+        return 'mergeObjects';
+      default: {
+        const _exhaustive: never = expr;
+        return _exhaustive;
+      }
+    }
+  }
+  assertType<(expr: MongoAggExpr) => string>(exhaustiveSwitch);
+});
+
+test('MongoAggExprVisitor requires all 11 methods', () => {
+  type Complete = MongoAggExprVisitor<string>;
+
+  expectTypeOf<Complete>().toHaveProperty('fieldRef');
+  expectTypeOf<Complete>().toHaveProperty('literal');
+  expectTypeOf<Complete>().toHaveProperty('operator');
+  expectTypeOf<Complete>().toHaveProperty('accumulator');
+  expectTypeOf<Complete>().toHaveProperty('cond');
+  expectTypeOf<Complete>().toHaveProperty('switch_');
+  expectTypeOf<Complete>().toHaveProperty('filter');
+  expectTypeOf<Complete>().toHaveProperty('map');
+  expectTypeOf<Complete>().toHaveProperty('reduce');
+  expectTypeOf<Complete>().toHaveProperty('let_');
+  expectTypeOf<Complete>().toHaveProperty('mergeObjects');
+
+  // @ts-expect-error - missing 'fieldRef' method
+  assertType<MongoAggExprVisitor<string>>({
+    literal: () => '',
+    operator: () => '',
+    accumulator: () => '',
+    cond: () => '',
+    switch_: () => '',
+    filter: () => '',
+    map: () => '',
+    reduce: () => '',
+    let_: () => '',
+    mergeObjects: () => '',
+  });
+});
+
+test('MongoAggExprRewriter accepts empty object (all optional)', () => {
+  assertType<MongoAggExprRewriter>({});
+});
+
+test('rewriter hooks return MongoAggExpr', () => {
+  const rewriter: MongoAggExprRewriter = {
+    fieldRef: (expr): MongoAggExpr => expr,
+    literal: (expr): MongoAggExpr => expr,
+    operator: (expr): MongoAggExpr => expr,
+  };
+  assertType<MongoAggExprRewriter>(rewriter);
+});
+
+test('accept returns R for any visitor R', () => {
+  const ref = MongoAggFieldRef.of('x');
+  const visitor: MongoAggExprVisitor<number> = {
+    fieldRef: () => 1,
+    literal: () => 2,
+    operator: () => 3,
+    accumulator: () => 4,
+    cond: () => 5,
+    switch_: () => 6,
+    filter: () => 7,
+    map: () => 8,
+    reduce: () => 9,
+    let_: () => 10,
+    mergeObjects: () => 11,
+  };
+  expectTypeOf(ref.accept(visitor)).toBeNumber();
+});
+
+test('rewrite returns MongoAggExpr', () => {
+  const ref = MongoAggFieldRef.of('x');
+  expectTypeOf(ref.rewrite({})).toEqualTypeOf<MongoAggExpr>();
+});
+
+test('MongoAggOperator.args accepts both forms', () => {
+  expectTypeOf<MongoAggOperator['args']>().toEqualTypeOf<
+    MongoAggExpr | ReadonlyArray<MongoAggExpr>
+  >();
+});
+
+test('MongoAggAccumulator.arg accepts MongoAggExpr or null', () => {
+  expectTypeOf<MongoAggAccumulator['arg']>().toEqualTypeOf<MongoAggExpr | null>();
+});
+
+test('MongoAggSwitch.branches is ReadonlyArray', () => {
+  expectTypeOf<MongoAggSwitch['branches']>().toExtend<ReadonlyArray<unknown>>();
+});
+
+test('MongoAggLet.vars is Readonly<Record>', () => {
+  expectTypeOf<MongoAggLet['vars']>().toExtend<Readonly<Record<string, MongoAggExpr>>>();
+});
+
+test('MongoAggMergeObjects.exprs is ReadonlyArray', () => {
+  expectTypeOf<MongoAggMergeObjects['exprs']>().toExtend<ReadonlyArray<MongoAggExpr>>();
+});