Add API to provide a selection range, and return the updated selection after formatting.

Fix #93.

BUG=https://github.com/dart-lang/dart_style/issues/93
R=brianwilkerson@google.com

Review URL: https://chromiumcodereview.appspot.com//822273004
diff --git a/benchmark/benchmark.dart b/benchmark/benchmark.dart
index a468261..3ab9d48 100644
--- a/benchmark/benchmark.dart
+++ b/benchmark/benchmark.dart
@@ -60,7 +60,7 @@
 
   // Sanity check to make sure the output is what we expect and to make sure
   // the VM doesn't optimize "dead" code away.
-  if (result.length != 29796) {
+  if (result.length != 29791) {
     print("Incorrect output (length ${result.length}):\n$result");
     exit(1);
   }
diff --git a/lib/src/chunk.dart b/lib/src/chunk.dart
index ffa96db..f4e46da 100644
--- a/lib/src/chunk.dart
+++ b/lib/src/chunk.dart
@@ -6,6 +6,43 @@
 
 import 'debug.dart';
 
+/// Tracks where a selection start or end point may appear in some piece of
+/// text.
+abstract class Selection {
+  /// The chunk of text.
+  String get text;
+
+  /// The offset from the beginning of [text] where the selection starts, or
+  /// `null` if the selection does not start within this chunk.
+  int get selectionStart => _selectionStart;
+  int _selectionStart;
+
+  /// The offset from the beginning of [text] where the selection ends, or
+  /// `null` if the selection does not start within this chunk.
+  int get selectionEnd => _selectionEnd;
+  int _selectionEnd;
+
+  /// Sets [selectionStart] to be [start] characters into [text].
+  void startSelection(int start) {
+    _selectionStart = start;
+  }
+
+  /// Sets [selectionStart] to be [fromEnd] characters from the end of [text].
+  void startSelectionFromEnd(int fromEnd) {
+    _selectionStart = text.length - fromEnd;
+  }
+
+  /// Sets [selectionEnd] to be [end] characters into [text].
+  void endSelection(int end) {
+    _selectionEnd = end;
+  }
+
+  /// Sets [selectionEnd] to be [fromEnd] characters from the end of [text].
+  void endSelectionFromEnd(int fromEnd) {
+    _selectionEnd = text.length - fromEnd;
+  }
+}
+
 /// A chunk of non-breaking output text terminated by a hard or soft newline.
 ///
 /// Chunks are created by [LineWriter] and fed into [LineSplitter]. Each
@@ -27,7 +64,7 @@
 ///
 /// A split controls the leading spacing of the subsequent line, both
 /// block-based [indent] and expression-wrapping-based [nesting].
-class Chunk {
+class Chunk extends Selection {
   /// The literal text output for the chunk.
   String get text => _text;
   String _text;
@@ -279,3 +316,29 @@
     return result + ")";
   }
 }
+
+/// A comment in the source, with a bit of information about the surrounding
+/// whitespace.
+class SourceComment extends Selection {
+  /// The text of the comment, including `//`, `/*`, and `*/`.
+  final String text;
+
+  /// The number of newlines between the comment or token preceding this comment
+  /// and the beginning of this one.
+  ///
+  /// Will be zero if the comment is a trailing one.
+  final int linesBefore;
+
+  /// Whether this comment is a line comment.
+  final bool isLineComment;
+
+  /// Whether this comment starts at column one in the source.
+  ///
+  /// Comments that start at the start of the line will not be indented in the
+  /// output. This way, commented out chunks of code do not get erroneously
+  /// re-indented.
+  final bool isStartOfLine;
+
+  SourceComment(this.text, this.linesBefore,
+      {this.isLineComment, this.isStartOfLine});
+}
diff --git a/lib/src/dart_formatter.dart b/lib/src/dart_formatter.dart
index 480f40d..5b59faa 100644
--- a/lib/src/dart_formatter.dart
+++ b/lib/src/dart_formatter.dart
@@ -12,6 +12,58 @@
 import 'error_listener.dart';
 import 'source_visitor.dart';
 
+/// Describes a chunk of source code that is to be formatted or has been
+/// formatted.
+class SourceCode {
+  /// The [uri] where the source code is from.
+  ///
+  /// Used in error messages if the code cannot be parsed.
+  final String uri;
+
+  /// The Dart source code text.
+  final String text;
+
+  /// Whether the source is a compilation unit or a bare statement.
+  final bool isCompilationUnit;
+
+  /// The offset in [text] where the selection begins, or `null` if there is
+  /// no selection.
+  final int selectionStart;
+
+  /// The number of selected characters or `null` if there is no selection.
+  final int selectionLength;
+
+  SourceCode(this.text,
+      {this.uri, this.isCompilationUnit: true, this.selectionStart,
+      this.selectionLength}) {
+    // Must either provide both selection bounds or neither.
+    if ((selectionStart == null) != (selectionLength == null)) {
+      throw new ArgumentError(
+          "Is selectionStart is provided, selectionLength must be too.");
+    }
+
+    if (selectionStart != null) {
+      if (selectionStart < 0) {
+        throw new ArgumentError("selectionStart must be non-negative.");
+      }
+
+      if (selectionStart > text.length) {
+        throw new ArgumentError("selectionStart must be within text.");
+      }
+    }
+
+    if (selectionLength != null) {
+      if (selectionLength < 0) {
+        throw new ArgumentError("selectionLength must be non-negative.");
+      }
+
+      if (selectionStart + selectionLength > text.length) {
+        throw new ArgumentError("selectionLength must end within text.");
+      }
+    }
+  }
+}
+
 /// Dart source code formatter.
 class DartFormatter {
   /// The string that newlines should use.
@@ -38,7 +90,7 @@
   DartFormatter({this.lineEnding, int pageWidth, this.indent: 0})
       : this.pageWidth = (pageWidth == null) ? 80 : pageWidth;
 
-  /// Format the given [source] string containing an entire Dart compilation
+  /// Formats the given [source] string containing an entire Dart compilation
   /// unit.
   ///
   /// If [uri] is given, it is a [String] or [Uri] used to identify the file
@@ -54,20 +106,25 @@
       throw new ArgumentError("uri must be `null`, a Uri, or a String.");
     }
 
-    return _format(source, uri: uri, isCompilationUnit: true);
+    return formatSource(
+        new SourceCode(source, uri: uri, isCompilationUnit: true)).text;
   }
 
