feat: allow re-direction of ReferenceExpr (#521)

.projenrc.ts

@@ -21,7 +21,7 @@ class GitHooksPreCommitComponent extends TextFile {
   }
 }
 
-const MIN_CDK_VERSION = "2.28.1";
+const MIN_CDK_VERSION = "2.43.1";
 
 /**
  * Projen does not currently support a way to set `*` for deerDependency versions.
@@ -288,8 +288,8 @@ const testApp = project.addTask("test:app", {
 
 project.testTask.spawn(typeCheck);
 project.testTask.spawn(testFast);
-project.testTask.spawn(testRuntime);
 project.testTask.spawn(testApp);
+project.testTask.spawn(testRuntime);
 project.testTask.spawn(project.tasks.tryFind("eslint")!);
 
 project.addPackageIgnore("/test-app");

src/event-bridge/event-bus.ts

@@ -14,7 +14,6 @@ import {
   PropAssignExpr,
   StringLiteralExpr,
 } from "../expression";
-import { Function } from "../function";
 import { NativePreWarmContext, PrewarmClients } from "../function-prewarm";
 import {
   isArrayLiteralExpr,
@@ -445,8 +444,8 @@ abstract class EventBusBase<in Evnt extends Event, OutEvnt extends Evnt = Evnt>
         },
       },
       native: {
-        bind: (context: Function<any, any>) => {
-          this.resource.grantPutEventsTo(context.resource);
+        bind: (context) => {
+          this.resource.grantPutEventsTo(context);
         },
         preWarm: (prewarmContext: NativePreWarmContext) => {
           prewarmContext.getOrInit(PrewarmClients.EventBridge);

src/function.ts

@@ -30,10 +30,10 @@ import { Construct } from "constructs";
 import esbuild from "esbuild";
 import ts from "typescript";
 import { ApiGatewayVtlIntegration } from "./api";
-import type { AppSyncVtlIntegration } from "./appsync";
+import { AppsyncResolver, AppSyncVtlIntegration } from "./appsync";
 import { ASL, ASLGraph } from "./asl";
 import { BindFunctionName, RegisterFunctionName } from "./compile";
-import { IntegrationInvocation } from "./declaration";
+import { FunctionLike, IntegrationInvocation } from "./declaration";
 import { ErrorCodes, formatErrorMessage, SynthError } from "./error-code";
 import {
   IEventBus,
@@ -309,8 +309,8 @@ abstract class FunctionBase<in Payload, Out>
       /**
        * Wire up permissions for this function to be called by the calling function
        */
