feat(formula): support lambda in function register (#4298)

Co-authored-by: Wenzhao Hu <wzhudev@gmail.com>
dream-num · Jan 10, 2025 · a1b7a16 · a1b7a16
1 parent 48080f6
commit a1b7a16
Show file tree

Hide file tree

Showing 8 changed files with 64 additions and 11 deletions.
diff --git a/packages/engine-formula/src/engine/ast-node/function-node.ts b/packages/engine-formula/src/engine/ast-node/function-node.ts
@@ -23,7 +23,7 @@ import type {
     BaseReferenceObject,
     FunctionVariantType,
     NodeValueType } from '../reference-object/base-reference-object';
-import type { FormulaFunctionValueType } from '../value-object/primitive-object';
+import type { FormulaFunctionResultValueType } from '../value-object/primitive-object';
 import { Inject, Injector } from '@univerjs/core';
 import { AstNodePromiseType } from '../../basics/common';
 import { ErrorType } from '../../basics/error-type';
@@ -206,7 +206,7 @@ export class FunctionNode extends BaseAstNode {
     /**
      * Transform the result of a custom function to a NodeValueType.
      */
-    private _handleCustomResult(resultVariantCustom: FormulaFunctionValueType): NodeValueType {
+    private _handleCustomResult(resultVariantCustom: FormulaFunctionResultValueType): NodeValueType {
         if (typeof resultVariantCustom !== 'object' || resultVariantCustom == null) {
             return ValueObjectFactory.create(resultVariantCustom);
         }
@@ -240,6 +240,11 @@ export class FunctionNode extends BaseAstNode {
             if (variant.isArray()) {
                 return (variant as ArrayValueObject).toValue();
             }
+
+            if (variant.isLambda()) {
+                return variant;
+            }
+
             return variant.getValue();
         });
     }
@@ -257,7 +262,7 @@ export class FunctionNode extends BaseAstNode {
         if (this._functionExecutor.isCustom()) {
             const resultVariantCustom = this._functionExecutor.calculateCustom(
                 ...this._mapVariantsToValues(variants)
-            ) as FormulaFunctionValueType;
+            ) as FormulaFunctionResultValueType;
 
             resultVariant = this._handleCustomResult(resultVariantCustom);
         } else {

diff --git a/packages/engine-formula/src/engine/value-object/lambda-value-object.ts b/packages/engine-formula/src/engine/value-object/lambda-value-object.ts
@@ -20,10 +20,12 @@ import type { BaseAstNode } from '../ast-node/base-ast-node';
 import type { LambdaParameterNode } from '../ast-node/lambda-parameter-node';
 import type { Interpreter } from '../interpreter/interpreter';
 import type { BaseReferenceObject, FunctionVariantType } from '../reference-object/base-reference-object';
+import type { PrimitiveValueType } from './primitive-object';
 import { ErrorType } from '../../basics/error-type';
 import { DEFAULT_TOKEN_TYPE_LAMBDA_RUNTIME_PARAMETER } from '../../basics/token-type';
 import { AsyncObject } from '../reference-object/base-reference-object';
 import { generateExecuteAstNodeData } from '../utils/ast-node-tool';
+import { ValueObjectFactory } from './array-value-object';
 import { BaseValueObject, ErrorValueObject } from './base-value-object';
 
 function getRootLexerHasValueNode(node: Nullable<BaseAstNode>): Nullable<BaseAstNode> {
@@ -110,6 +112,16 @@ export class LambdaValueObjectObject extends BaseValueObject {
         return value;
     }
 
+    /**
+     * Execute custom lambda function, handle basic types
+     * @param variants
+     */
+    executeCustom(...variants: PrimitiveValueType[]) {
+        // Create base value object from primitive value, then execute
+        const baseValueObjects = variants.map((variant) => ValueObjectFactory.create(variant));
+        return this.execute(...baseValueObjects);
+    }
+
     private _setLambdaNodeValue(node: Nullable<BaseAstNode>) {
         if (!node) {
             return;

diff --git a/packages/engine-formula/src/engine/value-object/primitive-object.ts b/packages/engine-formula/src/engine/value-object/primitive-object.ts
@@ -27,7 +27,8 @@ import { BaseValueObject, ErrorValueObject } from './base-value-object';
 
 export type PrimitiveValueType = string | boolean | number | null;
 
-export type FormulaFunctionValueType = PrimitiveValueType | PrimitiveValueType[][];
+export type FormulaFunctionValueType = PrimitiveValueType | PrimitiveValueType[][] | BaseValueObject;
+export type FormulaFunctionResultValueType = PrimitiveValueType | PrimitiveValueType[][];
 
 export class NullValueObject extends BaseValueObject {
     private static _instance: NullValueObject;

diff --git a/packages/engine-formula/src/functions/base-function.ts b/packages/engine-formula/src/functions/base-function.ts
@@ -19,7 +19,7 @@ import type { IFunctionNames } from '../basics/function';
 import type { BaseReferenceObject, FunctionVariantType, NodeValueType } from '../engine/reference-object/base-reference-object';
 import type { ArrayBinarySearchType } from '../engine/utils/compare';
 import type { ArrayValueObject } from '../engine/value-object/array-value-object';
-import type { FormulaFunctionValueType } from '../engine/value-object/primitive-object';
+import type { FormulaFunctionResultValueType, FormulaFunctionValueType } from '../engine/value-object/primitive-object';
 import type { FormulaDataModel } from '../models/formula-data.model';
 import type { IDefinedNameMapItem } from '../services/defined-names.service';
 import { ErrorType } from '../basics/error-type';
@@ -193,7 +193,7 @@ export class BaseFunction {
 
     calculateCustom(
         ...arg: Array<FormulaFunctionValueType>
-    ): FormulaFunctionValueType | Promise<FormulaFunctionValueType> {
+    ): FormulaFunctionResultValueType | Promise<FormulaFunctionResultValueType> {
         return null;
     }
 

diff --git a/packages/engine-formula/src/index.ts b/packages/engine-formula/src/index.ts
@@ -99,7 +99,7 @@ export { handleRefStringInfo } from './engine/utils/reference';
 export { generateStringWithSequence, type ISequenceNode, sequenceNodeType } from './engine/utils/sequence';
 export { ArrayValueObject, ValueObjectFactory } from './engine/value-object/array-value-object';
 export { BaseValueObject, ErrorValueObject } from './engine/value-object/base-value-object';
-export type { FormulaFunctionValueType, PrimitiveValueType } from './engine/value-object/primitive-object';
+export type { FormulaFunctionResultValueType, FormulaFunctionValueType, PrimitiveValueType } from './engine/value-object/primitive-object';
 export { BooleanValueObject, NullValueObject, NumberValueObject, StringValueObject } from './engine/value-object/primitive-object';
 export { functionArray } from './functions/array/function-map';
 export { FUNCTION_NAMES_ARRAY } from './functions/array/function-names';
@@ -167,3 +167,4 @@ export { ENGINE_FORMULA_CYCLE_REFERENCE_COUNT, ENGINE_FORMULA_PLUGIN_CONFIG_KEY,
 
 export { generateRandomDependencyTreeId } from './engine/dependency/formula-dependency';
 export { DependencyManagerBaseService } from './services/dependency-manager.service';
+export { LambdaValueObjectObject } from './engine/value-object/lambda-value-object';
diff --git a/packages/facade/src/apis/__tests__/facade.spec.ts b/packages/facade/src/apis/__tests__/facade.spec.ts
@@ -15,6 +15,7 @@
  */
 
 import type { ICellData, Injector, Nullable } from '@univerjs/core';
+import type { LambdaValueObjectObject, PrimitiveValueType } from '@univerjs/engine-formula';
 import type {
     ColumnHeaderLayout,
     RenderComponentType,
@@ -24,8 +25,8 @@ import type {
     SpreadsheetRowHeader,
 } from '@univerjs/engine-render';
 import type { FUniver } from '../everything';
-import { ICommandService, IUniverInstanceService } from '@univerjs/core';
 
+import { ICommandService, IUniverInstanceService } from '@univerjs/core';
 import { RegisterFunctionMutation, SetFormulaCalculationStartMutation } from '@univerjs/engine-formula';
 import { IRenderManagerService } from '@univerjs/engine-render';
 import { SetRangeValuesCommand, SetRangeValuesMutation, SetStyleCommand } from '@univerjs/sheets';
@@ -178,11 +179,22 @@ describe('Test FUniver', () => {
 
     it('Function registerFunction', () => {
         const funcionName = 'CUSTOMSUM';
+
         const functionsDisposable = univerAPI.registerFunction({
             calculate: [
                 [function (...variants) {
                     let sum = 0;
 
+                    const last = variants[variants.length - 1] as LambdaValueObjectObject;
+
+                    if (last.isLambda()) {
+                        variants.pop();
+
+                        const variantsList = variants.map((variant) => Array.isArray(variant) ? variant[0][0] : variant) as PrimitiveValueType[];
+
+                        sum += +last.executeCustom(...variantsList).getValue();
+                    }
+
                     for (const variant of variants) {
                         sum += Number(variant) || 0;
                     }

diff --git a/packages/sheets-formula/src/facade/f-formula.ts b/packages/sheets-formula/src/facade/f-formula.ts
@@ -58,6 +58,28 @@ export interface IFFormulaSheetsMixin {
      * univerAPI.getActiveWorkbook().getActiveSheet().getRange('A2').setValue({ f: '=DISCOUNT(A1, 20)' });
      * // A2 will display: 80
      * ```
+     * @example
+     * ```typescript
+     * // Registered formulas support lambda functions
+     * univerAPI.getFormula().registerFunction('CUSTOMSUM', (...variants) => {
+     *      let sum = 0;
+     *
+     *      const last = variants[variants.length - 1];
+     *      if (last.isLambda && last.isLambda()) {
+     *          variants.pop();
+     *
+     *          const variantsList = variants.map((variant) => Array.isArray(variant) ? variant[0][0]: variant);
+     *
+     *          sum += last.executeCustom(...variantsList).getValue();
+     *      }
+     *
+     *      for (const variant of variants) {
+     *          sum += Number(variant) || 0;
+     *      }
+     *
+     *      return sum;
+     * }, 'Adds its arguments');
+     * ```
      */
     registerFunction(name: string, func: IRegisterFunction, description?: string): IDisposable;
 

diff --git a/packages/sheets-formula/src/services/register-function.service.ts b/packages/sheets-formula/src/services/register-function.service.ts
@@ -15,19 +15,19 @@
  */
 
 import type { IDisposable, ILocales } from '@univerjs/core';
-import type { FormulaFunctionValueType, IFunctionInfo } from '@univerjs/engine-formula';
+import type { FormulaFunctionResultValueType, FormulaFunctionValueType, IFunctionInfo } from '@univerjs/engine-formula';
 import { createIdentifier, Disposable, DisposableCollection, Inject, LocaleService, Optional, toDisposable } from '@univerjs/core';
 import { AsyncCustomFunction, CustomFunction, FunctionType, IFunctionService } from '@univerjs/engine-formula';
 import { IDescriptionService } from './description.service';
 import { IRemoteRegisterFunctionService } from './remote/remote-register-function.service';
 
 export type IRegisterFunction = (
     ...arg: Array<FormulaFunctionValueType>
-) => FormulaFunctionValueType;
+) => FormulaFunctionResultValueType;
 
 export type IRegisterAsyncFunction = (
     ...arg: Array<FormulaFunctionValueType>
-) => Promise<FormulaFunctionValueType>;
+) => Promise<FormulaFunctionResultValueType>;
 
 // [[Function, FunctionName, Description]]
 export type IRegisterFunctionList = [[IRegisterFunction, string, string?]];