-  /// Format the given [source] string containing a single Dart statement.
+  /// Formats the given [source] string containing a single Dart statement.
   String formatStatement(String source) {
-    return _format(source, isCompilationUnit: false);
+    return formatSource(new SourceCode(source, isCompilationUnit: false)).text;
   }
 
-  String _format(String source, {String uri, bool isCompilationUnit}) {
+  /// Formats the given [source].
+  ///
+  /// Returns a new [SourceCode] containing the formatted code and the resulting
+  /// selection, if any.
+  SourceCode formatSource(SourceCode source) {
     var errorListener = new ErrorListener();
 
     // Tokenize the source.
-    var reader = new CharSequenceReader(source);
-    var stringSource = new StringSource(source, uri);
+    var reader = new CharSequenceReader(source.text);
+    var stringSource = new StringSource(source.text, source.uri);
     var scanner = new Scanner(stringSource, reader, errorListener);
     var startToken = scanner.tokenize();
     var lineInfo = new LineInfo(scanner.lineStarts);
@@ -78,7 +135,7 @@
       // If the first newline is "\r\n", use that. Otherwise, use "\n".
       if (scanner.lineStarts.length > 1 &&
           scanner.lineStarts[1] >= 2 &&
-          source[scanner.lineStarts[1] - 2] == '\r') {
+          source.text[scanner.lineStarts[1] - 2] == '\r') {
         lineEnding = "\r\n";
       } else {
         lineEnding = "\n";
@@ -92,7 +149,7 @@
     parser.parseAsync = true;
 
     var node;
-    if (isCompilationUnit) {
+    if (source.isCompilationUnit) {
       node = parser.parseCompilationUnit(startToken);
     } else {
       node = parser.parseStatement(startToken);
@@ -101,14 +158,7 @@
     errorListener.throwIfErrors();
 
     // Format it.
-    var buffer = new StringBuffer();
-    var visitor = new SourceVisitor(this, lineInfo, source, buffer);
-
-    visitor.run(node);
-
-    // Be a good citizen, end with a newline.
-    if (isCompilationUnit) buffer.write(lineEnding);
-
-    return buffer.toString();
+    var visitor = new SourceVisitor(this, lineInfo, source);
+    return visitor.run(node);
   }
 }
diff --git a/lib/src/formatter_exception.dart b/lib/src/formatter_exception.dart
index a70670f..7ccc353 100644
--- a/lib/src/formatter_exception.dart
+++ b/lib/src/formatter_exception.dart
@@ -32,4 +32,6 @@
 
     return buffer.toString();
   }
+
+  String toString() => message();
 }
diff --git a/lib/src/line_splitter.dart b/lib/src/line_splitter.dart
index 205fdfc..476d211 100644
--- a/lib/src/line_splitter.dart
+++ b/lib/src/line_splitter.dart
@@ -9,6 +9,7 @@
 import 'chunk.dart';
 import 'debug.dart';
 import 'line_prefix.dart';
+import 'line_writer.dart';
 
 /// The number of spaces in a single level of indentation.
 const spacesPerIndent = 2;
@@ -94,19 +95,38 @@
   ///
   /// It will determine how best to split it into multiple lines of output and
   /// return a single string that may contain one or more newline characters.
-  void apply(StringBuffer buffer) {
+  ///
+  /// Returns a two-element list. The first element will be an [int] indicating
+  /// where in [buffer] the selection start point should appear if it was
+  /// contained in the formatted list of chunks. Otherwise it will be `null`.
+  /// Likewise, the second element will be non-`null` if the selection endpoint
+  /// is within the list of chunks.
+  List<int> apply(StringBuffer buffer) {
     if (debugFormatter) dumpLine(_chunks, _indent);
 
     var splits = _findBestSplits(new LinePrefix());
+    var selection = [null, null];
 
     // Write each chunk and the split after it.
     buffer.write(" " * (_indent * spacesPerIndent));
-    for (var i = 0; i < _chunks.length - 1; i++) {
+    for (var i = 0; i < _chunks.length; i++) {
       var chunk = _chunks[i];
 
+      // If this chunk contains one of the selection markers, tell the writer
+      // where it ended up in the final output.
+      if (chunk.selectionStart != null) {
+        selection[0] = buffer.length + chunk.selectionStart;
+      }
+
+      if (chunk.selectionEnd != null) {
+        selection[1] = buffer.length + chunk.selectionEnd;
+      }
+
       buffer.write(chunk.text);
 
-      if (splits.shouldSplitAt(i)) {
+      if (i == _chunks.length - 1) {
+        // Don't write trailing whitespace after the last chunk.
+      } else if (splits.shouldSplitAt(i)) {
         buffer.write(_lineEnding);
         if (chunk.isDouble) buffer.write(_lineEnding);
 
@@ -120,8 +140,7 @@
       }
     }
 
-    // Write the final chunk without any trailing newlines.
-    buffer.write(_chunks.last.text);
+    return selection;
   }
 
   /// Finds the best set of splits to apply to the remainder of the line
diff --git a/lib/src/line_writer.dart b/lib/src/line_writer.dart
index 2a69ae5..87bfb48 100644
--- a/lib/src/line_writer.dart
+++ b/lib/src/line_writer.dart
@@ -11,32 +11,6 @@
 import 'multisplit.dart';
 import 'whitespace.dart';
 
-/// A comment in the source, with a bit of information about the surrounding
-/// whitespace.
-class SourceComment {
-  /// The text of the comment, including `//`, `/*`, and `*/`.
-  final String text;
-
-  /// The number of newlines between the comment or token preceding this comment
-  /// and the beginning of this one.
-  ///
-  /// Will be zero if the comment is a trailing one.
-  final int linesBefore;
-
-  /// Whether this comment is a line comment.
-  final bool isLineComment;
-
-  /// Whether this comment starts at column one in the source.
-  ///
-  /// Comments that start at the start of the line will not be indented in the
-  /// output. This way, commented out chunks of code do not get erroneously
-  /// re-indented.
-  final bool isStartOfLine;
-
-  SourceComment(this.text, this.linesBefore,
-      {this.isLineComment, this.isStartOfLine});
-}
-
 /// Takes the incremental serialized output of [SourceVisitor]--the source text
 /// along with any comments and preserved whitespace--and produces a coherent
 /// series of [Chunk]s which can then be split into physical lines.
@@ -47,7 +21,9 @@
 class LineWriter {
   final DartFormatter _formatter;
 
-  final StringBuffer _buffer;
+  final SourceCode _source;
+
+  final _buffer = new StringBuffer();
 
   final _chunks = <Chunk>[];
 
@@ -161,6 +137,18 @@
     return _chunks.length;
   }
 
+  /// The offset in [_buffer] where the selection starts in the formatted code.
+  ///
+  /// This will be `null` if there is no selection or the writer hasn't reached
+  /// the beginning of the selection yet.
+  int _selectionStart;
+
+  /// The length in [_buffer] of the selection in the formatted code.
+  ///
+  /// This will be `null` if there is no selection or the writer hasn't reached
+  /// the end of the selection yet.
+  int _selectionLength;
+
   /// Whether there is pending whitespace that depends on the number of
   /// newlines in the source.
   ///
@@ -174,7 +162,7 @@
   /// The number of characters of code that can fit in a single line.
   int get pageWidth => _formatter.pageWidth;
 
-  LineWriter(this._formatter, this._buffer) {
+  LineWriter(this._formatter, this._source) {
     indent(_formatter.indent);
     _beginningIndent = _formatter.indent;
   }
@@ -298,6 +286,14 @@
 
       _writeText(comment.text);
 
+      if (comment.selectionStart != null) {
+        startSelectionFromEnd(comment.text.length - comment.selectionStart);
+      }
+
+      if (comment.selectionEnd != null) {
+        endSelectionFromEnd(comment.text.length - comment.selectionEnd);
+      }
+
       // Make sure there is at least one newline after a line comment and allow
       // one or two after a block comment that has nothing after it.
       var linesAfter;
@@ -452,9 +448,51 @@
     _nesting--;
   }
 
-  /// Finish writing the last line.
-  void end() {
+  /// Marks the selection starting point as occurring [fromEnd] characters to
+  /// the left of the end of what's currently been written.
+  ///
+  /// It counts backwards from the end because this is called *after* the chunk
+  /// of text containing the selection has been output.
+  void startSelectionFromEnd(int fromEnd) {
+    assert(_chunks.isNotEmpty);
+    _chunks.last.startSelectionFromEnd(fromEnd);
+  }
+
+  /// Marks the selection ending point as occurring [fromEnd] characters to the
+  /// left of the end of what's currently been written.
+  ///
+  /// It counts backwards from the end because this is called *after* the chunk
+  /// of text containing the selection has been output.
+  void endSelectionFromEnd(int fromEnd) {
+    assert(_chunks.isNotEmpty);
+    _chunks.last.endSelectionFromEnd(fromEnd);
+  }
+
+  /// Finishes writing and returns a [SourceCode] containing the final output
+  /// and updated selection, if any.
+  SourceCode end() {
     if (_chunks.isNotEmpty) _completeLine();
+
+    // Be a good citizen, end with a newline.
+    if (_source.isCompilationUnit) _buffer.write(_formatter.lineEnding);
+
+    // If we haven't hit the beginning and/or end of the selection yet, they
+    // must be at the very end of the code.
+    if (_source.selectionStart != null) {
+      if (_selectionStart == null) {
+        _selectionStart = _buffer.length;
+      }
+
+      if (_selectionLength == null) {
+        _selectionLength = _buffer.length - _selectionStart;
+      }
+    }
+
+    return new SourceCode(_buffer.toString(),
+        uri: _source.uri,
+        isCompilationUnit: _source.isCompilationUnit,
+        selectionStart: _selectionStart,
+        selectionLength: _selectionLength);
   }
 
   /// Writes the current pending [Whitespace] to the output, if any.
@@ -648,9 +686,12 @@
       _buffer.write(_formatter.lineEnding);
     }
 
-    var splitter = new LineSplitter(_formatter.lineEnding,
-        _formatter.pageWidth, _chunks, _spans, _beginningIndent);
-    splitter.apply(_buffer);
+    var splitter = new LineSplitter(_formatter.lineEnding, _formatter.pageWidth,
+        _chunks, _spans, _beginningIndent);
+    var selection = splitter.apply(_buffer);
+
+    if (selection[0] != null) _selectionStart = selection[0];
+    if (selection[1] != null) _selectionLength = selection[1] - _selectionStart;
   }
 
   /// Handles multisplits when a hard line occurs.
diff --git a/lib/src/source_visitor.dart b/lib/src/source_visitor.dart
index e0780ff..7906fe4 100644
--- a/lib/src/source_visitor.dart
+++ b/lib/src/source_visitor.dart
@@ -21,28 +21,38 @@
   /// Cached line info for calculating blank lines.
   LineInfo _lineInfo;
 
-  /// The source being formatted (used in interpolation handling)
-  final String _source;
+  /// The source being formatted.
+  final SourceCode _source;
+
+  /// `true` if the visitor has written past the beginning of the selection in
+  /// the original source text.
+  bool _passedSelectionStart = false;
+
+  /// `true` if the visitor has written past the end of the selection in the
+  /// original source text.
+  bool _passedSelectionEnd = false;
 
   /// Initialize a newly created visitor to write source code representing
   /// the visited nodes to the given [writer].
-  SourceVisitor(DartFormatter formatter, this._lineInfo, this._source,
-      StringBuffer outputBuffer)
-      : _writer = new LineWriter(formatter, outputBuffer);
+  SourceVisitor(DartFormatter formatter, this._lineInfo, SourceCode source)
+      : _source = source,
+        _writer = new LineWriter(formatter, source);
 
-  /// Run the visitor on [node], writing all of the formatted output to the
-  /// output buffer.
+  /// Runs the visitor on [node], formatting its contents.
+  ///
+  /// Returns a [SourceCode] containing the resulting formatted source and
+  /// updated selection, if any.
   ///
   /// This is the only method that should be called externally. Everything else
   /// is effectively private.
-  void run(AstNode node) {
+  SourceCode run(AstNode node) {
     visit(node);
 
     // Output trailing comments.
     writePrecedingCommentsAndNewlines(node.endToken.next);
 
-    // Finish off the last line.
-    _writer.end();
+    // Finish writing and return the complete result.
+    return _writer.end();
   }
 
   visitAdjacentStrings(AdjacentStrings node) {
@@ -1073,7 +1083,8 @@
     // formatter ensures it gets a newline after it. Since the script tag must
     // come at the top of the file, we don't have to worry about preceding
     // comments or whitespace.
-    _writer.write(node.scriptTag.lexeme.trim());
+    _writeText(node.scriptTag.lexeme.trim(), node.offset);
+
     oneOrTwoNewlines();
   }
 
@@ -1097,7 +1108,7 @@
     // comments are written first.
     writePrecedingCommentsAndNewlines(node.literal);
 
-    _writeStringLiteral(node.literal.lexeme);
+    _writeStringLiteral(node.literal.lexeme, node.offset);
   }
 
   visitStringInterpolation(StringInterpolation node) {
@@ -1109,7 +1120,8 @@
     // contents of interpolated strings. Instead, it treats the entire thing as
     // a single (possibly multi-line) chunk of text.
      _writeStringLiteral(
-        _source.substring(node.beginToken.offset, node.endToken.end));
+        _source.text.substring(node.beginToken.offset, node.endToken.end),
+        node.offset);
   }
 
   visitSuperConstructorInvocation(SuperConstructorInvocation node) {
@@ -1552,15 +1564,18 @@
   ///
   /// Splits multiline strings into separate chunks so that the line splitter
   /// can handle them correctly.
-  void _writeStringLiteral(String string) {
+  void _writeStringLiteral(String string, int offset) {
     // Split each line of a multiline string into separate chunks.
     var lines = string.split("\n");
 
-    _writer.write(lines.first);
+    _writeText(lines.first, offset);
+    offset += lines.first.length;
 
     for (var line in lines.skip(1)) {
       _writer.writeWhitespace(Whitespace.newlineFlushLeft);
-      _writer.write(line);
+      offset++;
+      _writeText(line, offset);
+      offset += line.length;
     }
   }
 
@@ -1623,7 +1638,7 @@
 
     if (before != null) before();
 
-    _writer.write(token.lexeme);
+    _writeText(token.lexeme, token.offset);
 
     if (after != null) after();
   }
@@ -1661,10 +1676,20 @@
         previousLine = commentLine;
       }
 
-      comments.add(new SourceComment(comment.toString().trim(),
+      var sourceComment = new SourceComment(comment.toString().trim(),
           commentLine - previousLine,
           isLineComment: comment.type == TokenType.SINGLE_LINE_COMMENT,
-          isStartOfLine: _startColumn(comment) == 1));
+          isStartOfLine: _startColumn(comment) == 1);
+
+      // If this comment contains either of the selection endpoints, mark them
+      // in the comment.
+      var start = _getSelectionStartWithin(comment.offset, comment.length);
+      if (start != null) sourceComment.startSelection(start);
+
+      var end = _getSelectionEndWithin(comment.offset, comment.length);
+      if (end != null) sourceComment.endSelection(end);
+
+      comments.add(sourceComment);
 
       previousLine = _endLine(comment);
       comment = comment.next;
@@ -1673,6 +1698,80 @@
     _writer.writeComments(comments, tokenLine - previousLine, token.lexeme);
   }
 
+  /// Write [text] to the current chunk, given that it starts at [offset] in
+  /// the original source.
+  ///
+  /// Also outputs the selection endpoints if needed.
+  void _writeText(String text, int offset) {
+    _writer.write(text);
+
+    // If this text contains either of the selection endpoints, mark them in
+    // the chunk.
+    var start = _getSelectionStartWithin(offset, text.length);
+    if (start != null) {
+      _writer.startSelectionFromEnd(text.length - start);
+    }
+
+    var end = _getSelectionEndWithin(offset, text.length);
+    if (end != null) {
+      _writer.endSelectionFromEnd(text.length - end);
+    }
+  }
+
+  /// Returns the number of characters past [offset] in the source where the
+  /// selection start appears if it appears before `offset + length`.
+  ///
+  /// Returns `null` if the selection start has already been processed or is
+  /// not within that range.
+  int _getSelectionStartWithin(int offset, int length) {
+    // If there is no selection, do nothing.
+    if (_source.selectionStart == null) return null;
+
+    // If we've already passed it, don't consider it again.
+    if (_passedSelectionStart) return null;
+
+    var start = _source.selectionStart - offset;
+
+    // If it started in whitespace before this text, push it forward to the
+    // beginning of the non-whitespace text.
+    if (start < 0) start = 0;
+
+    // If we haven't reached it yet, don't consider it.
+    if (start >= length) return null;
+
+    // We found it.
+    _passedSelectionStart = true;
+
+    return start;
+  }
+
+  /// Returns the number of characters past [offset] in the source where the
+  /// selection endpoint appears if it appears before `offset + length`.
+  ///
+  /// Returns `null` if the selection endpoint has already been processed or is
+  /// not within that range.
+  int _getSelectionEndWithin(int offset, int length) {
+    // If there is no selection, do nothing.
+    if (_source.selectionLength == null) return null;
+
+    // If we've already passed it, don't consider it again.
+    if (_passedSelectionEnd) return null;
+
+    var end = _source.selectionStart + _source.selectionLength - offset;
+
+    // If it started in whitespace before this text, push it forward to the
+    // beginning of the non-whitespace text.
+    if (end < 0) end = 0;
+
+    // If we haven't reached it yet, don't consider it.
+    if (end >= length) return null;
+
+    // We found it.
+    _passedSelectionEnd = true;
+
+    return end;
+  }
+
   /// Gets the 1-based line number that the beginning of [token] lies on.
   int _startLine(Token token) => _lineInfo.getLocation(token.offset).lineNumber;
 
diff --git a/test/formatter_test.dart b/test/formatter_test.dart
index 13c0dad..e35c6d3 100644
--- a/test/formatter_test.dart
+++ b/test/formatter_test.dart
@@ -20,6 +20,7 @@
 
   testDirectory("comments");
   testDirectory("regression");
+  testDirectory("selections");
   testDirectory("splitting");
   testDirectory("whitespace");
 
@@ -47,6 +48,44 @@
     }
   });
 
