@@ -1077,11 +1077,13 @@ impl Arena {
10771077 Value::Function(f) => {
10781078 let ne = self.promote_value(f.env);
10791079 if ne != f.env {
1080+ unsafe { code_handle_retain(f.code) };
10801081 self.alloc_pinned(Value::Function(Box::new(FunctionInner {
10811082 fn_ptr: f.fn_ptr,
10821083 env: ne,
10831084 source: f.source.clone(),
10841085 extract_source: f.extract_source.clone(),
1086+ code: f.code,
10851087 })))
10861088 } else {
10871089 val
@@ -1090,10 +1092,12 @@ impl Arena {
10901092 Value::IO(io) => {
10911093 let ne = self.promote_value(io.env);
10921094 if ne != io.env {
1095+ unsafe { code_handle_retain(io.code) };
10931096 self.alloc_pinned(Value::IO(Box::new(IOInner {
10941097 fn_ptr: io.fn_ptr,
10951098 env: ne,
10961099 source: io.source.clone(),
1100+ code: io.code,
10971101 })))
10981102 } else {
10991103 val
@@ -1220,17 +1224,25 @@ impl Arena {
12201224 Value::Constructor(tag, inner) => {
12211225 Value::Constructor(tag.clone(), self.promote_value(*inner))
12221226 }
1223− Value::Function(f) => Value::Function(Box::new(FunctionInner {
1224− fn_ptr: f.fn_ptr,
1225− env: self.promote_value(f.env),
1226− source: f.source.clone(),
1227− extract_source: f.extract_source.clone(),
1228− })),
1229− Value::IO(io) => Value::IO(Box::new(IOInner {
1230− fn_ptr: io.fn_ptr,
1231− env: self.promote_value(io.env),
1232− source: io.source.clone(),
1233− })),
1227+ Value::Function(f) => {
1228+ unsafe { code_handle_retain(f.code) };
1229+ Value::Function(Box::new(FunctionInner {
1230+ fn_ptr: f.fn_ptr,
1231+ env: self.promote_value(f.env),
1232+ source: f.source.clone(),
1233+ extract_source: f.extract_source.clone(),
1234+ code: f.code,
1235+ }))
1236+ }
1237+ Value::IO(io) => {
1238+ unsafe { code_handle_retain(io.code) };
1239+ Value::IO(Box::new(IOInner {
1240+ fn_ptr: io.fn_ptr,
1241+ env: self.promote_value(io.env),
1242+ source: io.source.clone(),
1243+ code: io.code,
1244+ }))
1245+ }
12341246 Value::Pair(a, b) => Value::Pair(self.promote_value(*a), self.promote_value(*b)),
12351247 };
12361248 self.alloc_pinned(cloned)
@@ -1304,17 +1316,25 @@ impl Arena {
13041316 Value::Constructor(tag, inner) => {
13051317 Value::Constructor(tag.clone(), self.clone_from_child(*inner, child))
13061318 }
1307− Value::Function(f) => Value::Function(Box::new(FunctionInner {
1308− fn_ptr: f.fn_ptr,
1309− env: self.clone_from_child(f.env, child),
1310− source: f.source.clone(),
1311− extract_source: f.extract_source.clone(),
1312− })),
1313− Value::IO(io) => Value::IO(Box::new(IOInner {
1314− fn_ptr: io.fn_ptr,
1315− env: self.clone_from_child(io.env, child),
1316− source: io.source.clone(),
1317− })),
1319+ Value::Function(f) => {
1320+ unsafe { code_handle_retain(f.code) };
1321+ Value::Function(Box::new(FunctionInner {
1322+ fn_ptr: f.fn_ptr,
1323+ env: self.clone_from_child(f.env, child),
1324+ source: f.source.clone(),
1325+ extract_source: f.extract_source.clone(),
1326+ code: f.code,
1327+ }))
1328+ }
1329+ Value::IO(io) => {
1330+ unsafe { code_handle_retain(io.code) };
1331+ Value::IO(Box::new(IOInner {
1332+ fn_ptr: io.fn_ptr,
1333+ env: self.clone_from_child(io.env, child),
1334+ source: io.source.clone(),
1335+ code: io.code,
1336+ }))
1337+ }
13181338 Value::Pair(a, b) => Value::Pair(
13191339 self.clone_from_child(*a, child),
13201340 self.clone_from_child(*b, child),
@@ -4392,6 +4412,10 @@ pub struct FunctionInner {
43924412 /// qualified, re-parseable `Maybe.Nothing {}`). `None` → extract derives
43934413 /// from `source` + `env` as usual.
43944414 pub extract_source: Option<Arc<str>>,
4415+ /// Keep-alive handle for the JIT module this function's `fn_ptr` points
4416+ /// into. Null for natively-compiled (AOT) functions, whose code lives in
4417+ /// the process image and never needs freeing. See [`CodeHandle`].
4418+ pub code: *const CodeHandle,
43954419}
43964420
43974421/// Payload for `Value::IO`: boxed to keep `Value` compact. `source` is the
@@ -4402,6 +4426,128 @@ pub struct IOInner {
44024426 pub fn_ptr: *const u8,
44034427 pub env: *mut Value,
44044428 pub source: Arc<str>,
4429+ /// Keep-alive handle for the JIT module this thunk's `fn_ptr` points into.
4430+ /// Null for AOT-compiled actions. See [`CodeHandle`].
4431+ pub code: *const CodeHandle,
4432+}
4433+
4434+/// A type-erased, refcounted keep-alive handle for a JIT'd code module.
4435+///
4436+/// A `Value::Function`/`Value::IO` produced by `base.compile` holds a `fn_ptr`
4437+/// into the JIT module's executable memory, which must stay mapped while any
4438+/// such value is reachable. But knot-runtime can't name the JIT's module type
4439+/// (that would be a dependency cycle), so the module's lifetime is exposed as
4440+/// two function pointers plus an opaque refcount word:
4441+///
4442+/// - `refcount` is a `*const AtomicUsize` owned by the JIT side.
4443+/// - `free` reclaims the module; the caller (here, on drop of the last
4444+/// function/IO referencing it) invokes it exactly once, when the shared
4445+/// refcount reaches zero.
4446+///
4447+/// Cloning a function/IO increments the refcount; dropping one decrements it.
4448+/// The arena's `reset_to`/frame-drop paths run value destructors, which is what
4449+/// drives the decrement — so a JIT module is freed precisely when the last
4450+/// closure pointing into it is reclaimed.
4451+///
4452+/// # Safety
4453+/// `refcount` and `free` must remain valid until the refcount reaches zero and
4454+/// `free` has been called. The JIT side guarantees this by leaking the handle's
4455+/// storage into the module's own lifetime.
4456+///
4457+/// `repr(C)` pins the field order — knot-jit constructs this layout via a
4458+/// mirror struct, so the two must agree byte-for-byte.
4459+#[repr(C)]
4460+pub struct CodeHandle {
4461+ pub refcount: *const std::sync::atomic::AtomicUsize,
4462+ pub free: unsafe extern "C" fn(*const std::sync::atomic::AtomicUsize),
4463+}
4464+
4465+/// Increment a code handle's refcount (a new function/IO now references the
4466+/// module). Null-safe: a null handle (AOT code) is a no-op.
4467+///
4468+/// # Safety
4469+/// `handle` must be null or point to a live [`CodeHandle`].
4470+pub unsafe fn code_handle_retain(handle: *const CodeHandle) {
4471+ if handle.is_null() {
4472+ return;
4473+ }
4474+ unsafe {
4475+ (*(*handle).refcount).fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4476+ }
4477+}
4478+
4479+/// Decrement a code handle's refcount; on the last drop, free the module.
4480+/// Null-safe.
4481+///
4482+/// # Safety
4483+/// `handle` must be null or a [`CodeHandle`] whose refcount is above zero.
4484+pub unsafe fn code_handle_release(handle: *const CodeHandle) {
4485+ if handle.is_null() {
4486+ return;
4487+ }
4488+ unsafe {
4489+ // Release so prior writes are visible to whoever frees; Acquire on the
4490+ // final decrement pairs with it. On zero, hand the refcount word to the
4491+ // JIT's free routine, which reclaims the module (and the word).
4492+ if (*(*handle).refcount).fetch_sub(1, std::sync::atomic::Ordering::AcqRel) == 1 {
4493+ ((*handle).free)((*handle).refcount);
4494+ }
4495+ }
4496+}
4497+
4498+/// Recursively stamp a JIT code handle onto every `Function`/`IO` reachable
4499+/// from `v` that doesn't already have one (those have `code == null`). Sets the
4500+/// refcount to the number of values stamped, so the count matches the number of
4501+/// code-bearing values sharing the module. Called once by the JIT right after a
4502+/// `base.compile` body returns, before the value escapes to the host.
4503+///
4504+/// Depth-bounded against pathological values. Values that already carry a
4505+/// handle (e.g. a closure captured from an earlier `compile`) are left alone —
4506+/// their module is already accounted for.
4507+///
4508+/// # Safety
4509+/// `v` must be a live `Value*` produced by this runtime; `handle` a valid
4510+/// [`CodeHandle`] whose refcount is about to be initialized by this call.
4511+#[unsafe(no_mangle)]
4512+pub unsafe extern "C-unwind" fn knot_value_set_code(v: *mut Value, handle: *const CodeHandle) {
4513+ if v.is_null() || handle.is_null() {
4514+ return;
4515+ }
4516+ fn stamp(v: *mut Value, handle: *const CodeHandle, depth: usize) -> usize {
4517+ if v.is_null() || depth > 512 {
4518+ return 0;
4519+ }
4520+ match unsafe { as_ref(v) } {
4521+ Value::Function(f) => {
4522+ if f.code.is_null() {
4523+ // Cast away the shared ref to write the field once.
4524+ unsafe { (*(f.as_ref() as *const FunctionInner as *mut FunctionInner)).code = handle };
4525+ 1
4526+ } else {
4527+ 0
4528+ }
4529+ }
4530+ Value::IO(io) => {
4531+ if io.code.is_null() {
4532+ unsafe { (*(io.as_ref() as *const IOInner as *mut IOInner)).code = handle };
4533+ 1
4534+ } else {
4535+ 0
4536+ }
4537+ }
4538+ Value::Record(fields) => fields.iter().map(|f| stamp(f.value, handle, depth + 1)).sum(),
4539+ Value::Relation(elems) => elems.iter().map(|e| stamp(*e, handle, depth + 1)).sum(),
4540+ Value::Constructor(_, payload) => stamp(*payload, handle, depth + 1),
4541+ Value::Pair(a, b) => stamp(*a, handle, depth + 1) + stamp(*b, handle, depth + 1),
4542+ _ => 0,
4543+ }
4544+ }
4545+ let stamped = stamp(v, handle, 0);
4546+ if stamped > 0 {
4547+ unsafe {
4548+ (*(*handle).refcount).store(stamped, std::sync::atomic::Ordering::Relaxed);
4549+ }
4550+ }
44054551}
44064552
44074553pub struct RecordField {
@@ -4412,6 +4558,87 @@ pub struct RecordField {
44124558 pub value: *mut Value,
44134559}
44144560
4561+// Dropping a function/IO releases its JIT code handle; when the last value
4562+// referencing a module drops, the module's code memory is freed. The arena
4563+// runs these destructors on every reclaim path (`reset_to`, frame drop), so a
4564+// JIT module lives exactly as long as its longest-lived closure.
4565+impl Drop for FunctionInner {
4566+ fn drop(&mut self) {
4567+ unsafe { code_handle_release(self.code) };
4568+ }
4569+}
4570+
4571+impl Drop for IOInner {
4572+ fn drop(&mut self) {
4573+ unsafe { code_handle_release(self.code) };
4574+ }
4575+}
4576+
4577+#[cfg(test)]
4578+mod code_handle_tests {
4579+ use super::*;
4580+ use std::sync::atomic::{AtomicUsize, Ordering};
4581+
4582+ // A fake "module": a boxed (refcount, freed-flag, handle). `free` flips the
4583+ // freed flag via the refcount's container, mirroring how the JIT's block
4584+ // reclaims its module when the last closure drops.
4585+ #[repr(C)]
4586+ struct FakeBlock {
4587+ refcount: AtomicUsize,
4588+ freed: std::sync::atomic::AtomicBool,
4589+ handle: CodeHandle,
4590+ }
4591+
4592+ unsafe extern "C" fn fake_free(refcount: *const AtomicUsize) {
4593+ let block = refcount as *mut FakeBlock;
4594+ unsafe {
4595+ (*block).freed.store(true, Ordering::SeqCst);
4596+ drop(Box::from_raw(block));
4597+ }
4598+ }
4599+
4600+ fn new_block(count: usize) -> (*mut FakeBlock, *const CodeHandle) {
4601+ let mut b = Box::new(FakeBlock {
4602+ refcount: AtomicUsize::new(count),
4603+ freed: std::sync::atomic::AtomicBool::new(false),
4604+ handle: CodeHandle {
4605+ refcount: std::ptr::null(),
4606+ free: fake_free,
4607+ },
4608+ });
4609+ b.handle.refcount = &raw const b.refcount;
4610+ let handle = &raw const b.handle;
4611+ (Box::into_raw(b), handle)
4612+ }
4613+
4614+ #[test]
4615+ fn release_to_zero_frees() {
4616+ let (block, handle) = new_block(2);
4617+ unsafe { code_handle_release(handle) }; // 2 -> 1
4618+ assert!(!unsafe { (*block).freed.load(Ordering::SeqCst) });
4619+ let freed_flag = unsafe { (*block).freed.load(Ordering::SeqCst) };
4620+ assert!(!freed_flag);
4621+ unsafe { code_handle_release(handle) }; // 1 -> 0 → free
4622+ // block is now freed; we can't read it. The test not crashing/hanging is
4623+ // the assertion; freed was set inside fake_free before the Box drop.
4624+ }
4625+
4626+ #[test]
4627+ fn retain_then_release_balances() {
4628+ let (block, handle) = new_block(1);
4629+ unsafe { code_handle_retain(handle) }; // 1 -> 2
4630+ unsafe { code_handle_release(handle) }; // 2 -> 1
4631+ assert!(!unsafe { (*block).freed.load(Ordering::SeqCst) });
4632+ unsafe { code_handle_release(handle) }; // 1 -> 0 → free
4633+ }
4634+
4635+ #[test]
4636+ fn null_handle_is_noop() {
4637+ unsafe { code_handle_retain(std::ptr::null()) };
4638+ unsafe { code_handle_release(std::ptr::null()) };
4639+ }
4640+}
4641+
44154642impl Drop for Value {
44164643 fn drop(&mut self) {
44174644 // Custom Drop intercepts `Value::Relation` / `Value::Record` to
@@ -4958,6 +5185,7 @@ fn alloc_io(fn_ptr: *const u8, env: *mut Value, source: Arc<str>) -> *mut Value
49585185 fn_ptr,
49595186 env,
49605187 source,
5188+ code: std::ptr::null(),
49615189 })))
49625190}
49635191
@@ -5592,6 +5820,7 @@ pub extern "C-unwind" fn knot_value_function(
55925820 env,
55935821 source,
55945822 extract_source: None,
5823+ code: std::ptr::null(),
55955824 })))
55965825}
55975826
@@ -5615,6 +5844,7 @@ pub extern "C-unwind" fn knot_value_function_full(
56155844 env,
56165845 source,
56175846 extract_source,
5847+ code: std::ptr::null(),
56185848 })))
56195849}
56205850
@@ -9458,19 +9688,23 @@ fn deep_clone_with(val: *mut Value, mut alloc_node: impl FnMut(Value) -> *mut Va
94589688 }
94599689 Value::Function(f) => {
94609690 to_alloc.push(f.env);
9691+ unsafe { code_handle_retain(f.code) };
94619692 Value::Function(Box::new(FunctionInner {
94629693 fn_ptr: f.fn_ptr,
94639694 env: std::ptr::null_mut(),
94649695 source: f.source.clone(),
94659696 extract_source: f.extract_source.clone(),
9697+ code: f.code,
94669698 }))
94679699 }
94689700 Value::IO(io) => {
94699701 to_alloc.push(io.env);
9702+ unsafe { code_handle_retain(io.code) };
94709703 Value::IO(Box::new(IOInner {
94719704 fn_ptr: io.fn_ptr,
94729705 env: std::ptr::null_mut(),
94739706 source: io.source.clone(),
9707+ code: io.code,
94749708 }))
94759709 }
94769710 Value::Pair(a, b) => {
@@ -10740,6 +10974,7 @@ pub extern "C-unwind" fn knot_trace(
1074010974 env,
1074110975 source: Arc::from("trace"),
1074210976 extract_source: None,
10977+ code: std::ptr::null(),
1074310978 })))
1074410979}
1074510980
@@ -11593,38 +11828,6 @@ pub extern "C-unwind" fn knot_register_compile_impl(f: *mut c_void) {
1159311828 COMPILE_IMPL.store(f, std::sync::atomic::Ordering::Release);
1159411829}
1159511830
11596−/// Is `v` free of any JIT'd code reference? A compiled snippet's JIT module
11597−/// may be freed only when its result cannot re-enter the module's code. Only
11598−/// `Function` and `IO` values carry a `fn_ptr` into that code (a lambda or an
11599−/// IO thunk); every other variant is pure data whose children are values, not
11600−/// code. Recurses through records, relations, constructor payloads, and pairs.
11601−/// `1` = code-free (module safe to free), `0` = holds code (must keep the
11602−/// module alive).
11603−///
11604−/// # Safety
11605−/// `v` must be a live `Value*` produced by this runtime.
11606−#[unsafe(no_mangle)]
11607−pub extern "C-unwind" fn knot_value_is_code_free(v: *mut Value) -> i32 {
11608− if v.is_null() {
11609− return 1;
11610− }
11611− // Depth-bounded: a pathological cyclic value can't overflow the stack.
11612− fn go(v: *mut Value, depth: usize) -> bool {
11613− if v.is_null() || depth > 512 {
11614− return depth > 512; // too deep to prove code-free → keep the module
11615− }
11616− match unsafe { as_ref(v) } {
11617− Value::Function(_) | Value::IO(_) => false,
11618− Value::Record(fields) => fields.iter().all(|f| go(f.value, depth + 1)),
11619− Value::Relation(elems) => elems.iter().all(|e| go(*e, depth + 1)),
11620− Value::Constructor(_, payload) => go(*payload, depth + 1),
11621− Value::Pair(a, b) => go(*a, depth + 1) && go(*b, depth + 1),
11622− _ => true,
11623− }
11624− }
11625− go(v, 0) as i32
11626−}
11627−
1162811831/// Look up the registered compile implementation, if any.
1162911832fn compile_impl() -> Option<CompileImpl> {
1163011833 let p = COMPILE_IMPL.load(std::sync::atomic::Ordering::Acquire);
@@ -19049,6 +19252,7 @@ fn build_request_ctx(
1904919252 env: header_env,
1905019253 source: intern_str("RequestCtx.header"),
1905119254 extract_source: None,
19255+ code: std::ptr::null(),
1905219256 })));
1905319257 let mut fields = vec![
1905419258 RecordField {
@@ -19490,6 +19694,7 @@ pub extern "C-unwind" fn knot_constraint_register(
1949019694
1949119695// ── Atomic (transactions) ─────────────────────────────────────────
1949219696
19697+
1949319698#[unsafe(no_mangle)]
1949419699pub extern "C-unwind" fn knot_atomic_begin(db: *mut c_void) {
1949519700 let _guard = write_lock_guard();
@@ -24312,6 +24517,7 @@ mod extract_source_tests {
2431224517 env,
2431324518 source: intern_str(source),
2431424519 extract_source: None,
24520+ code: std::ptr::null(),
2431524521 })))
2431624522 }
2431724523 fn func_full(source: &str, env: *mut Value, xsrc: &str) -> *mut Value {
@@ -24320,6 +24526,7 @@ mod extract_source_tests {
2432024526 env,
2432124527 source: intern_str(source),
2432224528 extract_source: Some(intern_str(xsrc)),
24529+ code: std::ptr::null(),
2432324530 })))
2432424531 }
2432524532
@@ -24465,6 +24672,7 @@ mod extract_source_tests {
2446524672 fn_ptr: std::ptr::null(),
2446624673 env,
2446724674 source: intern_str(source),
24675+ code: std::ptr::null(),
2446824676 })))
2446924677 }
2447024678 #[test]
@@ -24514,6 +24722,7 @@ mod show_tests {
2451424722 env,
2451524723 source: intern_str(source),
2451624724 extract_source: None,
24725+ code: std::ptr::null(),
2451724726 })))
2451824727 }
2451924728 fn io(source: &str) -> *mut Value {
@@ -24521,6 +24730,7 @@ mod show_tests {
2452124730 fn_ptr: std::ptr::null(),
2452224731 env: std::ptr::null_mut(),
2452324732 source: intern_str(source),
24733+ code: std::ptr::null(),
2452424734 })))
2452524735 }
2452624736
@@ -24558,6 +24768,7 @@ mod show_tests {
2455824768 env: alloc(Value::Unit),
2455924769 source: intern_str("Nothing"),
2456024770 extract_source: Some(intern_str("Maybe.Nothing {}")),
24771+ code: std::ptr::null(),
2456124772 })));
2456224773 assert_eq!(format_value(f), "Nothing");
2456324774 }