feat: remove logic in compile.ts to flatten references to expressions

src/api.ts

@@ -39,7 +39,7 @@ import {
   isSpreadAssignExpr,
   isFunctionLike,
 } from "./guards";
-import { Integration, IntegrationImpl, isIntegration } from "./integration";
+import { IntegrationImpl, tryFindIntegration } from "./integration";
 import { validateFunctionLike } from "./reflect";
 import { Stmt } from "./statement";
 import { AnyFunction, singletonConstruct } from "./util";
@@ -732,12 +732,13 @@ export class APIGatewayVTL extends VTL {
         return this.exprToJson(expr.expr);
       }
     } else if (isCallExpr(expr)) {
-      if (isReferenceExpr(expr.expr) || isThisExpr(expr.expr)) {
+      const integration = tryFindIntegration(expr.expr);
+      if (integration) {
+        const serviceCall = new IntegrationImpl(integration);
+        return this.integrate(serviceCall, expr);
+      } else if (isReferenceExpr(expr.expr) || isThisExpr(expr.expr)) {
         const ref = expr.expr.ref();
-        if (isIntegration<Integration>(ref)) {
-          const serviceCall = new IntegrationImpl(ref);
-          return this.integrate(serviceCall, expr);
-        } else if (ref === Number) {
+        if (ref === Number) {
           return this.exprToJson(expr.args[0]);
         } else {
           throw new SynthError(

src/appsync.ts

@@ -466,14 +466,18 @@ function synthesizeFunctions(api: appsync.GraphqlApi, decl: FunctionLike) {
         // we find the range of nodes to hoist so that we avoid visiting the middle nodes.
         // The start node is the first node in the integration pattern (integ, await, or promise)
         // The end is always the integration.
+
         const end = getIntegrationExprFromIntegrationCallPattern(node);
 
-        const updatedChild = visitSpecificChildren(node, [end], (expr) =>
-          normalizeAST(expr, hoist)
-        );
-        // when we find an integration call,
-        // if it is nested, hoist it up (create variable, add above, replace expr with variable)
-        return hoist && doHoist(node) ? hoist(updatedChild) : updatedChild;
+        if (end) {
+          const updatedChild = visitSpecificChildren(node, [end], (expr) =>
+            normalizeAST(expr, hoist)
+          );
+
+          // when we find an integration call,
+          // if it is nested, hoist it up (create variable, add above, replace expr with variable)
+          return hoist && doHoist(node) ? hoist(updatedChild) : updatedChild;
+        }
       } else if (isVariableStmt(node) && node.declList.decls.length > 1) {
         /**
          * Flatten variable declarations into multiple variable statements.

src/asl.ts

@@ -110,10 +110,10 @@ import {
   isQuasiString,
 } from "./guards";
 import {
-  Integration,
   IntegrationImpl,
   isIntegration,
   isIntegrationCallPattern,
+  tryFindIntegration,
 } from "./integration";
 import { FunctionlessNode } from "./node";
 import {
@@ -1647,75 +1647,66 @@ export class ASL {
         };
       });
     } else if (isCallExpr(expr)) {
-      if (isReferenceExpr(expr.expr)) {
-        const ref = expr.expr.ref();
-        if (isIntegration<Integration>(ref)) {
-          const serviceCall = new IntegrationImpl(ref);
-          const integStates = serviceCall.asl(expr, this);
+      const integration = tryFindIntegration(expr.expr);
+      if (integration) {
+        const serviceCall = new IntegrationImpl(integration);
+        const integStates = serviceCall.asl(expr, this);
 
+        if (
+          ASLGraph.isLiteralValue(integStates) ||
+          ASLGraph.isJsonPath(integStates)
+        ) {
+          return integStates;
+        }
+
+        const updateState = (state: ASLGraph.NodeState): ASLGraph.NodeState => {
+          const throwOrPass = this.throw(expr);
           if (
-            ASLGraph.isLiteralValue(integStates) ||
-            ASLGraph.isJsonPath(integStates)
+            throwOrPass?.Next &&
+            (isTaskState(state) ||
+              isMapTaskState(state) ||
+              isParallelTaskState(state))
           ) {
-            return integStates;
+            return {
+              ...state,
+              Catch: [
+                {
+                  ErrorEquals: ["States.ALL"],
+                  Next: throwOrPass.Next,
+                  ResultPath: throwOrPass.ResultPath,
+                },
+              ],
+            };
+          } else {
+            return state;
           }
+        };
 
-          const updateState = (
-            state: ASLGraph.NodeState
-          ): ASLGraph.NodeState => {
-            const throwOrPass = this.throw(expr);
-            if (
-              throwOrPass?.Next &&
-              (isTaskState(state) ||
-                isMapTaskState(state) ||
-                isParallelTaskState(state))
-            ) {
-              return {
-                ...state,
-                Catch: [
-                  {
-                    ErrorEquals: ["States.ALL"],
-                    Next: throwOrPass.Next,
-                    ResultPath: throwOrPass.ResultPath,
-                  },
-                ],
-              };
-            } else {
-              return state;
-            }
-          };
-
-          const updateStates = (
-            states: ASLGraph.NodeState | ASLGraph.SubState
-          ): ASLGraph.NodeState | ASLGraph.SubState => {
-            return ASLGraph.isSubState(states)
-              ? {
-                  ...states,
-                  states: Object.fromEntries(
-                    Object.entries(states.states ?? {}).map(
-                      ([stateName, state]) => {
-                        if (ASLGraph.isSubState(state)) {
-                          return [stateName, updateStates(state)];
-                        } else {
-                          return [stateName, updateState(state)];
-                        }
+        const updateStates = (
+          states: ASLGraph.NodeState | ASLGraph.SubState
+        ): ASLGraph.NodeState | ASLGraph.SubState => {
+          return ASLGraph.isSubState(states)
+            ? {
+                ...states,
+                states: Object.fromEntries(
+                  Object.entries(states.states ?? {}).map(
+                    ([stateName, state]) => {
+                      if (ASLGraph.isSubState(state)) {
+                        return [stateName, updateStates(state)];
+                      } else {
+                        return [stateName, updateState(state)];
                       }
-                    )
-                  ),
-                }
-              : updateState(states);
-          };
+                    }
+                  )
+                ),
+              }
+            : updateState(states);
+        };
 
-          return {
-            ...integStates,
-            ...updateStates(integStates),
-          };
-        } else {
-          throw new SynthError(
-            ErrorCodes.Unexpected_Error,
-            "Called references are expected to be an integration."
-          );
-        }
+        return {
+          ...integStates,
+          ...updateStates(integStates),
+        };
       } else if (isMapOrForEach(expr)) {
         const throwTransition = this.throw(expr);
 

src/checker.ts

@@ -402,7 +402,8 @@ export function makeFunctionlessChecker(
         return updatedSymbol ? isSymbolOutOfScope(updatedSymbol, scope) : false;
       } else if (
         ts.isVariableDeclaration(symbol.valueDeclaration) ||
-        ts.isClassDeclaration(symbol.valueDeclaration)
+        ts.isClassDeclaration(symbol.valueDeclaration) ||
+        ts.isParameter(symbol.valueDeclaration)
       ) {
         return !hasParent(symbol.valueDeclaration, scope);
       } else if (ts.isBindingElement(symbol.valueDeclaration)) {

src/compile.ts

@@ -400,18 +400,6 @@ export function compile(
           } else if (node.text === "null") {
             return newExpr(NodeKind.NullLiteralExpr, []);
           }
-          if (checker.isIntegrationNode(node)) {
-            // if this is a reference to a Table or Lambda, retain it
-            const _ref = checker.getOutOfScopeValueNode(node, scope);
-            if (_ref) {
-              return ref(_ref);
-            } else {
-              throw new SynthError(
-                ErrorCodes.Unable_to_find_reference_out_of_application_function,
-                `Unable to find reference out of application function: ${node.getText()}`
-              );
-            }
-          }
 
           /**
            * If the identifier is not within the closure, we attempt to enclose the reference in its own closure.
@@ -423,28 +411,13 @@ export function compile(
            * return { value: () => val };
            */
           if (checker.isIdentifierOutOfScope(node, scope)) {
-            const _ref = checker.getOutOfScopeValueNode(node, scope);
-            if (_ref) {
-              return ref(_ref);
-            }
+            return ref(node);
           }
 
           return newExpr(NodeKind.Identifier, [
             ts.factory.createStringLiteral(node.text),
           ]);
         } else if (ts.isPropertyAccessExpression(node)) {
-          if (checker.isIntegrationNode(node)) {
-            // if this is a reference to a Table or Lambda, retain it
-            const _ref = checker.getOutOfScopeValueNode(node, scope);
-            if (_ref) {
-              return ref(_ref);
-            } else {
-              throw new SynthError(
-                ErrorCodes.Unable_to_find_reference_out_of_application_function,
-                `Unable to find reference out of application function: ${node.getText()}`
-              );
-            }
-          }
           return newExpr(NodeKind.PropAccessExpr, [
             toExpr(node.expression, scope),
             toExpr(node.name, scope),

src/expression.ts

@@ -187,11 +187,13 @@ export class Argument extends BaseExpr<NodeKind.Argument, CallExpr | NewExpr> {
   }
 }
 
-export class CallExpr extends BaseExpr<NodeKind.CallExpr> {
-  constructor(
-    readonly expr: Expr | SuperKeyword | ImportKeyword,
-    readonly args: Argument[]
-  ) {
+export class CallExpr<
+  E extends Expr | SuperKeyword | ImportKeyword =
+    | Expr
+    | SuperKeyword
+    | ImportKeyword
+> extends BaseExpr<NodeKind.CallExpr> {
+  constructor(readonly expr: E, readonly args: Argument[]) {
     super(NodeKind.CallExpr, arguments);
   }
 }

src/function.ts

@@ -31,7 +31,7 @@ import { ApiGatewayVtlIntegration } from "./api";
 import type { AppSyncVtlIntegration } from "./appsync";
 import { ASL, ASLGraph } from "./asl";
 import { BindFunctionName, RegisterFunctionName } from "./compile";
-import { FunctionLike, IntegrationInvocation } from "./declaration";
+import { IntegrationInvocation } from "./declaration";
 import { ErrorCodes, formatErrorMessage, SynthError } from "./error-code";
 import {
   IEventBus,
@@ -49,19 +49,16 @@ import {
   PrewarmProps,
 } from "./function-prewarm";
 import {
+  findDeepIntegrations,
   Integration,
-  IntegrationCallExpr,
   IntegrationImpl,
   INTEGRATION_TYPE_KEYS,
   isIntegration,
-  isIntegrationCallExpr,
 } from "./integration";
-import { FunctionlessNode } from "./node";
 import { ReflectionSymbols, validateFunctionLike } from "./reflect";
 import { isStepFunction } from "./step-function";
 import { isTable } from "./table";
 import { AnyAsyncFunction, AnyFunction } from "./util";
-import { visitEachChild } from "./visit";
 
 export function isFunction<Payload = any, Output = any>(
   a: any
@@ -659,7 +656,7 @@ export class Function<
 
     super(_resource);
 
-    const integrations = getInvokedIntegrations(ast).map(
+    const integrations = findDeepIntegrations(ast).map(
       (i) =>
         <IntegrationInvocation>{
           args: i.args,
@@ -723,19 +720,6 @@ export class Function<
   }
 }
 
-function getInvokedIntegrations(ast: FunctionLike): IntegrationCallExpr[] {
-  const nodes: IntegrationCallExpr[] = [];
-  visitEachChild(ast, function visit(node: FunctionlessNode): FunctionlessNode {
-    if (isIntegrationCallExpr(node)) {
-      nodes.push(node);
-    }
-
-    return visitEachChild(node, visit);
-  });
-
-  return nodes;
-}
-
 /**
  * A {@link Function} which wraps a CDK function.
  *