From a88ae21fcdb9e475b59c933352f0b6a5ccffa414 Mon Sep 17 00:00:00 2001
From: "Hernan G. Gonzalez" <camus.sonico@gmail.com>
Date: Tue, 9 Jan 2024 08:09:09 -0300
Subject: [PATCH] Fix page selection

---
 src/db.rs       |  3 ++-
 src/db/table.rs | 16 +++++++---------
 src/page.rs     |  6 +++++-
 3 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/src/db.rs b/src/db.rs
index 06b0322..3e95931 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -37,6 +37,7 @@ impl SQL for SQLiteFile {
     }
 
     fn page_at(&self, idx: i64) -> Result<Page> {
+        let idx = idx - 1;
         let offset = idx * self.head.page_size() as i64;
         decode::read_page(self, offset.into(), 0)
     }
@@ -44,7 +45,7 @@ impl SQL for SQLiteFile {
     fn table(&self, name: &str) -> Result<Table> {
         let schema = self.schema()?;
         let desc = schema.table_named(name)?.clone();
-        let root = self.page_at(desc.root - 1)?;
+        let root = self.page_at(desc.root)?;
         Ok(Table::new(self, desc, root))
     }
 }
diff --git a/src/db/table.rs b/src/db/table.rs
index bea9ecb..15c857a 100644
--- a/src/db/table.rs
+++ b/src/db/table.rs
@@ -56,26 +56,24 @@ impl<'a> Table<'a> {
         match kind {
             page::Kind::TableLeaf => {
                 let leave = self.root.into_leaf().unwrap();
-                Box::new(vec![leave].into_iter())
+                let iter = std::iter::once(leave);
+                Box::new(iter)
             }
             page::Kind::TableInterior => {
                 let interior = self.root.into_interior().unwrap();
+                let rhs = std::iter::once(interior.rhs())
+                    .flatten()
+                    .flat_map(|i| self.db.page_at(i as i64));
                 let iter = interior
                     .cells()
                     .into_iter()
                     .flat_map(|c| self.db.page_at(c.lhs as i64))
+                    .chain(rhs)
                     .flat_map(|p| p.into_leaf());
                 Box::new(iter)
             }
-            _ => Box::new(Vec::new().into_iter()),
+            _ => Box::new(std::iter::empty()),
         }
-        // if let Ok(leave) = self.root.into_leaf() {
-        //     return vec![leave].into_iter();
-        // }
-
-        // let Ok(interior) = self.root.into_interior() else {
-        //     return vec![].into_iter();
-        // };
     }
 
     pub fn select(
diff --git a/src/page.rs b/src/page.rs
index 4863f84..e608bfd 100644
--- a/src/page.rs
+++ b/src/page.rs
@@ -5,7 +5,7 @@ mod kind;
 use crate::offset::{self, Offset};
 pub use crate::page::kind::Kind;
 use crate::Result;
-use anyhow::{bail, ensure};
+use anyhow::{bail, ensure, Context};
 pub use cell::{Column, TableInteriorCell, TableLeafCell};
 pub use header::Header;
 
@@ -54,6 +54,10 @@ impl TableInteriorPage {
             .map(|r| r.1)
             .collect()
     }
+
+    pub fn rhs(&self) -> Result<u32> {
+        self.0.head.right_leave.context("Missing right leave")
+    }
 }
 
 pub mod decode {