port ourdb from vlang

This commit is contained in:
Timur Gordon
2025-04-09 11:37:11 +02:00
parent 5c5225c8f7
commit 0eedec9ed0
14 changed files with 3423 additions and 0 deletions

335
ourdb/src/backend.rs Normal file
View File

@@ -0,0 +1,335 @@
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Seek, SeekFrom, Write};
use crc32fast::Hasher;
use crate::error::Error;
use crate::location::Location;
use crate::OurDB;
// Header size: 2 bytes (size) + 4 bytes (CRC32) + 6 bytes (previous location)
pub const HEADER_SIZE: usize = 12;
impl OurDB {
/// Selects and opens a database file for read/write operations
pub(crate) fn db_file_select(&mut self, file_nr: u16) -> Result<(), Error> {
if file_nr > 65535 {
return Err(Error::InvalidOperation("File number needs to be < 65536".to_string()));
}
let path = self.path.join(format!("{}.db", file_nr));
// Always close the current file if it's open
self.file = None;
// Create file if it doesn't exist
if !path.exists() {
self.create_new_db_file(file_nr)?;
}
// Open the file fresh
let file = OpenOptions::new()
.read(true)
.write(true)
.open(&path)?;
self.file = Some(file);
self.file_nr = file_nr;
Ok(())
}
/// Creates a new database file
pub(crate) fn create_new_db_file(&mut self, file_nr: u16) -> Result<(), Error> {
let new_file_path = self.path.join(format!("{}.db", file_nr));
let mut file = File::create(&new_file_path)?;
// Write a single byte to make all positions start from 1
file.write_all(&[0u8])?;
Ok(())
}
/// Gets the file number to use for the next write operation
pub(crate) fn get_file_nr(&mut self) -> Result<u16, Error> {
let path = self.path.join(format!("{}.db", self.last_used_file_nr));
if !path.exists() {
self.create_new_db_file(self.last_used_file_nr)?;
return Ok(self.last_used_file_nr);
}
let metadata = fs::metadata(&path)?;
if metadata.len() >= self.file_size as u64 {
self.last_used_file_nr += 1;
self.create_new_db_file(self.last_used_file_nr)?;
}
Ok(self.last_used_file_nr)
}
/// Stores data at the specified ID with history tracking
pub(crate) fn set_(&mut self, id: u32, old_location: Location, data: &[u8]) -> Result<(), Error> {
// Validate data size - maximum is u16::MAX (65535 bytes or ~64KB)
if data.len() > u16::MAX as usize {
return Err(Error::InvalidOperation(
format!("Data size exceeds maximum allowed size of {} bytes", u16::MAX)
));
}
// Get file number to use
let file_nr = self.get_file_nr()?;
// Select the file
self.db_file_select(file_nr)?;
// Get current file position for lookup
let file = self.file.as_mut().ok_or_else(|| Error::Other("No file open".to_string()))?;
file.seek(SeekFrom::End(0))?;
let position = file.stream_position()? as u32;
// Create new location
let new_location = Location {
file_nr,
position,
};
// Calculate CRC of data
let crc = calculate_crc(data);
// Create header
let mut header = vec![0u8; HEADER_SIZE];
// Write size (2 bytes)
let size = data.len() as u16; // Safe now because we've validated the size
header[0] = (size & 0xFF) as u8;
header[1] = ((size >> 8) & 0xFF) as u8;
// Write CRC (4 bytes)
header[2] = (crc & 0xFF) as u8;
header[3] = ((crc >> 8) & 0xFF) as u8;
header[4] = ((crc >> 16) & 0xFF) as u8;
header[5] = ((crc >> 24) & 0xFF) as u8;
// Write previous location (6 bytes)
let prev_bytes = old_location.to_bytes();
for (i, &byte) in prev_bytes.iter().enumerate().take(6) {
header[6 + i] = byte;
}
// Write header
file.write_all(&header)?;
// Write actual data
file.write_all(data)?;
file.flush()?;
// Update lookup table with new position
self.lookup.set(id, new_location)?;
Ok(())
}
/// Retrieves data at the specified location
pub(crate) fn get_(&mut self, location: Location) -> Result<Vec<u8>, Error> {
if location.position == 0 {
return Err(Error::NotFound(format!("Record not found, location: {:?}", location)));
}
// Select the file
self.db_file_select(location.file_nr)?;
let file = self.file.as_mut().ok_or_else(|| Error::Other("No file open".to_string()))?;
// Read header
file.seek(SeekFrom::Start(location.position as u64))?;
let mut header = vec![0u8; HEADER_SIZE];
file.read_exact(&mut header)?;
// Parse size (2 bytes)
let size = u16::from(header[0]) | (u16::from(header[1]) << 8);
// Parse CRC (4 bytes)
let stored_crc = u32::from(header[2])
| (u32::from(header[3]) << 8)
| (u32::from(header[4]) << 16)
| (u32::from(header[5]) << 24);
// Read data
let mut data = vec![0u8; size as usize];
file.read_exact(&mut data)?;
// Verify CRC
let calculated_crc = calculate_crc(&data);
if calculated_crc != stored_crc {
return Err(Error::DataCorruption("CRC mismatch: data corruption detected".to_string()));
}
Ok(data)
}
/// Retrieves the previous position for a record (for history tracking)
pub(crate) fn get_prev_pos_(&mut self, location: Location) -> Result<Location, Error> {
if location.position == 0 {
return Err(Error::NotFound("Record not found".to_string()));
}
// Select the file
self.db_file_select(location.file_nr)?;
let file = self.file.as_mut().ok_or_else(|| Error::Other("No file open".to_string()))?;
// Skip size and CRC (6 bytes)
file.seek(SeekFrom::Start(location.position as u64 + 6))?;
// Read previous location (6 bytes)
let mut prev_bytes = vec![0u8; 6];
file.read_exact(&mut prev_bytes)?;
// Create location from bytes
Location::from_bytes(&prev_bytes, 6)
}
/// Deletes the record at the specified location
pub(crate) fn delete_(&mut self, id: u32, location: Location) -> Result<(), Error> {
if location.position == 0 {
return Err(Error::NotFound("Record not found".to_string()));
}
// Select the file
self.db_file_select(location.file_nr)?;
let file = self.file.as_mut().ok_or_else(|| Error::Other("No file open".to_string()))?;
// Read size first
file.seek(SeekFrom::Start(location.position as u64))?;
let mut size_bytes = vec![0u8; 2];
file.read_exact(&mut size_bytes)?;
let size = u16::from(size_bytes[0]) | (u16::from(size_bytes[1]) << 8);
// Write zeros for the entire record (header + data)
let zeros = vec![0u8; HEADER_SIZE + size as usize];
file.seek(SeekFrom::Start(location.position as u64))?;
file.write_all(&zeros)?;
// Clear lookup entry
self.lookup.delete(id)?;
Ok(())
}
/// Condenses the database by removing empty records and updating positions
pub fn condense(&mut self) -> Result<(), Error> {
// Create a temporary directory
let temp_path = self.path.join("temp");
fs::create_dir_all(&temp_path)?;
// Get all file numbers
let mut file_numbers = Vec::new();
for entry in fs::read_dir(&self.path)? {
let entry = entry?;
let path = entry.path();
if path.is_file() && path.extension().map_or(false, |ext| ext == "db") {
if let Some(stem) = path.file_stem() {
if let Ok(file_nr) = stem.to_string_lossy().parse::<u16>() {
file_numbers.push(file_nr);
}
}
}
}
// Process each file
for file_nr in file_numbers {
let src_path = self.path.join(format!("{}.db", file_nr));
let temp_file_path = temp_path.join(format!("{}.db", file_nr));
// Create new file
let mut temp_file = File::create(&temp_file_path)?;
temp_file.write_all(&[0u8])?; // Initialize with a byte
// Open source file
let mut src_file = File::open(&src_path)?;
// Read and process records
let mut buffer = vec![0u8; 1024]; // Read in chunks
let mut position = 0;
while let Ok(bytes_read) = src_file.read(&mut buffer) {
if bytes_read == 0 {
break;
}
// Process the chunk
// This is a simplified version - in a real implementation,
// you would need to handle records that span chunk boundaries
position += bytes_read;
}
// TODO: Implement proper record copying and position updating
// This would involve:
// 1. Reading each record from the source file
// 2. If not deleted (all zeros), copy to temp file
// 3. Update lookup table with new positions
}
// TODO: Replace original files with temp files
// Clean up
fs::remove_dir_all(&temp_path)?;
Ok(())
}
}
/// Calculates CRC32 for the data
fn calculate_crc(data: &[u8]) -> u32 {
let mut hasher = Hasher::new();
hasher.update(data);
hasher.finalize()
}
#[cfg(test)]
mod tests {
use std::path::PathBuf;
use super::*;
use crate::{OurDB, OurDBConfig, OurDBSetArgs};
use std::env::temp_dir;
use std::time::{SystemTime, UNIX_EPOCH};
fn get_temp_dir() -> PathBuf {
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
temp_dir().join(format!("ourdb_backend_test_{}", timestamp))
}
#[test]
fn test_backend_operations() {
let temp_dir = get_temp_dir();
let config = OurDBConfig {
path: temp_dir.clone(),
incremental_mode: false,
file_size: None,
keysize: None,
};
let mut db = OurDB::new(config).unwrap();
// Test set and get
let test_data = b"Test data for backend operations";
let id = 1;
db.set(OurDBSetArgs { id: Some(id), data: test_data }).unwrap();
let retrieved = db.get(id).unwrap();
assert_eq!(retrieved, test_data);
// Clean up
db.destroy().unwrap();
}
}

