memkit_bevy/

lib.rs

//! # memkit-bevy
//!
//! Bevy ECS integration for memkit.
//!
//! Provides automatic frame lifecycle management, resource integration, and
//! optional GPU memory coordination for Bevy applications.
//!
//! ## Features
//!
//! - **Automatic frame lifecycle** — Frame begin/end tied to Bevy schedule
//! - **Resource integration** — `MkAllocator` as a Bevy `Resource`
//! - **System ordering** — Proper allocation timing via system sets
//! - **Optional GPU support** — Enable `gpu` feature for memkit-co integration
//! - **Prelude compatibility** — Enable `bevy_prelude` for full Bevy prelude
//!
//! ## Quick Start
//!
//! ```rust,ignore
//! use bevy::prelude::*;
//! use memkit_bevy::MkPlugin;
//!
//! fn main() {
//!     App::new()
//!         .add_plugins(DefaultPlugins)
//!         .add_plugins(MkPlugin::default())
//!         .add_systems(Update, my_system)
//!         .run();
//! }
//!
//! fn my_system(alloc: Res<MkAllocatorRes>) {
//!     // Frame automatically managed - allocations valid until frame end
//!     let data = alloc.frame_alloc::<[f32; 4]>();
//!     // ...
//! }
//! ```
//!
//! ## Feature Flags
//!
//! - `bevy_prelude` — Enables full Bevy prelude re-exports
//! - `gpu` — Enables GPU memory coordination via memkit-co

use bevy_app::{App, Plugin, First, Last};
use bevy_ecs::prelude::*;
use memkit::{MkAllocator, MkConfig};

#[cfg(feature = "gpu")]
use memkit_co::bevy::BevyGpuCoordinator;
#[cfg(feature = "gpu")]
use memkit_gpu::DummyBackend;

// Re-export bevy prelude if feature enabled
#[cfg(feature = "bevy_prelude")]
pub use bevy::prelude::*;

// ============================================================================
// System Sets
// ============================================================================

/// System set for memkit frame lifecycle.
#[derive(SystemSet, Debug, Clone, PartialEq, Eq, Hash)]
pub enum MkSystemSet {
    /// Runs at the very start of each frame (calls `begin_frame`).
    FrameBegin,
    /// Runs at the very end of each frame (calls `end_frame`).
    FrameEnd,
}

// ============================================================================
// Resources
// ============================================================================

/// Bevy resource wrapping the memkit allocator.
///
/// Use this to access frame allocations in your systems:
///
/// ```rust,ignore
/// fn my_system(alloc: Res<MkAllocatorRes>) {
///     let data = alloc.frame_alloc::<MyStruct>();
///     // data is valid until end of frame
/// }
/// ```
#[derive(Resource)]
pub struct MkAllocatorRes {
    allocator: MkAllocator,
    frame_active: bool,
}

impl MkAllocatorRes {
    /// Create a new allocator resource with the given config.
    pub fn new(config: MkConfig) -> Self {
        Self {
            allocator: MkAllocator::new(config),
            frame_active: false,
        }
    }

    /// Get a reference to the underlying allocator.
    pub fn allocator(&self) -> &MkAllocator {
        &self.allocator
    }

    /// Check if a frame is currently active.
    pub fn is_frame_active(&self) -> bool {
        self.frame_active
    }

    /// Allocate a value in the current frame's arena.
    ///
    /// Returns a raw pointer to uninitialized memory. You must initialize it.
    ///
    /// # Panics
    /// Panics if called outside of an active frame.
    pub fn frame_alloc<T>(&self) -> *mut T {
        assert!(self.frame_active, "frame_alloc called outside active frame");
        self.allocator.frame_alloc::<T>()
    }

    /// Allocate a slice in the current frame's arena.
    ///
    /// Returns an `MkFrameSlice` wrapper if successful.
    pub fn frame_slice<T>(&self, len: usize) -> Option<memkit::MkFrameSlice<'_, T>> {
        assert!(self.frame_active, "frame_slice called outside active frame");
        self.allocator.frame_slice::<T>(len)
    }

    /// Allocate and initialize a value in the frame arena.
    ///
    /// Returns an `MkFrameBox` wrapper if successful.
    pub fn frame_box<T>(&self, value: T) -> Option<memkit::MkFrameBox<'_, T>> {
        assert!(self.frame_active, "frame_box called outside active frame");
        self.allocator.frame_box(value)
    }
}

impl std::ops::Deref for MkAllocatorRes {
    type Target = MkAllocator;

    fn deref(&self) -> &Self::Target {
        &self.allocator
    }
}

// ============================================================================
// GPU Resource (optional)
// ============================================================================

