xlog_runtime/

relation.rs

//! Relation management for GPU-based Datalog execution
//!
//! This module provides [`RelationStore`], a container for managing named relations
//! stored as GPU buffers. It provides CRUD operations for relations during query
//! execution.

use std::collections::HashMap;
use std::sync::Arc;

use xlog_core::{Schema, XlogError};
use xlog_cuda::{CudaBuffer, CudaKernelProvider};

/// Storage for named relations as GPU buffers
///
/// `RelationStore` manages a collection of named relations, each stored as a
/// [`CudaBuffer`]. It provides CRUD operations for relation management during
/// query execution.
///
/// # Thread Safety
///
/// This implementation is NOT thread-safe. It is designed for single-threaded
/// runtime execution in the MVP.
///
/// # Example
///
/// ```ignore
/// use xlog_runtime::RelationStore;
/// use xlog_cuda::CudaBuffer;
/// use xlog_core::Schema;
///
/// let mut store = RelationStore::new(provider);
///
/// // Add a relation
/// store.put("edge", buffer);
///
/// // Check if relation exists
/// if store.contains("edge") {
///     let edge = store.get("edge").unwrap();
/// }
///
/// // Remove a relation
/// let removed = store.remove("edge");
/// ```
pub struct RelationStore {
    /// CUDA kernel provider for GPU allocations
    provider: Arc<CudaKernelProvider>,
    /// Map of relation names to GPU buffers
    relations: HashMap<String, VersionedCudaBuffer>,
    /// Monotonic semantic-mutation epoch used by staged transactions.
    ///
    /// Per-name versions are not sufficient to detect clear-and-reinsert ABA:
    /// a replacement can return to version one while owning different device
    /// allocations. This epoch never changes for read-only access or reserve.
    mutation_epoch: u64,
}

struct VersionedCudaBuffer {
    buffer: CudaBuffer,
    version: u64,
}

impl RelationStore {
    /// Create a new empty relation store
    pub fn new(provider: Arc<CudaKernelProvider>) -> Self {
        Self {
            provider,
            relations: HashMap::new(),
            mutation_epoch: 0,
        }
    }

    fn advance_mutation_epoch(&mut self) {
        self.mutation_epoch = self
            .mutation_epoch
            .checked_add(1)
            .expect("relation-store mutation epoch exhausted");
    }

    /// Snapshot the semantic store revision for an optimistic transaction.
    pub(crate) fn mutation_epoch(&self) -> u64 {
        self.mutation_epoch
    }

    /// Get a reference to a relation by name
    ///
    /// # Arguments
    /// * `name` - The name of the relation
    ///
    /// # Returns
    /// `Some(&CudaBuffer)` if the relation exists, `None` otherwise
    pub fn get(&self, name: &str) -> Option<&CudaBuffer> {
        self.relations.get(name).map(|e| &e.buffer)
    }

    /// Get a mutable reference to a relation by name
    ///
    /// # Arguments
    /// * `name` - The name of the relation
    ///
    /// # Returns
    /// `Some(&mut CudaBuffer)` if the relation exists, `None` otherwise
    pub fn get_mut(&mut self, name: &str) -> Option<&mut CudaBuffer> {
        if self.relations.contains_key(name) {
            self.advance_mutation_epoch();
        }
        self.relations.get_mut(name).map(|e| {
            // Any mutable access may change the contents; bump the version so cached
            // indexes can be invalidated conservatively.
            e.version = e.version.saturating_add(1);
            &mut e.buffer
        })
    }

    /// Get a relation by name along with its current version.
    pub fn get_with_version(&self, name: &str) -> Option<(&CudaBuffer, u64)> {
        self.relations.get(name).map(|e| (&e.buffer, e.version))
    }

    /// Get the current version for a relation.
    pub fn version(&self, name: &str) -> Option<u64> {
        self.relations.get(name).map(|e| e.version)
    }

    /// Store a relation with the given name
    ///
    /// If a relation with the same name already exists, it will be replaced.
    ///
    /// # Arguments
    /// * `name` - The name of the relation
    /// * `buffer` - The GPU buffer containing the relation data
    pub fn put(&mut self, name: &str, buffer: CudaBuffer) {
        self.put_owned(name.to_string(), buffer);
    }