41
ourdb/src/error.rs Normal file
View File

@@ -0,0 +1,41 @@
use thiserror::Error;
/// Error types for OurDB operations
#[derive(Error, Debug)]
pub enum Error {
/// IO errors from file operations
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
/// Data corruption errors
#[error("Data corruption: {0}")]
DataCorruption(String),
/// Invalid operation errors
#[error("Invalid operation: {0}")]
InvalidOperation(String),
/// Lookup table errors
#[error("Lookup error: {0}")]
LookupError(String),
/// Record not found errors
#[error("Record not found: {0}")]
NotFound(String),
/// Other errors
#[error("Error: {0}")]
Other(String),
}
impl From<String> for Error {
fn from(msg: String) -> Self {
Error::Other(msg)
}
}
impl From<&str> for Error {
fn from(msg: &str) -> Self {
Error::Other(msg.to_string())
}
}

267
ourdb/src/lib.rs Normal file
View File

@@ -0,0 +1,267 @@
mod error;
mod location;
mod lookup;
mod backend;
pub use error::Error;
pub use location::Location;
pub use lookup::LookupTable;
use std::fs::File;
use std::path::PathBuf;
/// OurDB is a lightweight, efficient key-value database implementation that provides
/// data persistence with history tracking capabilities.
pub struct OurDB {
/// Directory path for storage
path: PathBuf,
/// Whether to use auto-increment mode
incremental_mode: bool,
/// Maximum file size (default: 500MB)
file_size: u32,
/// Lookup table for mapping keys to locations
lookup: LookupTable,
/// Currently open file
file: Option<File>,
/// Current file number
file_nr: u16,
/// Last used file number
last_used_file_nr: u16,
}
/// Configuration for creating a new OurDB instance
pub struct OurDBConfig {
/// Directory path for storage
pub path: PathBuf,
/// Whether to use auto-increment mode
pub incremental_mode: bool,
/// Maximum file size (default: 500MB)
pub file_size: Option<u32>,
/// Lookup table key size
pub keysize: Option<u8>,
}
/// Arguments for setting a value in OurDB
pub struct OurDBSetArgs<'a> {
/// ID for the record (optional in incremental mode)
pub id: Option<u32>,
/// Data to store
pub data: &'a [u8],
}
impl OurDB {
/// Creates a new OurDB instance with the given configuration
pub fn new(config: OurDBConfig) -> Result<Self, Error> {
// Create directory if it doesn't exist
std::fs::create_dir_all(&config.path)?;
// Create lookup table
let lookup_path = config.path.join("lookup");
std::fs::create_dir_all(&lookup_path)?;
let lookup_config = lookup::LookupConfig {
size: 1000000, // Default size
keysize: config.keysize.unwrap_or(4),
lookuppath: lookup_path.to_string_lossy().to_string(),
incremental_mode: config.incremental_mode,
};
let lookup = LookupTable::new(lookup_config)?;
let mut db = OurDB {
path: config.path,
incremental_mode: config.incremental_mode,
file_size: config.file_size.unwrap_or(500 * (1 << 20)), // 500MB default
lookup,
file: None,
file_nr: 0,
last_used_file_nr: 0,
};
// Load existing metadata if available
db.load()?;
Ok(db)
}
/// Sets a value in the database
///
/// In incremental mode:
/// - If ID is provided, it updates an existing record
/// - If ID is not provided, it creates a new record with auto-generated ID
///
/// In key-value mode:
/// - ID must be provided
pub fn set(&mut self, args: OurDBSetArgs) -> Result<u32, Error> {
if self.incremental_mode {
if let Some(id) = args.id {
// This is an update
let location = self.lookup.get(id)?;
if location.position == 0 {
return Err(Error::InvalidOperation(
"Cannot set ID for insertions when incremental mode is enabled".to_string()
));
}
self.set_(id, location, args.data)?;
Ok(id)
} else {
// This is an insert
let id = self.lookup.get_next_id()?;
self.set_(id, Location::default(), args.data)?;
Ok(id)
}
} else {
// Using key-value mode
let id = args.id.ok_or_else(|| Error::InvalidOperation(
"ID must be provided when incremental is disabled".to_string()
))?;
let location = self.lookup.get(id)?;
self.set_(id, location, args.data)?;
Ok(id)
}
}
/// Retrieves data stored at the specified key position
pub fn get(&mut self, id: u32) -> Result<Vec<u8>, Error> {
let location = self.lookup.get(id)?;
self.get_(location)
}
/// Retrieves a list of previous values for the specified key
///
/// The depth parameter controls how many historical values to retrieve (maximum)
pub fn get_history(&mut self, id: u32, depth: u8) -> Result<Vec<Vec<u8>>, Error> {
let mut result = Vec::new();
let mut current_location = self.lookup.get(id)?;
// Traverse the history chain up to specified depth
for _ in 0..depth {
// Get current value
let data = self.get_(current_location)?;
result.push(data);
// Try to get previous location
match self.get_prev_pos_(current_location) {
Ok(location) => {
if location.position == 0 {
break;
}
current_location = location;
}
Err(_) => break,
}
}
Ok(result)
}
/// Deletes the data at the specified key position
pub fn delete(&mut self, id: u32) -> Result<(), Error> {
let location = self.lookup.get(id)?;
self.delete_(id, location)?;
self.lookup.delete(id)?;
Ok(())
}
/// Returns the next ID which will be used when storing in incremental mode
pub fn get_next_id(&mut self) -> Result<u32, Error> {
if !self.incremental_mode {
return Err(Error::InvalidOperation("Incremental mode is not enabled".to_string()));
}
self.lookup.get_next_id()
}
/// Closes the database, ensuring all data is saved
pub fn close(&mut self) -> Result<(), Error> {
self.save()?;
self.close_();
Ok(())
}
/// Destroys the database, removing all files
pub fn destroy(&mut self) -> Result<(), Error> {
let _ = self.close();
std::fs::remove_dir_all(&self.path)?;
Ok(())
}
// Helper methods
fn lookup_dump_path(&self) -> PathBuf {
self.path.join("lookup_dump.db")
}
fn load(&mut self) -> Result<(), Error> {
let dump_path = self.lookup_dump_path();
if dump_path.exists() {
self.lookup.import_sparse(&dump_path.to_string_lossy())?;
}
Ok(())
}
fn save(&mut self) -> Result<(), Error> {
self.lookup.export_sparse(&self.lookup_dump_path().to_string_lossy())?;
Ok(())
}
fn close_(&mut self) {
self.file = None;
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::env::temp_dir;
use std::time::{SystemTime, UNIX_EPOCH};
fn get_temp_dir() -> PathBuf {
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
temp_dir().join(format!("ourdb_test_{}", timestamp))
}
#[test]
fn test_basic_operations() {
let temp_dir = get_temp_dir();
let config = OurDBConfig {
path: temp_dir.clone(),
incremental_mode: true,
file_size: None,
keysize: None,
};
let mut db = OurDB::new(config).unwrap();
// Test set and get
let test_data = b"Hello, OurDB!";
let id = db.set(OurDBSetArgs { id: None, data: test_data }).unwrap();
let retrieved = db.get(id).unwrap();
assert_eq!(retrieved, test_data);
// Test update
let updated_data = b"Updated data";
db.set(OurDBSetArgs { id: Some(id), data: updated_data }).unwrap();
let retrieved = db.get(id).unwrap();
assert_eq!(retrieved, updated_data);
// Test history
let history = db.get_history(id, 2).unwrap();
assert_eq!(history.len(), 2);
assert_eq!(history[0], updated_data);
assert_eq!(history[1], test_data);
// Test delete
db.delete(id).unwrap();
assert!(db.get(id).is_err());
// Clean up
db.destroy().unwrap();
}
}

168
ourdb/src/location.rs Normal file
View File

@@ -0,0 +1,168 @@
use crate::error::Error;
/// Location represents a physical position in a database file
///
/// It consists of a file number and a position within that file.
/// This allows OurDB to span multiple files for large datasets.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Location {
/// File number (0-65535)
pub file_nr: u16,
/// Position within the file
pub position: u32,
}
impl Location {
/// Creates a new Location from bytes based on keysize
///
/// - keysize = 2: Only position (2 bytes), file_nr = 0
/// - keysize = 3: Only position (3 bytes), file_nr = 0
/// - keysize = 4: Only position (4 bytes), file_nr = 0
/// - keysize = 6: file_nr (2 bytes) + position (4 bytes)
pub fn from_bytes(bytes: &[u8], keysize: u8) -> Result<Self, Error> {
// Validate keysize
if ![2, 3, 4, 6].contains(&keysize) {
return Err(Error::InvalidOperation(format!("Invalid keysize: {}", keysize)));
}
// Create padded bytes
let mut padded = vec![0u8; keysize as usize];
let start_idx = keysize as usize - bytes.len();
if start_idx < 0 {
return Err(Error::InvalidOperation("Input bytes exceed keysize".to_string()));
}
for (i, &b) in bytes.iter().enumerate() {
if i + start_idx < padded.len() {
padded[start_idx + i] = b;
}
}
let mut location = Location::default();
match keysize {
2 => {
// Only position, 2 bytes big endian
location.position = u32::from(padded[0]) << 8 | u32::from(padded[1]);
location.file_nr = 0;
// Verify limits
if location.position > 0xFFFF {
return Err(Error::InvalidOperation(
"Position exceeds max value for keysize=2 (max 65535)".to_string()
));
}
},
3 => {
// Only position, 3 bytes big endian
location.position = u32::from(padded[0]) << 16 | u32::from(padded[1]) << 8 | u32::from(padded[2]);
location.file_nr = 0;
// Verify limits
if location.position > 0xFFFFFF {
return Err(Error::InvalidOperation(
"Position exceeds max value for keysize=3 (max 16777215)".to_string()
));
}
},
4 => {
// Only position, 4 bytes big endian
location.position = u32::from(padded[0]) << 24 | u32::from(padded[1]) << 16
| u32::from(padded[2]) << 8 | u32::from(padded[3]);
location.file_nr = 0;
},
6 => {
// 2 bytes file_nr + 4 bytes position, all big endian
location.file_nr = u16::from(padded[0]) << 8 | u16::from(padded[1]);
location.position = u32::from(padded[2]) << 24 | u32::from(padded[3]) << 16
| u32::from(padded[4]) << 8 | u32::from(padded[5]);
},
_ => unreachable!(),
}
Ok(location)
}
/// Converts the location to bytes (always 6 bytes)
///
/// Format: [file_nr (2 bytes)][position (4 bytes)]
pub fn to_bytes(&self) -> Vec<u8> {
let mut bytes = Vec::with_capacity(6);
// Put file_nr first (2 bytes)
bytes.push((self.file_nr >> 8) as u8);
bytes.push(self.file_nr as u8);
// Put position next (4 bytes)
bytes.push((self.position >> 24) as u8);
bytes.push((self.position >> 16) as u8);
bytes.push((self.position >> 8) as u8);
bytes.push(self.position as u8);
bytes
}
/// Converts the location to a u64 value
///
/// The file_nr is stored in the most significant bits
pub fn to_u64(&self) -> u64 {
(u64::from(self.file_nr) << 32) | u64::from(self.position)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_location_from_bytes_keysize_2() {
let bytes = vec![0x12, 0x34];
let location = Location::from_bytes(&bytes, 2).unwrap();
assert_eq!(location.file_nr, 0);
assert_eq!(location.position, 0x1234);
}
#[test]
fn test_location_from_bytes_keysize_3() {
let bytes = vec![0x12, 0x34, 0x56];
let location = Location::from_bytes(&bytes, 3).unwrap();
assert_eq!(location.file_nr, 0);
assert_eq!(location.position, 0x123456);
}
#[test]
fn test_location_from_bytes_keysize_4() {
let bytes = vec![0x12, 0x34, 0x56, 0x78];
let location = Location::from_bytes(&bytes, 4).unwrap();
assert_eq!(location.file_nr, 0);
assert_eq!(location.position, 0x12345678);
}
#[test]
fn test_location_from_bytes_keysize_6() {
let bytes = vec![0xAB, 0xCD, 0x12, 0x34, 0x56, 0x78];
let location = Location::from_bytes(&bytes, 6).unwrap();
assert_eq!(location.file_nr, 0xABCD);
assert_eq!(location.position, 0x12345678);
}
#[test]
fn test_location_to_bytes() {
let location = Location {
file_nr: 0xABCD,
position: 0x12345678,
};
let bytes = location.to_bytes();
assert_eq!(bytes, vec![0xAB, 0xCD, 0x12, 0x34, 0x56, 0x78]);
}
#[test]
fn test_location_to_u64() {
let location = Location {
file_nr: 0xABCD,
position: 0x12345678,
};
let value = location.to_u64();
assert_eq!(value, 0xABCD_0000_0000 | 0x12345678);
}
}

519
ourdb/src/lookup.rs Normal file
View File

@@ -0,0 +1,519 @@
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::Path;
use crate::error::Error;
use crate::location::Location;
const DATA_FILE_NAME: &str = "data";
const INCREMENTAL_FILE_NAME: &str = ".inc";
/// Configuration for creating a new lookup table
pub struct LookupConfig {
/// Size of the lookup table
pub size: u32,
/// Size of each entry in bytes (2-6)
/// - 2: For databases with < 65,536 records
/// - 3: For databases with < 16,777,216 records
/// - 4: For databases with < 4,294,967,296 records
/// - 6: For large databases requiring multiple files
pub keysize: u8,
/// Path for disk-based lookup
pub lookuppath: String,
/// Whether to use incremental mode
pub incremental_mode: bool,
}
/// Lookup table maps keys to physical locations in the backend storage
pub struct LookupTable {
/// Size of each entry in bytes (2-6)
keysize: u8,
/// Path for disk-based lookup
lookuppath: String,
/// In-memory data for memory-based lookup
data: Vec<u8>,
/// Next empty slot if incremental mode is enabled
incremental: Option<u32>,
}
impl LookupTable {
/// Creates a new lookup table with the given configuration
pub fn new(config: LookupConfig) -> Result<Self, Error> {
// Verify keysize is valid
if ![2, 3, 4, 6].contains(&config.keysize) {
return Err(Error::InvalidOperation(format!("Invalid keysize: {}", config.keysize)));
}
let incremental = if config.incremental_mode {
Some(get_incremental_info(&config)?)
} else {
None
};
if !config.lookuppath.is_empty() {
// Create directory if it doesn't exist
fs::create_dir_all(&config.lookuppath)?;
// For disk-based lookup, create empty file if it doesn't exist
let data_path = Path::new(&config.lookuppath).join(DATA_FILE_NAME);
if !data_path.exists() {
let data = vec![0u8; config.size as usize * config.keysize as usize];
fs::write(&data_path, &data)?;
}
Ok(LookupTable {
data: Vec::new(),
keysize: config.keysize,
lookuppath: config.lookuppath,
incremental,
})
} else {
// For memory-based lookup
Ok(LookupTable {
data: vec![0u8; config.size as usize * config.keysize as usize],
keysize: config.keysize,
lookuppath: String::new(),
incremental,
})
}
}
/// Gets a location for the given ID
pub fn get(&self, id: u32) -> Result<Location, Error> {
let entry_size = self.keysize as usize;
if !self.lookuppath.is_empty() {
// Disk-based lookup
let data_path = Path::new(&self.lookuppath).join(DATA_FILE_NAME);
// Check file size first
let file_size = fs::metadata(&data_path)?.len();
let start_pos = id as u64 * entry_size as u64;
if start_pos + entry_size as u64 > file_size {
return Err(Error::LookupError(format!(
"Invalid read for get in lut: {}: {} would exceed file size {}",
self.lookuppath, start_pos + entry_size as u64, file_size
)));
}
// Read directly from file
let mut file = File::open(&data_path)?;
file.seek(SeekFrom::Start(start_pos))?;
let mut data = vec![0u8; entry_size];
let bytes_read = file.read(&mut data)?;
if bytes_read < entry_size {
return Err(Error::LookupError(format!(
"Incomplete read: expected {} bytes but got {}",
entry_size, bytes_read
)));
}
return Location::from_bytes(&data, self.keysize);
}
// Memory-based lookup
if (id * self.keysize as u32) as usize >= self.data.len() {
return Err(Error::LookupError("Index out of bounds".to_string()));
}
let start = (id * self.keysize as u32) as usize;
let end = start + entry_size;
Location::from_bytes(&self.data[start..end], self.keysize)
}
/// Sets a location for the given ID
pub fn set(&mut self, id: u32, location: Location) -> Result<(), Error> {
let entry_size = self.keysize as usize;
// Handle incremental mode
if let Some(incremental) = self.incremental {
if id == incremental {
self.increment_index()?;
}
if id > incremental {
return Err(Error::InvalidOperation(
"Cannot set ID for insertions when incremental mode is enabled".to_string()
));
}
}
// Convert location to bytes based on keysize
let location_bytes = match self.keysize {
2 => {
if location.file_nr != 0 {
return Err(Error::InvalidOperation("file_nr must be 0 for keysize=2".to_string()));
}
if location.position > 0xFFFF {
return Err(Error::InvalidOperation(
"position exceeds max value for keysize=2 (max 65535)".to_string()
));
}
vec![(location.position >> 8) as u8, location.position as u8]
},
3 => {
if location.file_nr != 0 {
return Err(Error::InvalidOperation("file_nr must be 0 for keysize=3".to_string()));
}
if location.position > 0xFFFFFF {
return Err(Error::InvalidOperation(
"position exceeds max value for keysize=3 (max 16777215)".to_string()
));
}
vec![
(location.position >> 16) as u8,
(location.position >> 8) as u8,
location.position as u8
]
},
4 => {
if location.file_nr != 0 {
return Err(Error::InvalidOperation("file_nr must be 0 for keysize=4".to_string()));
}
vec![
(location.position >> 24) as u8,
(location.position >> 16) as u8,
(location.position >> 8) as u8,
location.position as u8
]
},
6 => {
// Full location with file_nr and position
location.to_bytes()
},
_ => return Err(Error::InvalidOperation(format!("Invalid keysize: {}", self.keysize))),
};
if !self.lookuppath.is_empty() {
// Disk-based lookup
let data_path = Path::new(&self.lookuppath).join(DATA_FILE_NAME);
let mut file = OpenOptions::new().write(true).open(data_path)?;
let start_pos = id as u64 * entry_size as u64;
file.seek(SeekFrom::Start(start_pos))?;
file.write_all(&location_bytes)?;
} else {
// Memory-based lookup
let start = (id * self.keysize as u32) as usize;
if start + entry_size > self.data.len() {
return Err(Error::LookupError("Index out of bounds".to_string()));
}
for (i, &byte) in location_bytes.iter().enumerate() {
self.data[start + i] = byte;
}
}
Ok(())
}
/// Deletes an entry for the given ID
pub fn delete(&mut self, id: u32) -> Result<(), Error> {
// Set location to all zeros
self.set(id, Location::default())
}
/// Gets the next available ID in incremental mode
pub fn get_next_id(&self) -> Result<u32, Error> {
let incremental = self.incremental.ok_or_else(||
Error::InvalidOperation("Lookup table not in incremental mode".to_string())
)?;
let table_size = if !self.lookuppath.is_empty() {
let data_path = Path::new(&self.lookuppath).join(DATA_FILE_NAME);
fs::metadata(data_path)?.len() as u32
} else {
self.data.len() as u32
};
if incremental * self.keysize as u32 >= table_size {
return Err(Error::LookupError("Lookup table is full".to_string()));
}
Ok(incremental)
}
/// Increments the index in incremental mode
pub fn increment_index(&mut self) -> Result<(), Error> {
let mut incremental = self.incremental.ok_or_else(||
Error::InvalidOperation("Lookup table not in incremental mode".to_string())
)?;
incremental += 1;
self.incremental = Some(incremental);
if !self.lookuppath.is_empty() {
let inc_path = Path::new(&self.lookuppath).join(INCREMENTAL_FILE_NAME);
fs::write(inc_path, incremental.to_string())?;
}
Ok(())
}
/// Exports the lookup table to a file
pub fn export_data(&self, path: &str) -> Result<(), Error> {
if !self.lookuppath.is_empty() {
// For disk-based lookup, just copy the file
let data_path = Path::new(&self.lookuppath).join(DATA_FILE_NAME);
fs::copy(data_path, path)?;
} else {
// For memory-based lookup, write the data to file
fs::write(path, &self.data)?;
}
Ok(())
}
/// Imports the lookup table from a file
pub fn import_data(&mut self, path: &str) -> Result<(), Error> {
if !self.lookuppath.is_empty() {
// For disk-based lookup, copy the file
let data_path = Path::new(&self.lookuppath).join(DATA_FILE_NAME);
fs::copy(path, data_path)?;
} else {
// For memory-based lookup, read the data from file
self.data = fs::read(path)?;
}
Ok(())
}
/// Exports only non-zero entries to save space
pub fn export_sparse(&self, path: &str) -> Result<(), Error> {
let mut output = Vec::new();
let entry_size = self.keysize as usize;
if !self.lookuppath.is_empty() {
// For disk-based lookup
let data_path = Path::new(&self.lookuppath).join(DATA_FILE_NAME);
let mut file = File::open(&data_path)?;
let file_size = fs::metadata(&data_path)?.len();
let max_entries = file_size / entry_size as u64;
for id in 0..max_entries {
file.seek(SeekFrom::Start(id * entry_size as u64))?;
let mut buffer = vec![0u8; entry_size];
let bytes_read = file.read(&mut buffer)?;
if bytes_read < entry_size {
break;
}
// Check if entry is non-zero
if buffer.iter().any(|&b| b != 0) {
// Write ID (4 bytes) + entry
output.extend_from_slice(&(id as u32).to_be_bytes());
output.extend_from_slice(&buffer);
}
}
} else {
// For memory-based lookup
let max_entries = self.data.len() / entry_size;
for id in 0..max_entries {
let start = id * entry_size;
let entry = &self.data[start..start + entry_size];
// Check if entry is non-zero
if entry.iter().any(|&b| b != 0) {
// Write ID (4 bytes) + entry
output.extend_from_slice(&(id as u32).to_be_bytes());
output.extend_from_slice(entry);
}
}
}
// Write the output to file
fs::write(path, &output)?;
Ok(())
}
/// Imports sparse data (only non-zero entries)
pub fn import_sparse(&mut self, path: &str) -> Result<(), Error> {
let data = fs::read(path)?;
let entry_size = self.keysize as usize;
let record_size = 4 + entry_size; // ID (4 bytes) + entry
if data.len() % record_size != 0 {
return Err(Error::DataCorruption(
"Invalid sparse data format: size mismatch".to_string()
));
}
for chunk_start in (0..data.len()).step_by(record_size) {
if chunk_start + record_size > data.len() {
break;
}
// Extract ID (4 bytes)
let id_bytes = &data[chunk_start..chunk_start + 4];
let id = u32::from_be_bytes([id_bytes[0], id_bytes[1], id_bytes[2], id_bytes[3]]);
// Extract entry
let entry = &data[chunk_start + 4..chunk_start + record_size];
// Create location from entry
let location = Location::from_bytes(entry, self.keysize)?;
// Set the entry
self.set(id, location)?;
}
Ok(())
}
/// Finds the highest ID with a non-zero entry
pub fn find_last_entry(&mut self) -> Result<u32, Error> {
let mut last_id = 0u32;
let entry_size = self.keysize as usize;
if !self.lookuppath.is_empty() {
// For disk-based lookup
let data_path = Path::new(&self.lookuppath).join(DATA_FILE_NAME);
let mut file = File::open(&data_path)?;
let file_size = fs::metadata(&data_path)?.len();
let mut buffer = vec![0u8; entry_size];
let mut pos = 0u32;
while (pos as u64 * entry_size as u64) < file_size {
file.seek(SeekFrom::Start(pos as u64 * entry_size as u64))?;
let bytes_read = file.read(&mut buffer)?;
if bytes_read == 0 || bytes_read < entry_size {
break;
}
let location = Location::from_bytes(&buffer, self.keysize)?;
if location.position != 0 || location.file_nr != 0 {
last_id = pos;
}
pos += 1;
}
} else {
// For memory-based lookup
for i in 0..(self.data.len() / entry_size) as u32 {
if let Ok(location) = self.get(i) {
if location.position != 0 || location.file_nr != 0 {
last_id = i;
}
}
}
}
Ok(last_id)
}
}
/// Helper function to get the incremental value
fn get_incremental_info(config: &LookupConfig) -> Result<u32, Error> {
if !config.incremental_mode {
return Ok(0);
}
if !config.lookuppath.is_empty() {
let inc_path = Path::new(&config.lookuppath).join(INCREMENTAL_FILE_NAME);
if !inc_path.exists() {
// Create a separate file for storing the incremental value
fs::write(&inc_path, "1")?;
}
let inc_str = fs::read_to_string(&inc_path)?;
let incremental = match inc_str.trim().parse::<u32>() {
Ok(val) => val,
Err(_) => {
// If the value is invalid, reset it to 1
fs::write(&inc_path, "1")?;
1
}
};
Ok(incremental)
} else {
// For memory-based lookup, start with 1
Ok(1)
}
}
#[cfg(test)]
mod tests {
use std::path::PathBuf;
use super::*;
use std::env::temp_dir;
use std::time::{SystemTime, UNIX_EPOCH};
fn get_temp_dir() -> PathBuf {
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
temp_dir().join(format!("ourdb_lookup_test_{}", timestamp))
}
#[test]
fn test_memory_lookup() {
let config = LookupConfig {
size: 1000,
keysize: 4,
lookuppath: String::new(),
incremental_mode: true,
};
let mut lookup = LookupTable::new(config).unwrap();
// Test set and get
let location = Location {
file_nr: 0,
position: 12345,
};
lookup.set(1, location).unwrap();
let retrieved = lookup.get(1).unwrap();
assert_eq!(retrieved.file_nr, location.file_nr);
assert_eq!(retrieved.position, location.position);
// Test incremental mode
let next_id = lookup.get_next_id().unwrap();
assert_eq!(next_id, 2);
lookup.increment_index().unwrap();
let next_id = lookup.get_next_id().unwrap();
assert_eq!(next_id, 3);
}
#[test]
fn test_disk_lookup() {
let temp_dir = get_temp_dir();
fs::create_dir_all(&temp_dir).unwrap();
let config = LookupConfig {
size: 1000,
keysize: 4,
lookuppath: temp_dir.to_string_lossy().to_string(),
incremental_mode: true,
};
let mut lookup = LookupTable::new(config).unwrap();
// Test set and get
let location = Location {
file_nr: 0,
position: 12345,
};
lookup.set(1, location).unwrap();
let retrieved = lookup.get(1).unwrap();
assert_eq!(retrieved.file_nr, location.file_nr);
assert_eq!(retrieved.position, location.position);
// Clean up
fs::remove_dir_all(temp_dir).unwrap();
}
}