#[cfg(feature = "gpu")]
/// Bevy resource for GPU memory coordination.
///
/// This uses the thread-safe `BevyGpuCoordinator` from `memkit-co`.
/// Only available when the `gpu` feature is enabled.
///
/// # Example
///
/// ```rust,ignore
/// fn my_gpu_system(gpu: Res<MkGpuRes>) {
///     let buffer = gpu.upload_slice(&[1.0f32, 2.0, 3.0]).unwrap();
/// }
/// ```
#[derive(Resource)]
pub struct MkGpuRes {
    coordinator: BevyGpuCoordinator<DummyBackend>,
}

#[cfg(feature = "gpu")]
impl MkGpuRes {
    /// Create a new GPU resource.
    pub fn new() -> Self {
        Self {
            coordinator: BevyGpuCoordinator::new(DummyBackend::new()),
        }
    }

    /// Get a reference to the coordinator.
    pub fn coordinator(&self) -> &BevyGpuCoordinator<DummyBackend> {
        &self.coordinator
    }
}

#[cfg(feature = "gpu")]
impl Default for MkGpuRes {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(feature = "gpu")]
impl std::ops::Deref for MkGpuRes {
    type Target = BevyGpuCoordinator<DummyBackend>;

    fn deref(&self) -> &Self::Target {
        &self.coordinator
    }
}

// ============================================================================
// Plugin
// ============================================================================

/// Memkit plugin for Bevy.
///
/// Adds automatic frame lifecycle management to your Bevy app.
/// Frame begin is called at the start of each frame (in `First`),
/// and frame end is called at the end (in `Last`).
///
/// # Example
///
/// ```rust,ignore
/// App::new()
///     .add_plugins(DefaultPlugins)
///     .add_plugins(MkPlugin::default())
///     .run();
/// ```
pub struct MkPlugin {
    config: MkConfig,
    #[cfg(feature = "gpu")]
    enable_gpu: bool,
}

impl Default for MkPlugin {
    fn default() -> Self {
        Self {
            config: MkConfig::default(),
            #[cfg(feature = "gpu")]
            enable_gpu: true,
        }
    }
}

impl MkPlugin {
    /// Create a new plugin with custom allocator configuration.
    pub fn with_config(config: MkConfig) -> Self {
        Self {
            config,
            #[cfg(feature = "gpu")]
            enable_gpu: true,
        }
    }

    /// Create a plugin with specific arena size.
    pub fn with_arena_size(size: usize) -> Self {
        Self::with_config(MkConfig {
            frame_arena_size: size,
            ..Default::default()
        })
    }

    #[cfg(feature = "gpu")]
    /// Disable GPU coordination even when the feature is enabled.
    pub fn without_gpu(mut self) -> Self {
        self.enable_gpu = false;
        self
    }
}

impl Plugin for MkPlugin {
    fn build(&self, app: &mut App) {
        // Insert allocator resource
        app.insert_resource(MkAllocatorRes::new(self.config.clone()));

        // Configure system sets
        app.configure_sets(First, MkSystemSet::FrameBegin);
        app.configure_sets(Last, MkSystemSet::FrameEnd);

        // Add frame lifecycle systems
        app.add_systems(First, frame_begin_system.in_set(MkSystemSet::FrameBegin));
        app.add_systems(Last, frame_end_system.in_set(MkSystemSet::FrameEnd));

        // Add GPU coordination if enabled
        #[cfg(feature = "gpu")]
        if self.enable_gpu {
            app.insert_resource(MkGpuRes::new());
            app.add_systems(First, gpu_frame_begin_system.after(MkSystemSet::FrameBegin));
            app.add_systems(Last, gpu_frame_end_system.before(MkSystemSet::FrameEnd));
        }
    }
}

// ============================================================================
// Systems
// ============================================================================

/// System that begins a new allocation frame.
fn frame_begin_system(mut alloc: ResMut<MkAllocatorRes>) {
    alloc.allocator.begin_frame();
    alloc.frame_active = true;
}

/// System that ends the current allocation frame.
fn frame_end_system(mut alloc: ResMut<MkAllocatorRes>) {
    alloc.frame_active = false;
    alloc.allocator.end_frame();
}

#[cfg(feature = "gpu")]
/// System that begins GPU frame coordination.
fn gpu_frame_begin_system(gpu: Res<MkGpuRes>) {
    gpu.coordinator.begin_frame();
}

#[cfg(feature = "gpu")]
/// System that ends GPU frame coordination.
fn gpu_frame_end_system(gpu: Res<MkGpuRes>) {
    let _ = gpu.coordinator.end_frame();
}

// ============================================================================
// Convenience Exports
// ============================================================================

/// Prelude module for common imports.
pub mod prelude {
    pub use super::{MkPlugin, MkAllocatorRes, MkSystemSet};
    pub use memkit::{MkConfig, MkFrameBox, MkFrameSlice};
    
    #[cfg(feature = "gpu")]
    pub use super::MkGpuRes;
}

// For backwards compatibility
pub use MkAllocatorRes as MkAllocatorResource;