feat(custom-function): improve custom function registration to correc…

…tly know when a function updates GitOrigin-RevId: d4c13c5f8eb71eac73184e55f9a0fb5ea382c6a1
plasmicapp · Dec 18, 2024 · 40d5b8a · 40d5b8a
1 parent 34bf9b3
commit 40d5b8a
Show file tree

Hide file tree

Showing 10 changed files with 170 additions and 24 deletions.
diff --git a/platform/wab/src/wab/client/components/canvas/canvas-rendering.ts b/platform/wab/src/wab/client/components/canvas/canvas-rendering.ts
@@ -221,6 +221,7 @@ import {
   ComponentDataQuery,
   CompositeExpr,
   CustomCode,
+  CustomFunctionExpr,
   DataSourceOpExpr,
   EventHandler,
   Expr,
@@ -2075,6 +2076,9 @@ function computeTplComponentArgs(
           ctx.viewCtx.canvasCtx.win()
         )
       )
+      .when(CustomFunctionExpr, (_expr) =>
+        unexpected(`Cannot evaluate CustomFunctionExpr as a component arg`)
+      )
       .result();
   };
 
@@ -2632,6 +2636,7 @@ function evalTagAttrExprToString(
         TplRef,
         QueryInvalidationExpr,
         CompositeExpr,
+        CustomFunctionExpr,
       ],
       (_) => {
         assert(false, "Unexpected expr type");

diff --git a/platform/wab/src/wab/client/components/sidebar-tabs/PropEditorRow.tsx b/platform/wab/src/wab/client/components/sidebar-tabs/PropEditorRow.tsx
@@ -110,6 +110,7 @@ import {
   CollectionExpr,
   CompositeExpr,
   CustomCode,
+  CustomFunctionExpr,
   DataSourceOpExpr,
   EventHandler,
   Expr,
@@ -304,6 +305,7 @@ export function extractLitFromMaybeRenderable(
         CompositeExpr,
         MapExpr,
         CollectionExpr,
+        CustomFunctionExpr,
       ],
       (_expr): [any, boolean] => [_expr, true]
     )

diff --git a/platform/wab/src/wab/client/studio-ctx/StudioCtx.tsx b/platform/wab/src/wab/client/studio-ctx/StudioCtx.tsx
@@ -3189,6 +3189,10 @@ export class StudioCtx extends WithDbCtx {
     ];
   }
 
