WIP - Update support

Author: PThorpe92

core/storage/btree.rs

@@ -820,7 +820,27 @@ impl BTreeCursor {
                         (self.find_cell(page, int_key), page.page_type())
                     };
 
-                    // TODO: if overwrite drop cell
+                    // check if we have an exact match on the cell at cell_idx:
+                    // (`update` support/insert record to same existing rowid)
+                    if cell_idx < page.get_contents().cell_count() {
+                        // parse the existing cell to get the stored key
+                        if let BTreeCell::TableLeafCell(tbl_leaf) = page.get_contents().cell_get(
+                            cell_idx,
+                            self.pager.clone(),
+                            payload_overflow_threshold_max(page_type, self.usable_space() as u16),
+                            payload_overflow_threshold_min(page_type, self.usable_space() as u16),
+                            self.usable_space(),
+                        )? {
+                            if tbl_leaf._rowid == int_key {
+                                tracing::trace!("insert_into_page: found exact match with cell_idx={cell_idx}, overwriting");
+                                // we are "overwriting in place" the old cell, not removing the pointers
+                                self.overwrite_cell(page.clone(), cell_idx, record)?;
+                                // advance to the next cell
+                                self.move_to(SeekKey::TableRowId(int_key), SeekOp::GT)?;
+                                return Ok(CursorResult::Ok(()));
+                            }
+                        }
+                    }
 
                     // insert cell
                     let mut cell_payload: Vec<u8> = Vec::new();
