rename crates/ to engrams/, bindings/ to receptors/
- crates/ → engrams/ (Rust engrams live here) - bindings/ → receptors/ (cross-language access points into the graph) - Cargo.toml workspace paths updated
This commit is contained in:
Will Anderson
@@ -0,0 +1,496 @@
|
||||
/// JNI bindings for engram-core.
|
||||
///
|
||||
/// These functions expose the Engram API to Kotlin/JVM callers via Java Native Interface.
|
||||
/// The convention is:
|
||||
///
|
||||
/// Java_<package>_<class>_<method>
|
||||
/// → Java_ai_neuron_engram_EngramDb_<method>
|
||||
///
|
||||
/// The `EngramDb` handle is stored as a Java `long` (native pointer). The Kotlin
|
||||
/// wrapper class casts it to/from `Long` and keeps it private.
|
||||
///
|
||||
/// # Memory model
|
||||
/// - `open` allocates an `EngramHandle` on the Rust heap and returns its address as `jlong`.
|
||||
/// - `close` takes that `jlong`, reconstructs the Box, and drops it.
|
||||
/// - All other methods borrow the handle via `&*ptr`.
|
||||
///
|
||||
/// # JSON wire format
|
||||
/// Nodes and results are passed as JSON strings to avoid bespoke JNI object marshalling.
|
||||
/// The Kotlin layer converts between the data classes and JSON.
|
||||
use engram_core::{ActivatedNode, EngramDb, MemoryTier, Node, NodeType, ScoredNode};
|
||||
use jni::objects::{JClass, JString};
|
||||
use jni::sys::{jfloatArray, jint, jlong, jstring};
|
||||
use jni::JNIEnv;
|
||||
use std::path::Path;
|
||||
use uuid::Uuid;
|
||||
|
||||
// ── Handle ────────────────────────────────────────────────────────────────────
|
||||
|
||||
struct EngramHandle {
|
||||
db: EngramDb,
|
||||
}
|
||||
|
||||
// ── Helper macros ─────────────────────────────────────────────────────────────
|
||||
|
||||
macro_rules! handle_ref {
|
||||
($handle:expr) => {
|
||||
unsafe { &*($handle as *const EngramHandle) }
|
||||
};
|
||||
}
|
||||
|
||||
// ── JNI methods: lifecycle ────────────────────────────────────────────────────
|
||||
|
||||
/// Open an engram database and return a native handle as jlong.
|
||||
///
|
||||
/// Kotlin: `external fun open(path: String): Long`
|
||||
#[no_mangle]
|
||||
pub extern "system" fn Java_ai_neuron_engram_EngramDb_open(
|
||||
mut env: JNIEnv,
|
||||
_class: JClass,
|
||||
path: JString,
|
||||
) -> jlong {
|
||||
let path_str: String = match env.get_string(&path) {
|
||||
Ok(s) => s.into(),
|
||||
Err(_) => return 0,
|
||||
};
|
||||
match EngramDb::open(Path::new(&path_str)) {
|
||||
Ok(db) => {
|
||||
let handle = Box::new(EngramHandle { db });
|
||||
Box::into_raw(handle) as jlong
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = env.throw_new("java/lang/RuntimeException", e.to_string());
|
||||
0
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Close and free a database handle.
|
||||
///
|
||||
/// Kotlin: `external fun close(handle: Long)`
|
||||
#[no_mangle]
|
||||
pub extern "system" fn Java_ai_neuron_engram_EngramDb_close(
|
||||
_env: JNIEnv,
|
||||
_class: JClass,
|
||||
handle: jlong,
|
||||
) {
|
||||
if handle != 0 {
|
||||
unsafe {
|
||||
drop(Box::from_raw(handle as *mut EngramHandle));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── JNI methods: statistics ───────────────────────────────────────────────────
|
||||
|
||||
/// Kotlin: `external fun nodeCount(handle: Long): Long`
|
||||
#[no_mangle]
|
||||
pub extern "system" fn Java_ai_neuron_engram_EngramDb_nodeCount(
|
||||
mut env: JNIEnv,
|
||||
_class: JClass,
|
||||
handle: jlong,
|
||||
) -> jlong {
|
||||
let h = handle_ref!(handle);
|
||||
match h.db.node_count() {
|
||||
Ok(n) => n as jlong,
|
||||
Err(e) => {
|
||||
let _ = env.throw_new("java/lang/RuntimeException", e.to_string());
|
||||
-1
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Kotlin: `external fun edgeCount(handle: Long): Long`
|
||||
#[no_mangle]
|
||||
pub extern "system" fn Java_ai_neuron_engram_EngramDb_edgeCount(
|
||||
mut env: JNIEnv,
|
||||
_class: JClass,
|
||||
handle: jlong,
|
||||
) -> jlong {
|
||||
let h = handle_ref!(handle);
|
||||
match h.db.edge_count() {
|
||||
Ok(n) => n as jlong,
|
||||
Err(e) => {
|
||||
let _ = env.throw_new("java/lang/RuntimeException", e.to_string());
|
||||
-1
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── JNI methods: nodes ────────────────────────────────────────────────────────
|
||||
|
||||
/// Store a node from JSON and return the assigned UUID string.
|
||||
///
|
||||
/// Kotlin: `external fun putNode(handle: Long, nodeJson: String): String`
|
||||
#[no_mangle]
|
||||
pub extern "system" fn Java_ai_neuron_engram_EngramDb_putNode(
|
||||
mut env: JNIEnv,
|
||||
_class: JClass,
|
||||
handle: jlong,
|
||||
node_json: JString,
|
||||
) -> jstring {
|
||||
let json: String = match env.get_string(&node_json) {
|
||||
Ok(s) => s.into(),
|
||||
Err(_) => return std::ptr::null_mut(),
|
||||
};
|
||||
|
||||
let node = match node_from_json(&json) {
|
||||
Some(n) => n,
|
||||
None => {
|
||||
let _ = env.throw_new("java/lang/IllegalArgumentException", "Invalid node JSON");
|
||||
return std::ptr::null_mut();
|
||||
}
|
||||
};
|
||||
|
||||
let h = handle_ref!(handle);
|
||||
match h.db.put_node(node) {
|
||||
Ok(id) => {
|
||||
let id_str = id.to_string();
|
||||
env.new_string(&id_str)
|
||||
.map(|s| s.into_raw())
|
||||
.unwrap_or(std::ptr::null_mut())
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = env.throw_new("java/lang/RuntimeException", e.to_string());
|
||||
std::ptr::null_mut()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Retrieve a node by UUID, returned as JSON, or null if not found.
|
||||
///
|
||||
/// Kotlin: `external fun getNode(handle: Long, id: String): String?`
|
||||
#[no_mangle]
|
||||
pub extern "system" fn Java_ai_neuron_engram_EngramDb_getNode(
|
||||
mut env: JNIEnv,
|
||||
_class: JClass,
|
||||
handle: jlong,
|
||||
id: JString,
|
||||
) -> jstring {
|
||||
let id_str: String = match env.get_string(&id) {
|
||||
Ok(s) => s.into(),
|
||||
Err(_) => return std::ptr::null_mut(),
|
||||
};
|
||||
let uuid = match id_str.parse::<Uuid>() {
|
||||
Ok(u) => u,
|
||||
Err(_) => return std::ptr::null_mut(),
|
||||
};
|
||||
|
||||
let h = handle_ref!(handle);
|
||||
match h.db.get_node(uuid) {
|
||||
Ok(Some(node)) => {
|
||||
let json = node_to_json(&node);
|
||||
env.new_string(&json)
|
||||
.map(|s| s.into_raw())
|
||||
.unwrap_or(std::ptr::null_mut())
|
||||
}
|
||||
Ok(None) => std::ptr::null_mut(),
|
||||
Err(e) => {
|
||||
let _ = env.throw_new("java/lang/RuntimeException", e.to_string());
|
||||
std::ptr::null_mut()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── JNI methods: search ───────────────────────────────────────────────────────
|
||||
|
||||
/// Search for similar nodes by embedding vector.
|
||||
/// Returns a JSON array of scored nodes.
|
||||
///
|
||||
/// Kotlin: `external fun searchEmbedding(handle: Long, embedding: FloatArray, limit: Int): String`
|
||||
#[no_mangle]
|
||||
pub extern "system" fn Java_ai_neuron_engram_EngramDb_searchEmbedding(
|
||||
mut env: JNIEnv,
|
||||
_class: JClass,
|
||||
handle: jlong,
|
||||
embedding: jfloatArray,
|
||||
limit: jint,
|
||||
) -> jstring {
|
||||
let emb = match float_array_from_jni(&mut env, embedding) {
|
||||
Some(v) => v,
|
||||
None => return std::ptr::null_mut(),
|
||||
};
|
||||
|
||||
let h = handle_ref!(handle);
|
||||
match h.db.search_embedding(&emb, limit as usize) {
|
||||
Ok(results) => {
|
||||
let json = scored_nodes_to_json(&results);
|
||||
env.new_string(&json)
|
||||
.map(|s| s.into_raw())
|
||||
.unwrap_or(std::ptr::null_mut())
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = env.throw_new("java/lang/RuntimeException", e.to_string());
|
||||
std::ptr::null_mut()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Run spreading activation.
|
||||
/// `seeds_json` is a JSON array of UUID strings.
|
||||
/// Returns a JSON array of activated nodes.
|
||||
///
|
||||
/// Kotlin: `external fun activate(handle: Long, seedsJson: String, queryEmbedding: FloatArray, maxDepth: Int, limit: Int): String`
|
||||
#[no_mangle]
|
||||
pub extern "system" fn Java_ai_neuron_engram_EngramDb_activate(
|
||||
mut env: JNIEnv,
|
||||
_class: JClass,
|
||||
handle: jlong,
|
||||
seeds_json: JString,
|
||||
query_embedding: jfloatArray,
|
||||
max_depth: jint,
|
||||
limit: jint,
|
||||
) -> jstring {
|
||||
let seeds_str: String = match env.get_string(&seeds_json) {
|
||||
Ok(s) => s.into(),
|
||||
Err(_) => return std::ptr::null_mut(),
|
||||
};
|
||||
|
||||
let seeds: Vec<Uuid> = parse_uuid_array(&seeds_str);
|
||||
let query_emb = match float_array_from_jni(&mut env, query_embedding) {
|
||||
Some(v) => v,
|
||||
None => return std::ptr::null_mut(),
|
||||
};
|
||||
|
||||
let h = handle_ref!(handle);
|
||||
match h.db.activate(&seeds, &query_emb, max_depth as u8, limit as usize) {
|
||||
Ok(results) => {
|
||||
let json = activated_nodes_to_json(&results);
|
||||
env.new_string(&json)
|
||||
.map(|s| s.into_raw())
|
||||
.unwrap_or(std::ptr::null_mut())
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = env.throw_new("java/lang/RuntimeException", e.to_string());
|
||||
std::ptr::null_mut()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── JNI methods: salience ─────────────────────────────────────────────────────
|
||||
|
||||
/// Touch a node (increment activation count and update salience).
|
||||
///
|
||||
/// Kotlin: `external fun touch(handle: Long, id: String)`
|
||||
#[no_mangle]
|
||||
pub extern "system" fn Java_ai_neuron_engram_EngramDb_touch(
|
||||
mut env: JNIEnv,
|
||||
_class: JClass,
|
||||
handle: jlong,
|
||||
id: JString,
|
||||
) {
|
||||
let id_str: String = match env.get_string(&id) {
|
||||
Ok(s) => s.into(),
|
||||
Err(_) => return,
|
||||
};
|
||||
let uuid = match id_str.parse::<Uuid>() {
|
||||
Ok(u) => u,
|
||||
Err(_) => return,
|
||||
};
|
||||
let h = handle_ref!(handle);
|
||||
if let Err(e) = h.db.touch(uuid) {
|
||||
let _ = env.throw_new("java/lang/RuntimeException", e.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply salience decay. Returns the number of nodes updated.
|
||||
///
|
||||
/// Kotlin: `external fun decay(handle: Long, factor: Float): Int`
|
||||
#[no_mangle]
|
||||
pub extern "system" fn Java_ai_neuron_engram_EngramDb_decay(
|
||||
mut env: JNIEnv,
|
||||
_class: JClass,
|
||||
handle: jlong,
|
||||
factor: f32,
|
||||
) -> jint {
|
||||
let h = handle_ref!(handle);
|
||||
match h.db.decay(factor) {
|
||||
Ok(n) => n as jint,
|
||||
Err(e) => {
|
||||
let _ = env.throw_new("java/lang/RuntimeException", e.to_string());
|
||||
-1
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── JNI helpers ───────────────────────────────────────────────────────────────
|
||||
|
||||
fn float_array_from_jni(env: &mut JNIEnv, arr: jfloatArray) -> Option<Vec<f32>> {
|
||||
if arr.is_null() {
|
||||
return None;
|
||||
}
|
||||
let arr_obj = unsafe { jni::objects::JFloatArray::from_raw(arr) };
|
||||
let len = env.get_array_length(&arr_obj).ok()? as usize;
|
||||
let mut buf = vec![0f32; len];
|
||||
env.get_float_array_region(&arr_obj, 0, &mut buf).ok()?;
|
||||
Some(buf)
|
||||
}
|
||||
|
||||
fn parse_uuid_array(json: &str) -> Vec<Uuid> {
|
||||
// Minimal parser: `["uuid1","uuid2",...]`
|
||||
json.trim_matches(|c| c == '[' || c == ']')
|
||||
.split(',')
|
||||
.filter_map(|s| {
|
||||
let s = s.trim().trim_matches('"');
|
||||
s.parse::<Uuid>().ok()
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
// ── JSON helpers ──────────────────────────────────────────────────────────────
|
||||
|
||||
fn node_from_json(json: &str) -> Option<Node> {
|
||||
let content = extract_string_field(json, "content").unwrap_or_default();
|
||||
let node_type_str =
|
||||
extract_string_field(json, "node_type").unwrap_or_else(|| "Memory".into());
|
||||
let tier_str = extract_string_field(json, "tier").unwrap_or_else(|| "Episodic".into());
|
||||
let importance: f32 = extract_f32_field(json, "importance").unwrap_or(0.5);
|
||||
let embedding = extract_f32_array(json, "embedding").unwrap_or_default();
|
||||
|
||||
let node_type = match node_type_str.as_str() {
|
||||
"Concept" => NodeType::Concept,
|
||||
"Event" => NodeType::Event,
|
||||
"Entity" => NodeType::Entity,
|
||||
"Process" => NodeType::Process,
|
||||
"InternalState" => NodeType::InternalState,
|
||||
_ => NodeType::Memory,
|
||||
};
|
||||
let tier = match tier_str.as_str() {
|
||||
"Working" => MemoryTier::Working,
|
||||
"Semantic" => MemoryTier::Semantic,
|
||||
"Procedural" => MemoryTier::Procedural,
|
||||
_ => MemoryTier::Episodic,
|
||||
};
|
||||
Some(Node::new(node_type, embedding, content.into_bytes(), tier, importance))
|
||||
}
|
||||
|
||||
fn node_to_json(node: &Node) -> String {
|
||||
let content = String::from_utf8_lossy(&node.content)
|
||||
.replace('\\', "\\\\")
|
||||
.replace('"', "\\\"");
|
||||
let emb_str = node
|
||||
.embedding
|
||||
.iter()
|
||||
.map(|f| format!("{:.6}", f))
|
||||
.collect::<Vec<_>>()
|
||||
.join(",");
|
||||
format!(
|
||||
r#"{{"id":"{}","node_type":"{:?}","tier":"{:?}","content":"{}","salience":{:.6},"importance":{:.6},"activation_count":{},"embedding":[{}]}}"#,
|
||||
node.id, node.node_type, node.tier, content, node.salience, node.importance,
|
||||
node.activation_count, emb_str,
|
||||
)
|
||||
}
|
||||
|
||||
fn scored_nodes_to_json(nodes: &[ScoredNode]) -> String {
|
||||
let items: Vec<String> = nodes
|
||||
.iter()
|
||||
.map(|s| format!(r#"{{"node":{},"score":{:.6}}}"#, node_to_json(&s.node), s.score))
|
||||
.collect();
|
||||
format!("[{}]", items.join(","))
|
||||
}
|
||||
|
||||
fn activated_nodes_to_json(nodes: &[ActivatedNode]) -> String {
|
||||
let items: Vec<String> = nodes
|
||||
.iter()
|
||||
.map(|a| {
|
||||
format!(
|
||||
r#"{{"node":{},"activation_strength":{:.6},"hops":{}}}"#,
|
||||
node_to_json(&a.node), a.activation_strength, a.hops,
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
format!("[{}]", items.join(","))
|
||||
}
|
||||
|
||||
fn extract_string_field(json: &str, key: &str) -> Option<String> {
|
||||
let needle = format!("\"{}\":", key);
|
||||
let start = json.find(&needle)? + needle.len();
|
||||
let rest = json[start..].trim_start();
|
||||
if !rest.starts_with('"') {
|
||||
return None;
|
||||
}
|
||||
let inner = &rest[1..];
|
||||
let end = inner.find('"')?;
|
||||
Some(inner[..end].to_string())
|
||||
}
|
||||
|
||||
fn extract_f32_field(json: &str, key: &str) -> Option<f32> {
|
||||
let needle = format!("\"{}\":", key);
|
||||
let start = json.find(&needle)? + needle.len();
|
||||
let rest = json[start..].trim_start();
|
||||
let end = rest
|
||||
.find(|c: char| c == ',' || c == '}')
|
||||
.unwrap_or(rest.len());
|
||||
rest[..end].trim().parse::<f32>().ok()
|
||||
}
|
||||
|
||||
fn extract_f32_array(json: &str, key: &str) -> Option<Vec<f32>> {
|
||||
let needle = format!("\"{}\":", key);
|
||||
let start = json.find(&needle)? + needle.len();
|
||||
let rest = json[start..].trim_start();
|
||||
if !rest.starts_with('[') {
|
||||
return None;
|
||||
}
|
||||
let end = rest.find(']')?;
|
||||
let inner = &rest[1..end];
|
||||
Some(
|
||||
inner
|
||||
.split(',')
|
||||
.filter_map(|s| s.trim().parse::<f32>().ok())
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
// ── Tests ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn node_json_roundtrip() {
|
||||
let node = Node::new(
|
||||
NodeType::Memory,
|
||||
vec![0.1, 0.2, 0.3],
|
||||
b"test content".to_vec(),
|
||||
MemoryTier::Episodic,
|
||||
0.8,
|
||||
);
|
||||
let json = node_to_json(&node);
|
||||
assert!(json.contains("test content"));
|
||||
assert!(json.contains("Memory"));
|
||||
assert!(json.contains("Episodic"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_uuid_array_valid() {
|
||||
let uuids = parse_uuid_array(r#"["550e8400-e29b-41d4-a716-446655440000"]"#);
|
||||
assert_eq!(uuids.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_uuid_array_empty() {
|
||||
let uuids = parse_uuid_array("[]");
|
||||
assert_eq!(uuids.len(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_string_field_works() {
|
||||
let json = r#"{"content":"hello","type":"Memory"}"#;
|
||||
assert_eq!(extract_string_field(json, "content"), Some("hello".into()));
|
||||
assert_eq!(extract_string_field(json, "type"), Some("Memory".into()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_f32_array_works() {
|
||||
let json = r#"{"embedding":[0.1,0.2,0.3]}"#;
|
||||
let arr = extract_f32_array(json, "embedding").unwrap();
|
||||
assert_eq!(arr.len(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn activated_nodes_json_is_array() {
|
||||
let json = activated_nodes_to_json(&[]);
|
||||
assert_eq!(json, "[]");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user