Repository: diesel-rs/diesel
File: diesel/src/sqlite/connection/statement_iterator.rs (Rust)

Code before:
use std::cell::RefCell;
use std::rc::Rc;

use super::row::{PrivateSqliteRow, SqliteRow};
use super::stmt::StatementUse;
use crate::result::QueryResult;

#[allow(missing_debug_implementations)]
pub struct StatementIterator<'stmt, 'query> {
    inner: PrivateStatementIterator<'stmt, 'query>,
    column_names: Option<Rc<[Option<String>]>>,
    field_count: usize,
}

enum PrivateStatementIterator<'stmt, 'query> {
    NotStarted(StatementUse<'stmt, 'query>),
    Started(Rc<RefCell<PrivateSqliteRow<'stmt, 'query>>>),
    TemporaryEmpty,
}

impl<'stmt, 'query> StatementIterator<'stmt, 'query> {
    pub fn new(stmt: StatementUse<'stmt, 'query>) -> StatementIterator<'stmt, 'query> {
        Self {
            inner: PrivateStatementIterator::NotStarted(stmt),
            column_names: None,
            field_count: 0,
        }
    }
}

impl<'stmt, 'query> Iterator for StatementIterator<'stmt, 'query> {
    type Item = QueryResult<SqliteRow<'stmt, 'query>>;

    fn next(&mut self) -> Option<Self::Item> {
        use PrivateStatementIterator::{NotStarted, Started, TemporaryEmpty};

        match std::mem::replace(&mut self.inner, TemporaryEmpty) {
            NotStarted(stmt) => match stmt.step() {
                Err(e) => Some(Err(e)),
                Ok(None) => None,
                Ok(Some(stmt)) => {
                    let field_count = stmt.column_count() as usize;
                    self.field_count = field_count;
                    let inner = Rc::new(RefCell::new(PrivateSqliteRow::Direct(stmt)));
                    self.inner = Started(inner.clone());
                    Some(Ok(SqliteRow { inner, field_count }))

Review comment (bug):
Curious what this kind of duplication handling is needed for. Is this a SQLite or SQLite driver bug?

Code after:
use std::cell::RefCell;
use std::rc::Rc;

use super::row::{PrivateSqliteRow, SqliteRow};
use super::stmt::StatementUse;
use crate::result::QueryResult;

#[allow(missing_debug_implementations)]
pub struct StatementIterator<'stmt, 'query> {
    inner: PrivateStatementIterator<'stmt, 'query>,
    column_names: Option<Rc<[Option<String>]>>,
    field_count: usize,
}

impl<'stmt, 'query> StatementIterator<'stmt, 'query> {
    #[cold]
    fn handle_duplicated_row_case(
        outer_last_row: &mut Rc<RefCell<PrivateSqliteRow<'stmt, 'query>>>,
        column_names: &mut Option<Rc<[Option<String>]>>,
        field_count: usize,
    ) -> Option<QueryResult<SqliteRow<'stmt, 'query>>> {
        // We don't own the statement. There is another existing reference, likly because
        // a user stored the row in some long time container before calling next another time
        // In this case we copy out the current values into a temporary store and advance
        // the statement iterator internally afterwards
        let last_row = {
            let mut last_row = match outer_last_row.try_borrow_mut() {
                Ok(o) => o,
                Err(_e) => {
                    return Some(Err(crate::result::Error::DeserializationError(
                                    "Failed to reborrow row. Try to release any `SqliteField` or `SqliteValue` \
                                     that exists at this point"
                                        .into(),
                                )));
                }
            };
            let last_row = &mut *last_row;
            let duplicated = last_row.duplicate(column_names);
            std::mem::replace(last_row, duplicated)
        };
        if let PrivateSqliteRow::Direct(mut stmt) = last_row {
            let res = unsafe {
                // This is actually safe here as we've already
                // performed one step. For the first step we would have
                // used `PrivateStatementIterator::NotStarted` where we don't
                // have access to `PrivateSqliteRow` at all