@@ -2201,6 +2221,183 @@ impl BTreeCursor {
     pub fn table_id(&self) -> usize {
         self.root_page
     }
+
+    pub fn overwrite_cell(
+        &mut self,
+        page_ref: PageRef,
+        cell_idx: usize,
+        record: &Record,
+    ) -> Result<CursorResult<()>> {
+        // build the new payload
+        let page_type = page_ref.get().contents.as_ref().unwrap().page_type();
+        let mut new_payload = Vec::new();
+        fill_cell_payload(
+            page_type,
+            self.rowid.get(),
+            &mut new_payload,
+            record,
+            self.usable_space() as u16,
+            self.pager.clone(),
+        );
+
+        // figure out old cell offset & size
+        let (old_offset, old_local_size) = {
+            let page = page_ref.get().contents.as_ref().unwrap();
+            page.cell_get_raw_region(
+                cell_idx,
+                payload_overflow_threshold_max(page_type, self.usable_space() as u16),
+                payload_overflow_threshold_min(page_type, self.usable_space() as u16),
+                self.usable_space(),
+            )
+        };
+
+        // if it all fits in local space and old_local_size is enough, do an in-place copy
+        //  (like sqlite btreeOverwriteContent local portion)
+        if new_payload.len() <= old_local_size {
+            self.overwrite_content(
+                page_ref.clone(),
+                old_offset,
+                &new_payload,
+                0,
+                new_payload.len(),
+            )?;
+            // if there's leftover local space (old_local_size > new_payload.len()), zero it or free it
+            let remaining = old_local_size - new_payload.len();
+            if remaining > 0 {
+                let buf = page_ref.get().contents.as_mut().unwrap().as_ptr();
+                for i in 0..remaining {
+                    buf[old_offset + new_payload.len() + i] = 0;
+                }
+            }
+            Ok(CursorResult::Ok(()))
+        } else {
+            // otherwise, we have to handle overflow (or page-splitting)... overflow for simplicity
+            self.overwrite_overflow_cell(
+                page_ref,
+                cell_idx,
+                old_offset,
+                old_local_size,
+                &new_payload,
+            )
+        }
+    }
+
+    pub fn overwrite_content(
+        &mut self,
+        page_ref: PageRef,
+        dest_offset: usize,
+        new_payload: &[u8],
+        src_offset: usize,
+        amount: usize,
+    ) -> Result<CursorResult<()>> {
+        return_if_locked!(page_ref);
+        page_ref.set_dirty();
+        self.pager.add_dirty(page_ref.get().id);
+
+        let buf = page_ref.get().contents.as_mut().unwrap().as_ptr();
+
+        // if new_payload doesn't have enough data, we fill with zeros
+        let n_data = new_payload.len().saturating_sub(src_offset);
+        if n_data == 0 {
+            // everything is zeros
+            for i in 0..amount {
+                if buf[dest_offset + i] != 0 {
+                    // do page-write if needed
+                    // ...
+                    buf[dest_offset + i] = 0;
+                }
+            }
+        } else {
+            let copy_len = n_data.min(amount);
+            // copy the overlapping portion
+            buf[dest_offset..dest_offset + copy_len]
+                .copy_from_slice(&new_payload[src_offset..src_offset + copy_len]);
+
+            // if copy_len < amount => fill remainder with 0
+            if copy_len < amount {
+                for i in copy_len..amount {
+                    buf[dest_offset + i] = 0;
+                }
+            }
+        }
+        Ok(CursorResult::Ok(()))
+    }
+
+    fn overwrite_overflow_cell(
+        &mut self,
+        page_ref: PageRef,
+        cell_idx: usize,
+        old_offset: usize,
+        old_local_size: usize,
+        new_payload: &[u8],
+    ) -> Result<CursorResult<()>> {
+        // possibly defragment
+        let page = page_ref.get().contents.as_mut().unwrap();
+        defragment_page(page, self.usable_space() as u16);
+
+        // re-check how many bytes we can store locally
+        let (local_needed, _) = compute_local_needed(new_payload.len(), self.usable_space() as u16);
+        let local_area_size = old_local_size;
+
+        let local_bytes = (local_needed.min(old_local_size)).saturating_sub(4);
+
+        // overwrite that many bytes
+        self.overwrite_content(page_ref.clone(), old_offset, new_payload, 0, local_bytes)?;
+
+        // if local_bytes < local_area_size => zero out leftover
+        if local_bytes < local_area_size {
+            let leftover = local_area_size - local_bytes;
+            // zero the leftover area (for now)
+            for i in 0..leftover {
+                page.as_ptr()[old_offset + local_bytes + i] = 0;
+            }
+        }
+
+        let remainder_start = if local_bytes < 4 { 4 } else { local_bytes };
+
+        let remainder = &new_payload[remainder_start..];
+        // now we must store new_payload[local_bytes..] in overflow
+        if !remainder.is_empty() {
+            let overflow_page = allocate_overflow_page(self.pager.clone());
+            let arr = overflow_page.get().id.to_be_bytes();
+            let buf = page.as_ptr();
+            let pointer_offset = old_offset + local_bytes;
+            if pointer_offset + 4 > buf.len() {
+                panic!(
+                    "Overflow pointer exceeds buffer! pointer_offset: {}, buf.len(): {}",
+                    pointer_offset,
+                    buf.len()
+                );
+            }
+            buf[pointer_offset..pointer_offset + 4].copy_from_slice(&arr);
+        }
+        // no remainder => store 0 as next overflow pointer
+        let pointer_offset = old_offset + local_bytes.saturating_sub(4);
+        for b in &mut page.as_ptr()[pointer_offset..pointer_offset + 4] {
+            *b = 0;
+        }
+        // zero out the rest of the local area
+        Ok(CursorResult::Ok(()))
+    }
+}
+
+// Compute how many bytes of a new payload to store locally on a b-tree table leaf page.
+// returns (local_size, needs_overflow).
+fn compute_local_needed(total_payload_size: usize, usable_space: u16) -> (usize, bool) {
+    // Compute the min/max thresholds.
+    let max_local = payload_overflow_threshold_max(PageType::TableLeaf, usable_space);
+    let min_local = payload_overflow_threshold_min(PageType::TableLeaf, usable_space);
+
+    // if the entire record fits within max_local, store it all on page.
+    if total_payload_size <= max_local {
+        return (total_payload_size, false);
+    }
+    let c = (usable_space as isize - 4).max(1) as usize;
+    let mut local = min_local + ((total_payload_size - min_local) % (c - 1));
+    if local > max_local {
+        local = min_local;
+    }
+    (local, true)
 }
 
 impl PageStack {

core/translate/mod.rs

@@ -23,6 +23,7 @@ pub(crate) mod result_row;
 pub(crate) mod select;
 pub(crate) mod subquery;
 pub(crate) mod transaction;
+pub(crate) mod update;
 
 use crate::schema::Schema;
 use crate::storage::pager::Pager;
@@ -40,6 +41,7 @@ use std::fmt::Display;
 use std::rc::{Rc, Weak};
 use std::sync::{Arc, Mutex};
 use transaction::{translate_tx_begin, translate_tx_commit};
+use update::translate_update;
 
 /// Translate SQL statement into bytecode program.
 pub fn translate(
@@ -108,7 +110,7 @@ pub fn translate(
         ast::Stmt::Rollback { .. } => bail_parse_error!("ROLLBACK not supported yet"),
         ast::Stmt::Savepoint(_) => bail_parse_error!("SAVEPOINT not supported yet"),
         ast::Stmt::Select(select) => translate_select(query_mode, schema, *select, syms)?,
-        ast::Stmt::Update { .. } => bail_parse_error!("UPDATE not supported yet"),
+        ast::Stmt::Update(update) => translate_update(query_mode, schema, &update, syms)?,
         ast::Stmt::Vacuum(_, _) => bail_parse_error!("VACUUM not supported yet"),
         ast::Stmt::Insert(insert) => {
             let Insert {