From 3b9f5aa51178b3abe3f7d589fe5dbd49a8130375 Mon Sep 17 00:00:00 2001
From: Pere Diaz Bou <pere-altea@hotmail.com>
Date: Fri, 26 Jul 2024 14:52:38 +0200
Subject: [PATCH] core: implement exists

---
 core/btree.rs          | 72 +++++++++++++++++++++++++++++++++---------
 core/pseudo.rs         |  2 +-
 core/sqlite3_ondisk.rs |  1 -
 core/types.rs          |  2 +-
 core/vdbe/mod.rs       |  7 ++--
 core/vdbe/sorter.rs    |  2 +-
 6 files changed, 64 insertions(+), 22 deletions(-)

diff --git a/core/btree.rs b/core/btree.rs
index fcd74b78c..a0388f269 100644
--- a/core/btree.rs
+++ b/core/btree.rs
@@ -232,20 +232,9 @@ impl BTreeCursor {
             OwnedValue::Integer(i) => *i as u64,
             _ => unreachable!("btree tables are indexed by integers!"),
         };
-        let mut cell_idx = 0;
-        for cell in &page.cells {
-            match cell {
-                BTreeCell::TableLeafCell(cell) => {
-                    if int_key <= cell._rowid {
-                        break;
-                    }
-                }
-                _ => todo!(),
-            }
-            cell_idx += 1;
-        }
+        let cell_idx = find_cell(page, int_key);
 
-        // if overwrite drop cell
+        // TODO: if overwrite drop cell
 
         // insert cell
         let mut payload: Vec<u8> = Vec::new();
@@ -571,7 +560,60 @@ impl Cursor for BTreeCursor {
         self.null_flag
     }
 
-    fn exists(&mut self, key: &OwnedValue) -> Result<bool> {
-        Ok(false)
+    fn exists(&mut self, key: &OwnedValue) -> Result<CursorResult<bool>> {
+        let int_key = match key {
+            OwnedValue::Integer(i) => i,
+            _ => unreachable!("btree tables are indexed by integers!"),
+        };
+        match self.move_to(*int_key as u64)? {
+            CursorResult::Ok(_) => {}
+            CursorResult::IO => return Ok(CursorResult::IO),
+        };
+        let mem_page = {
+            let mem_page = self.page.borrow();
+            let mem_page = mem_page.as_ref().unwrap();
+            mem_page.clone()
+        };
+        let page_idx = mem_page.page_idx;
+        let page_ref = self.pager.read_page(page_idx)?;
+        let page = page_ref.borrow();
+        if page.is_locked() {
+            return Ok(CursorResult::IO);
+        }
+
+        let page = page.contents.read().unwrap();
+        let page = page.as_ref().unwrap();
+
+        // find cell
+        let int_key = match key {
+            OwnedValue::Integer(i) => *i as u64,
+            _ => unreachable!("btree tables are indexed by integers!"),
+        };
+        let cell_idx = find_cell(page, int_key);
+        if cell_idx >= page.cells.len() {
+            Ok(CursorResult::Ok(false))
+        } else {
+            let equals = match &page.cells[cell_idx] {
+                BTreeCell::TableLeafCell(l) => l._rowid == int_key,
+                _ => unreachable!(),
+            };
+            Ok(CursorResult::Ok(equals))
+        }
     }
 }
+
+fn find_cell(page: &BTreePage, int_key: u64) -> usize {
+    let mut cell_idx = 0;
+    for cell in &page.cells {
+        match cell {
+            BTreeCell::TableLeafCell(cell) => {
+                if int_key <= cell._rowid {
+                    break;
+                }
+            }
+            _ => todo!(),
+        }
+        cell_idx += 1;
+    }
+    cell_idx
+}
diff --git a/core/pseudo.rs b/core/pseudo.rs
index ef35a9dc6..04881597d 100644
--- a/core/pseudo.rs
+++ b/core/pseudo.rs
@@ -64,7 +64,7 @@ impl Cursor for PseudoCursor {
         // Do nothing
     }
 
-    fn exists(&mut self, key: &OwnedValue) -> Result<bool> {
+    fn exists(&mut self, key: &OwnedValue) -> Result<CursorResult<bool>> {
         let _ = key;
         todo!()
     }
diff --git a/core/sqlite3_ondisk.rs b/core/sqlite3_ondisk.rs
index d2ff1cac4..d4834f465 100644
--- a/core/sqlite3_ondisk.rs
+++ b/core/sqlite3_ondisk.rs
@@ -334,7 +334,6 @@ pub fn begin_write_btree_page(pager: &Pager, page: &Rc<RefCell<Page>>) -> Result
             }
         })
     };
-    dbg!(buffer.borrow().len());
     let c = Rc::new(Completion::Write(WriteCompletion::new(write_complete)));
     page_source.write(page.id, buffer.clone(), c)?;
     Ok(())
diff --git a/core/types.rs b/core/types.rs
index 024541eae..2eb6c9ef9 100644
--- a/core/types.rs
+++ b/core/types.rs
@@ -371,7 +371,7 @@ pub trait Cursor {
     fn rowid(&self) -> Result<Option<u64>>;
     fn record(&self) -> Result<Ref<Option<OwnedRecord>>>;
     fn insert(&mut self, key: &OwnedValue, record: &OwnedRecord) -> Result<CursorResult<()>>;
-    fn exists(&mut self, key: &OwnedValue) -> Result<bool>;
+    fn exists(&mut self, key: &OwnedValue) -> Result<CursorResult<bool>>;
     fn set_null_flag(&mut self, flag: bool);
     fn get_null_flag(&self) -> bool;
 }
diff --git a/core/vdbe/mod.rs b/core/vdbe/mod.rs
index 6a2209acb..2fc2f8e17 100644
--- a/core/vdbe/mod.rs
+++ b/core/vdbe/mod.rs
@@ -1346,10 +1346,11 @@ impl Program {
                 } => {
                     let cursor = cursors.get_mut(cursor).unwrap();
                     match cursor.exists(&state.registers[*rowid_reg])? {
-                        true => state.pc += 1,
-                        false => state.pc = *target_pc,
+                        CursorResult::Ok(true) => state.pc += 1,
+                        CursorResult::Ok(false) => state.pc = *target_pc,
+                        CursorResult::IO => return Ok(StepResult::IO),
                     };
-                } // TODO(pere): how is not exists implemented? We probably need to traverse keys my pointing cursor.
+                }
                 // this cursor may be reused for next insert
                 // Update: tablemoveto is used to travers on not exists, on insert depending on flags if nonseek it traverses again.
                 // If not there might be some optimizations obviously.
diff --git a/core/vdbe/sorter.rs b/core/vdbe/sorter.rs
index 51bdc6b29..7ec6753c5 100644
--- a/core/vdbe/sorter.rs
+++ b/core/vdbe/sorter.rs
@@ -95,7 +95,7 @@ impl Cursor for Sorter {
         todo!();
     }
 
-    fn exists(&mut self, key: &OwnedValue) -> Result<bool> {
+    fn exists(&mut self, key: &OwnedValue) -> Result<CursorResult<bool>> {
         let _ = key;
         todo!()
     }