jump through some hoops to keep statements appropriately spaced, and comments attached to their owners

src/Bundle.js

@@ -132,6 +132,8 @@ export default class Bundle {
 		// Determine export mode - 'default', 'named', 'none'
 		let exportMode = this.getExportMode( options.exports );
 
+		let previousMargin = 0;
+
 		// Apply new names and add to the output bundle
 		this.statements.forEach( statement => {
 			let replacements = {};
@@ -146,7 +148,7 @@ export default class Bundle {
 					}
 				});
 
-			const source = statement._source.clone();
+			const source = statement._source.clone().trim();
 
 			// modify exports as necessary
 			if ( /^Export/.test( statement.type ) ) {
@@ -174,9 +176,40 @@ export default class Bundle {
 				}
 			}
 
-
 			replaceIdentifiers( statement, source, replacements );
-			magicString.addSource( source );
+
+			// add leading comments
+			if ( statement._leadingComments.length ) {
+				const commentBlock = statement._leadingComments.map( comment => {
+					return comment.block ?
+						`/*${comment.text}*/` :
+						`//${comment.text}`;
+				}).join( '\n' );
+
+				magicString.addSource( new MagicString( commentBlock ) );
+			}
+
+			// add margin
+			const margin = Math.max( statement._margin[0], previousMargin );
+			const newLines = new Array( margin ).join( '\n' );
+
+			// add the statement itself
+			magicString.addSource({
+				content: source,
+				separator: newLines
+			});
+
+			// add trailing comments
+			const comment = statement._trailingComment;
+			if ( comment ) {
+				const commentBlock = comment.block ?
+					` /*${comment.text}*/` :
+					` //${comment.text}`;
+
+				magicString.append( commentBlock );
+			}
+
+			previousMargin = statement._margin[1];
 		});
 
 		// prepend bundle with internal namespaces

src/Module.js

@@ -21,18 +21,19 @@ export default class Module {
 		});
 
 		this.suggestedNames = {};
+		this.comments = [];
 
 		try {
 			this.ast = parse( code, {
 				ecmaVersion: 6,
-				sourceType: 'module'
+				sourceType: 'module',
+				onComment: ( block, text, start, end ) => this.comments.push({ block, text, start, end })
 			});
 		} catch ( err ) {
 			err.file = path;
 			throw err;
 		}
 
-
 		this.analyse();
 	}
 
