rewrite exported vars, to keep exports live

src/Bundle.js

@@ -32,6 +32,10 @@ export default class Bundle {
 		};
 
 		this.entryModule = null;
+
+		this.varExports = blank();
+		this.toExport = null;
+
 		this.modulePromises = blank();
 		this.statements = [];
 		this.externalModules = [];
@@ -179,21 +183,59 @@ export default class Bundle {
 
 		let previousMargin = 0;
 
+		// If we have named exports from the bundle, and those exports
+		// are assigned to *within* the bundle, we may need to rewrite e.g.
+		//
+		//   export let count = 0;
+		//   export function incr () { count++ }
+		//
+		// might become...
+		//
+		//   exports.count = 0;
+		//   function incr () {
+		//     exports.count += 1;
+		//   }
+		//   exports.incr = incr;
+		//
+		// TODO This doesn't apply if the bundle is exported as ES6!
+		let allBundleExports = blank();
+
+		keys( this.entryModule.exports ).forEach( key => {
+			const exportDeclaration = this.entryModule.exports[ key ];
+
+			const originalDeclaration = this.entryModule.findDeclaration( exportDeclaration.localName );
+
+			if ( originalDeclaration && originalDeclaration.type === 'VariableDeclaration' ) {
+				const canonicalName = this.entryModule.getCanonicalName( exportDeclaration.localName );
+
+				allBundleExports[ canonicalName ] = `exports.${key}`;
+				this.varExports[ key ] = true;
+			}
+		});
+
+		// since we're rewriting variable exports, we want to
+		// ensure we don't try and export them again at the bottom
+		this.toExport = keys( this.entryModule.exports )
+			.filter( key => !this.varExports[ key ] );
+
 		// Apply new names and add to the output bundle
 		this.statements.forEach( statement => {
 			let replacements = blank();
+			let bundleExports = blank();
 
 			keys( statement.dependsOn )
 				.concat( keys( statement.defines ) )
 				.forEach( name => {
 					const canonicalName = statement.module.getCanonicalName( name );
 
-					if ( name !== canonicalName ) {
+					if ( allBundleExports[ canonicalName ] ) {
+						bundleExports[ name ] = replacements[ name ] = allBundleExports[ canonicalName ];
+					} else if ( name !== canonicalName ) {
 						replacements[ name ] = canonicalName;
 					}
 				});
 
-			const source = statement.replaceIdentifiers( replacements );
+			const source = statement.replaceIdentifiers( replacements, bundleExports );
 
 			// modify exports as necessary
 			if ( statement.isExportDeclaration ) {

src/ExternalModule.js

@@ -37,4 +37,8 @@ export default class ExternalModule {
 			this.suggestedNames[ exportName ] = suggestion;
 		}
 	}
+
+	findDefiningStatement () {
+		return null;
+	}
 }

src/Module.js

@@ -178,6 +178,32 @@ export default class Module {
 		});
 	}
 
+	findDeclaration ( localName ) {
+		const importDeclaration = this.imports[ localName ];
+
+		// name was defined by another module
+		if ( importDeclaration ) {
+			const module = importDeclaration.module;
+
+			if ( module.isExternal ) return null;
+
+			const exportDeclaration = module.exports[ importDeclaration.name ];
+			return module.findDeclaration( exportDeclaration.localName );
+		}
+
+		// name was defined by this module, if any
+		let i = this.statements.length;
+		while ( i-- ) {
+			const statement = this.statements[i];
+			const declaration = this.statements[i].scope.declarations[ localName ];
+			if ( declaration ) {
+				return declaration;
+			}
+		}
+
+		return null;
+	}
+
 	getCanonicalName ( localName ) {
 		if ( this.suggestedNames[ localName ] ) {
 			localName = this.suggestedNames[ localName ];

src/Statement.js

@@ -37,24 +37,6 @@ export default class Statement {
 
 		let scope = this.scope;
 
-		function addToScope ( declarator ) {
-			var name = declarator.id.name;
-			scope.add( name, false );
-
-			if ( !scope.parent ) {
-				statement.defines[ name ] = true;
-			}
-		}
-
-		function addToBlockScope ( declarator ) {
-			var name = declarator.id.name;
-			scope.add( name, true );
-
-			if ( !scope.parent ) {
-				statement.defines[ name ] = true;
-			}
-		}
-
 		walk( this.node, {
 			enter ( node ) {
 				let newScope;
@@ -65,20 +47,22 @@ export default class Statement {
 					case 'FunctionExpression':
 					case 'FunctionDeclaration':
 					case 'ArrowFunctionExpression':
-						let names = node.params.map( getName );
-
 						if ( node.type === 'FunctionDeclaration' ) {
-							addToScope( node );
-						} else if ( node.type === 'FunctionExpression' && node.id ) {
-							names.push( node.id.name );
+							scope.addDeclaration( node.id.name, node );
 						}
 
 						newScope = new Scope({
 							parent: scope,
-							params: names, // TODO rest params?
+							params: node.params, // TODO rest params?
 							block: false
 						});
 
+						// named function expressions - the name is considered
+						// part of the function's scope
+						if ( node.type === 'FunctionExpression' && node.id ) {
+							newScope.addDeclaration( node.id.name, node );
+						}
+
 						break;
 
 					case 'BlockStatement':
@@ -92,18 +76,20 @@ export default class Statement {
 					case 'CatchClause':
 						newScope = new Scope({
 							parent: scope,
-							params: [ node.param.name ],
+							params: [ node.param ],
 							block: true
 						});
 
 						break;
 
 					case 'VariableDeclaration':
-						node.declarations.forEach( node.kind === 'let' ? addToBlockScope : addToScope ); // TODO const?
+						node.declarations.forEach( declarator => {
+							scope.addDeclaration( declarator.id.name, node );
+						});
 						break;
 
 					case 'ClassDeclaration':
-						addToScope( node );
+						scope.addDeclaration( node.id.name, node );
 						break;
 				}
 
@@ -132,6 +118,10 @@ export default class Statement {
 				}
 			});
 		}
+
+		keys( scope.declarations ).forEach( name => {
+			statement.defines[ name ] = true;
+		});
 	}
 
 	checkForReads ( scope, node, parent ) {
@@ -247,7 +237,9 @@ export default class Statement {
 			});
 	}
 
-	replaceIdentifiers ( names ) {
+	replaceIdentifiers ( names, bundleExports ) {
+		const module = this.module;
+
 		const magicString = this.magicString.clone();
 		const replacementStack = [ names ];
 		const nameList = keys( names );
@@ -258,24 +250,60 @@ export default class Statement {
 			deshadowList.push( replacement.split( '.' )[0] );
 		});
 
-		if ( nameList.length > 0 ) {
+		if ( nameList.length > 0 || keys( bundleExports ).length ) {
+			let topLevel = true;
+
 			walk( this.node, {
 				enter ( node, parent ) {
+					if ( node._skip ) return this.skip();
+
+					// special case - variable declarations that need to be rewritten
+					// as bundle exports
+					if ( topLevel ) {
+						if ( node.type === 'VariableDeclaration' ) {
+							// if this contains a single declarator, and it's one that
+							// needs to be rewritten, we replace the whole lot
+							const name = node.declarations[0].id.name;
+							if ( node.declarations.length === 1 && bundleExports[ name ] ) {
+								magicString.overwrite( node.start, node.declarations[0].id.end, bundleExports[ name ] );
+								node.declarations[0].id._skip = true;
+							}
+
+							// otherwise, we insert the `exports.foo = foo` after the declaration
+							else {
+								const exportInitialisers = node.declarations
+									.map( declarator => declarator.id.name )
+									.filter( name => !!bundleExports[ name ] )
+									.map( name => `\n${bundleExports[name]} = ${name};` )
+									.join( '' );
+
+								// TODO clean this up
+								try {
+									magicString.insert( node.end, exportInitialisers );
+								} catch ( err ) {
+									magicString.append( exportInitialisers );
+								}
+							}
+						}
+					}
+
 					const scope = node._scope;
 
 					if ( scope ) {
+						topLevel = false;
+
 						let newNames = blank();
 						let hasReplacements;
 
 						keys( names ).forEach( key => {
-							if ( !~scope.names.indexOf( key ) ) {
+							if ( !scope.declarations[ key ] ) {
 								newNames[ key ] = names[ key ];
 								hasReplacements = true;
 							}
 						});
 
 						deshadowList.forEach( name => {
-							if ( ~scope.names.indexOf( name ) ) {
+							if ( ~scope.declarations[ name ] ) {
 								newNames[ name ] = name + '$$'; // TODO better mechanism
 								hasReplacements = true;
 							}
@@ -289,7 +317,7 @@ export default class Statement {
 						replacementStack.push( newNames );
 					}
 
-					// We want to rewrite identifiers (that aren't property names)
+					// We want to rewrite identifiers (that aren't property names etc)
 					if ( node.type !== 'Identifier' ) return;
 					if ( parent.type === 'MemberExpression' && !parent.computed && node !== parent.object ) return;
 					if ( parent.type === 'Property' && node !== parent.value ) return;

src/ast/Scope.js

@@ -1,29 +1,59 @@
+import { blank } from '../utils/object';
+
+const blockDeclarations = {
+	'const': true,
+	'let': true
+};
+
 export default class Scope {
 	constructor ( options ) {
 		options = options || {};
 
 		this.parent = options.parent;
 		this.depth = this.parent ? this.parent.depth + 1 : 0;
-		this.names = options.params || [];
+		this.declarations = blank();
 		this.isBlockScope = !!options.block;
+
+		if ( options.params ) {
+			options.params.forEach( param => {
+				this.declarations[ param.name ] = param;
+			});
+		}
 	}
 
-	add ( name, isBlockDeclaration ) {
+	// add ( name, isBlockDeclaration ) {
+	// 	if ( !isBlockDeclaration && this.isBlockScope ) {
+	// 		// it's a `var` or function declaration, and this
+	// 		// is a block scope, so we need to go up
+	// 		this.parent.add( name, isBlockDeclaration );
+	// 	} else {
+	// 		this.names.push( name );
+	// 	}
+	// }
+
+	addDeclaration ( name, declaration ) {
+		const isBlockDeclaration = declaration.type === 'VariableDeclaration' && blockDeclarations[ declaration.kind ];
+
 		if ( !isBlockDeclaration && this.isBlockScope ) {
 			// it's a `var` or function declaration, and this
 			// is a block scope, so we need to go up
-			this.parent.add( name, isBlockDeclaration );
+			this.parent.addDeclaration( name, declaration );
 		} else {
-			this.names.push( name );
+			this.declarations[ name ] = declaration;
 		}
 	}
 
+	getDeclaration ( name ) {
+		return this.declarations[ name ] ||
+		       this.parent && this.parent.getDeclaration( name );
+	}
+
 	contains ( name ) {
-		return !!this.findDefiningScope( name );
+		return !!this.getDeclaration( name );
 	}
 
 	findDefiningScope ( name ) {
-		if ( ~this.names.indexOf( name ) ) {
+		if ( !!this.declarations[ name ] ) {
 			return this;
 		}
 

src/finalisers/amd.js

@@ -22,7 +22,7 @@ export default function amd ( bundle, magicString, exportMode, options ) {
 	if ( exportMode === 'default' ) {
 		exportBlock = `return ${bundle.entryModule.getCanonicalName('default')};`;
 	} else {
-		exportBlock = Object.keys( exports ).map( name => {
+		exportBlock = bundle.toExport.map( name => {
 			return `exports.${name} = ${exports[name].localName};`;
 		}).join( '\n' );
 	}

src/finalisers/cjs.js

@@ -27,7 +27,7 @@ export default function cjs ( bundle, magicString, exportMode ) {
 	if ( exportMode === 'default' && bundle.entryModule.exports.default ) {
 		exportBlock = `module.exports = ${bundle.entryModule.getCanonicalName('default')};`;
 	} else if ( exportMode === 'named' ) {
-		exportBlock = keys( bundle.entryModule.exports )
+		exportBlock = bundle.toExport
 			.map( key => {
 				const specifier = bundle.entryModule.exports[ key ];
 				const name = bundle.entryModule.getCanonicalName( specifier.localName );

src/finalisers/iife.js

@@ -27,6 +27,8 @@ export default function iife ( bundle, magicString, exportMode, options ) {
 		magicString.append( `\n\nreturn ${bundle.entryModule.getCanonicalName('default')};` );
 	}
 
+	// TODO named exports
+
 	return magicString
 		.indent()
 		.prepend( intro )

src/finalisers/umd.js

@@ -50,7 +50,7 @@ export default function umd ( bundle, magicString, exportMode, options ) {
 		const canonicalName = bundle.entryModule.getCanonicalName( 'default' );
 		exportBlock = `return ${canonicalName};`;
 	} else {
-		exportBlock = Object.keys( exports ).map( name => {
+		exportBlock = bundle.toExport.map( name => {
 			const canonicalName = bundle.entryModule.getCanonicalName( exports[ name ].localName );
 			return `exports.${name} = ${canonicalName};`;
 		}).join( '\n' );

test/function/assignment-to-exports/_config.js

@@ -0,0 +1,10 @@
+var assert = require( 'assert' );
+
+module.exports = {
+	description: 'exports are kept up-to-date',
+	exports: function ( exports ) {
+		assert.equal( exports.count, 0 );
+		exports.incr();
+		assert.equal( exports.count, 1 );
+	}
+};

test/function/assignment-to-exports/main.js

@@ -0,0 +1,5 @@
+export var count = 0;
+
+export function incr () {
+	count += 1;
+}