Skip to main content

loong_kernel/
memory.rs

1use std::{collections::BTreeMap, sync::Arc};
2
3use async_trait::async_trait;
4
5// Re-export data types from contracts
6pub use loong_contracts::{
7    MemoryCoreOutcome, MemoryCoreRequest, MemoryExtensionOutcome, MemoryExtensionRequest,
8    MemoryTier,
9};
10
11use crate::errors::MemoryPlaneError;
12
13#[async_trait]
14pub trait CoreMemoryAdapter: Send + Sync {
15    fn name(&self) -> &str;
16
17    async fn execute_core_memory(
18        &self,
19        request: MemoryCoreRequest,
20    ) -> Result<MemoryCoreOutcome, MemoryPlaneError>;
21}
22
23#[async_trait]
24pub trait MemoryExtensionAdapter: Send + Sync {
25    fn name(&self) -> &str;
26
27    async fn execute_memory_extension(
28        &self,
29        request: MemoryExtensionRequest,
30        core: &(dyn CoreMemoryAdapter + Sync),
31    ) -> Result<MemoryExtensionOutcome, MemoryPlaneError>;
32}
33
34#[derive(Default)]
35pub struct MemoryPlane {
36    core_adapters: BTreeMap<String, Arc<dyn CoreMemoryAdapter>>,
37    extension_adapters: BTreeMap<String, Arc<dyn MemoryExtensionAdapter>>,
38    default_core_adapter: Option<String>,
39}
40
41impl MemoryPlane {
42    #[must_use]
43    pub fn new() -> Self {
44        Self {
45            core_adapters: BTreeMap::new(),
46            extension_adapters: BTreeMap::new(),
47            default_core_adapter: None,
48        }
49    }
50
51    pub fn register_core_adapter<A: CoreMemoryAdapter + 'static>(&mut self, adapter: A) {
52        let name = adapter.name().to_owned();
53        if self.default_core_adapter.is_none() {
54            self.default_core_adapter = Some(name.clone());
55        }
56        self.core_adapters.insert(name, Arc::new(adapter));
57    }
58
59    pub fn register_extension_adapter<A: MemoryExtensionAdapter + 'static>(&mut self, adapter: A) {
60        let name = adapter.name().to_owned();
61        self.extension_adapters.insert(name, Arc::new(adapter));
62    }
63
64    pub fn set_default_core_adapter(&mut self, name: &str) -> Result<(), MemoryPlaneError> {
65        if !self.core_adapters.contains_key(name) {
66            return Err(MemoryPlaneError::CoreAdapterNotFound(name.to_owned()));
67        }
68        self.default_core_adapter = Some(name.to_owned());
69        Ok(())
70    }
71
72    #[must_use]
73    pub fn default_core_adapter_name(&self) -> Option<&str> {
74        self.default_core_adapter.as_deref()
75    }
76
77    pub async fn execute_core(
78        &self,
79        core_name: Option<&str>,
80        request: MemoryCoreRequest,
81    ) -> Result<MemoryCoreOutcome, MemoryPlaneError> {
82        let resolved_name = if let Some(name) = core_name {
83            name
84        } else {
85            self.default_core_adapter
86                .as_deref()
87                .ok_or(MemoryPlaneError::NoDefaultCoreAdapter)?
88        };
89
90        let adapter = self
91            .core_adapters
92            .get(resolved_name)
93            .ok_or(MemoryPlaneError::CoreAdapterNotFound(
94                resolved_name.to_owned(),
95            ))?
96            .clone();
97
98        return adapter.execute_core_memory(request).await;
99    }
100
101    pub async fn execute_extension(
102        &self,
103        extension_name: &str,
104        core_name: Option<&str>,
105        request: MemoryExtensionRequest,
106    ) -> Result<MemoryExtensionOutcome, MemoryPlaneError> {
107        let extension = self
108            .extension_adapters
109            .get(extension_name)
110            .ok_or_else(|| MemoryPlaneError::ExtensionNotFound(extension_name.to_owned()))?
111            .clone();
112
113        let resolved_core_name = if let Some(name) = core_name {
114            name
115        } else {
116            self.default_core_adapter
117                .as_deref()
118                .ok_or(MemoryPlaneError::NoDefaultCoreAdapter)?
119        };
120
121        let core = self
122            .core_adapters
123            .get(resolved_core_name)
124            .ok_or(MemoryPlaneError::CoreAdapterNotFound(
125                resolved_core_name.to_owned(),
126            ))?
127            .clone();
128
129        return extension
130            .execute_memory_extension(request, core.as_ref())
131            .await;
132    }
133}