    /// Store a relation using an already-owned name.
    ///
    /// This has the same replacement and versioning behavior as [`Self::put`],
    /// but avoids allocating a second copy of a name that the caller already
    /// owns. Call [`Self::try_reserve_relations`] before a sequence of inserts
    /// when that sequence must not grow the relation map.
    pub fn put_owned(&mut self, name: String, buffer: CudaBuffer) {
        self.advance_mutation_epoch();
        let version = self
            .relations
            .get(name.as_str())
            .map(|e| e.version.saturating_add(1))
            .unwrap_or(1);
        self.relations
            .insert(name, VersionedCudaBuffer { buffer, version });
    }

    /// Reserve map capacity for `additional` new relation names.
    ///
    /// This operation is fallible and changes capacity only: relation contents
    /// and versions remain unchanged. Callers should count only names that are
    /// not already present before choosing `additional`.
    pub fn try_reserve_relations(&mut self, additional: usize) -> xlog_core::Result<()> {
        self.relations.try_reserve(additional).map_err(|error| {
            XlogError::Execution(format!(
                "Failed to reserve capacity for {additional} relation entries: {error}"
            ))
        })
    }

    /// Get a relation by name, or insert an empty buffer with the given schema
    ///
    /// This is useful for semi-naive evaluation where delta relations may not
    /// exist yet on the first iteration. If the relation doesn't exist, an empty
    /// buffer with the given schema is inserted into the store.
    ///
    /// # Arguments
    /// * `name` - The name of the relation
    /// * `schema` - The schema to use if creating an empty buffer
    ///
    /// # Returns
    /// A reference to the existing buffer, or the newly inserted empty buffer
    pub fn get_or_insert_empty(
        &mut self,
        name: &str,
        schema: &Schema,
    ) -> xlog_core::Result<&CudaBuffer> {
        if !self.relations.contains_key(name) {
            let buffer = self.provider.create_empty_buffer(schema.clone())?;
            self.relations
                .insert(name.to_string(), VersionedCudaBuffer { buffer, version: 1 });
            self.advance_mutation_epoch();
        }
        Ok(&self
            .relations
            .get(name)
            .expect("Relation must exist after insertion")
            .buffer)
    }

    /// Get a mutable reference to a relation, or insert an empty buffer with the given schema
    ///
    /// This is useful for semi-naive evaluation where delta relations may not
    /// exist yet on the first iteration. If the relation doesn't exist, an empty
    /// buffer with the given schema is inserted into the store.
    ///
    /// # Arguments
    /// * `name` - The name of the relation
    /// * `schema` - The schema to use if creating an empty buffer
    ///
    /// # Returns
    /// A mutable reference to the existing buffer, or the newly inserted empty buffer
    pub fn get_or_insert_empty_mut(
        &mut self,
        name: &str,
        schema: &Schema,
    ) -> xlog_core::Result<&mut CudaBuffer> {
        if !self.relations.contains_key(name) {
            let buffer = self.provider.create_empty_buffer(schema.clone())?;
            self.relations
                .insert(name.to_string(), VersionedCudaBuffer { buffer, version: 1 });
        }
        self.advance_mutation_epoch();
        let entry = self
            .relations
            .get_mut(name)
            .expect("Relation must exist after insertion");
        entry.version = entry.version.saturating_add(1);
        Ok(&mut entry.buffer)
    }

    /// Check if a relation exists in the store
    ///
    /// # Arguments
    /// * `name` - The name of the relation
    ///
    /// # Returns
    /// `true` if the relation exists, `false` otherwise
    pub fn contains(&self, name: &str) -> bool {
        self.relations.contains_key(name)
    }

    /// Remove a relation from the store
    ///
    /// # Arguments
    /// * `name` - The name of the relation
    ///
    /// # Returns
    /// `Some(CudaBuffer)` if the relation existed, `None` otherwise
    pub fn remove(&mut self, name: &str) -> Option<CudaBuffer> {
        let removed = self.relations.remove(name).map(|e| e.buffer);
        if removed.is_some() {
            self.advance_mutation_epoch();
        }
        removed
    }

