Introduce btree::CellWriter

CellWriter allocates space from the tail of unallocated space. This will
be used for balancing logic as well.

src/btree.rs

@@ -19,7 +19,6 @@ use crate::pager::MemPage;
 use crate::pager::PageBuffer;
 use crate::pager::PageBufferMut;
 use crate::pager::PageId;
-use crate::pager::TemporaryPage;
 use crate::utils::parse_varint;
 use crate::utils::u64_to_i64;
 
@@ -159,15 +158,6 @@ impl<'a> BtreePageHeaderMut<'a> {
         )
     }
 
-    pub fn from_tmp(page: &MemPage, tmp: &'a mut TemporaryPage) -> Self {
-        // SAFETY: TemporaryPage is always more than 512 bytes.
-        Self(
-            (&mut tmp[page.header_offset..page.header_offset + BTREE_PAGE_HEADER_MAX_SIZE])
-                .try_into()
-                .unwrap(),
-        )
-    }
-
     pub fn set_page_type(&mut self, page_type: BtreePageType) {
         self.0[0] = page_type.0;
     }
@@ -583,6 +573,56 @@ pub fn allocate_from_freeblocks(
     None
 }
 
+/// Insert a cell to the page.
+///
+/// The cell is allocated from just before the cell content area.
+pub struct CellWriter<'a, 'page> {
+    page: &'a MemPage,
+    buffer: &'a mut PageBufferMut<'page>,
+    header_size: u8,
+    cell_content_area_offset: usize,
+}
+
+impl<'a, 'page> CellWriter<'a, 'page> {
+    /// Initialize [CellWriter].
+    pub fn new(
+        page: &'a MemPage,
+        buffer: &'a mut PageBufferMut<'page>,
+        header_size: u8,
+        cell_content_area_offset: usize,
+    ) -> Self {
+        Self {
+            page,
+            buffer,
+            header_size,
+            cell_content_area_offset,
+        }
+    }
+
+    /// Write a new cell.
+    ///
+    /// Even if a cell of idx exists, it is overwritten.
+    ///
+    /// NOTE: This does not check if the cell is actually fit in the page.
+    pub fn write(&mut self, idx: u16, cell: &[u8]) {
+        let new_cell_content_area_offset = self.cell_content_area_offset - cell.len();
+        let cell_pointer_offset = cell_pointer_offset(self.page, idx, self.header_size);
+        set_u16(
+            self.buffer,
+            cell_pointer_offset,
+            new_cell_content_area_offset as u16,
+        );
+        self.buffer[new_cell_content_area_offset..self.cell_content_area_offset]
+            .copy_from_slice(cell);
+        self.cell_content_area_offset = new_cell_content_area_offset;
+    }
+
+    #[inline(always)]
+    pub fn cell_content_area_offset(&self) -> usize {
+        self.cell_content_area_offset
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -903,4 +943,61 @@ mod tests {
         );
         assert_eq!(page_header.fragmented_free_bytes(), 3);
     }
+
+    #[test]
+    fn test_cell_writer() {
+        let pager = create_empty_pager(&[], 2 * 4096);
+        let page_type = BtreePageType(BTREE_PAGE_TYPE_LEAF_TABLE);
+        let header_size = page_type.header_size();
+
+        // Tests for page 1.
+        let (page_id, page) = pager.allocate_page().unwrap();
+        assert_eq!(page_id, 1);
+        let mut buffer = pager.make_page_mut(&page).unwrap();
+        let mut cell_writer = CellWriter::new(&page, &mut buffer, header_size, 4096);
+        cell_writer.write(0, &[1; 100]);
+        cell_writer.write(2, &[3; 100]);
+        cell_writer.write(1, &[2; 100]);
+        assert_eq!(cell_writer.cell_content_area_offset(), 4096 - 300);
+        assert_eq!(
+            get_cell_offset(&page, &buffer, 0, header_size).unwrap(),
+            3996
+        );
+        assert_eq!(&buffer[3996..4096], &[1; 100]);
+        assert_eq!(
+            get_cell_offset(&page, &buffer, 2, header_size).unwrap(),
+            3896
+        );
+        assert_eq!(&buffer[3896..3996], &[3; 100]);
+        assert_eq!(
+            get_cell_offset(&page, &buffer, 1, header_size).unwrap(),
+            3796
+        );
+        assert_eq!(&buffer[3796..3896], &[2; 100]);
+
+        // Test for page non-one.
+        let (page_id, page) = pager.allocate_page().unwrap();
+        assert_ne!(page_id, 1);
+        let mut buffer = pager.make_page_mut(&page).unwrap();
+        let mut cell_writer = CellWriter::new(&page, &mut buffer, header_size, 3000);
+        cell_writer.write(0, &[1; 100]);
+        cell_writer.write(2, &[3; 100]);
+        cell_writer.write(1, &[2; 100]);
+        assert_eq!(cell_writer.cell_content_area_offset(), 3000 - 300);
+        assert_eq!(
+            get_cell_offset(&page, &buffer, 0, header_size).unwrap(),
+            2900
+        );
+        assert_eq!(&buffer[2900..3000], &[1; 100]);
+        assert_eq!(
+            get_cell_offset(&page, &buffer, 2, header_size).unwrap(),
+            2800
+        );
+        assert_eq!(&buffer[2800..2900], &[3; 100]);
+        assert_eq!(
+            get_cell_offset(&page, &buffer, 1, header_size).unwrap(),
+            2700
+        );
+        assert_eq!(&buffer[2700..2800], &[2; 100]);
+    }
 }

