From 2df922d3e9988d15de66538841521cc26f97f9d4 Mon Sep 17 00:00:00 2001
From: Matt Keeter <matt.j.keeter@gmail.com>
Date: Wed, 6 Nov 2024 08:35:44 -0500
Subject: [PATCH 1/3] Helper functions to convert JitFunctions to/from
 GenericVmFunction

---
 fidget/src/core/eval/mod.rs |  1 +
 fidget/src/jit/mod.rs       | 24 ++++++++++++++++++------
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/fidget/src/core/eval/mod.rs b/fidget/src/core/eval/mod.rs
index 7d18ef75..6be39ea6 100644
--- a/fidget/src/core/eval/mod.rs
+++ b/fidget/src/core/eval/mod.rs
@@ -3,6 +3,7 @@ use crate::{
     context::{Context, Node},
     types::{Grad, Interval},
     var::VarMap,
+    vm::GenericVmFunction,
     Error,
 };
 
diff --git a/fidget/src/jit/mod.rs b/fidget/src/jit/mod.rs
index 6d0da4fb..525b5b4d 100644
--- a/fidget/src/jit/mod.rs
+++ b/fidget/src/jit/mod.rs
@@ -926,6 +926,24 @@ impl RenderHints for JitFunction {
     }
 }
 
+impl MathFunction for JitFunction {
+    fn new(ctx: &Context, nodes: &[Node]) -> Result<Self, Error> {
+        GenericVmFunction::new(ctx, nodes).map(JitFunction)
+    }
+}
+
+impl From<GenericVmFunction<REGISTER_LIMIT>> for JitFunction {
+    fn from(v: GenericVmFunction<REGISTER_LIMIT>) -> Self {
+        Self(v)
+    }
+}
+
+impl From<JitFunction> for GenericVmFunction<REGISTER_LIMIT> {
+    fn from(v: JitFunction) -> Self {
+        v.0.clone()
+    }
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 // Selects the calling convention based on platform; this is forward-looking for
@@ -1278,12 +1296,6 @@ impl BulkEvaluator for JitGradSliceEval {
     }
 }
 
-impl MathFunction for JitFunction {
-    fn new(ctx: &Context, nodes: &[Node]) -> Result<Self, Error> {
-        GenericVmFunction::new(ctx, nodes).map(JitFunction)
-    }
-}
-
 /// A [`Shape`](crate::shape::Shape) which uses the JIT evaluator
 pub type JitShape = crate::shape::Shape<JitFunction>;
 

From e92d09d506eb8468915bb1059320fe4d2951bef3 Mon Sep 17 00:00:00 2001
From: Matt Keeter <matt.j.keeter@gmail.com>
Date: Wed, 6 Nov 2024 08:37:56 -0500
Subject: [PATCH 2/3] Helper function to resize and get register count

---
 fidget/src/core/eval/mod.rs |  1 -
 fidget/src/core/vm/mod.rs   | 14 ++++++++++++--
 fidget/src/jit/mod.rs       |  6 +++---
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/fidget/src/core/eval/mod.rs b/fidget/src/core/eval/mod.rs
index 6be39ea6..7d18ef75 100644
--- a/fidget/src/core/eval/mod.rs
+++ b/fidget/src/core/eval/mod.rs
@@ -3,7 +3,6 @@ use crate::{
     context::{Context, Node},
     types::{Grad, Interval},
     var::VarMap,
-    vm::GenericVmFunction,
     Error,
 };
 
diff --git a/fidget/src/core/vm/mod.rs b/fidget/src/core/vm/mod.rs
index c95be7b1..b574023e 100644
--- a/fidget/src/core/vm/mod.rs
+++ b/fidget/src/core/vm/mod.rs
@@ -135,6 +135,17 @@ impl<const N: usize> GenericVmFunction<N> {
     pub fn choice_count(&self) -> usize {
         self.0.choice_count()
     }
+
+    /// Simplifies the function with the given trace and a new register count
+    pub fn simplify_with<const M: usize>(
+        &self,
+        trace: &VmTrace,
+        storage: VmData<M>,
+        workspace: &mut VmWorkspace<M>,
+    ) -> Result<GenericVmFunction<M>, Error> {
+        let d = self.0.simplify::<M>(trace.as_slice(), workspace, storage)?;
+        Ok(GenericVmFunction(Arc::new(d)))
+    }
 }
 
 impl<const N: usize> Function for GenericVmFunction<N> {
@@ -166,8 +177,7 @@ impl<const N: usize> Function for GenericVmFunction<N> {
         storage: Self::Storage,
         workspace: &mut Self::Workspace,
     ) -> Result<Self, Error> {
-        let d = self.0.simplify(trace.as_slice(), workspace, storage)?;
-        Ok(Self(Arc::new(d)))
+        self.simplify_with(trace, storage, workspace)
     }
 
     fn recycle(self) -> Option<Self::Storage> {
diff --git a/fidget/src/jit/mod.rs b/fidget/src/jit/mod.rs
index 525b5b4d..2893281d 100644
--- a/fidget/src/jit/mod.rs
+++ b/fidget/src/jit/mod.rs
@@ -938,9 +938,9 @@ impl From<GenericVmFunction<REGISTER_LIMIT>> for JitFunction {
     }
 }
 
-impl From<JitFunction> for GenericVmFunction<REGISTER_LIMIT> {
-    fn from(v: JitFunction) -> Self {
-        v.0.clone()
+impl<'a> From<&'a JitFunction> for &'a GenericVmFunction<REGISTER_LIMIT> {
+    fn from(v: &'a JitFunction) -> Self {
+        &v.0
     }
 }
 

From 739f7ce81f5cc55426319c1ed0d34d04c096bed2 Mon Sep 17 00:00:00 2001
From: Matt Keeter <matt.j.keeter@gmail.com>
Date: Wed, 6 Nov 2024 08:45:51 -0500
Subject: [PATCH 3/3] Update CHANGELOG and bump version

---
 CHANGELOG.md      | 6 ++++++
 Cargo.lock        | 2 +-
 fidget/Cargo.toml | 2 +-
 3 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 02f2ce98..9a937a97 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,9 @@
+# 0.3.4 (unreleased)
+- Add `GenericVmFunction::simplify_with` to simultaneously simplify a function
+  and pick a new register count for the resulting tape
+- Add bidirectional conversions between `JitFunction` and `GenericVmFunction`
+  (exposing the inner data member)
+
 # 0.3.3
 - `Function` and evaluator types now produce multiple outputs
     - `MathFunction::new` now takes a slice of nodes, instead of a single node
diff --git a/Cargo.lock b/Cargo.lock
index ffa15734..75e64e90 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -871,7 +871,7 @@ dependencies = [
 
 [[package]]
 name = "fidget"
-version = "0.3.3"
+version = "0.3.4"
 dependencies = [
  "approx",
  "arrayvec",
diff --git a/fidget/Cargo.toml b/fidget/Cargo.toml
index 47f03bcb..7cbd6bad 100644
--- a/fidget/Cargo.toml
+++ b/fidget/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "fidget"
-version = "0.3.3"
+version = "0.3.4"
 edition = "2021"
 license = "MPL-2.0"
 repository = "https://github.com/mkeeter/fidget"