    /// Clear all relations from the store
    ///
    /// This removes all stored relations. The GPU memory will be freed
    /// when the CudaBuffer instances are dropped.
    pub fn clear(&mut self) {
        if !self.relations.is_empty() {
            self.relations.clear();
            self.advance_mutation_epoch();
        }
    }

    /// Get the number of relations in the store
    pub fn len(&self) -> usize {
        self.relations.len()
    }

    /// Check if the store is empty
    pub fn is_empty(&self) -> bool {
        self.relations.is_empty()
    }

    /// Get an iterator over relation names
    pub fn names(&self) -> impl Iterator<Item = &str> {
        self.relations.keys().map(|s| s.as_str())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;
    use xlog_core::{MemoryBudget, ScalarType};
    use xlog_cuda::{CudaDevice, CudaKernelProvider, GpuMemoryManager};

    fn setup_provider() -> Option<Arc<CudaKernelProvider>> {
        let device = match CudaDevice::new(0) {
            Ok(d) => Arc::new(d),
            Err(e) => {
                eprintln!("Skipping: CUDA runtime unavailable: {}", e);
                return None;
            }
        };
        let memory = Arc::new(GpuMemoryManager::new(
            device.clone(),
            MemoryBudget::with_limit(1024 * 1024 * 1024),
        ));
        CudaKernelProvider::new(device, memory).ok().map(Arc::new)
    }

    fn setup_store() -> Option<(RelationStore, Arc<CudaKernelProvider>)> {
        let provider = setup_provider()?;
        let store = RelationStore::new(provider.clone());
        Some((store, provider))
    }

    fn test_schema() -> Schema {
        Schema::new(vec![
            ("a".to_string(), ScalarType::U32),
            ("b".to_string(), ScalarType::U64),
        ])
    }

    fn device_row_count(provider: &CudaKernelProvider, buffer: &CudaBuffer) -> u32 {
        let mut host_rows = [0u32];
        provider
            .device()
            .inner()
            .dtoh_sync_copy_into(buffer.num_rows_device(), &mut host_rows)
            .expect("dtoh row count");
        host_rows[0]
    }

    fn make_buffer(provider: &CudaKernelProvider, schema: Schema, rows: usize) -> CudaBuffer {
        if schema.arity() == 0 {
            if rows == 0 {
                return provider.create_empty_buffer(schema).expect("empty buffer");
            }
            let rows_u32 = u32::try_from(rows).expect("row count fits u32");
            let mut d_num_rows = provider.memory().alloc::<u32>(1).expect("alloc");
            provider
                .device()
                .inner()
                .htod_sync_copy_into(&[rows_u32], &mut d_num_rows)
                .expect("htod row count");
            return CudaBuffer::from_columns(Vec::new(), rows as u64, d_num_rows, schema);
        }
        if rows == 0 {
            return provider.create_empty_buffer(schema).expect("empty buffer");
        }
        let mut columns: Vec<Vec<u8>> = Vec::with_capacity(schema.arity());
        for col_idx in 0..schema.arity() {
            let size = schema
                .column_type(col_idx)
                .map(|t| t.size_bytes())
                .unwrap_or(4);
            columns.push(vec![0u8; rows * size]);
        }
        let slices: Vec<&[u8]> = columns.iter().map(|c| c.as_slice()).collect();
        provider
            .create_buffer_from_slices(&slices, schema)
            .expect("buffer")
    }

    #[test]
    fn mutation_epoch_detects_clear_reinsert_aba_without_counting_reserve() {
        let Some((mut store, provider)) = setup_store() else {
            return;
        };
        let schema = test_schema();
        assert_eq!(store.mutation_epoch(), 0);
        store
            .try_reserve_relations(4)
            .expect("reserve must not be a semantic mutation");
        assert_eq!(store.mutation_epoch(), 0);

        store.put("rows", make_buffer(&provider, schema.clone(), 1));
        let after_first_put = store.mutation_epoch();
        assert_eq!(after_first_put, 1);
        assert_eq!(store.version("rows"), Some(1));

        store.clear();
        let after_clear = store.mutation_epoch();
        assert!(after_clear > after_first_put);
        store.put("rows", make_buffer(&provider, schema, 1));
        assert_eq!(store.version("rows"), Some(1));
        assert!(
            store.mutation_epoch() > after_clear,
            "clear-and-reinsert must not recreate the original transaction revision"
        );
    }

    #[test]
    fn test_new_store_is_empty() {
        let Some((store, _provider)) = setup_store() else {
            return;
        };
        assert!(store.is_empty());
        assert_eq!(store.len(), 0);
    }

    #[test]
    fn test_put_and_get() {
        let Some((mut store, provider)) = setup_store() else {
            return;
        };
        let buffer = provider
            .create_empty_buffer(Schema::new(vec![]))
            .expect("empty");

        store.put("test_rel", buffer);

        assert!(store.contains("test_rel"));
        assert!(!store.is_empty());
        assert_eq!(store.len(), 1);

        let retrieved = store.get("test_rel");
        assert!(retrieved.is_some());
    }

    #[test]
    fn test_get_nonexistent() {
        let Some((store, _provider)) = setup_store() else {
            return;
        };
        assert!(store.get("nonexistent").is_none());
    }

    #[test]
    fn test_contains() {
        let Some((mut store, provider)) = setup_store() else {
            return;
        };

        assert!(!store.contains("test"));

        store.put(
            "test",
            provider
                .create_empty_buffer(Schema::new(vec![]))
                .expect("empty"),
        );

        assert!(store.contains("test"));
        assert!(!store.contains("other"));
    }

    #[test]
    fn test_remove() {
        let Some((mut store, provider)) = setup_store() else {
            return;
        };
        store.put(
            "test",
            provider
                .create_empty_buffer(Schema::new(vec![]))
                .expect("empty"),
        );

        assert!(store.contains("test"));

        let removed = store.remove("test");
        assert!(removed.is_some());
        assert!(!store.contains("test"));
        assert!(store.is_empty());
    }

    #[test]
    fn test_remove_nonexistent() {
        let Some((mut store, _provider)) = setup_store() else {
            return;
        };
        let removed = store.remove("nonexistent");
        assert!(removed.is_none());
    }

    #[test]
    fn test_clear() {
        let Some((mut store, provider)) = setup_store() else {
            return;
        };
        let empty = provider
            .create_empty_buffer(Schema::new(vec![]))
            .expect("empty");
        store.put("rel1", empty);
        store.put(
            "rel2",
            provider
                .create_empty_buffer(Schema::new(vec![]))
                .expect("empty"),
        );
        store.put(
            "rel3",
            provider
                .create_empty_buffer(Schema::new(vec![]))
                .expect("empty"),
        );

        assert_eq!(store.len(), 3);

        store.clear();

        assert!(store.is_empty());
        assert_eq!(store.len(), 0);
    }

    #[test]
    fn test_get_or_insert_empty_existing() {
        let Some((mut store, provider)) = setup_store() else {
            return;
        };
        let schema = test_schema();

        let buffer = make_buffer(&provider, schema.clone(), 100);
        store.put("existing", buffer);

        let result = store.get_or_insert_empty("existing", &schema).unwrap();
        assert_eq!(device_row_count(&provider, result), 100);
        assert_eq!(result.schema(), &schema);
        assert_eq!(store.len(), 1);
    }

    #[test]
    fn test_get_or_insert_empty_nonexistent() {
        let Some((mut store, provider)) = setup_store() else {
            return;
        };
        let schema = test_schema();

        assert!(store.is_empty());

        let result = store.get_or_insert_empty("nonexistent", &schema).unwrap();
        assert_eq!(device_row_count(&provider, result), 0);
        assert_eq!(result.schema(), &schema);
        assert!(result.is_empty());

        assert!(store.contains("nonexistent"));
        assert_eq!(store.len(), 1);
    }

    #[test]
    fn test_get_mut() {
        let Some((mut store, provider)) = setup_store() else {
            return;
        };
        let buffer = make_buffer(&provider, Schema::new(vec![]), 10);
        store.put("test", buffer);

        {
            let buf_mut = store.get_mut("test").unwrap();
            buf_mut.set_row_capacity(50);
            provider
                .device()
                .inner()
                .htod_sync_copy_into(&[50u32], buf_mut.num_rows_device_mut())
                .expect("htod row count");
        }

        assert_eq!(device_row_count(&provider, store.get("test").unwrap()), 50);
    }

    #[test]
    fn test_get_mut_nonexistent() {
        let Some((mut store, _provider)) = setup_store() else {
            return;
        };
        assert!(store.get_mut("nonexistent").is_none());
    }

    #[test]
    fn test_get_or_insert_empty_mut() {
        let Some((mut store, provider)) = setup_store() else {
            return;
        };
        let schema = test_schema();

        {
            let buf_mut = store.get_or_insert_empty_mut("new_rel", &schema).unwrap();
            assert_eq!(device_row_count(&provider, buf_mut), 0);
            buf_mut.set_row_capacity(42);
            provider
                .device()
                .inner()
                .htod_sync_copy_into(&[42u32], buf_mut.num_rows_device_mut())
                .expect("htod row count");
        }

        assert!(store.contains("new_rel"));
        assert_eq!(
            device_row_count(&provider, store.get("new_rel").unwrap()),
            42
        );
    }

    #[test]
    fn test_put_replaces_existing() {
        let Some((mut store, provider)) = setup_store() else {
            return;
        };

        let buffer1 = make_buffer(&provider, Schema::new(vec![]), 10);
        let buffer2 = make_buffer(&provider, Schema::new(vec![]), 20);

        store.put("test", buffer1);
        assert_eq!(device_row_count(&provider, store.get("test").unwrap()), 10);

        store.put("test", buffer2);
        assert_eq!(device_row_count(&provider, store.get("test").unwrap()), 20);
        assert_eq!(store.len(), 1);
    }

    #[test]
    fn reserved_owned_put_reuses_capacity_and_versions_replacements() {
        let Some((mut store, provider)) = setup_store() else {
            return;
        };
        let capacity_before = store.relations.capacity();

        store
            .try_reserve_relations(1)
            .expect("reserve one relation entry");
        let reserved_capacity = store.relations.capacity();
        assert!(reserved_capacity > capacity_before);

        store.put_owned(
            "owned_relation".to_string(),
            make_buffer(&provider, Schema::new(vec![]), 10),
        );
        assert_eq!(store.relations.capacity(), reserved_capacity);
        assert_eq!(store.version("owned_relation"), Some(1));

        store.put_owned(
            "owned_relation".to_string(),
            make_buffer(&provider, Schema::new(vec![]), 20),
        );
        assert_eq!(store.relations.capacity(), reserved_capacity);
        assert_eq!(store.version("owned_relation"), Some(2));
        assert_eq!(
            device_row_count(&provider, store.get("owned_relation").unwrap()),
            20
        );
    }

    #[test]
    fn test_names_iterator() {
        let Some((mut store, provider)) = setup_store() else {
            return;
        };
        store.put(
            "alpha",
            provider
                .create_empty_buffer(Schema::new(vec![]))
                .expect("empty"),
        );
        store.put(
            "beta",
            provider
                .create_empty_buffer(Schema::new(vec![]))
                .expect("empty"),
        );
        store.put(
            "gamma",
            provider
                .create_empty_buffer(Schema::new(vec![]))
                .expect("empty"),
        );

        let mut names: Vec<&str> = store.names().collect();
        names.sort();

        assert_eq!(names, vec!["alpha", "beta", "gamma"]);
    }

    #[test]
    fn test_multiple_operations() {
        let Some((mut store, provider)) = setup_store() else {
            return;
        };

        let empty = provider
            .create_empty_buffer(Schema::new(vec![]))
            .expect("empty");
        store.put("a", empty);
        store.put(
            "b",
            provider
                .create_empty_buffer(Schema::new(vec![]))
                .expect("empty"),
        );
        store.put(
            "c",
            provider
                .create_empty_buffer(Schema::new(vec![]))
                .expect("empty"),
        );
        assert_eq!(store.len(), 3);

        store.remove("b");
        assert_eq!(store.len(), 2);
        assert!(!store.contains("b"));

        store.put("a", make_buffer(&provider, Schema::new(vec![]), 50));
        assert_eq!(store.len(), 2);
        assert_eq!(device_row_count(&provider, store.get("a").unwrap()), 50);

        store.clear();
        assert!(store.is_empty());
    }
}