src/cursor.rs

@@ -26,6 +26,7 @@ use crate::btree::BtreeContext;
 use crate::btree::BtreePageHeader;
 use crate::btree::BtreePageHeaderMut;
 use crate::btree::BtreePageType;
+use crate::btree::CellWriter;
 use crate::btree::FreeblockIterator;
 use crate::btree::IndexCellKeyParser;
 use crate::btree::PayloadInfo;
@@ -540,23 +541,30 @@ impl<'a> BtreeCursor<'a> {
                             // TODO: Optimization: Move cell block of the size of a freeblock if
                             // there is only a few freeblocks. See defragmentPage() of SQLite.
                             let mut tmp_page = self.pager.allocate_tmp_page();
+                            buffer.swap(&mut tmp_page);
                             // Copy database header if the page is page 1.
-                            tmp_page[..self.current_page.mem.header_offset]
-                                .copy_from_slice(&buffer[..self.current_page.mem.header_offset]);
-                            let mut new_page_header =
-                                BtreePageHeaderMut::from_tmp(&self.current_page.mem, &mut tmp_page);
-                            new_page_header.set_page_type(page_type);
-                            new_page_header.set_first_freeblock_offset(0);
-                            new_page_header.clear_fragmented_free_bytes();
+                            buffer[..self.current_page.mem.header_offset]
+                                .copy_from_slice(&tmp_page[..self.current_page.mem.header_offset]);
+                            let mut page_header =
+                                BtreePageHeaderMut::from_page(&self.current_page.mem, &mut buffer);
+                            page_header.set_page_type(page_type);
+                            page_header.set_first_freeblock_offset(0);
+                            page_header.clear_fragmented_free_bytes();
                             // cell_content_area_offset and n_cells will be set later.
                             cell_content_area_offset = self.btree_ctx.usable_size as usize;
                             // Copy reserved area.
-                            tmp_page[cell_content_area_offset..]
-                                .copy_from_slice(&buffer[cell_content_area_offset..]);
+                            buffer[cell_content_area_offset..]
+                                .copy_from_slice(&tmp_page[cell_content_area_offset..]);
+                            let mut cell_writer = CellWriter::new(
+                                &self.current_page.mem,
+                                &mut buffer,
+                                header_size,
+                                cell_content_area_offset,
+                            );
                             for i in 0..self.current_page.n_cells {
                                 let offset = get_cell_offset(
                                     &self.current_page.mem,
-                                    &buffer,
+                                    &tmp_page,
                                     i,
                                     header_size,
                                 )
@@ -565,41 +573,34 @@ impl<'a> BtreeCursor<'a> {
                                 })?;
                                 // Compute cell size.
                                 let (payload_size, payload_size_length) =
-                                    parse_varint(&buffer[offset..])
+                                    parse_varint(&tmp_page[offset..])
                                         .ok_or_else(|| anyhow::anyhow!("parse payload size"))?;
                                 let key_length =
-                                    len_varint_buffer(&buffer[offset + payload_size_length..]);
+                                    len_varint_buffer(&tmp_page[offset + payload_size_length..]);
                                 let n_local =
                                     if payload_size <= self.btree_ctx.max_local(is_table) as u64 {
                                         payload_size as u16
                                     } else {
                                         self.btree_ctx.n_local(is_table, payload_size as i32)
                                     };
                                 let cell_size = payload_size_length + key_length + n_local as usize;
-                                cell_content_area_offset -= cell_size;
-
-                                let cell_pointer_offset =
-                                    cell_pointer_offset(&self.current_page.mem, i, header_size);
-                                set_u16(
-                                    &mut tmp_page,
-                                    cell_pointer_offset,
-                                    cell_content_area_offset as u16,
-                                );
-                                tmp_page[cell_content_area_offset
-                                    ..cell_content_area_offset + cell_size]
-                                    .copy_from_slice(&buffer[offset..offset + cell_size]);
+
+                                // There must be enough size of unallocated space before
+                                // cell_content_area_offset because all cells have fit in the page
+                                // before defragmentation.
+                                cell_writer.write(i, &tmp_page[offset..offset + cell_size]);
                             }
+                            cell_content_area_offset = cell_writer.cell_content_area_offset();
                             // unallocated_space_offset does not change because n_cell is not
                             // changed.
                             assert!(
                                 unallocated_space_offset + 2
                                     <= cell_content_area_offset - cell_size as usize
                             );
                             // Clear the unallocated space.
-                            tmp_page[unallocated_space_offset + 2
+                            buffer[unallocated_space_offset + 2
                                 ..cell_content_area_offset - cell_size as usize]
                                 .fill(0);
-                            buffer.swap(&mut tmp_page);
                         }
 
                         // 3. Allocate space from unallocated space.