Skip to main content

wasm_dbms_memory/
table_registry.rs

1// Rust guideline compliant 2026-02-28
2
3mod autoincrement_ledger;
4mod free_segments_ledger;
5mod index_ledger;
6mod page_ledger;
7mod raw_record;
8mod record_address;
9mod table_reader;
10mod write_at;
11
12use wasm_dbms_api::prelude::{Encode, MemoryResult, PageOffset, Value};
13
14pub use self::autoincrement_ledger::AutoincrementLedger;
15use self::free_segments_ledger::FreeSegmentsLedger;
16pub use self::index_ledger::{IndexLedger, IndexTreeWalker};
17use self::page_ledger::PageLedger;
18use self::raw_record::RawRecord;
19pub use self::record_address::RecordAddress;
20pub use self::table_reader::{NextRecord, TableReader};
21use self::write_at::WriteAt;
22use crate::{MemoryAccess, TableRegistryPage, align_up};
23
24/// The table registry takes care of storing the records for each table,
25/// using the [`FreeSegmentsLedger`] and [`PageLedger`] to derive exactly where to read/write.
26///
27/// A registry is generic over a record, which must implement [`Encode`].
28///
29/// The CRUD operations provided by the table registry do NOT perform any logical checks,
30/// but just allow to read/write records from/to memory.
31/// So CRUD checks must be performed by a higher layer, prior to calling these methods.
32pub struct TableRegistry {
33    free_segments_ledger: FreeSegmentsLedger,
34    pub(crate) page_ledger: PageLedger,
35    index_ledger: IndexLedger,
36    auto_increment_ledger: Option<AutoincrementLedger>,
37}
38
39impl TableRegistry {
40    /// Loads the table registry from memory.
41    pub fn load(table_pages: TableRegistryPage, mm: &mut impl MemoryAccess) -> MemoryResult<Self> {
42        Ok(Self {
43            free_segments_ledger: FreeSegmentsLedger::load(table_pages.free_segments_page, mm)?,
44            page_ledger: PageLedger::load(table_pages.pages_list_page, mm)?,
45            index_ledger: IndexLedger::load(table_pages.index_registry_page, mm)?,
46            auto_increment_ledger: if let Some(page) = table_pages.autoincrement_registry_page {
47                Some(AutoincrementLedger::load(page, mm)?)
48            } else {
49                None
50            },
51        })
52    }
53
54    /// Inserts a new record into the table registry.
55    ///
56    /// Returns the address where the record was inserted, which can be used to read it back or to update/delete it.
57    ///
58    /// NOTE: this function does NOT make any logical checks on the record being inserted.
59    pub fn insert<E>(
60        &mut self,
61        record: E,
62        mm: &mut impl MemoryAccess,
63    ) -> MemoryResult<RecordAddress>
64    where
65        E: Encode,
66    {
67        // get position to write the record
68        let raw_record = RawRecord::new(record);
69        let write_at = self.get_write_position(&raw_record, mm)?;
70
71        // align insert to RawRecord<E> alignment (includes the 2-byte header)
72        let aligned_offset = align_up::<RawRecord<E>>(write_at.offset() as usize) as PageOffset;
73
74        // write record
75        mm.write_at(write_at.page(), aligned_offset, &raw_record)?;
76
77        let pointer = RecordAddress {
78            page: write_at.page(),
79            offset: aligned_offset,
80        };
81
82        // commit post-write actions
83        self.post_write(write_at, &raw_record, mm)?;
84
85        Ok(pointer)
86    }
87
88    /// Creates a [`TableReader`] to read records from the table registry.
89    ///
90    /// Use [`TableReader::try_next`] to read records one by one.
91    pub fn read<'a, E, MA>(&'a self, mm: &'a mut MA) -> TableReader<'a, E, MA>
92    where
93        E: Encode,
94        MA: MemoryAccess,
95    {
96        TableReader::new(&self.page_ledger, mm)
97    }
98
99    /// Reads a single record at the given address.
100    pub fn read_at<E, MA>(&self, address: RecordAddress, mm: &mut MA) -> MemoryResult<E>
101    where
102        E: Encode,
103        MA: MemoryAccess,
104    {
105        let raw_record: RawRecord<E> = mm.read_at(address.page, address.offset)?;
106        Ok(raw_record.data)
107    }
108
109    /// Deletes a record at the given page and offset.
110    ///
111    /// The space occupied by the record is marked as free and zeroed.
112    pub fn delete(
113        &mut self,
114        record: impl Encode,
115        address: RecordAddress,
116        mm: &mut impl MemoryAccess,
117    ) -> MemoryResult<()> {
118        let raw_record = RawRecord::new(record);
119
120        // zero the record in memory
121        mm.zero(address.page, address.offset, &raw_record)?;
122
123        // insert a free segment for the deleted record
124        self.free_segments_ledger
125            .insert_free_segment(address.page, address.offset, &raw_record, mm)
126    }
127
128    /// Updates a record at the given page and offset.
129    ///
130    /// The [`RecordAddress`] of the new record is returned, which can be different from the old one if the record was reallocated.
131    ///
132    /// The logic is the following:
133    ///
134    /// 1. If the new record has exactly the same size of the old record, overwrite it in place.
135    /// 2. If the new record does not fit, delete the old record and insert the new record.
136    pub fn update(
137        &mut self,
138        new_record: impl Encode,
139        old_record: impl Encode,
140        old_address: RecordAddress,
141        mm: &mut impl MemoryAccess,
142    ) -> MemoryResult<RecordAddress> {
143        if new_record.size() == old_record.size() {
144            self.update_in_place(new_record, old_address, mm)
145        } else {
146            self.update_by_realloc(new_record, old_record, old_address, mm)
147        }
148    }
149
150    /// Get a reference to the index ledger, allowing to read the indexes.
151    pub fn index_ledger(&self) -> &IndexLedger {
152        &self.index_ledger
153    }
154
155    /// Get a mutable reference to the index ledger, allowing to modify the indexes.
156    pub fn index_ledger_mut(&mut self) -> &mut IndexLedger {
157        &mut self.index_ledger
158    }
159
160    /// Get next value for an autoincrement column of the given type, and increment it in the ledger.
161    pub fn next_autoincrement(
162        &mut self,
163        column_name: &str,
164        mm: &mut impl MemoryAccess,
165    ) -> MemoryResult<Option<Value>> {
166        if let Some(ledger) = &mut self.auto_increment_ledger {
167            ledger.next(column_name, mm).map(Some)
168        } else {
169            Ok(None)
170        }
171    }
172
173    /// Update a [`RawRecord`] in place at the given page and offset.
174    ///
175    /// The [`RecordAddress`] of the record is returned, which is the same as the old one.
176    ///
177    /// This must be used IF AND ONLY if the new record has the SAME size as the old record.
178    fn update_in_place(
179        &mut self,
180        record: impl Encode,
181        address: RecordAddress,
182        mm: &mut impl MemoryAccess,
183    ) -> MemoryResult<RecordAddress> {
184        let raw_record = RawRecord::new(record);
185        mm.write_at(address.page, address.offset, &raw_record)?;
186
187        Ok(address)
188    }
189
190    /// Updates a record by reallocating it.
191    ///
192    /// The old record is deleted and the new record is inserted.
193    ///
194    /// The [`RecordAddress`] of the new record is returned, which can be different from the old one.
195    fn update_by_realloc(
196        &mut self,
197        new_record: impl Encode,
198        old_record: impl Encode,
199        old_address: RecordAddress,
200        mm: &mut impl MemoryAccess,
201    ) -> MemoryResult<RecordAddress> {
202        // delete old record
203        self.delete(old_record, old_address, mm)?;
204
205        // insert new record
206        self.insert(new_record, mm)
207    }
208
209    /// Gets the position where to write a record of the given size.
210    fn get_write_position<E>(
211        &mut self,
212        record: &RawRecord<E>,
213        mm: &mut impl MemoryAccess,
214    ) -> MemoryResult<WriteAt>
215    where
216        E: Encode,
217    {
218        // check if there is a free segment that can hold the record
219        if let Some(segment) = self
220            .free_segments_ledger
221            .find_reusable_segment(record, mm)?
222        {
223            return Ok(WriteAt::ReusedSegment(segment));
224        }
225
226        // otherwise, write at the end of the table
227        self.page_ledger
228            .get_page_and_offset_for_record(record, mm)
229            .map(|(page, offset)| WriteAt::End(page, offset))
230    }
231
232    /// Commits the post-write actions after writing a record at the given position.
233    ///
234    /// - If the record was a [`WriteAt::ReusedSegment`], the free segment is marked as used.
235    /// - If the record was a [`WriteAt::End`], the page ledger is updated.
236    fn post_write<E>(
237        &mut self,
238        write_at: WriteAt,
239        record: &RawRecord<E>,
240        mm: &mut impl MemoryAccess,
241    ) -> MemoryResult<()>
242    where
243        E: Encode,
244    {
245        match write_at {
246            WriteAt::ReusedSegment(free_segment) => {
247                // mark segment as used
248                self.free_segments_ledger
249                    .commit_reused_space(record, free_segment, mm)
250            }
251            WriteAt::End(page, ..) => {
252                // update page ledger
253                self.page_ledger.commit(page, record, mm)
254            }
255        }
256    }
257}
258
259/// Test utilities shared across the table_registry submodules.
260#[cfg(test)]
261pub(crate) mod test_utils {
262    use wasm_dbms_api::prelude::{
263        DEFAULT_ALIGNMENT, DataSize, DecodeError, Encode, MSize, MemoryError, MemoryResult,
264        PageOffset,
265    };
266
267    /// A simple user struct for testing purposes (no macro dependencies).
268    #[derive(Debug, Clone, PartialEq, Eq)]
269    pub struct User {
270        pub id: u32,
271        pub name: String,
272        pub email: String,
273        pub age: u32,
274    }
275
276    impl Encode for User {
277        const SIZE: DataSize = DataSize::Dynamic;
278
279        const ALIGNMENT: PageOffset = DEFAULT_ALIGNMENT;
280
281        fn encode(&'_ self) -> std::borrow::Cow<'_, [u8]> {
282            let mut buf = Vec::new();
283            // id: 4 bytes
284            buf.extend_from_slice(&self.id.to_le_bytes());
285            // name length: 2 bytes + name bytes
286            buf.extend_from_slice(&(self.name.len() as u16).to_le_bytes());
287            buf.extend_from_slice(self.name.as_bytes());
288            // email length: 2 bytes + email bytes
289            buf.extend_from_slice(&(self.email.len() as u16).to_le_bytes());
290            buf.extend_from_slice(self.email.as_bytes());
291            // age: 4 bytes
292            buf.extend_from_slice(&self.age.to_le_bytes());
293            std::borrow::Cow::Owned(buf)
294        }
295
296        fn decode(data: std::borrow::Cow<[u8]>) -> MemoryResult<Self>
297        where
298            Self: Sized,
299        {
300            if data.len() < 12 {
301                return Err(MemoryError::DecodeError(DecodeError::TooShort));
302            }
303            let mut offset = 0;
304            // id
305            let id = u32::from_le_bytes(data[offset..offset + 4].try_into().unwrap());
306            offset += 4;
307            // name
308            let name_len =
309                u16::from_le_bytes(data[offset..offset + 2].try_into().unwrap()) as usize;
310            offset += 2;
311            let name = String::from_utf8_lossy(&data[offset..offset + name_len]).to_string();
312            offset += name_len;
313            // email
314            let email_len =
315                u16::from_le_bytes(data[offset..offset + 2].try_into().unwrap()) as usize;
316            offset += 2;
317            let email = String::from_utf8_lossy(&data[offset..offset + email_len]).to_string();
318            offset += email_len;
319            // age
320            let age = u32::from_le_bytes(data[offset..offset + 4].try_into().unwrap());
321
322            Ok(User {
323                id,
324                name,
325                email,
326                age,
327            })
328        }
329
330        fn size(&self) -> MSize {
331            (4 + 2 + self.name.len() + 2 + self.email.len() + 4) as MSize
332        }
333    }
334}
335
336#[cfg(test)]
337mod tests {
338
339    use self::test_utils::User;
340    use super::free_segments_ledger::FreeSegment;
341    use super::table_reader::NextRecord;
342    use super::*;
343    use crate::{HeapMemoryProvider, MemoryManager};
344
345    #[test]
346    fn test_should_create_table_registry() {
347        let mut mm = MemoryManager::init(HeapMemoryProvider::default());
348        let page_ledger_page = mm.allocate_page().expect("failed to get page");
349        let free_segments_page = mm.allocate_page().expect("failed to get page");
350        let index_registry_page = mm.allocate_page().expect("failed to get page");
351        let autoincrement_page = mm.allocate_page().expect("failed to get page");
352        let table_pages = TableRegistryPage {
353            pages_list_page: page_ledger_page,
354            free_segments_page,
355            index_registry_page,
356            autoincrement_registry_page: Some(autoincrement_page),
357        };
358
359        let registry: MemoryResult<TableRegistry> = TableRegistry::load(table_pages, &mut mm);
360        assert!(registry.is_ok());
361    }
362
363    #[test]
364    fn test_should_get_write_at_end() {
365        let mut mm = MemoryManager::init(HeapMemoryProvider::default());
366        let mut registry = registry(&mut mm);
367
368        let record = RawRecord::new(User {
369            id: 1,
370            name: "Test".to_string(),
371            email: "new_user@example.com".to_string(),
372            age: 25,
373        });
374        let write_at = registry
375            .get_write_position(&record, &mut mm)
376            .expect("failed to get write at");
377
378        assert!(matches!(write_at, WriteAt::End(_, 0)));
379    }
380
381    #[test]
382    fn test_should_get_write_at_free_segment() {
383        let mut mm = MemoryManager::init(HeapMemoryProvider::default());
384        let mut registry = registry(&mut mm);
385
386        let record = RawRecord::new(User {
387            id: 1,
388            name: "Test".to_string(),
389            email: "new_user@example.com".to_string(),
390            age: 25,
391        });
392        // allocate a page to insert a free segment
393        let (page, _) = registry
394            .page_ledger
395            .get_page_and_offset_for_record(&record, &mut mm)
396            .expect("failed to get page and offset");
397        registry
398            .page_ledger
399            .commit(page, &record, &mut mm)
400            .expect("failed to commit page ledger");
401        // insert data about a free segment
402        registry
403            .free_segments_ledger
404            .insert_free_segment(page, 256, &record, &mut mm)
405            .expect("failed to insert free segment");
406
407        let write_at = registry
408            .get_write_position(&record, &mut mm)
409            .expect("failed to get write at");
410
411        let reused_segment = match write_at {
412            WriteAt::ReusedSegment(segment) => segment.segment,
413            _ => panic!("expected reused segment"),
414        };
415
416        assert_eq!(
417            reused_segment,
418            FreeSegment {
419                page,
420                offset: 256,
421                size: 64, // padded size
422            }
423        );
424    }
425
426    #[test]
427    fn test_should_insert_record_into_table_registry() {
428        let mut mm = MemoryManager::init(HeapMemoryProvider::default());
429        let mut registry = registry(&mut mm);
430
431        let record = User {
432            id: 1,
433            name: "Test".to_string(),
434            email: "new_user@example.com".to_string(),
435            age: 25,
436        };
437
438        // insert record
439        assert!(registry.insert(record, &mut mm).is_ok());
440    }
441
442    #[test]
443    fn test_should_manage_to_insert_users_to_exceed_one_page() {
444        let mut mm = MemoryManager::init(HeapMemoryProvider::default());
445        let mut registry = registry(&mut mm);
446
447        for id in 0..4000 {
448            let record = User {
449                id,
450                name: format!("User {}", id),
451                email: "new_user@example.com".to_string(),
452                age: 20 + id,
453            };
454            registry
455                .insert(record, &mut mm)
456                .expect("failed to insert record");
457        }
458    }
459
460    #[test]
461    fn test_should_delete_record() {
462        let mut mm = MemoryManager::init(HeapMemoryProvider::default());
463        let mut registry = registry(&mut mm);
464
465        let record = User {
466            id: 1,
467            name: "Test".to_string(),
468            email: "new_user@example.com".to_string(),
469            age: 25,
470        };
471
472        // insert record
473        registry
474            .insert(record.clone(), &mut mm)
475            .expect("failed to insert");
476
477        // find where it was written
478        let mut reader = registry.read(&mut mm);
479        let next_record: NextRecord<User> = reader
480            .try_next()
481            .expect("failed to read")
482            .expect("no record");
483        let page = next_record.page;
484        let offset = next_record.offset;
485        let record = next_record.record;
486        let raw_user = RawRecord::new(record.clone());
487        let raw_user_size = raw_user.size();
488
489        // delete record
490        assert!(
491            registry
492                .delete(record, RecordAddress { page, offset }, &mut mm)
493                .is_ok()
494        );
495
496        // should have been deleted
497        let mut reader = registry.read::<User, _>(&mut mm);
498        assert!(reader.try_next().expect("failed to read").is_none());
499
500        // should have a free segment
501        let free_segment = registry
502            .free_segments_ledger
503            .find_reusable_segment(
504                &User {
505                    id: 2,
506                    name: "Test".to_string(),
507                    email: "new_user@example.com".to_string(),
508                    age: 25,
509                },
510                &mut mm,
511            )
512            .expect("failed to find free segment")
513            .expect("could not find the free segment after free")
514            .segment;
515        assert_eq!(free_segment.page, page);
516        assert_eq!(free_segment.offset, offset);
517        assert_eq!(free_segment.size, 64); // padded
518
519        // should have zeroed the memory
520        let mut buffer = vec![0u8; raw_user_size as usize];
521        mm.read_at_raw(page, offset, &mut buffer)
522            .expect("failed to read memory");
523        assert!(buffer.iter().all(|&b| b == 0));
524    }
525
526    #[test]
527    fn test_read_at_returns_record_at_address() {
528        let mut mm = MemoryManager::init(HeapMemoryProvider::default());
529        let mut registry = registry(&mut mm);
530        let record = User {
531            id: 1,
532            name: "Alice".to_string(),
533            email: "alice@example.com".to_string(),
534            age: 30,
535        };
536
537        let address = registry
538            .insert(record.clone(), &mut mm)
539            .expect("failed to insert record");
540
541        let stored: User = registry
542            .read_at(address, &mut mm)
543            .expect("failed to read record");
544        assert_eq!(stored, record);
545    }
546
547    #[test]
548    fn test_read_at_after_update_returns_updated_record() {
549        let mut mm = MemoryManager::init(HeapMemoryProvider::default());
550        let mut registry = registry(&mut mm);
551        let old_record = User {
552            id: 1,
553            name: "Alice".to_string(),
554            email: "alice@example.com".to_string(),
555            age: 30,
556        };
557        let new_record = User {
558            id: 1,
559            name: "Alice Updated".to_string(),
560            email: "alice.updated@example.com".to_string(),
561            age: 31,
562        };
563
564        let old_address = registry
565            .insert(old_record.clone(), &mut mm)
566            .expect("failed to insert record");
567        let new_address = registry
568            .update(new_record.clone(), old_record, old_address, &mut mm)
569            .expect("failed to update record");
570
571        let stored: User = registry
572            .read_at(new_address, &mut mm)
573            .expect("failed to read updated record");
574        assert_eq!(stored, new_record);
575    }
576
577    #[test]
578    fn test_should_update_record_in_place() {
579        let mut mm = MemoryManager::init(HeapMemoryProvider::default());
580        let mut registry = registry(&mut mm);
581
582        let old_record = User {
583            id: 1,
584            name: "John".to_string(),
585            email: "new_user@example.com".to_string(),
586            age: 28,
587        };
588        let new_record = User {
589            id: 1,
590            name: "Mark".to_string(), // same length as "John"
591            email: "new_user@example.com".to_string(),
592            age: 30,
593        };
594
595        // insert old record
596        registry
597            .insert(old_record.clone(), &mut mm)
598            .expect("failed to insert");
599
600        // find where it was written
601        let mut reader = registry.read::<User, _>(&mut mm);
602        let next_record = reader
603            .try_next()
604            .expect("failed to read")
605            .expect("no record");
606        let page = next_record.page;
607        let offset = next_record.offset;
608
609        // update in place
610        let old_address = RecordAddress { page, offset };
611        let new_location = registry
612            .update(
613                new_record.clone(),
614                next_record.record.clone(),
615                old_address,
616                &mut mm,
617            )
618            .expect("failed to update record");
619        assert_eq!(new_location, old_address); // should be same address
620
621        // read back the record
622        let mut reader = registry.read::<User, _>(&mut mm);
623        let next_record = reader
624            .try_next()
625            .expect("failed to read")
626            .expect("no record");
627        assert_eq!(next_record.page, page); // should be same page
628        assert_eq!(next_record.offset, offset); // should be same offset
629        assert_eq!(next_record.record, new_record);
630    }
631
632    #[test]
633    fn test_should_update_record_reallocating() {
634        let mut mm = MemoryManager::init(HeapMemoryProvider::default());
635        let mut registry = registry(&mut mm);
636
637        let old_record = User {
638            id: 1,
639            name: "John".to_string(),
640            email: "new_user@example.com".to_string(),
641            age: 28,
642        };
643        // this user creates a record with same size as old_record to avoid reusing the free segment
644        let extra_user = User {
645            id: 2,
646            name: "Extra".to_string(),
647            email: "new_user@example.com".to_string(),
648            age: 25,
649        };
650        let new_record = User {
651            id: 1,
652            name: "Alexanderejruwgjowergjioewrgjioewrigjewriogjweoirgjiowerjgoiwerjiogewirogjowejrgiwer".to_string(), // must exceed padding
653            email: "new_user@example.com".to_string(),
654            age: 30,
655        };
656
657        // insert old record
658        registry
659            .insert(old_record.clone(), &mut mm)
660            .expect("failed to insert");
661        // insert extra record to avoid reusing the free segment
662        registry
663            .insert(extra_user.clone(), &mut mm)
664            .expect("failed to insert extra user");
665
666        // find where it was written
667        let mut reader = registry.read::<User, _>(&mut mm);
668        let old_record_from_db = reader
669            .try_next()
670            .expect("failed to read")
671            .expect("no record");
672        assert_eq!(old_record_from_db.record, old_record);
673        let page = old_record_from_db.page;
674        let offset = old_record_from_db.offset;
675
676        // update by reallocating
677        let old_address = RecordAddress { page, offset };
678        let new_location = registry
679            .update(
680                new_record.clone(),
681                old_record_from_db.record.clone(),
682                old_address,
683                &mut mm,
684            )
685            .expect("failed to update record");
686        assert_ne!(new_location, old_address); // should be different page
687
688        // read back the record
689        let mut reader = registry.read::<User, _>(&mut mm);
690
691        // find extra record first
692        let _ = reader
693            .try_next()
694            .expect("failed to read")
695            .expect("no record");
696
697        let updated_record = reader
698            .try_next()
699            .expect("failed to read")
700            .expect("no record");
701        assert_ne!(updated_record.offset, offset); // should be different offset
702        assert_eq!(updated_record.record, new_record);
703    }
704
705    #[test]
706    fn test_should_insert_delete_insert_many() {
707        const COUNT: u32 = 1_000;
708        let mut mm = MemoryManager::init(HeapMemoryProvider::default());
709        let mut registry = registry(&mut mm);
710        for id in 0..COUNT {
711            let record = User {
712                id,
713                name: format!("User {id}"),
714                email: format!("user_{id}@example.com"),
715                age: 20,
716            };
717
718            // insert record
719            registry
720                .insert(record.clone(), &mut mm)
721                .expect("failed to insert");
722        }
723
724        // delete odd records
725        for id in (0..COUNT).filter(|id| id % 2 == 1) {
726            let record = User {
727                id,
728                name: format!("User {id}"),
729                email: format!("user_{id}@example.com"),
730                age: 20,
731            };
732            // find where it was written
733            let mut reader = registry.read::<User, _>(&mut mm);
734            let mut deleted = false;
735            while let Some(next_record) = reader.try_next().expect("failed to read") {
736                if next_record.record.id == id {
737                    registry
738                        .delete(
739                            record.clone(),
740                            RecordAddress {
741                                page: next_record.page,
742                                offset: next_record.offset,
743                            },
744                            &mut mm,
745                        )
746                        .expect("failed to delete");
747                    deleted = true;
748                    break;
749                }
750            }
751            assert!(deleted, "record with id {} was not found", id);
752        }
753
754        // now delete also the others
755        for id in (0..COUNT).filter(|id| id % 2 == 0) {
756            let record = User {
757                id,
758                name: format!("User {id}"),
759                email: format!("user_{id}@example.com"),
760                age: 20,
761            };
762            // find where it was written
763            let mut reader = registry.read::<User, _>(&mut mm);
764            let mut deleted = false;
765            while let Some(next_record) = reader.try_next().expect("failed to read") {
766                if next_record.record.id == id {
767                    registry
768                        .delete(
769                            record.clone(),
770                            RecordAddress {
771                                page: next_record.page,
772                                offset: next_record.offset,
773                            },
774                            &mut mm,
775                        )
776                        .expect("failed to delete");
777                    deleted = true;
778                    break;
779                }
780            }
781            assert!(deleted, "record with id {} was not found", id);
782        }
783
784        // insert back
785        for id in 0..COUNT {
786            let record = User {
787                id,
788                name: format!("User {id}"),
789                email: format!("user_{id}@example.com"),
790                age: 20,
791            };
792
793            // insert record
794            registry
795                .insert(record.clone(), &mut mm)
796                .expect("failed to insert");
797        }
798    }
799
800    #[test]
801    fn test_should_reduce_free_segment_size_with_padding() {
802        let mut mm = MemoryManager::init(HeapMemoryProvider::default());
803        let mut registry = registry(&mut mm);
804
805        // first insert a user
806        let long_name = vec!['A'; 1024].into_iter().collect::<String>();
807        let record = User {
808            id: 1,
809            name: "Test User".to_string(),
810            email: long_name,
811            age: 30,
812        };
813        registry
814            .insert(record.clone(), &mut mm)
815            .expect("failed to insert");
816        // get record page
817        let mut reader = registry.read::<User, _>(&mut mm);
818        let next_record = reader
819            .try_next()
820            .expect("failed to read")
821            .expect("no record");
822        // delete user
823        registry
824            .delete(
825                next_record.record,
826                RecordAddress {
827                    page: next_record.page,
828                    offset: next_record.offset,
829                },
830                &mut mm,
831            )
832            .expect("failed to delete");
833
834        // get the free segment
835        let raw_record = RawRecord::new(record.clone());
836        let free_segment = registry
837            .free_segments_ledger
838            .find_reusable_segment(&raw_record, &mut mm)
839            .expect("failed to find reusable segment")
840            .expect("could not find the free segment after free")
841            .segment;
842        // size should be at least 1024
843        assert!(free_segment.size >= 1024);
844        let previous_size = free_segment.size;
845
846        // now insert a small user at 0
847        let small_record = User {
848            id: 2,
849            name: "Bob The Builder".to_string(),
850            email: "bob@hotmail.com".to_string(),
851            age: 22,
852        };
853        registry
854            .insert(small_record.clone(), &mut mm)
855            .expect("failed to insert small user");
856
857        // get free segment
858        let free_segment_after = registry
859            .free_segments_ledger
860            .find_reusable_segment(&small_record, &mut mm)
861            .expect("failed to find reusable segment")
862            .expect("could not find the free segment after inserting small user")
863            .segment;
864
865        // size should be reduced
866        assert_eq!(
867            free_segment_after.offset, 64,
868            "expected offset to be 64, but had: {}",
869            free_segment_after.offset
870        ); // which is the padding
871        assert_eq!(
872            free_segment_after.size,
873            previous_size - 64,
874            "Expected free segment to have size: {} but got: {}",
875            previous_size - 64,
876            free_segment_after.size
877        );
878    }
879
880    fn registry(mm: &mut MemoryManager<HeapMemoryProvider>) -> TableRegistry {
881        let page_ledger_page = mm.allocate_page().expect("failed to get page");
882        let free_segments_page = mm.allocate_page().expect("failed to get page");
883        let index_registry_page = mm.allocate_page().expect("failed to get page");
884        let autoincrement_page = mm.allocate_page().expect("failed to get page");
885        let table_pages = TableRegistryPage {
886            pages_list_page: page_ledger_page,
887            free_segments_page,
888            index_registry_page,
889            autoincrement_registry_page: Some(autoincrement_page),
890        };
891
892        TableRegistry::load(table_pages, mm).expect("failed to load")
893    }
894
895    /// Creates a [`TableRegistry`] with a properly initialized autoincrement ledger
896    /// via [`SchemaRegistry::register_table`].
897    fn registry_with_autoincrement(mm: &mut MemoryManager<HeapMemoryProvider>) -> TableRegistry {
898        use crate::SchemaRegistry;
899
900        let mut schema = SchemaRegistry::load(mm).expect("failed to load schema");
901        let pages = schema
902            .register_table::<AutoincUser>(mm)
903            .expect("failed to register table");
904        TableRegistry::load(pages, mm).expect("failed to load")
905    }
906
907    /// Creates a [`TableRegistry`] without an autoincrement ledger.
908    fn registry_without_autoincrement(mm: &mut MemoryManager<HeapMemoryProvider>) -> TableRegistry {
909        let page_ledger_page = mm.allocate_page().expect("failed to get page");
910        let free_segments_page = mm.allocate_page().expect("failed to get page");
911        let index_registry_page = mm.allocate_page().expect("failed to get page");
912        let table_pages = TableRegistryPage {
913            pages_list_page: page_ledger_page,
914            free_segments_page,
915            index_registry_page,
916            autoincrement_registry_page: None,
917        };
918
919        TableRegistry::load(table_pages, mm).expect("failed to load")
920    }
921
922    // -- AutoincUser mock: a table with an autoincrement Uint32 column --
923
924    use candid::CandidType;
925    use serde::{Deserialize, Serialize};
926    use wasm_dbms_api::prelude::{
927        ColumnDef, DbmsResult, IndexDef, InsertRecord, NoForeignFetcher, TableColumns, TableRecord,
928        TableSchema, UpdateRecord,
929    };
930
931    #[derive(Clone, CandidType)]
932    struct AutoincUser;
933
934    impl Encode for AutoincUser {
935        const SIZE: wasm_dbms_api::prelude::DataSize = wasm_dbms_api::prelude::DataSize::Dynamic;
936        const ALIGNMENT: PageOffset = wasm_dbms_api::prelude::DEFAULT_ALIGNMENT;
937
938        fn encode(&'_ self) -> std::borrow::Cow<'_, [u8]> {
939            std::borrow::Cow::Owned(vec![])
940        }
941
942        fn decode(_data: std::borrow::Cow<[u8]>) -> MemoryResult<Self>
943        where
944            Self: Sized,
945        {
946            Ok(Self)
947        }
948
949        fn size(&self) -> wasm_dbms_api::prelude::MSize {
950            0
951        }
952    }
953
954    #[derive(Clone, CandidType, Deserialize)]
955    struct AutoincUserRecord;
956
957    impl TableRecord for AutoincUserRecord {
958        type Schema = AutoincUser;
959
960        fn from_values(_values: TableColumns) -> Self {
961            Self
962        }
963
964        fn to_values(&self) -> Vec<(ColumnDef, Value)> {
965            vec![]
966        }
967    }
968
969    #[derive(Clone, CandidType, Serialize)]
970    struct AutoincUserInsert;
971
972    impl InsertRecord for AutoincUserInsert {
973        type Record = AutoincUserRecord;
974        type Schema = AutoincUser;
975
976        fn from_values(_values: &[(ColumnDef, Value)]) -> DbmsResult<Self> {
977            Ok(Self)
978        }
979
980        fn into_values(self) -> Vec<(ColumnDef, Value)> {
981            vec![]
982        }
983
984        fn into_record(self) -> Self::Schema {
985            AutoincUser
986        }
987    }
988
989    #[derive(Clone, CandidType, Serialize)]
990    struct AutoincUserUpdate;
991
992    impl UpdateRecord for AutoincUserUpdate {
993        type Record = AutoincUserRecord;
994        type Schema = AutoincUser;
995
996        fn from_values(
997            _values: &[(ColumnDef, Value)],
998            _where_clause: Option<wasm_dbms_api::prelude::Filter>,
999        ) -> Self {
1000            Self
1001        }
1002
1003        fn update_values(&self) -> Vec<(ColumnDef, Value)> {
1004            vec![]
1005        }
1006
1007        fn where_clause(&self) -> Option<wasm_dbms_api::prelude::Filter> {
1008            None
1009        }
1010    }
1011
1012    impl TableSchema for AutoincUser {
1013        type Record = AutoincUserRecord;
1014        type Insert = AutoincUserInsert;
1015        type Update = AutoincUserUpdate;
1016        type ForeignFetcher = NoForeignFetcher;
1017
1018        fn table_name() -> &'static str {
1019            "autoinc_users"
1020        }
1021
1022        fn columns() -> &'static [ColumnDef] {
1023            use wasm_dbms_api::prelude::DataTypeKind;
1024
1025            &[ColumnDef {
1026                name: "id",
1027                data_type: DataTypeKind::Uint32,
1028                auto_increment: true,
1029                nullable: false,
1030                primary_key: true,
1031                unique: true,
1032                foreign_key: None,
1033            }]
1034        }
1035
1036        fn primary_key() -> &'static str {
1037            "id"
1038        }
1039
1040        fn indexes() -> &'static [IndexDef] {
1041            &[IndexDef(&["id"])]
1042        }
1043
1044        fn to_values(self) -> Vec<(ColumnDef, Value)> {
1045            vec![]
1046        }
1047
1048        fn sanitizer(
1049            _column_name: &'static str,
1050        ) -> Option<Box<dyn wasm_dbms_api::prelude::Sanitize>> {
1051            None
1052        }
1053
1054        fn validator(
1055            _column_name: &'static str,
1056        ) -> Option<Box<dyn wasm_dbms_api::prelude::Validate>> {
1057            None
1058        }
1059    }
1060
1061    // -- next_autoincrement tests --
1062
1063    #[test]
1064    fn test_next_autoincrement_returns_sequential_values() {
1065        let mut mm = MemoryManager::init(HeapMemoryProvider::default());
1066        let mut registry = registry_with_autoincrement(&mut mm);
1067
1068        let v1 = registry
1069            .next_autoincrement("id", &mut mm)
1070            .expect("failed")
1071            .expect("expected Some");
1072        let v2 = registry
1073            .next_autoincrement("id", &mut mm)
1074            .expect("failed")
1075            .expect("expected Some");
1076        let v3 = registry
1077            .next_autoincrement("id", &mut mm)
1078            .expect("failed")
1079            .expect("expected Some");
1080
1081        assert_eq!(v1, Value::Uint32(1u32.into()));
1082        assert_eq!(v2, Value::Uint32(2u32.into()));
1083        assert_eq!(v3, Value::Uint32(3u32.into()));
1084    }
1085
1086    #[test]
1087    fn test_next_autoincrement_returns_none_without_ledger() {
1088        let mut mm = MemoryManager::init(HeapMemoryProvider::default());
1089        let mut registry = registry_without_autoincrement(&mut mm);
1090
1091        let result = registry.next_autoincrement("id", &mut mm).expect("failed");
1092        assert!(result.is_none());
1093    }
1094
1095    #[test]
1096    fn test_next_autoincrement_persists_across_reload() {
1097        let mut mm = MemoryManager::init(HeapMemoryProvider::default());
1098
1099        use crate::SchemaRegistry;
1100        let mut schema = SchemaRegistry::load(&mut mm).expect("failed to load schema");
1101        let pages = schema
1102            .register_table::<AutoincUser>(&mut mm)
1103            .expect("failed to register table");
1104
1105        // advance 5 times
1106        let mut registry = TableRegistry::load(pages, &mut mm).expect("failed to load");
1107        for _ in 0..5 {
1108            let _ = registry
1109                .next_autoincrement("id", &mut mm)
1110                .expect("next failed");
1111        }
1112
1113        // reload the registry from the same pages
1114        let mut reloaded = TableRegistry::load(pages, &mut mm).expect("failed to reload");
1115        let value = reloaded
1116            .next_autoincrement("id", &mut mm)
1117            .expect("failed")
1118            .expect("expected Some");
1119        assert_eq!(value, Value::Uint32(6u32.into()));
1120    }
1121
1122    #[test]
1123    fn test_next_autoincrement_overflow_returns_error() {
1124        let mut mm = MemoryManager::init(HeapMemoryProvider::default());
1125
1126        // manually set up a Uint8 autoincrement to hit overflow quickly
1127        let page_ledger_page = mm.allocate_page().expect("failed to get page");
1128        let free_segments_page = mm.allocate_page().expect("failed to get page");
1129        let index_registry_page = mm.allocate_page().expect("failed to get page");
1130        let autoinc_page = mm.allocate_page().expect("failed to get page");
1131
1132        // Use the Uint8AutoincTable from the autoincrement_ledger tests — we replicate the
1133        // TableSchema inline since it's in a sibling test module.
1134        // Instead, just init the ledger page directly with a Uint8 value.
1135        {
1136            let mut registry_data = super::autoincrement_ledger::AutoincrementLedger::init::<
1137                Uint8AutoincSchema,
1138            >(autoinc_page, &mut mm)
1139            .expect("failed to init autoinc ledger");
1140
1141            // advance to 255
1142            for _ in 0..255 {
1143                let _ = registry_data.next("val", &mut mm).expect("next failed");
1144            }
1145        }
1146
1147        // init index ledger
1148        IndexLedger::init(index_registry_page, &[], &mut mm).expect("failed to init index ledger");
1149
1150        let table_pages = TableRegistryPage {
1151            pages_list_page: page_ledger_page,
1152            free_segments_page,
1153            index_registry_page,
1154            autoincrement_registry_page: Some(autoinc_page),
1155        };
1156
1157        let mut registry = TableRegistry::load(table_pages, &mut mm).expect("failed to load");
1158        let result = registry.next_autoincrement("val", &mut mm);
1159        assert!(result.is_err());
1160        assert!(matches!(
1161            result.unwrap_err(),
1162            wasm_dbms_api::prelude::MemoryError::AutoincrementOverflow(_)
1163        ));
1164    }
1165
1166    // Minimal Uint8 autoincrement schema for overflow test
1167
1168    #[derive(Clone, CandidType)]
1169    struct Uint8AutoincSchema;
1170
1171    impl Encode for Uint8AutoincSchema {
1172        const SIZE: wasm_dbms_api::prelude::DataSize = wasm_dbms_api::prelude::DataSize::Dynamic;
1173        const ALIGNMENT: PageOffset = wasm_dbms_api::prelude::DEFAULT_ALIGNMENT;
1174
1175        fn encode(&'_ self) -> std::borrow::Cow<'_, [u8]> {
1176            std::borrow::Cow::Owned(vec![])
1177        }
1178
1179        fn decode(_data: std::borrow::Cow<[u8]>) -> MemoryResult<Self>
1180        where
1181            Self: Sized,
1182        {
1183            Ok(Self)
1184        }
1185
1186        fn size(&self) -> wasm_dbms_api::prelude::MSize {
1187            0
1188        }
1189    }
1190
1191    #[derive(Clone, CandidType, Deserialize)]
1192    struct Uint8AutoincSchemaRecord;
1193
1194    impl TableRecord for Uint8AutoincSchemaRecord {
1195        type Schema = Uint8AutoincSchema;
1196
1197        fn from_values(_values: TableColumns) -> Self {
1198            Self
1199        }
1200
1201        fn to_values(&self) -> Vec<(ColumnDef, Value)> {
1202            vec![]
1203        }
1204    }
1205
1206    #[derive(Clone, CandidType, Serialize)]
1207    struct Uint8AutoincSchemaInsert;
1208
1209    impl InsertRecord for Uint8AutoincSchemaInsert {
1210        type Record = Uint8AutoincSchemaRecord;
1211        type Schema = Uint8AutoincSchema;
1212
1213        fn from_values(_values: &[(ColumnDef, Value)]) -> DbmsResult<Self> {
1214            Ok(Self)
1215        }
1216
1217        fn into_values(self) -> Vec<(ColumnDef, Value)> {
1218            vec![]
1219        }
1220
1221        fn into_record(self) -> Self::Schema {
1222            Uint8AutoincSchema
1223        }
1224    }
1225
1226    #[derive(Clone, CandidType, Serialize)]
1227    struct Uint8AutoincSchemaUpdate;
1228
1229    impl UpdateRecord for Uint8AutoincSchemaUpdate {
1230        type Record = Uint8AutoincSchemaRecord;
1231        type Schema = Uint8AutoincSchema;
1232
1233        fn from_values(
1234            _values: &[(ColumnDef, Value)],
1235            _where_clause: Option<wasm_dbms_api::prelude::Filter>,
1236        ) -> Self {
1237            Self
1238        }
1239
1240        fn update_values(&self) -> Vec<(ColumnDef, Value)> {
1241            vec![]
1242        }
1243
1244        fn where_clause(&self) -> Option<wasm_dbms_api::prelude::Filter> {
1245            None
1246        }
1247    }
1248
1249    impl TableSchema for Uint8AutoincSchema {
1250        type Record = Uint8AutoincSchemaRecord;
1251        type Insert = Uint8AutoincSchemaInsert;
1252        type Update = Uint8AutoincSchemaUpdate;
1253        type ForeignFetcher = NoForeignFetcher;
1254
1255        fn table_name() -> &'static str {
1256            "uint8_autoinc_schema"
1257        }
1258
1259        fn columns() -> &'static [ColumnDef] {
1260            use wasm_dbms_api::prelude::DataTypeKind;
1261
1262            &[ColumnDef {
1263                name: "val",
1264                data_type: DataTypeKind::Uint8,
1265                auto_increment: true,
1266                nullable: false,
1267                primary_key: true,
1268                unique: true,
1269                foreign_key: None,
1270            }]
1271        }
1272
1273        fn primary_key() -> &'static str {
1274            "val"
1275        }
1276
1277        fn indexes() -> &'static [IndexDef] {
1278            &[IndexDef(&["val"])]
1279        }
1280
1281        fn to_values(self) -> Vec<(ColumnDef, Value)> {
1282            vec![]
1283        }
1284
1285        fn sanitizer(
1286            _column_name: &'static str,
1287        ) -> Option<Box<dyn wasm_dbms_api::prelude::Sanitize>> {
1288            None
1289        }
1290
1291        fn validator(
1292            _column_name: &'static str,
1293        ) -> Option<Box<dyn wasm_dbms_api::prelude::Validate>> {
1294            None
1295        }
1296    }
1297
1298    use wasm_dbms_api::prelude::{DataSize, DecodeError, MSize, MemoryError};
1299
1300    /// A fixed-size record for regression testing (issue #80).
1301    ///
1302    /// Layout: u64 (8) + u64 (8) + u8 (1) = 17 bytes, all fixed-size.
1303    #[derive(Debug, Clone, PartialEq, Eq)]
1304    struct FixedSizeRecord {
1305        id: u64,
1306        timestamp: u64,
1307        tag: u8,
1308    }
1309
1310    impl Encode for FixedSizeRecord {
1311        const SIZE: DataSize = DataSize::Fixed(17);
1312        const ALIGNMENT: PageOffset = 17;
1313
1314        fn encode(&'_ self) -> std::borrow::Cow<'_, [u8]> {
1315            let mut buf = Vec::with_capacity(17);
1316            buf.extend_from_slice(&self.id.to_le_bytes());
1317            buf.extend_from_slice(&self.timestamp.to_le_bytes());
1318            buf.push(self.tag);
1319            std::borrow::Cow::Owned(buf)
1320        }
1321
1322        fn decode(data: std::borrow::Cow<[u8]>) -> MemoryResult<Self>
1323        where
1324            Self: Sized,
1325        {
1326            if data.len() < 17 {
1327                return Err(MemoryError::DecodeError(DecodeError::TooShort));
1328            }
1329            let id = u64::from_le_bytes(data[0..8].try_into().unwrap());
1330            let timestamp = u64::from_le_bytes(data[8..16].try_into().unwrap());
1331            let tag = data[16];
1332            Ok(Self { id, timestamp, tag })
1333        }
1334
1335        fn size(&self) -> MSize {
1336            17
1337        }
1338    }
1339
1340    #[test]
1341    fn test_should_insert_multiple_fixed_size_records() {
1342        let mut mm = MemoryManager::init(HeapMemoryProvider::default());
1343        let mut registry = registry(&mut mm);
1344
1345        for i in 0..10 {
1346            let record = FixedSizeRecord {
1347                id: i,
1348                timestamp: 1000 + i,
1349                tag: (i % 2) as u8,
1350            };
1351            registry
1352                .insert(record, &mut mm)
1353                .unwrap_or_else(|e| panic!("failed to insert record {i}: {e}"));
1354        }
1355
1356        // verify all records can be read back
1357        let mut reader = registry.read::<FixedSizeRecord, _>(&mut mm);
1358        let mut count = 0;
1359        while let Some(next) = reader.try_next().expect("failed to read") {
1360            assert_eq!(next.record.id, count);
1361            count += 1;
1362        }
1363        assert_eq!(count, 10);
1364    }
1365}