@@ -41,6 +42,10 @@ export default class Module {
 		this.imports = {};
 		this.exports = {};
 
+		let commentIndex = 0;
+		let previousNode;
+		let previousEnd = 0;
+
 		this.ast.body.forEach( node => {
 			let source;
 
@@ -371,27 +376,6 @@ export default class Module {
 			}
 
 			// include everything else
-
-
-
-			// // modify exports as necessary
-			// if ( /^Export/.test( statement.type ) ) {
-			// 	// remove `export` from `export var foo = 42`
-			// 	if ( statement.type === 'ExportNamedDeclaration' && statement.declaration.type === 'VariableDeclaration' ) {
-			// 		statement._source.remove( statement.start, statement.declaration.start );
-			// 	}
-			//
-			// 	// remove `export` from `export class Foo {...}` or `export default Foo`
-			// 	else if ( statement.declaration.id ) {
-			// 		statement._source.remove( statement.start, statement.declaration.start );
-			// 	}
-			//
-			// 	// declare variables for expressions
-			// 	else {
-			// 		statement._source.overwrite( statement.start, statement.declaration.start, `var TODO = ` );
-			// 	}
-			// }
-
 			return this.expandStatement( statement )
 				.then( statements => {
 					allStatements.push.apply( allStatements, statements );

src/ast/analyse.js

@@ -27,20 +27,55 @@ export default function analyse ( ast, magicString, module ) {
 	}
 
 	// first we need to generate comprehensive scope info
-	let previous = 0;
+	let previousStatement = null;
+	let commentIndex = 0;
 
 	ast.body.forEach( statement => {
+		currentTopLevelStatement = statement; // so we can attach scoping info
+
 		Object.defineProperties( statement, {
-			_defines:   { value: {} },
-			_modifies:  { value: {} },
-			_dependsOn: { value: {} },
-			_included:  { value: false, writable: true },
-			_module:    { value: module },
-			_source:    { value: magicString.snip( previous, statement.end ) } // TODO don't use snip, it's a waste of memory
+			_defines:          { value: {} },
+			_modifies:         { value: {} },
+			_dependsOn:        { value: {} },
+			_included:         { value: false, writable: true },
+			_module:           { value: module },
+			_source:           { value: magicString.snip( statement.start, statement.end ) }, // TODO don't use snip, it's a waste of memory
+			_margin:           { value: [ 0, 0 ] },
+			_leadingComments:  { value: [] },
+			_trailingComment:  { value: null, writable: true },
 		});
 
-		previous = statement.end;
-		currentTopLevelStatement = statement; // so we can attach scoping info
+		let trailing = !!previousStatement;
+
+		// attach leading comment
+		do {
+			const comment = module.comments[ commentIndex ];
+
+			if ( !comment || ( comment.end > statement.start ) ) break;
+
+			// attach any trailing comment to the previous statement
+			if ( trailing && !/\n/.test( magicString.slice( previousStatement.end, comment.start ) ) ) {
+				previousStatement._trailingComment = comment;
+			}
+
+			// then attach leading comments to this statement
+			else {
+				statement._leadingComments.push( comment );
+			}
+
+			commentIndex += 1;
+			trailing = false;
+		} while ( module.comments[ commentIndex ] );
+
+		// determine margin
+		const previousEnd = previousStatement ? ( previousStatement._trailingComment || previousStatement ).end : 0;
+		const start = ( statement._leadingComments[0] || statement ).start;
+
+		const gap = magicString.original.slice( previousEnd, start );
+		const margin = gap.split( '\n' ).length;
+
+		if ( previousStatement ) previousStatement._margin[1] = margin;
+		statement._margin[0] = margin;
 
 		walk( statement, {
 			enter ( node ) {
@@ -109,6 +144,8 @@ export default function analyse ( ast, magicString, module ) {
 				}
 			}
 		});
+
+		previousStatement = statement;
 	});
 
 	// then, we need to find which top-level dependencies this statement has,

test/form/self-contained-bundle/_config.js

@@ -0,0 +1,3 @@
+module.exports = {
+	description: 'self-contained bundle'
+};

test/form/self-contained-bundle/_expected/amd.js

@@ -0,0 +1,19 @@
+define(function () { 'use strict';
+
+	// comment before 1
+
+	console.log( 1 );
+	console.log( 2 ); // comment alongside 2
+
+	function bar () {
+		return 42;
+	}
+
+	function foo () {
+		return bar();
+	}
+
+	foo();
+	console.log( 3 );
+
+});

test/form/self-contained-bundle/_expected/cjs.js

@@ -0,0 +1,17 @@
+'use strict';
+
+// comment before 1
+
+console.log( 1 );
+console.log( 2 ); // comment alongside 2
+
+function bar () {
+	return 42;
+}
+
+function foo () {
+	return bar();
+}
+
+foo();
+console.log( 3 );

test/form/self-contained-bundle/_expected/es6.js

@@ -0,0 +1,15 @@
+// comment before 1
+
+console.log( 1 );
+console.log( 2 ); // comment alongside 2
+
+function bar () {
+	return 42;
+}
+
+function foo () {
+	return bar();
+}
+
+foo();
+console.log( 3 );

test/form/self-contained-bundle/_expected/iife.js

@@ -0,0 +1,19 @@
+(function () { 'use strict';
+
+	// comment before 1
+
+	console.log( 1 );
+	console.log( 2 ); // comment alongside 2
+
+	function bar () {
+		return 42;
+	}
+
+	function foo () {
+		return bar();
+	}
+
+	foo();
+	console.log( 3 );
+
+})();

test/form/self-contained-bundle/_expected/umd.js

@@ -0,0 +1,23 @@
+(function (global, factory) {
+	typeof exports === 'object' && typeof module !== 'undefined' ? factory() :
+	typeof define === 'function' && define.amd ? define(factory) :
+	factory();
+}(this, function () { 'use strict';
+
+	// comment before 1
+
+	console.log( 1 );
+	console.log( 2 ); // comment alongside 2
+
+	function bar () {
+		return 42;
+	}
+
+	function foo () {
+		return bar();
+	}
+
+	foo();
+	console.log( 3 );
+
+}));

test/form/self-contained-bundle/foo.js

@@ -0,0 +1,11 @@
+export default function foo () {
+	return bar();
+}
+
+function bar () {
+	return 42;
+}
+
+function baz () {
+	return 'this should be excluded';
+}

test/form/self-contained-bundle/main.js

@@ -0,0 +1,8 @@
+import foo from './foo';
+
+// comment before 1
+
+console.log( 1 );
+console.log( 2 ); // comment alongside 2
+foo();
+console.log( 3 );