From 01c9d0b2d347f13fd1a8be2833867a6ceed32ed8 Mon Sep 17 00:00:00 2001
From: Robert Hensing <robert@roberthensing.nl>
Date: Thu, 27 Jun 2024 16:01:40 +0200
Subject: [PATCH] feat: EvalState.weak_ref()

(cherry picked from commit d7a72c82ebfbfbb1b58fa15044b63648b2109260)
---
 rust/nix-expr/src/eval_state.rs | 57 +++++++++++++++++++++++++++++++--
 rust/nix-store/src/store.rs     |  2 ++
 2 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/rust/nix-expr/src/eval_state.rs b/rust/nix-expr/src/eval_state.rs
index 8596dc5..b1ada05 100644
--- a/rust/nix-expr/src/eval_state.rs
+++ b/rust/nix-expr/src/eval_state.rs
@@ -4,14 +4,14 @@ use anyhow::{bail, Result};
 use lazy_static::lazy_static;
 use nix_c_raw as raw;
 use nix_store::path::StorePath;
-use nix_store::store::Store;
+use nix_store::store::{Store, StoreWeak};
 use nix_util::context::Context;
 use nix_util::string_return::{callback_get_result_string, callback_get_result_string_data};
 use nix_util::{check_call, check_call_opt_key, result_string_init};
 use std::ffi::{c_char, CString};
 use std::os::raw::c_uint;
 use std::ptr::{null, null_mut, NonNull};
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 
 lazy_static! {
     static ref INIT: Result<()> = {
@@ -38,6 +38,26 @@ pub struct RealisedString {
     pub paths: Vec<StorePath>,
 }
 
+/// A [Weak] reference to an [EvalState]
+pub struct EvalStateWeak {
+    inner: Weak<EvalStateRef>,
+    store: StoreWeak,
+}
+impl EvalStateWeak {
+    /// Upgrade the weak reference to a proper [EvalState].
+    ///
+    /// If no normal reference to the [EvalState] is around anymore elsewhere, this fails by returning `None`.
+    pub fn upgrade(&self) -> Option<EvalState> {
+        self.inner.upgrade().and_then(|eval_state| {
+            self.store.upgrade().map(|store| EvalState {
+                eval_state: eval_state,
+                store: store,
+                context: Context::new(),
+            })
+        })
+    }
+}
+
 struct EvalStateRef {
     eval_state: NonNull<raw::EvalState>,
 }
@@ -106,6 +126,13 @@ impl EvalState {
     pub fn store(&self) -> &Store {
         &self.store
     }
+    pub fn weak_ref(&self) -> EvalStateWeak {
+        EvalStateWeak {
+            inner: Arc::downgrade(&self.eval_state),
+            store: self.store.weak_ref(),
+        }
+    }
+
     /// Parses and evaluates a Nix expression `expr`.
     ///
     /// Expressions can contain relative paths such as `./.` that are resolved relative to the given `path`.
@@ -481,6 +508,32 @@ mod tests {
         .unwrap();
     }
 
+    #[test]
+    fn weak_ref() {
+        gc_registering_current_thread(|| {
+            let store = Store::open("auto", HashMap::new()).unwrap();
+            let es = EvalState::new(store, []).unwrap();
+            let weak = es.weak_ref();
+            let _es = weak.upgrade().unwrap();
+        })
+        .unwrap();
+    }
+
+    #[test]
+    fn weak_ref_gone() {
+        gc_registering_current_thread(|| {
+            let weak = {
+                let store = Store::open("auto", HashMap::new()).unwrap();
+                let es = EvalState::new(store, []).unwrap();
+                es.weak_ref()
+            };
+            assert!(weak.upgrade().is_none());
+            assert!(weak.store.upgrade().is_none());
+            assert!(weak.inner.upgrade().is_none());
+        })
+        .unwrap();
+    }
+
     #[test]
     fn eval_state_lookup_path() {
         let import_expression = "import <test_file0> + import <test_file1>";
diff --git a/rust/nix-store/src/store.rs b/rust/nix-store/src/store.rs
index 8ca2031..631a6f5 100644
--- a/rust/nix-store/src/store.rs
+++ b/rust/nix-store/src/store.rs
@@ -39,6 +39,8 @@ pub struct StoreWeak {
 }
 impl StoreWeak {
     /// Upgrade the weak reference to a proper [Store].
+    ///
+    /// If no normal reference to the [Store] is around anymore elsewhere, this fails by returning `None`.
     pub fn upgrade(&self) -> Option<Store> {
         self.inner.upgrade().map(|inner| Store {
             inner,