+  group('SourceCode', () {
+    test('throws on negative start', () {
+      expect(() {
+        new SourceCode("12345;", selectionStart: -1, selectionLength: 0);
+      }, throwsArgumentError);
+    });
+
+    test('throws on out of bounds start', () {
+      expect(() {
+        new SourceCode("12345;", selectionStart: 7, selectionLength: 0);
+      }, throwsArgumentError);
+    });
+
+    test('throws on negative length', () {
+      expect(() {
+        new SourceCode("12345;", selectionStart: 1, selectionLength: -1);
+      }, throwsArgumentError);
+    });
+
+    test('throws on out of bounds length', () {
+      expect(() {
+        new SourceCode("12345;", selectionStart: 2, selectionLength: 5);
+      }, throwsArgumentError);
+    });
+
+    test('throws is start is null and length is not', () {
+      expect(() {
+        new SourceCode("12345;", selectionStart: 0);
+      }, throwsArgumentError);
+    });
+
+    test('throws is length is null and start is not', () {
+      expect(() {
+        new SourceCode("12345;", selectionLength: 1);
+      }, throwsArgumentError);
+    });
+  });
+
   test("adds newline to unit", () {
     expect(new DartFormatter().format("var x = 1;"),
         equals("var x = 1;\n"));
@@ -149,19 +188,52 @@
         }
 
         test(description, () {
+          var isCompilationUnit = p.extension(entry.path) == ".unit";
+
+          var inputCode = _extractSelection(input,
+              isCompilationUnit: isCompilationUnit);
+
+          var expected = _extractSelection(expectedOutput,
+              isCompilationUnit: isCompilationUnit);
+
           var formatter = new DartFormatter(
               pageWidth: pageWidth, indent: leadingIndent);
 
-          var result;
-          if (p.extension(entry.path) == ".stmt") {
-            result = formatter.formatStatement(input) + "\n";
-          } else {
-            result = formatter.format(input);
-          }
+          var actual = formatter.formatSource(inputCode);
 
-          expect(result, equals(expectedOutput));
+          // The test files always put a newline at the end of the expectation.
+          // Statements from the formatter (correctly) don't have that, so add
+          // one to line up with the expected result.
+          var actualText = actual.text;
+          if (!isCompilationUnit) actualText += "\n";
+
+          expect(actualText, equals(expected.text));
+          expect(actual.selectionStart, equals(expected.selectionStart));
+          expect(actual.selectionLength, equals(expected.selectionLength));
         });
       }
     });
   }
 }
