From 06da6c67578a350ebda87635377e408310dbacac Mon Sep 17 00:00:00 2001
From: Edward Amsden <edward@blackriversoft.com>
Date: Thu, 7 Dec 2023 00:10:19 -0600
Subject: [PATCH] pma: persist instance for Batteries

---
 rust/ares/src/jets/cold.rs | 50 ++++++++++++++++++++++++++++++++++++++
 rust/ares/src/noun.rs      | 15 ++++++++++++
 rust/ares/src/persist.rs   | 43 +++++++++++++++++++++++++++++++-
 3 files changed, 107 insertions(+), 1 deletion(-)

diff --git a/rust/ares/src/jets/cold.rs b/rust/ares/src/jets/cold.rs
index da9205c..4aa392f 100644
--- a/rust/ares/src/jets/cold.rs
+++ b/rust/ares/src/jets/cold.rs
@@ -5,6 +5,7 @@ use crate::noun::{Atom, DirectAtom, Noun, Slots, D, T};
 use crate::persist::{Persist, PMA};
 use std::ptr::copy_nonoverlapping;
 use std::ptr::null_mut;
+use std::mem::size_of;
 
 pub enum Error {
     NoParent,
@@ -32,6 +33,55 @@ struct BatteriesMem {
     parent_batteries: Batteries,
 }
 
+impl Persist for Batteries {
+    unsafe fn space_needed(&mut self, stack: &mut NockStack, pma: &PMA) -> usize {
+        let mut bytes = 0;
+        let mut batteries = *self;
+
+        loop {
+            if batteries.0.is_null() { break; }
+            if pma.contains(batteries.0, 1) { break; }
+            bytes += size_of::<BatteriesMem>();
+            bytes += (*batteries.0).battery.space_needed(stack, pma);
+            bytes += (*batteries.0).parent_axis.space_needed(stack, pma);
+            batteries = (*batteries.0).parent_batteries;
+        }
+        bytes
+    }
+
+    unsafe fn copy_to_buffer(
+        &mut self,           
+        stack: &mut NockStack,
+        pma: &PMA,
+        buffer: &mut *mut u8,
+    ) {
+        let mut dest = self;
+        loop {
+            if (*dest).0.is_null() { break; }
+            if pma.contains((*dest).0, 1) { break; }
+
+            let batteries_mem_ptr = *buffer as *mut BatteriesMem;
+            copy_nonoverlapping((*dest).0, batteries_mem_ptr, 1);
+            *buffer = batteries_mem_ptr.add(1) as *mut u8;
+
+            (*batteries_mem_ptr).battery.copy_to_buffer(stack, pma, buffer);
+            (*batteries_mem_ptr).parent_axis.copy_to_buffer(stack, pma, buffer);
+
+            (*dest).0 = batteries_mem_ptr;
+            dest = &mut (*(*dest).0).parent_batteries;
+        }
+    }
+
+    unsafe fn handle_to_u64(&self) -> u64 {
+        self.0 as u64
+    }
+
+    unsafe fn handle_from_u64(meta_handle: u64) -> Self {
+        Batteries(meta_handle as *mut BatteriesMem)
+    }
+
+}
+
 impl Preserve for Batteries {
     unsafe fn assert_in_stack(&self, stack: &NockStack) {
         if self.0.is_null() {
diff --git a/rust/ares/src/noun.rs b/rust/ares/src/noun.rs
index 0e4b5af..9936de0 100644
--- a/rust/ares/src/noun.rs
+++ b/rust/ares/src/noun.rs
@@ -876,6 +876,21 @@ impl Atom {
             *self
         }
     }
+
+    /** Make an atom from a raw u64
+     *
+     * # Safety 
+     *
+     * Note that the [u64] parameter is *not*, in general, the value of the atom!
+     *
+     * In particular, anything with the high bit set will be treated as a tagged pointer.
+     * This method is only to be used to restore an atom from the raw [u64] representation
+     * returned by [Noun::as_raw], and should only be used if we are sure the restored noun is in
+     * fact an atom.
+     */
+    pub unsafe fn from_raw(raw: u64) -> Atom {
+        Atom { raw }
+    }
 }
 
 impl fmt::Display for Atom {
diff --git a/rust/ares/src/persist.rs b/rust/ares/src/persist.rs
index 16c3354..42c0420 100644
--- a/rust/ares/src/persist.rs
+++ b/rust/ares/src/persist.rs
@@ -1,6 +1,6 @@
 use crate::jets::cold::Cold;
 use crate::mem::NockStack;
-use crate::noun::{Allocated, Cell, CellMemory, IndirectAtom, Noun};
+use crate::noun::{Allocated, Cell, CellMemory, IndirectAtom, Noun, Atom};
 use ares_pma::*;
 use either::Either::{Left, Right};
 use std::ffi::{c_void, CString};
@@ -199,6 +199,47 @@ unsafe fn unmark(a: Allocated) {
     a.set_metadata(metadata & !NOUN_MARKED);
 }
 
+impl Persist for Atom {
+    unsafe fn space_needed(&mut self, stack: &mut NockStack, pma: &PMA) -> usize {
+        if let Ok(indirect) = self.as_indirect() {
+            let count = indirect.raw_size();
+            if !pma.contains(indirect.to_raw_pointer(), count) {
+                if !mark(indirect.as_allocated()) {
+                    return count * size_of::<u64>();
+                }
+            }
+        }
+        0
+    }
+
+    unsafe fn copy_to_buffer(&mut self, stack: &mut NockStack, pma: &PMA, buffer: &mut *mut u8) {
+        if let Ok(mut indirect) = self.as_indirect() {
+            let count = indirect.raw_size();
+            if !pma.contains(indirect.to_raw_pointer(), count) {
+                if let Some(forward) = indirect.forwarding_pointer() {
+                    *self = forward.as_atom();
+                } else {
+                    let indirect_buffer_ptr = *buffer as *mut u64;
+                    copy_nonoverlapping(indirect.to_raw_pointer(), indirect_buffer_ptr, count);
+                    *buffer = indirect_buffer_ptr.add(count) as *mut u8;
+
+                    indirect.set_forwarding_pointer(indirect_buffer_ptr);
+
+                    *self = IndirectAtom::from_raw_pointer(indirect_buffer_ptr).as_atom();
+                }
+            }
+        }
+    }
+
+    unsafe fn handle_to_u64(&self) -> u64 {
+        self.as_noun().as_raw()
+    }
+
+    unsafe fn handle_from_u64(meta_handle: u64) -> Self {
+        Atom::from_raw(meta_handle)
+    }
+}
+
 impl Persist for Noun {
     unsafe fn space_needed(&mut self, stack: &mut NockStack, pma: &PMA) -> usize {
         let mut space = 0usize;