From f7966d7908f92bf213f543d0448cceeea71ea7d8 Mon Sep 17 00:00:00 2001
From: Jayaprabhakar Kadarkarai <jp@fizzbee.io>
Date: Wed, 1 May 2024 12:49:29 -0700
Subject: [PATCH] SetCustomPtrFunc to customize cloning pointers

---
 allocator.go              | 17 +++++++++++++++++
 clone.go                  |  6 ++++++
 structtype.go             |  8 ++++++++
 structtype_sample_test.go | 40 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 71 insertions(+)

diff --git a/allocator.go b/allocator.go
index af415f3..ee39fc9 100644
--- a/allocator.go
+++ b/allocator.go
@@ -318,6 +318,23 @@ func (a *Allocator) SetCustomFunc(t reflect.Type, fn Func) {
 	a.cachedCustomFuncTypes.Store(t, fn)
 }
 
+// SetCustomPtrFunc sets a custom clone function for type t.
+// If t is not pointer to struct, SetCustomFunc ignores t.
+//
+// If fn is nil, remove the custom clone function for type t.
+func (a *Allocator) SetCustomPtrFunc(t reflect.Type, fn Func) {
+        if fn == nil {
+                a.cachedCustomFuncTypes.Delete(t)
+                return
+        }
+
+        for t.Kind() != reflect.Ptr || t.Elem().Kind() != reflect.Struct {
+		return
+        }
+
+        a.cachedCustomFuncTypes.Store(t, fn)
+}
+
 func heapNew(pool unsafe.Pointer, t reflect.Type) reflect.Value {
 	return reflect.New(t)
 }
diff --git a/clone.go b/clone.go
index 18a3c73..c4d836d 100644
--- a/clone.go
+++ b/clone.go
@@ -190,6 +190,12 @@ func (state *cloneState) clonePtr(v reflect.Value) reflect.Value {
 
 	t := v.Type()
 
+	if fn, ok := state.allocator.cachedCustomFuncTypes.Load(t); ok {
+		nv := state.allocator.New(t)
+		fn.(Func)(state.allocator, v, nv.Elem())
+		return nv.Elem()
+	}
+
 	if state.allocator.isOpaquePointer(t) {
 		if v.CanInterface() {
 			return v
diff --git a/structtype.go b/structtype.go
index c6cce98..3a89bdf 100644
--- a/structtype.go
+++ b/structtype.go
@@ -151,6 +151,14 @@ func SetCustomFunc(t reflect.Type, fn Func) {
 	defaultAllocator.SetCustomFunc(t, fn)
 }
 
+// SetCustomPtrFunc sets a custom clone function for type t in heap allocator.
+// If t is not pointer to struct, SetCustomFunc ignores t.
+//
+// If fn is nil, remove the custom clone function for type t.
+func SetCustomPtrFunc(t reflect.Type, fn Func) {
+	defaultAllocator.SetCustomPtrFunc(t, fn)
+}
+
 // Init creates a new value of src.Type() and shadow copies all content from src.
 // If noCustomFunc is set to true, custom clone function will be ignored.
 //
diff --git a/structtype_sample_test.go b/structtype_sample_test.go
index 9e9a183..0a1c913 100644
--- a/structtype_sample_test.go
+++ b/structtype_sample_test.go
@@ -90,6 +90,46 @@ func ExampleSetCustomFunc() {
 	// [abc 123]
 }
 
+// ExampleSetCustomPtrFunc tests whether we can customize the cloning behavior of a pointer type.
+// In this example, the custom function reuses cached values.
+func ExampleSetCustomPtrFunc() {
+	type Data struct {
+		Name  string
+	Value int
+	}
+	refs := make(map[string]*Data)
+	// Filter nil values in Data when cloning old value.
+	SetCustomPtrFunc(reflect.TypeOf(&Data{}), func(allocator *Allocator, old, new reflect.Value) {
+		// The new is a zero value of MyStruct.
+		// We can get its address to update it.
+		value := new.Addr().Interface().(**Data)
+		oldRole := old.Interface().(*Data)
+
+		if cached, ok := refs[(*oldRole).Name]; ok {
+			*value = cached
+		} else {
+			*value = &Data{
+				Name:  (*oldRole).Name,
+				Value: (*oldRole).Value,
+			}
+			refs[(*value).Name] = *value
+		}
+	})
+
+	orig := &Data{
+		Name:  "abc",
+		Value: 123,
+	}
+	cloned1 := Clone(orig).(*Data)
+	cloned2 := Clone(orig).(*Data)
+	cloned1.Value = 456
+	orig.Value = -1
+	fmt.Println(*orig, *cloned1, *cloned2)
+
+	// Output:
+	// {abc -1} {abc 456} {abc 456}
+}
+
 func ExampleSetCustomFunc_partiallyClone() {
 	type T struct {
 		Value int