+
+/// Given a source string that contains ‹ and › to indicate a selection, returns
+/// a [SourceCode] with the text (with the selection markers removed) and the
+/// correct selection range.
+SourceCode _extractSelection(String source, {bool isCompilationUnit: false}) {
+  var start = source.indexOf("‹");
+  source = source.replaceAll("‹", "");
+
+  var end = source.indexOf("›");
+  source = source.replaceAll("›", "");
+
+  // If the selection end is after a trailing newline, there will be an extra
+  // newline *after* the "›", so remove it.
+  if (end != -1 && source.endsWith("\n\n")) {
+    source = source.substring(0, source.length - 1);
+  }
+
+  return new SourceCode(source,
+      isCompilationUnit: isCompilationUnit,
+      selectionStart: start == -1 ? null : start,
+      selectionLength: end == -1 ? null : end - start);
+}
diff --git a/test/selections/selections.stmt b/test/selections/selections.stmt
new file mode 100644
index 0000000..354bc83
--- /dev/null
+++ b/test/selections/selections.stmt
@@ -0,0 +1,73 @@
+40 columns                              |
+>>> start at beginning
+‹123›45;
+<<<
+‹123›45;
+>>> start at end
+12345;‹›
+<<<
+12345;‹›
+>>> zero length
+123‹›45;
+<<<
+123‹›45;
+>>> length at end
+12‹345;›
+<<<
+12‹345;›
+>>> unchanged
+f‹oo(a, ›b, c);
+<<<
+f‹oo(a, ›b, c);
+>>> includes added whitespace
+a+f‹irst+se›cond;
+<<<
+a + f‹irst + se›cond;
+>>> inside comment
+foo(   /* ‹ */  bar/*›*/);
+<<<
+foo(/* ‹ */ bar /*›*/);
+>>> in beginning of multi-line string literal
+  """f‹irs›t
+second""";
+<<<
+"""f‹irs›t
+second""";
+>>> in middle of multi-line string literal
+    """first
+se‹cond
+thi›rd
+fourth""";
+<<<
+"""first
+se‹cond
+thi›rd
+fourth""";
+>>> in end of multi-line string literal
+  """first
+sec‹ond"""  ;›
+<<<
+"""first
+sec‹ond""";›
+>>> in string interpolation
+foo(  "$fi‹rst",  "$sec›ond" );
+<<<
+foo("$fi‹rst", "$sec›ond");
+>>> in moved comment
+someMethod(argument /* long com‹ment that wraps */, other /* last com›ment */);
+<<<
+someMethod(
+    argument /* long com‹ment that wraps */,
+    other /* last com›ment */);
+>>> before comments
+1  ‹  /* */ +  ›  /* */ 2;
+<<<
+1 ‹/* */ + ›/* */ 2;
+>>> after comments
+1/* */  ‹  +/* */  ›   2;
+<<<
+1 /* */ ‹+ /* */ ›2;
+>>> between adjacent comments
+1/* */  ‹  /* */ › /* */ +  2;
+<<<
+1 /* */ ‹/* */ ›/* */ + 2;
\ No newline at end of file
diff --git a/test/selections/selections.unit b/test/selections/selections.unit
new file mode 100644
index 0000000..73fbddc
--- /dev/null
+++ b/test/selections/selections.unit
@@ -0,0 +1,69 @@
+40 columns                              |
+>>> inside script tag
+#!scr‹ip›t
+<<<
+#!scr‹ip›t
+>>> select entire file
+‹main(  )  {
+  body(   )  ;}›
+<<<
+‹main() {
+  body();
+}
+›
+>>> trailing comment
+   ma‹in() {}
+// com›ment
+<<<
+ma‹in() {}
+// com›ment
+>>> in discarded whitespace
+foo(  ‹  argument){  ›  }
+<<<
+foo(‹argument) {›}
+>>> in zero split whitespace
+main(){veryLongMethodCall(‹veryLongArgumentName);
+veryLongMethodCall(›veryLongArgumentName);
+}
+<<<
+main() {
+  veryLongMethodCall(
+      ‹veryLongArgumentName);
+  veryLongMethodCall(
+      ›veryLongArgumentName);
+}
+>>> in soft space split whitespace
+main() {shortCall(argument, ‹  argument);
+shortCall(argument,  › argument);
+}
+<<<
+main() {
+  shortCall(argument, ‹argument);
+  shortCall(argument, ›argument);
+}
+>>> in hard split whitespace
+foo() {body;  ‹  }
+bar() {body;  ›  }
+<<<
+foo() {
+  body;
+‹}
+bar() {
+  body;
+›}
+>>> across lines that get split separately
+foo() {
+
+
+  fir‹st();
+}
+
+bar() {sec›ond();}
+<<<
+foo() {
+  fir‹st();
+}
+
+bar() {
+  sec›ond();
+}