From 477744dcdbe5fc64ea6db106b4ad91a009af547d Mon Sep 17 00:00:00 2001
From: elftausend <76885970+elftausend@users.noreply.github.com>
Date: Sun, 24 Nov 2024 13:01:01 +0100
Subject: [PATCH] Update some parts of the Fork module

---
 Cargo.toml          |  2 +-
 src/modules/fork.rs | 35 +++++++++++++----------------------
 2 files changed, 14 insertions(+), 23 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 086cfbd2..48ead8a5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -53,7 +53,7 @@ min-cl = { git = "https://github.com/elftausend/min-cl", optional = true }
 [features]
 # default = ["cpu", "blas", "static-api", "macro", "cached", "autograd", "stack", "opencl", "fork", "graph", "untyped"]
 
-default = ["cpu", "cached", "autograd", "static-api", "blas", "macro", "stack"]
+default = ["cpu", "cached", "autograd", "static-api", "blas", "macro"]
 # default = ["no-std"]
 # default = ["opencl"]
 # default = ["untyped", "cpu", "lazy", "graph", "autograd", "fork", "serde", "json", "half", "cached", "static-api", "stack", "opencl", "nnapi"]
diff --git a/src/modules/fork.rs b/src/modules/fork.rs
index f42504ce..9a8da810 100644
--- a/src/modules/fork.rs
+++ b/src/modules/fork.rs
@@ -1,8 +1,5 @@
 use crate::{
-    impl_remove_layer, pass_down_add_operation, pass_down_exec_now, pass_down_replace_buf_module,
-    pass_down_tape_actions, AddLayer, Alloc, Buffer, Device, HasId, HasModules, IsShapeIndep,
-    Module, OnDropBuffer, OnNewBuffer, Parents, PtrType, Retrieve, RunModule, Setup, Shape, Unit,
-    UseGpuOrCpu, WrappedData, VERSION,
+    impl_remove_layer, pass_down_add_operation, pass_down_exec_now, pass_down_replace_buf_module, pass_down_tape_actions, AddLayer, Alloc, Buffer, Device, HasId, HasModules, IsBasePtr, IsShapeIndep, Module, OnNewBuffer, Parents, PtrType, Retrieve, RunModule, Setup, Shape, Unit, UseGpuOrCpu, WrappedData, VERSION
 };
 use core::cell::{Cell, RefCell};
 
@@ -28,22 +25,27 @@ pub struct Fork<Mods> {
 }
 
 impl<Mods: WrappedData> WrappedData for Fork<Mods> {
-    type Wrap<T: Unit, Base: HasId + PtrType> = Mods::Wrap<T, Base>;
+    type Wrap<'a, T: Unit, Base: IsBasePtr> = Mods::Wrap<'a, T, Base>;
 
     #[inline]
-    fn wrap_in_base<T: Unit, Base: HasId + PtrType>(&self, base: Base) -> Self::Wrap<T, Base> {
+    fn wrap_in_base<'a, T: Unit, Base: IsBasePtr>(&'a self, base: Base) -> Self::Wrap<'a, T, Base> {
         self.modules.wrap_in_base(base)
     }
 
     #[inline]
-    fn wrapped_as_base<T: Unit, Base: HasId + PtrType>(wrap: &Self::Wrap<T, Base>) -> &Base {
+    fn wrap_in_base_unbound<'a, T: Unit, Base: IsBasePtr>(&self, base: Base) -> Self::Wrap<'a, T, Base> {
+        self.modules.wrap_in_base_unbound(base)
+    }
+
+    #[inline]
+    fn wrapped_as_base<'a, 'b, T: Unit, Base: IsBasePtr>(wrap: &'b Self::Wrap<'a, T, Base>) -> &'b Base {
         Mods::wrapped_as_base(wrap)
     }
 
     #[inline]
-    fn wrapped_as_base_mut<T: Unit, Base: HasId + PtrType>(
-        wrap: &mut Self::Wrap<T, Base>,
-    ) -> &mut Base {
+    fn wrapped_as_base_mut<'a, 'b, T: Unit, Base: IsBasePtr>(
+        wrap: &'b mut Self::Wrap<'a, T, Base>,
+    ) -> &'b mut Base {
         Mods::wrapped_as_base_mut(wrap)
     }
 }
@@ -86,22 +88,11 @@ crate::pass_down_cursor!(Fork);
 pass_down_add_operation!(Fork);
 pass_down_exec_now!(Fork);
 
-impl<Mods: OnDropBuffer> OnDropBuffer for Fork<Mods> {
-    #[inline]
-    fn on_drop_buffer<T: Unit, D: crate::Device, S: crate::Shape>(
-        &self,
-        device: &D,
-        buf: &crate::Buffer<T, D, S>,
-    ) {
-        self.modules.on_drop_buffer(device, buf)
-    }
-}
-
 impl<'a, Mods: OnNewBuffer<'a, T, D, S>, T: Unit, D: Device, S: Shape> OnNewBuffer<'a, T, D, S>
     for Fork<Mods>
 {
     #[inline]
-    unsafe fn on_new_buffer(&self, device: &'a D, new_buf: &crate::Buffer<'a, T, D, S>) {
+    unsafe fn on_new_buffer(&'a self, device: &'a D, new_buf: &mut crate::Buffer<'a, T, D, S>) {
         self.modules.on_new_buffer(device, new_buf)
     }
 }