+  getRegisteredFunctionsMap() {
+    return this._ccRegistry.getRegisteredFunctionsMap();
+  }
+
   getRegisteredLibraries() {
     return [
       ...(swallow(() => this.codeComponentsRegistry.getRegisteredLibraries()) ??

diff --git a/platform/wab/src/wab/shared/cached-selectors.ts b/platform/wab/src/wab/shared/cached-selectors.ts
@@ -57,6 +57,7 @@ import {
   getNonVariantParams,
   getParamNames,
   isCodeComponent,
+  isPageComponent,
   isPlumeComponent,
   tryGetVariantGroupValueFromArg,
 } from "@/wab/shared/core/components";
@@ -536,6 +537,29 @@ export const allCustomFunctions = maybeComputedFn(function allCustomFunctions(
   return customFunctions;
 });
 
+export const findCustomFunctionUsages = maybeComputedFn(
+  function findCustomFunctionUsages(site: Site) {
+    return withoutNils(
+      site.components.map((component) => {
+        if (
+          !isPageComponent(component) ||
+          component.serverQueries.length === 0
+        ) {
+          return undefined;
+        }
+        return {
+          ownerComponent: component,
+          customFunctions: withoutNils(
+            component.serverQueries.map((serverQuery) => {
+              return serverQuery.op?.func;
+            })
+          ),
+        };
+      })
+    );
+  }
+);
+
 export const allCodeLibraries = maybeComputedFn(function allCodeLibraries(
   rootSite: Site
 ) {

diff --git a/platform/wab/src/wab/shared/code-components/code-components.ts b/platform/wab/src/wab/shared/code-components/code-components.ts
@@ -13,6 +13,7 @@ import {
   componentToReferencers,
   componentToTplComponents,
   computedProjectFlags,
+  findCustomFunctionUsages,
   flattenComponent,
 } from "@/wab/shared/cached-selectors";
 import {
@@ -2389,13 +2390,29 @@ export function registeredFunctionId(r: CustomFunctionRegistration) {
 }
 
 export function createCustomFunctionFromRegistration(
-  functionReg: CustomFunctionRegistration
+  functionReg: CustomFunctionRegistration,
+  existingFunction?: CustomFunction
 ) {
+  const existingParams = existingFunction?.params ?? [];
   return new CustomFunction({
     defaultExport: functionReg.meta.isDefaultExport ?? false,
     importName: functionReg.meta.name,
     importPath: functionReg.meta.importPath,
     namespace: functionReg.meta.namespace ?? null,
+    params:
+      functionReg.meta.params?.map((paramReg: string | BaseParam<any>) => {
+        const name = isString(paramReg) ? paramReg : paramReg.name;
+        const argType = isString(paramReg)
+          ? typeFactory.text()
+          : isArray(paramReg.type)
+          ? typeFactory.any()
+          : (convertTsToWabType(paramReg.type ?? "string") as ArgType["type"]);
+        const existingParam = existingParams.find((p) => p.argName === name);
+        if (existingParam && existingParam.type.name === argType.name) {
+          return existingParam;
+        }
+        return typeFactory.arg(name, argType);
+      }) ?? [],
   });
 }
 
@@ -4558,13 +4575,13 @@ async function upsertRegisteredFunctions(
           const updateableFields: Omit<
             CustomFunction,
             "importName" | "namespace" | "typeTag" | "uid"
-          > = pick(createCustomFunctionFromRegistration(functionReg), [
-            "defaultExport",
-            "importPath",
-          ]);
+          > = pick(
+            createCustomFunctionFromRegistration(functionReg, existing),
+            ["defaultExport", "importPath", "params"]
+          );
           if (
             Object.entries(updateableFields).some(
-              ([key, value]) => value !== existing[key]
+              ([key, value]) => !isEqual(value, existing[key])
             )
           ) {
             updatedFunctionRegs.push(functionReg);
@@ -4633,14 +4650,29 @@ async function upsertRegisteredFunctions(
           functionNamespaces.add(functionReg.meta.namespace);
         }
       }
+      const customFunctionUsages = findCustomFunctionUsages(site);
+
+      ctx.observeComponents(
+        customFunctionUsages
+          .filter((usage) => {
+            return usage.customFunctions.some(
+              (fn) =>
+                removedFunctions.has(fn) ||
+                updatedFunctionRegs.some(
+                  (reg) => registeredFunctionId(reg) === customFunctionId(fn)
+                )
+            );
+          })
+          .map((usage) => usage.ownerComponent)
+      );
 
       if (
         newFunctionRegs.length > 0 ||
         updatedFunctionRegs.length > 0 ||
         removedFunctions.size > 0
       ) {
         run(
-          await ctx.change<never>(
+          await ctx.change(
             ({ success: changeSuccess }) => {
               const newFunctions: CustomFunction[] = [];
               const updatedFunctions: CustomFunction[] = [];
@@ -4659,10 +4691,32 @@ async function upsertRegisteredFunctions(
                 const updateableFields: Omit<
                   CustomFunction,
                   "importName" | "namespace" | "typeTag" | "uid"
-                > = pick(createCustomFunctionFromRegistration(functionReg), [
-                  "defaultExport",
-                  "importPath",
-                ]);
+                > = pick(
+                  createCustomFunctionFromRegistration(functionReg, existing),
+                  ["defaultExport", "importPath", "params"]
+                );
+                customFunctionUsages.forEach((usage) => {
+                  usage.ownerComponent.serverQueries
+                    .filter((serverQuery) => {
+                      return serverQuery.op?.func === existing;
+                    })
+                    .map((serverQuery) => {
+                      removeWhere(
+                        serverQuery.op!.args,
+                        (arg) =>
+                          !updateableFields.params.find(
+                            (param) => param === arg.argType
+                          )
+                      );
+                    });
+                  removeWhere(
+                    usage.ownerComponent.serverQueries,
+                    (serverQuery) =>
+                      !!serverQuery.op?.func &&
+                      removedFunctions.has(serverQuery.op.func)
+                  );
+                });
+
                 Object.assign(existing, updateableFields);
                 updatedFunctions.push(existing);
               }
@@ -4674,6 +4728,14 @@ async function upsertRegisteredFunctions(
               removeWhere(site.customFunctions, (customFunction) =>
                 removedFunctions.has(customFunction)
               );
+              customFunctionUsages.forEach((usage) => {
+                removeWhere(
+                  usage.ownerComponent.serverQueries,
+                  (serverQuery) =>
+                    !!serverQuery.op?.func &&
+                    removedFunctions.has(serverQuery.op.func)
+                );
+              });
 
               fns.onUpdatedCustomFunctions?.({
                 newFunctions,

diff --git a/platform/wab/src/wab/shared/core/components.ts b/platform/wab/src/wab/shared/core/components.ts
@@ -342,6 +342,7 @@ export function mkComponent(obj: {
     metadata: {},
     states,
     dataQueries: [],
+    serverQueries: [],
     figmaMappings: obj.figmaMappings ?? [],
     alwaysAutoName: obj.alwaysAutoName ?? false,
     trapsFocus: obj.trapsFocus ?? false,
@@ -729,6 +730,7 @@ export function cloneComponent(
       oldToNewComponentQuery.set(componentDataQuery, cloned);
       return cloned;
     }),
+    serverQueries: fromComponent.serverQueries,
     figmaMappings: fromComponent.figmaMappings.map(
       (c) => new FigmaComponentMapping({ ...c })
     ),

diff --git a/platform/wab/src/wab/shared/core/exprs.ts b/platform/wab/src/wab/shared/core/exprs.ts
@@ -3,6 +3,7 @@ import {
   Component,
   CompositeExpr,
   CustomCode,
+  CustomFunctionExpr,
   DataSourceOpExpr,
   DataSourceTemplate,
   EventHandler,
@@ -102,7 +103,7 @@ import { pathToString } from "@/wab/shared/eval/expression-parser";
 import { maybeComputedFn } from "@/wab/shared/mobx-util";
 import { typeDisplayName } from "@/wab/shared/model/model-util";
 import { maybeConvertToIife } from "@/wab/shared/parser-utils";
-import L, { escapeRegExp, isString, mapValues, set } from "lodash";
+import L, { escapeRegExp, groupBy, isString, mapValues, set } from "lodash";
 
 export interface ExprCtx {
   component: Component | null;
@@ -162,6 +163,7 @@ export const summarizeExpr = (expr: Expr, exprCtx: ExprCtx): string =>
     )
     .when(QueryInvalidationExpr, (_expr) => `(query invalidations)`)
     .when(CompositeExpr, (_expr) => `(composite value)`)
+    .when(CustomFunctionExpr, (_expr) => `(custom function)`)
     .result();
 
 // Deep copy is necessary, since with shallow copies, the user could potentially
@@ -340,6 +342,14 @@ export function clone(_expr: Expr): Expr {
           ),
         })
     )
+    .when(
+      CustomFunctionExpr,
+      (expr) =>
+        new CustomFunctionExpr({
+          func: expr.func,
+          args: expr.args.map((arg) => clone(arg)),
+        })
+    )
     .result();
 }
 
@@ -740,6 +750,22 @@ const _asCode = maybeComputedFn(
   `.trim()
         );
       })
+      .when(CustomFunctionExpr, (expr) => {
+        const { func, args } = expr;
+        const argsMap = groupBy(args, (arg) => arg.argType.argName);
+        const orderedArgs =
+          func.params.map((param) => {
+            if (argsMap[param.argName]) {
+              return getRawCode(argsMap[param.argName][0].expr, exprCtx);
+            }
+            return undefined;
+          }) ?? [];
+        return code(
+          `$$${expr.func.namespace ? `.${expr.func.namespace}` : ""}.${
+            expr.func.importName
+          }(${orderedArgs.join(",")})`
+        );
+      })
       .result()
 );
 

diff --git a/platform/wab/src/wab/shared/core/sites.ts b/platform/wab/src/wab/shared/core/sites.ts
@@ -113,6 +113,7 @@ import {
   CompositeExpr,
   CustomCode,
   CustomFunction,
+  CustomFunctionExpr,
   DataSourceOpExpr,
   EventHandler,
   Expr,
@@ -453,6 +454,7 @@ export function cloneCustomFunction(
     importName: customFunction.importName,
     importPath: customFunction.importPath,
     namespace: customFunction.namespace,
+    params: [],
   });
 }
 
@@ -868,6 +870,9 @@ export function cloneSite(fromSite: Site) {
       .when(FunctionExpr, (functionExpr) =>
         fixGlobalRefForExpr(functionExpr.bodyExpr)
       )
+      .when(CustomFunctionExpr, (customFunctionExpr) =>
+        customFunctionExpr.args.forEach((arg) => fixGlobalRefForExpr(arg))
+      )
       .when([RenderExpr, VirtualRenderExpr], () => {})
       .result();
   };

diff --git a/platform/wab/src/wab/shared/core/tpls.ts b/platform/wab/src/wab/shared/core/tpls.ts
@@ -28,6 +28,7 @@ import {
   isKnownCollectionExpr,
   isKnownCompositeExpr,
   isKnownCustomCode,
+  isKnownCustomFunctionExpr,
   isKnownDataSourceOpExpr,
   isKnownEventHandler,
   isKnownExpr,
@@ -2360,6 +2361,10 @@ export function pushExprs(exprs: Expr[], expr: Expr | null | undefined) {
     Object.values(expr.substitutions).forEach((subExpr) =>
       pushExprs(exprs, subExpr)
     );
+  } else if (isKnownCustomFunctionExpr(expr)) {
+    for (const arg of expr.args) {
+      pushExprs(exprs, arg);
+    }
   }
 }
 
@@ -2440,6 +2445,12 @@ export function findExprsInComponent(component: Component) {
     }
   }
 
+  for (const query of component.serverQueries) {
+    if (query.op) {
+      pushExprs(componentExprs, query.op);
+    }
+  }
+
   const refs: ExprReference[] = componentExprs.map((expr) => ({
     expr,
   }));