Move logic for collecting tokens to token layer

packages/@aws-cdk/cdk/lib/cloudformation/stack-element.ts

@@ -1,5 +1,4 @@
 import { Construct, IConstruct, PATH_SEP } from "../core/construct";
-import { RESOLVE_OPTIONS } from "../core/tokens/options";
 
 const LOGICAL_ID_MD = 'aws:cdk:logicalId';
 
@@ -119,17 +118,22 @@ export abstract class StackElement extends Construct implements IDependable {
    * Automatically detect references in this StackElement
    */
   protected prepare() {
-    const options = RESOLVE_OPTIONS.push({ preProcess: (token, _) => { this.node.recordReference(token); return token; } });
     try {
-      // Execute for side effect of calling 'preProcess'.
-      // Note: it might be that the properties of the CFN object aren't valid. This will usually be preventatively
-      // caught in a construct's validate() and turned into a nicely descriptive error, but we're running prepare()
-      // before validate(). Swallow errors that occur because the CFN layer doesn't validate completely.
-      this.node.resolve(this.toCloudFormation());
+      // Note: it might be that the properties of the CFN object aren't valid.
+      // This will usually be preventatively caught in a construct's validate()
+      // and turned into a nicely descriptive error, but we're running prepare()
+      // before validate(). Swallow errors that occur because the CFN layer
+      // doesn't validate completely.
+      //
+      // This does make the assumption that the error will not be rectified,
+      // but the error will be thrown later on anyway. If the error doesn't
+      // get thrown down the line, we may miss references.
+      this.node.recordReference(...findTokens(this.toCloudFormation(), {
+        scope: this,
+        prefix: []
+      }));
     } catch (e) {
       if (e.type !== 'CfnSynthesisError') { throw e; }
-    } finally {
-      options.pop();
     }
   }
 }
@@ -145,6 +149,7 @@ export class Ref extends CfnReference {
   }
 }
 
+import { findTokens } from "../core/tokens/resolve";
 import { Stack } from "./stack";
 
 /**

packages/@aws-cdk/cdk/lib/core/construct.ts

@@ -418,9 +418,11 @@ export class ConstructNode {
   /**
    * Record a reference originating from this construct node
    */
-  public recordReference(ref: Token) {
-    if (ref.isReference) {
-      this.references.add(ref);
+  public recordReference(...refs: Token[]) {
+    for (const ref of refs) {
+      if (ref.isReference) {
+        this.references.add(ref);
+      }
     }
   }
 

packages/@aws-cdk/cdk/lib/core/tokens/options.ts

@@ -1,9 +1,9 @@
-import { ResolveContext, Token } from "./token";
+import { Token } from "./token";
 
 /**
  * Function used to preprocess Tokens before resolving
  */
-export type PreProcessFunc = (token: Token, context: ResolveContext) => Token;
+export type CollectFunc = (token: Token) => void;
 
 /**
  * Global options for resolve()
@@ -28,12 +28,12 @@ export class ResolveConfiguration {
     };
   }
 
-  public get preProcess(): PreProcessFunc {
+  public get collect(): CollectFunc | undefined {
     for (let i = this.options.length - 1; i >= 0; i--) {
-      const ret = this.options[i].preProcess;
+      const ret = this.options[i].collect;
       if (ret !== undefined) { return ret; }
     }
-    return noPreprocessFunction;
+    return undefined;
   }
 }
 
@@ -45,7 +45,7 @@ interface ResolveOptions {
   /**
    * What function to use to preprocess Tokens before resolving them
    */
-  preProcess?: PreProcessFunc;
+  collect?: CollectFunc;
 }
 
 const glob = global as any;
@@ -54,7 +54,3 @@ const glob = global as any;
  * Singleton instance of resolver options
  */
 export const RESOLVE_OPTIONS: ResolveConfiguration = glob.__cdkResolveOptions = glob.__cdkResolveOptions || new ResolveConfiguration();
-
-function noPreprocessFunction(x: Token, _: ResolveContext) {
-  return x;
-}

packages/@aws-cdk/cdk/lib/core/tokens/resolve.ts

@@ -1,6 +1,6 @@
 import { containsListToken, TOKEN_MAP } from "./encoding";
 import { RESOLVE_OPTIONS } from "./options";
-import { RESOLVE_METHOD, ResolveContext } from "./token";
+import { RESOLVE_METHOD, ResolveContext, Token } from "./token";
 import { unresolved } from "./unresolved";
 
 // This file should not be exported to consumers, resolving should happen through Construct.resolve()
@@ -80,8 +80,9 @@ export function resolve(obj: any, context: ResolveContext): any {
   //
 
   if (unresolved(obj)) {
-    const preProcess = RESOLVE_OPTIONS.preProcess;
-    const value = preProcess(obj, context)[RESOLVE_METHOD](context);
+    const collect = RESOLVE_OPTIONS.collect;
+    if (collect) { collect(obj); }
+    const value = obj[RESOLVE_METHOD](context);
     return resolve(value, context);
   }
 
@@ -116,6 +117,23 @@ export function resolve(obj: any, context: ResolveContext): any {
   return result;
 }
 
+/**
+ * Find all Tokens that are used in the given structure
+ */
+export function findTokens(obj: any, context: ResolveContext): Token[] {
+  const ret = new Array<Token>();
+
+  const options = RESOLVE_OPTIONS.push({ collect: ret.push.bind(ret) });
+  try {
+    // resolve() for side effect of calling 'preProcess', which adds to the
+    resolve(obj, context);
+  } finally {
+    options.pop();
+  }
+
+  return ret;
+}
+
 /**
  * Determine whether an object is a Construct
  *