-      bind: (context: Function<any, any>) => {
-        this.resource.grantInvoke(context.resource);
+      bind: (context) => {
+        this.resource.grantInvoke(context);
       },
       /**
        * Code that runs once per lambda invocation
@@ -706,10 +706,7 @@ export class Function<
           : {}),
         functionless_infer: Lazy.string({
           produce: () => {
-            // retrieve and bind all found native integrations. Will fail if the integration does not support native integration.
-            findAllIntegrations().forEach(({ integration, args }) => {
-              new IntegrationImpl(integration).native.bind(this, args);
-            });
+            inferIamPolicies(ast, _resource);
 
             return "DONE";
           },
@@ -721,7 +718,7 @@ export class Function<
 
     // retrieve all found native integrations. Will fail if the integration does not support native integration.
     // TODO: move this logic into the synthesis phase, see https://github.com/functionless/functionless/issues/476
-    const nativeIntegrationsPrewarm = findAllIntegrations().flatMap(
+    const nativeIntegrationsPrewarm = findAllIntegrations(ast).flatMap(
       ({ integration }) => {
         const native = new IntegrationImpl(integration).native;
         if (native.preWarm) {
@@ -785,19 +782,29 @@ export class Function<
               ),
             ],
     });
-
-    function findAllIntegrations() {
-      return findDeepIntegrations(ast).map(
-        (i) =>
-          <IntegrationInvocation>{
-            args: i.args,
-            integration: i.expr.ref(),
-          }
-      );
-    }
   }
 }
 
+export function inferIamPolicies(
+  decl: FunctionLike<AnyFunction>,
+  func: aws_lambda.IFunction
+) {
+  findAllIntegrations(decl).forEach(({ integration, args }) => {
+    const native = new IntegrationImpl(integration).native;
+    native.bind(func, args);
+  });
+}
+
+export function findAllIntegrations(decl: FunctionLike<AnyFunction>) {
+  return findDeepIntegrations(decl).map(
+    (i) =>
+      <IntegrationInvocation>{
+        args: i.args,
+        integration: i.expr.ref(),
+      }
+  );
+}
+
 /**
  * A {@link Function} which wraps a CDK function.
  *
@@ -947,7 +954,7 @@ export interface NativeIntegration<Func extends AnyFunction> {
    * @param context - The function invoking this function.
    * @param args - The functionless encoded AST form of the arguments passed to the integration.
    */
-  bind: (context: Function<any, any>, args: Expr[]) => void;
+  bind: (context: aws_lambda.IFunction, args: Expr[]) => void;
   /**
    * @param args The arguments passed to the integration function by the user.
    * @param preWarmContext contains singleton instances of client and other objects initialized outside of the native

src/integration.ts

@@ -18,7 +18,7 @@ import {
   isVariableDecl,
 } from "./guards";
 import { FunctionlessNode } from "./node";
-import { reflect } from "./reflect";
+import { reflect, resolveSubstitution } from "./reflect";
 import { AnyFunction, evalToConstant, isAnyFunction } from "./util";
 import { forEachChild } from "./visit";
 import { VTL } from "./vtl";
@@ -389,7 +389,9 @@ export function tryFindIntegration(
  * @returns a list of all the {@link Integration}s that the {@link node} could evaluate to.
  */
 export function tryFindIntegrations(node: FunctionlessNode): Integration[] {
-  return tryResolveReferences(node, undefined).filter(isIntegration);
+  return tryResolveReferences(node, undefined).filter((i) => {
+    return isIntegration(i);
+  });
 }
 
 /**
@@ -420,7 +422,7 @@ export function tryResolveReferences(
     return tryResolveReferences(node.parent, node.initializer).flatMap(
       (value) => {
         if (isIdentifier(node.name)) {
-          return [value?.[node.name.name]];
+          return [resolveSubstitution(value?.[node.name.name])];
         } else {
           throw new Error("should be impossible");
         }
@@ -438,7 +440,7 @@ export function tryResolveReferences(
         ? node.name.name
         : evalToConstant(node.element)?.constant;
       if (key !== undefined) {
-        return [(<any>expr)?.[key]];
+        return [resolveSubstitution((<any>expr)?.[key])];
       }
       return [];
     });

src/queue.ts

@@ -381,7 +381,7 @@ abstract class BaseQueue<Message>
     >({
       kind: "AWS.SQS.SendMessage",
       native: {
-        bind: (func) => this.resource.grantSendMessages(func.resource),
+        bind: (func) => this.resource.grantSendMessages(func),
         preWarm: (context) => context.getOrInit(SQSClient),
         call: async ([input], context) => {
           const sqs = context.getOrInit(SQSClient);
@@ -529,7 +529,7 @@ abstract class BaseQueue<Message>
     >({
       kind: "AWS.SQS.SendMessageBatch",
       native: {
-        bind: (func) => this.resource.grantSendMessages(func.resource),
+        bind: (func) => this.resource.grantSendMessages(func),
         preWarm: (context) => context.getOrInit(SQSClient),
         call: async ([input], context) => {
           const sqs: AWS.SQS = context.getOrInit<AWS.SQS>(SQSClient);
@@ -661,7 +661,7 @@ abstract class BaseQueue<Message>
       kind: "AWS.SQS.ReceiveMessage",
       // asl: (call, context) => {},
       native: {
-        bind: (func) => this.resource.grantConsumeMessages(func.resource),
+        bind: (func) => this.resource.grantConsumeMessages(func),
         preWarm: (context) => context.getOrInit(SQSClient),
         call: async ([input], context) => {
           const sqs: AWS.SQS = context.getOrInit(SQSClient);
@@ -807,7 +807,7 @@ abstract class BaseQueue<Message>
     this.purge = makeIntegration<"AWS.SQS.PurgeQueue", () => Promise<{}>>({
       kind: "AWS.SQS.PurgeQueue",
       native: {
-        bind: (func) => this.resource.grantPurge(func.resource),
+        bind: (func) => this.resource.grantPurge(func),
         preWarm: (context) => context.getOrInit(SQSClient),
         call: async ([], context) => {
           const updatedQueueUrl =

src/reflect.ts

@@ -6,12 +6,13 @@ import {
 } from "./declaration";
 import { Err } from "./error";
 import { ErrorCodes, SynthError } from "./error-code";
-import { isFunctionLike, isErr, isNewExpr } from "./guards";
+import { ReferenceExpr } from "./expression";
+import { isFunctionLike, isErr, isNewExpr, isReferenceExpr } from "./guards";
 import { tryResolveReferences } from "./integration";
 import type { FunctionlessNode } from "./node";
 import { parseSExpr } from "./s-expression";
 import { AnyAsyncFunction, AnyFunction } from "./util";
-import { forEachChild } from "./visit";
+import { forEachChild, visitEachChild } from "./visit";
 
 const Global: any = global;
 
@@ -67,14 +68,59 @@ export function reflect<F extends AnyFunction | AnyAsyncFunction>(
     if (!reflectCache.has(astCallback)) {
       reflectCache.set(
         astCallback,
-        parseSExpr(astCallback()) as FunctionLike<F> | Err
+        parseAst(astCallback()) as FunctionLike<F> | Err
       );
     }
     return reflectCache.get(astCallback) as FunctionLike<F> | Err | undefined;
   }
   return undefined;
 }
 
+function parseAst(expr: any): any {
+  const decl = parseSExpr(expr);
+  return visitEachChild(decl, function visit(node): any {
+    if (isReferenceExpr(node)) {
+      return new ReferenceExpr(
+        node.span,
+        node.name,
+        () => resolveSubstitution(node.ref()),
+        node.id,
+        node.referenceId
+      );
+    } else {
+      return visitEachChild(node, visit);
+    }
+  });
+}
+
+const substitutionSymbol = Symbol.for("functionless::substitutions");
+
+/**
+ * A global WeakMap for re-directing ReferenceExpr references.
+ *
+ * This hack is to support the fl-exp inverted model.
+ */
+const substitutions: WeakMap<any, any> = (Global[substitutionSymbol] =
+  Global[substitutionSymbol] ?? new WeakMap());
+
+/**
+ * Register a substitution. Any ReferenceExpr pointing to {@link from}
+ * will be substituted with {@link to} when ReferenceExpr is resolved.
+ */
+export function registerSubstitution(from: any, to: any) {
+  substitutions.set(from, to);
+}
+
+/**
+ * Resolve a value, potentially substituting it if a substitution has
+ * been registered for the {@link from} value. If no substitution
+ * has been registered, then the original value in the ReferenceExpr
+ * is returned.
+ */
+export function resolveSubstitution(from: any): any {
+  return substitutions.has(from) ? substitutions.get(from) : from;
+}
+
 export interface BoundFunctionComponents<F extends AnyFunction = AnyFunction> {
   /**
    * The target function