testdrive: fix --rewrite-results dropping bytes between commands

doc/developer/testdrive.md

@@ -243,8 +243,6 @@ Shuffle the list of tests before running them (using the value from --seed, if a
 
 #### `--rewrite-results`
 
-Note that this option is currently buggy!
-
 Automatically rewrite the testdrive file with the correct results when they are not as expected. Consider setting a lower `--default-max-tries` value too to get a result faster.
 
 ## Other options

src/testdrive/src/parser.rs

@@ -326,7 +326,7 @@ fn parse_version_constraint(
 fn parse_sql(line_reader: &mut LineReader) -> Result<SqlCommand, PosError> {
     let (_, line1) = line_reader.next().unwrap();
     let query = line1[1..].trim().to_owned();
-    let expected_start = line_reader.raw_pos;
+    let expected_start = line_reader.consumed_raw_pos;
     let line2 = slurp_one(line_reader);
     let line3 = slurp_one(line_reader);
     let mut column_names = None;
@@ -369,7 +369,7 @@ fn parse_sql(line_reader: &mut LineReader) -> Result<SqlCommand, PosError> {
     while let Some((pos, line)) = slurp_one(line_reader) {
         expected_rows.push(split_line(pos, &line)?)
     }
-    let expected_end = line_reader.raw_pos;
+    let expected_end = line_reader.consumed_raw_pos;
     Ok(SqlCommand {
         query,
         expected_output: SqlOutput::Full {
@@ -383,7 +383,7 @@ fn parse_sql(line_reader: &mut LineReader) -> Result<SqlCommand, PosError> {
 
 fn parse_explain_sql(line_reader: &mut LineReader) -> Result<SqlCommand, PosError> {
     let (_, line1) = line_reader.next().unwrap();
-    let expected_start = line_reader.raw_pos;
+    let expected_start = line_reader.consumed_raw_pos;
     // This is a bit of a hack to extract the next chunk of the file with
     // blank lines intact. Ideally the `LineReader` would expose the API we
     // need directly, but that would require a large refactor.
@@ -402,7 +402,7 @@ fn parse_explain_sql(line_reader: &mut LineReader) -> Result<SqlCommand, PosErro
     // We parsed the multiline expected_output directly using line_reader.inner
     // above.
     slurp_all(line_reader);
-    let expected_end = line_reader.raw_pos;
+    let expected_end = line_reader.consumed_raw_pos;
 
     Ok(SqlCommand {
         query: line1[1..].trim().to_owned(),
@@ -559,6 +559,13 @@ pub struct LineReader<'a> {
     pos: usize,
     pos_map: BTreeMap<usize, (usize, usize)>,
     raw_pos: usize,
+    // Position one byte past the end of the most recently *consumed* line —
+    // i.e. a line returned by an external call to `next()`. `peek()` reads a
+    // line from `inner` and advances `raw_pos`, but the line is not consumed
+    // until a caller takes it via `next()`. `consumed_raw_pos` therefore lags
+    // `raw_pos` by one peeked-but-not-consumed line, which is what callers
+    // need when slicing the input string at command boundaries.
+    consumed_raw_pos: usize,
 }
 
 impl<'a> LineReader<'a> {
@@ -572,12 +579,13 @@ impl<'a> LineReader<'a> {
             pos: 0,
             pos_map,
             raw_pos: 0,
+            consumed_raw_pos: 0,
         }
     }
 
     fn peek(&mut self) -> Option<&(usize, String)> {
         if self.next.is_none() {
-            self.next = Some(self.next())
+            self.next = Some(self.read_one())
         }
         self.next.as_ref().unwrap().as_ref()
     }
@@ -590,15 +598,8 @@ impl<'a> LineReader<'a> {
     fn push(&mut self, text: &String) {
         self.next = Some(Some((0usize, text.to_string())));
     }
-}
-
-impl<'a> Iterator for LineReader<'a> {
-    type Item = (usize, String);
 
-    fn next(&mut self) -> Option<Self::Item> {
-        if let Some(next) = self.next.take() {
-            return next;
-        }
+    fn read_one(&mut self) -> Option<(usize, String)> {
         if self.inner.is_empty() {
             return None;
         }
@@ -651,6 +652,27 @@ impl<'a> Iterator for LineReader<'a> {
     }
 }
 
+impl<'a> Iterator for LineReader<'a> {
+    type Item = (usize, String);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let item = if let Some(next) = self.next.take() {
+            next
+        } else {
+            self.read_one()
+        };
+        if item.is_some() {
+            // Sync `consumed_raw_pos` with `raw_pos` only when the line is
+            // actually returned to an external caller. `peek()` advances
+            // `raw_pos` via `read_one()` without consuming, so updating
+            // `consumed_raw_pos` here keeps it pointing at the end of the most
+            // recently consumed line — even if a later line has been peeked.
+            self.consumed_raw_pos = self.raw_pos;
+        }
+        item
+    }
+}
+
 fn is_sigil(c: Option<char>) -> bool {
     is_sql_sigil(c) || is_non_sql_sigil(c)
 }