Implement support for genericStage and genericFunction

Author: MarkDuckworth

packages/firestore/src/lite-api/expressions.ts

@@ -2175,6 +2175,29 @@ export class Constant extends Expr {
   }
 }
 
+/**
+ * Internal only
+ * @internal
+ * @private
+ */
+export class MapValue extends Expr {
+  constructor(private plainObject: Map<string, Expr>) {
+    super();
+  }
+
+  exprType: ExprType = 'Constant';
+
+  _readUserData(dataReader: UserDataReader): void {
+    this.plainObject.forEach(expr => {
+      expr._readUserData(dataReader);
+    });
+  }
+
+  _toProto(serializer: JsonProtoSerializer): ProtoValue {
+    return toMapValue(serializer, this.plainObject);
+  }
+}
+
 /**
  * @beta
  *
@@ -2214,6 +2237,17 @@ export class FirestoreFunction extends Expr {
   }
 }
 
+/**
+ * @beta
+ */
+export class GenericFunction
+  extends FirestoreFunction
+  implements FilterCondition, Accumulator
+{
+  accumulator: true = true;
+  filterable: true = true;
+}
+
 /**
  * @beta
  */
@@ -4098,6 +4132,28 @@ export function map(elements: Record<string, any>): MapFunction {
   return new MapFunction(result);
 }
 
+/**
+ * Internal use only
+ * Converts a plainObject to a mapValue in the proto representation,
+ * rather than a functionValue+map that is the result of the map(...) function.
+ * This behaves different than Constant.of(plainObject) because it
+ * traverses the input object, converts values in the object to expressions,
+ * and calls _readUserData on each of these expressions.
+ * @private
+ * @internal
+ * @param plainObject
+ */
+export function _mapValue(plainObject: Record<string, any>): MapValue {
+  const result: Map<string, Expr> = new Map<string, Expr>();
+  for (const key in plainObject) {
+    if (Object.prototype.hasOwnProperty.call(plainObject, key)) {
+      const value = plainObject[key];
+      result.set(key, valueToDefaultExpr(value));
+    }
+  }
+  return new MapValue(result);
+}
+
 export function array(elements: any[]): ArrayFunction {
   return new ArrayFunction(
     elements.map(element => valueToDefaultExpr(element))
@@ -7499,15 +7555,15 @@ export function timestampSub(
  * genericFunction("sum", [Field.of("price")]);
  * ```
  *
- * @param name The name of the user defined function.
+ * @param functionName The name of the user defined function.
  * @param params The arguments to pass to the function.
- * @return A new {@code Function} representing the function call.
+ * @return A new {@code GenericFucntion} representing the function call.
  */
 export function genericFunction(
-  name: string,
-  params: Expr[]
-): FirestoreFunction {
-  return new FirestoreFunction(name, params);
+  functionName: string,
+  params: any[]
+): GenericFunction {
+  return new GenericFunction(functionName, params.map(valueToDefaultExpr));
 }
 
 /**

packages/firestore/src/lite-api/pipeline.ts

@@ -34,6 +34,7 @@ import { isPlainObject } from '../util/input_validation';
 import { getDatastore } from './components';
 import { Firestore } from './database';
 import {
+  _mapValue,
   Accumulator,
   AccumulatorTarget,
   Constant,
@@ -672,12 +673,27 @@ export class Pipeline implements ProtoSerializable<ExecutePipelineRequest> {
    */
   genericStage(name: string, params: any[]): Pipeline {
     const copy = this.stages.map(s => s);
-    params.forEach(param => {
+
+    // Convert input values to Expressions.
+    // We treat objects as mapValues and arrays as arrayValues,
+    // this is unlike the default conversion for objects and arrays
+    // passed to an expression.
+    const expressionParams = params.map((value: any) => {
+      if (value instanceof Expr) {
+        return value;
+      } else if (isPlainObject(value)) {
+        return _mapValue(value);
+      } else {
+        return Constant.of(value);
+      }
+    });
+
+    expressionParams.forEach(param => {
       if (isReadableUserData(param)) {
         param._readUserData(this.userDataReader);
       }
     });
-    copy.push(new GenericStage(name, params));
+    copy.push(new GenericStage(name, expressionParams));
     return this.newPipeline(
       this._db,
       this.userDataReader,

packages/firestore/src/lite-api/stage.ts

@@ -384,14 +384,16 @@ export class Sort implements Stage {
  * @beta
  */
 export class GenericStage implements Stage {
-  constructor(public name: string, params: unknown[]) {}
+  constructor(public name: string, private params: Expr[]) {}
 
   /**
    * @internal
    * @private
    */
   _toProto(serializer: JsonProtoSerializer): ProtoStage {
-    // TODO support generic stage
-    return {};
+    return {
+      name: this.name,
+      args: this.params.map(o => o._toProto(serializer))
+    };
   }
 }