mark statements rather than expanding them

src/Bundle.js

@@ -84,7 +84,10 @@ export default class Bundle {
 					}
 				}
 
-				return entryModule.expandAllStatements( true );
+				return entryModule.markAllStatements( true );
+			})
+			.then( () => {
+				return this.markAllModifierStatements();
 			})
 			.then( () => {
 				this.statements = this.sort();
@@ -404,6 +407,34 @@ export default class Bundle {
 		return { code, map };
 	}
 
+	markAllModifierStatements () {
+		let settled = true;
+		let promises = [];
+
+		this.modules.forEach( module => {
+			module.statements.forEach( statement => {
+				if ( statement.isIncluded ) return;
+
+				keys( statement.modifies ).forEach( name => {
+					const definingStatement = module.definitions[ name ];
+					const exportDeclaration = module.exports[ name ];
+
+					const shouldMark = ( definingStatement && definingStatement.isIncluded ) ||
+					                   ( exportDeclaration && exportDeclaration.isUsed );
+
+					if ( shouldMark ) {
+						settled = false;
+						promises.push( statement.mark() );
+					}
+				});
+			});
+		});
+
+		return Promise.all( promises ).then( () => {
+			if ( !settled ) return this.markAllModifierStatements();
+		});
+	}
+
 	sort () {
 		let seen = {};
 		let ordered = [];

src/Module.js

@@ -220,7 +220,7 @@ export default class Module {
 		return { strongDependencies, weakDependencies };
 	}
 
-	define ( name ) {
+	mark ( name ) {
 		// shortcut cycles. TODO this won't work everywhere...
 		if ( this.definitionPromises[ name ] ) {
 			return emptyArrayPromise;
@@ -272,7 +272,7 @@ export default class Module {
 							this.bundle.internalNamespaceModules.push( module );
 						}
 
-						return module.expandAllStatements();
+						return module.markAllStatements();
 					}
 
 					const exportDeclaration = module.exports[ importDeclaration.name ];
@@ -281,7 +281,7 @@ export default class Module {
 						throw new Error( `Module ${module.id} does not export ${importDeclaration.name} (imported by ${this.id})` );
 					}
 
-					return module.define( exportDeclaration.localName );
+					return module.mark( exportDeclaration.localName );
 				});
 		}
 
@@ -289,14 +289,14 @@ export default class Module {
 		else if ( name === 'default' && this.exports.default.isDeclaration ) {
 			// We have something like `export default foo` - so we just start again,
 			// searching for `foo` instead of default
-			promise = this.define( this.exports.default.name );
+			promise = this.mark( this.exports.default.name );
 		}
 
 		else {
 			let statement;
 
 			statement = name === 'default' ? this.exports.default.statement : this.definitions[ name ];
-			promise = statement && !statement.isIncluded ? statement.expand() : emptyArrayPromise;
+			promise = statement && !statement.isIncluded ? statement.mark() : emptyArrayPromise;
 
 			// Special case - `export default foo; foo += 1` - need to be
 			// vigilant about maintaining the correct order of the export
@@ -331,22 +331,9 @@ export default class Module {
 		return this.definitionPromises[ name ];
 	}
 
-	expandAllStatements ( isEntryModule ) {
+	markAllStatements ( isEntryModule ) {
-		let allStatements = [];
-
 		return sequence( this.statements, statement => {
-			// A statement may have already been included, in which case we need to
+			if ( statement.isIncluded ) return; // TODO can this happen? probably not...
-			// curb rollup's enthusiasm and move it down here. It remains to be seen
-			// if this approach is bulletproof
-			if ( statement.isIncluded ) {
-				const index = allStatements.indexOf( statement );
-				if ( ~index ) {
-					allStatements.splice( index, 1 );
-					allStatements.push( statement );
-				}
-
-				return;
-			}
 
 			// skip import declarations...
 			if ( statement.isImportDeclaration ) {
@@ -356,10 +343,7 @@ export default class Module {
 					return this.bundle.fetchModule( statement.node.source.value, this.id )
 						.then( module => {
 							statement.module = module;
-							return module.expandAllStatements();
+							return module.markAllStatements();
-						})
-						.then( statements => {
-							allStatements.push.apply( allStatements, statements );
 						});
 				}
 
@@ -370,20 +354,14 @@ export default class Module {
 			if ( statement.node.type === 'ExportNamedDeclaration' && statement.node.specifiers.length ) {
 				// ...but ensure they are defined, if this is the entry module
 				if ( isEntryModule ) {
-					return statement.expand().then( statements => {
+					return statement.mark();
-						allStatements.push.apply( allStatements, statements );
-					});
 				}
 
 				return;
 			}
 
 			// include everything else
-			return statement.expand().then( statements => {
+			return statement.mark();
-				allStatements.push.apply( allStatements, statements );
-			});
-		}).then( () => {
-			return allStatements;
 		});
 	}
 

src/Statement.js

@@ -234,52 +234,16 @@ export default class Statement {
 		}
 	}
 
-	expand () {
+	mark () {
-		this.isIncluded = true; // prevent statement being included twice
+		if ( this.included ) return; // prevent infinite loops
+		this.isIncluded = true;
 
-		let result = [];
-
-		// We have a statement, and it hasn't been included yet. First, include
-		// the statements it depends on
 		const dependencies = Object.keys( this.dependsOn );
 
 		return sequence( dependencies, name => {
 			if ( this.defines[ name ] ) return; // TODO maybe exclude from `this.dependsOn` in the first place?
-
+			return this.module.mark( name );
-			return this.module.define( name ).then( definition => {
+		});
-				result.push.apply( result, definition );
-			});
-		})
-
-		// then include the statement itself
-			.then( () => {
-				result.push( this );
-			})
-
-		// then include any statements that could modify the
-		// thing(s) this statement defines
-			.then( () => {
-				return sequence( keys( this.defines ), name => {
-					const modifications = this.module.modifications[ name ];
-
-					if ( modifications ) {
-						return sequence( modifications, statement => {
-							if ( !statement.isIncluded ) {
-								return statement.expand()
-									.then( statements => {
-										result.push.apply( result, statements );
-									});
-							}
-						});
-					}
-				});
-			})
-
-		// the `result` is an array of all statements that need
-		// to be included if this one is
-			.then( () => {
-				return result;
-			});
 	}
 
 	replaceIdentifiers ( names, bundleExports ) {