From f67eb4ac3c65291c590775f1b56cd1caa5444efc Mon Sep 17 00:00:00 2001
From: Yusuke Hata <yusuke.hata@mirrativ.co.jp>
Date: Wed, 10 Feb 2021 20:43:50 +0900
Subject: [PATCH] image.NRGBA

---
 imagepool.go       |  42 ++++++++++++++
 imagepool_test.go  | 134 +++++++++++++++++++++++++++++++++++++++++++++
 mimagepool.go      |  60 +++++++++++++++++++-
 mimagepool_test.go | 134 +++++++++++++++++++++++++++++++++++++++++++++
 ref.go             |  38 +++++++++++++
 ref_test.go        |  18 ++++++
 6 files changed, 424 insertions(+), 2 deletions(-)

diff --git a/imagepool.go b/imagepool.go
index 7d45a52..bb724f5 100644
--- a/imagepool.go
+++ b/imagepool.go
@@ -99,6 +99,48 @@ func (b *ImageRGBAPool) Cap() int {
 	return cap(b.pool)
 }
 
+type ImageNRGBAPool struct {
+	ImageRGBAPool
+}
+
+func NewImageNRGBAPool(poolSize int, rect image.Rectangle, funcs ...optionFunc) *ImageNRGBAPool {
+	opt := newOption()
+	for _, fn := range funcs {
+		fn(opt)
+	}
+
+	b := new(ImageNRGBAPool)
+	b.pool = make(chan []uint8, poolSize)
+	b.init(rect)
+
+	if opt.preload {
+		b.preload(opt.preloadRate)
+	}
+	return b
+}
+
+func (b *ImageNRGBAPool) createImageNRGBARef(pix []uint8, pool *ImageNRGBAPool) *ImageNRGBARef {
+	ref := newImageNRGBARef(pix, &image.NRGBA{
+		Pix:    pix,
+		Stride: b.stride,
+		Rect:   b.rect,
+	}, pool)
+	ref.setFinalizer()
+	return ref
+}
+
+func (b *ImageNRGBAPool) GetRef() *ImageNRGBARef {
+	var pix []uint8
+	select {
+	case pix = <-b.pool:
+		// reuse exists pool
+	default:
+		// create []uint8
+		pix = make([]uint8, b.length)
+	}
+	return b.createImageNRGBARef(pix, b)
+}
+
 type ImageYCbCrPool struct {
 	pool     chan []uint8
 	rect     image.Rectangle
diff --git a/imagepool_test.go b/imagepool_test.go
index 9650cc7..a990bd5 100644
--- a/imagepool_test.go
+++ b/imagepool_test.go
@@ -168,6 +168,61 @@ func TestImageRGBAPoolBufSize(t *testing.T) {
 	})
 }
 
+func TestImageNRGBAPoolBufSize(t *testing.T) {
+	t.Run("getsamecap", func(tt *testing.T) {
+		rect := image.Rect(0, 0, 100, 100)
+		pool := NewImageNRGBAPool(10, rect)
+		img := image.NewNRGBA(rect)
+
+		d1 := pool.GetRef()
+		d2 := pool.GetRef()
+		if cap(d1.pix) != cap(img.Pix) {
+			tt.Errorf("buf size = %d", cap(img.Pix))
+		}
+		if cap(d2.pix) != cap(img.Pix) {
+			tt.Errorf("buf size = %d", cap(img.Pix))
+		}
+	})
+	t.Run("getput/smallcap", func(tt *testing.T) {
+		rect := image.Rect(0, 0, 100, 100)
+		pool := NewImageNRGBAPool(10, rect)
+		img := image.NewNRGBA(rect)
+
+		r1 := image.Rect(0, 0, 64, 36)
+		i1 := image.NewNRGBA(r1)
+		if pool.Put(i1.Pix) {
+			tt.Errorf("discard small pix")
+		}
+
+		d1 := pool.GetRef()
+		if cap(d1.pix) != cap(img.Pix) {
+			tt.Errorf("discard small pix = %d", cap(d1.pix))
+		}
+		if len(d1.pix) != len(img.Pix) {
+			tt.Errorf("discard small pix")
+		}
+	})
+	t.Run("getput/largecap", func(tt *testing.T) {
+		rect := image.Rect(0, 0, 100, 100)
+		pool := NewImageNRGBAPool(10, rect)
+		img := image.NewNRGBA(rect)
+
+		r1 := image.Rect(0, 0, 640, 360)
+		i1 := image.NewNRGBA(r1)
+		if pool.Put(i1.Pix) != true {
+			tt.Errorf("allow large pix")
+		}
+
+		d1 := pool.GetRef()
+		if cap(d1.pix) == cap(img.Pix) {
+			tt.Errorf("large pix ok")
+		}
+		if len(d1.pix) != len(img.Pix) {
+			tt.Errorf("same len")
+		}
+	})
+}
+
 func TestImageYCbCrPoolBufSize(t *testing.T) {
 	t.Run("getsamecap", func(tt *testing.T) {
 		rect := image.Rect(0, 0, 100, 100)
@@ -307,6 +362,76 @@ func TestImageRGBAPoolCapLen(t *testing.T) {
 	})
 }
 
+func TestImageNRGBAPoolCapLen(t *testing.T) {
+	t.Run("getput", func(tt *testing.T) {
+		r := image.Rect(0, 0, 16, 9)
+		p := NewImageNRGBAPool(10, r)
+		if 0 != p.Len() {
+			tt.Errorf("initial len 0")
+		}
+		if 10 != p.Cap() {
+			tt.Errorf("initial cap 10")
+		}
+
+		data := p.GetRef()
+		if 0 != p.Len() {
+			tt.Errorf("initial len 0")
+		}
+		if 10 != p.Cap() {
+			tt.Errorf("initial cap 10")
+		}
+		p.Put(data.pix)
+
+		if 1 != p.Len() {
+			tt.Errorf("put one")
+		}
+		if 10 != p.Cap() {
+			tt.Errorf("initial cap 10")
+		}
+
+		d1 := p.GetRef()
+		if 0 != p.Len() {
+			tt.Errorf("acquire pool")
+		}
+		p.Put(d1.pix)
+		if 1 != p.Len() {
+			tt.Errorf("release pool")
+		}
+	})
+	t.Run("maxcap", func(tt *testing.T) {
+		r := image.Rect(0, 0, 16, 9)
+		p := NewImageNRGBAPool(10, r)
+		s := make([]*ImageNRGBARef, 0)
+		for i := 0; i < 10; i += 1 {
+			r := p.GetRef()
+			s = append(s, r)
+		}
+		for _, r := range s {
+			p.Put(r.pix)
+		}
+
+		if 10 != p.Len() {
+			tt.Errorf("fill-ed pool: %d", p.Len())
+		}
+		if 10 != p.Cap() {
+			tt.Errorf("max capacity = 10")
+		}
+
+		i1 := image.NewNRGBA(r)
+		d1 := newImageNRGBARef(i1.Pix, i1, p)
+		i2 := image.NewNRGBA(r)
+		d2 := newImageNRGBARef(i2.Pix, i2, p)
+		p.Put(d1.pix)
+		p.Put(d2.pix)
+		if 10 != p.Len() {
+			tt.Errorf("fixed size pool")
+		}
+		if 10 != p.Cap() {
+			tt.Errorf("max capacity = 10")
+		}
+	})
+}
+
 func TestImageYCbCrPoolCapLen(t *testing.T) {
 	t.Run("getput", func(tt *testing.T) {
 		r := image.Rect(0, 0, 16, 9)
@@ -424,6 +549,15 @@ func TestImageRGBAPoolPreload(t *testing.T) {
 	}
 }
 
+func TestImageNRGBAPoolPreload(t *testing.T) {
+	r := image.Rect(0, 0, 100, 100)
+	p := NewImageNRGBAPool(12, r, Preload(true))
+	l := int(float64(12) * defaultPreloadRate)
+	if p.Len() != l {
+		t.Errorf("preloaded buffer = %d", p.Len())
+	}
+}
+
 func TestImageYCbCrPoolPreload(t *testing.T) {
 	r := image.Rect(0, 0, 100, 100)
 	p := NewImageYCbCrPool(12, r, image.YCbCrSubsampleRatio420, Preload(true))
diff --git a/mimagepool.go b/mimagepool.go
index 46a29a2..5a48afc 100644
--- a/mimagepool.go
+++ b/mimagepool.go
@@ -79,12 +79,11 @@ func NewMultiImageRGBAPool(funcs ...multiImageBufferPoolOptionFunc) *MultiImageR
 	}
 
 	tuples := uniqImagepoolTuple(mOpt.tuples)
-	poolFuncs := mOpt.poolFuncs
 
 	sortTuples(tuples)
 	pools := make([]*ImageRGBAPool, len(tuples))
 	for i, t := range tuples {
-		pools[i] = NewImageRGBAPool(t.poolSize, t.rect, poolFuncs...)
+		pools[i] = NewImageRGBAPool(t.poolSize, t.rect, mOpt.poolFuncs...)
 	}
 	return &MultiImageRGBAPool{
 		tuples: tuples,
@@ -125,6 +124,63 @@ func (b *MultiImageRGBAPool) Put(pix []uint8, r image.Rectangle) bool {
 	return false
 }
 
+type MultiImageNRGBAPool struct {
+	tuples []imagepoolTuple
+	pools  []*ImageNRGBAPool
+}
+
+func NewMultiImageNRGBAPool(funcs ...multiImageBufferPoolOptionFunc) *MultiImageNRGBAPool {
+	mOpt := newMultiImageBufferPoolOption()
+	for _, fn := range funcs {
+		fn(mOpt)
+	}
+
+	tuples := uniqImagepoolTuple(mOpt.tuples)
+
+	sortTuples(tuples)
+	pools := make([]*ImageNRGBAPool, len(tuples))
+	for i, t := range tuples {
+		pools[i] = NewImageNRGBAPool(t.poolSize, t.rect, mOpt.poolFuncs...)
+	}
+	return &MultiImageNRGBAPool{
+		tuples: tuples,
+		pools:  pools,
+	}
+}
+
+func (b *MultiImageNRGBAPool) find(r image.Rectangle) (*ImageNRGBAPool, bool) {
+	if r.Empty() {
+		return nil, false
+	}
+
+	for i, t := range b.tuples {
+		if rectIn(t.rect, r) {
+			return b.pools[i], true
+		}
+	}
+	return nil, false
+}
+
+func (b *MultiImageNRGBAPool) GetRef(r image.Rectangle) *ImageNRGBARef {
+	if pool, ok := b.find(r); ok {
+		return pool.GetRef()
+	}
+
+	pool := &ImageNRGBAPool{}
+	pool.init(r)
+
+	pix := make([]uint8, pool.length)
+	return pool.createImageNRGBARef(pix, b.pools[len(b.pools)-1])
+}
+
+func (b *MultiImageNRGBAPool) Put(pix []uint8, r image.Rectangle) bool {
+	if pool, ok := b.find(r); ok {
+		return pool.Put(pix)
+	}
+	// discard
+	return false
+}
+
 type MultiImageYCbCrPool struct {
 	tuples []imagepoolTuple
 	pools  []*ImageYCbCrPool
diff --git a/mimagepool_test.go b/mimagepool_test.go
index eca1761..421ba3e 100644
--- a/mimagepool_test.go
+++ b/mimagepool_test.go
@@ -140,6 +140,140 @@ func TestMultiImageRGBAPoolPutGet(t *testing.T) {
 	})
 }
 
+func TestMultiImageNRGBAPoolNew(t *testing.T) {
+	t.Run("sorted", func(tt *testing.T) {
+		mp := NewMultiImageNRGBAPool(
+			MultiImagePoolSize(10, image.Rect(0, 0, 640, 640)),
+			MultiImagePoolSize(10, image.Rect(0, 0, 640, 360)),
+			MultiImagePoolSize(10, image.Rect(0, 0, 1280, 720)),
+			MultiImagePoolSize(10, image.Rect(0, 0, 360, 640)),
+			MultiImagePoolSize(10, image.Rect(0, 0, 720, 1280)),
+		)
+		rects := make([]string, len(mp.pools))
+		for i, p := range mp.pools {
+			rects[i] = fmt.Sprintf("%dx%d", p.rect.Dx(), p.rect.Dy())
+		}
+		order := []string{
+			"360x640",
+			"640x360",
+			"640x640",
+			"720x1280",
+			"1280x720",
+		}
+		for i, s := range order {
+			if rects[i] != s {
+				tt.Errorf("sorted expect:pools[%d]=%s", i, s)
+			}
+		}
+	})
+	t.Run("preload", func(tt *testing.T) {
+		mp := NewMultiImageNRGBAPool(
+			MultiImagePoolSize(10, image.Rect(0, 0, 640, 360)),
+			MultiImagePoolSize(10, image.Rect(0, 0, 1280, 720)),
+			MultiImagePoolSize(10, image.Rect(0, 0, 360, 640)),
+			MultiImagePoolOption(
+				Preload(true),
+				PreloadRate(0.5),
+			),
+		)
+		if mp.pools[0].rect.Eq(image.Rect(0, 0, 360, 640)) != true {
+			tt.Errorf("sorted head")
+		}
+		if mp.pools[2].rect.Eq(image.Rect(0, 0, 1280, 720)) != true {
+			tt.Errorf("sorted tail")
+		}
+		for _, p := range mp.pools {
+			l := int(float64(p.Cap()) * 0.5)
+			if p.Len() != l {
+				tt.Errorf("preloaded %d", l)
+			}
+		}
+	})
+}
+
+func TestMultiImageNRGBAPoolPutGet(t *testing.T) {
+	t.Run("getput", func(tt *testing.T) {
+		mp := NewMultiImageNRGBAPool(
+			MultiImagePoolSize(10, image.Rect(0, 0, 640, 360)),
+			MultiImagePoolSize(10, image.Rect(0, 0, 1280, 720)),
+			MultiImagePoolSize(10, image.Rect(0, 0, 360, 640)),
+		)
+		d1 := mp.GetRef(image.Rect(0, 0, 100, 100)) // 100x100 < pools[0]
+		d2 := mp.GetRef(image.Rect(0, 0, 360, 120)) // 360x120 < pools[0]
+		d3 := mp.GetRef(image.Rect(0, 0, 360, 700)) // pools[1] < 360x700 < pools[2]
+		d4 := mp.GetRef(image.Rect(0, 0, 640, 320)) // 640x320 < pools[1]
+		d5 := mp.GetRef(image.Rect(0, 0, 768, 432)) // pools[1] < 768x432 < pools[2]
+		if mp.Put(d1.pix, d1.Img.Bounds()) != true {
+			tt.Errorf("release ok / free cap")
+		}
+		if mp.pools[0].Len() != 1 {
+			tt.Errorf("release pool[0] 100x100")
+		}
+		if mp.Put(d2.pix, d2.Img.Bounds()) != true {
+			tt.Errorf("release ok / free cap")
+		}
+		if mp.pools[0].Len() != 2 {
+			tt.Errorf("release pool[0] 360x120")
+		}
+		if mp.Put(d3.pix, d3.Img.Bounds()) != true {
+			tt.Errorf("release ok / free cap")
+		}
+		if mp.pools[2].Len() != 1 {
+			tt.Errorf("release pool[2] 360x700")
+		}
+		if mp.Put(d4.pix, d4.Img.Bounds()) != true {
+			tt.Errorf("release ok / free cap")
+		}
+		if mp.pools[1].Len() != 1 {
+			tt.Errorf("release pool[1] 640x320")
+		}
+		if mp.Put(d5.pix, d5.Img.Bounds()) != true {
+			tt.Errorf("release ok / free cap")
+		}
+		if mp.pools[2].Len() != 2 {
+			tt.Errorf("release pool[2] 768x432")
+		}
+	})
+	t.Run("getref", func(tt *testing.T) {
+		mp := NewMultiImageNRGBAPool(
+			MultiImagePoolSize(10, image.Rect(0, 0, 640, 360)),
+			MultiImagePoolSize(10, image.Rect(0, 0, 1280, 720)),
+			MultiImagePoolSize(10, image.Rect(0, 0, 360, 640)),
+		)
+		d1 := mp.GetRef(image.Rect(0, 0, 100, 100)) // 100x100 < pools[0]
+		d2 := mp.GetRef(image.Rect(0, 0, 360, 120)) // 360x120 < pools[0]
+		d3 := mp.GetRef(image.Rect(0, 0, 360, 700)) // pools[1] < 360x700 < pools[2]
+		d4 := mp.GetRef(image.Rect(0, 0, 640, 320)) // 640x320 < pools[1]
+		d5 := mp.GetRef(image.Rect(0, 0, 768, 432)) // pools[1] < 768x432 < pools[2]
+		d1.Release()
+		if mp.pools[0].Len() != 1 {
+			tt.Errorf("release pool[0] 100x100")
+		}
+		d2.Release()
+		if mp.pools[0].Len() != 2 {
+			tt.Errorf("release pool[0] 360x120")
+		}
+		d3.Release()
+		if mp.pools[2].Len() != 1 {
+			tt.Errorf("release pool[2] 360x700")
+		}
+		d4.Release()
+		if mp.pools[1].Len() != 1 {
+			tt.Errorf("release pool[1] 640x320")
+		}
+		d5.Release()
+		if mp.pools[2].Len() != 2 {
+			tt.Errorf("release pool[2] 768x432")
+		}
+
+		d6 := mp.GetRef(image.Rect(0, 0, 1920, 1080))
+		d6.Release()
+		if mp.pools[2].Len() == 2 {
+			tt.Errorf("put ok large pix")
+		}
+	})
+}
+
 func TestMultiImageYCbCrPoolNew(t *testing.T) {
 	t.Run("sorted", func(tt *testing.T) {
 		mp := NewMultiImageYCbCrPool(
diff --git a/ref.go b/ref.go
index 607dacd..d9e4e2f 100644
--- a/ref.go
+++ b/ref.go
@@ -211,6 +211,44 @@ func (b *ImageRGBARef) Release() {
 	}
 }
 
+type ImageNRGBARef struct {
+	Img    *image.NRGBA
+	pix    []uint8
+	pool   *ImageNRGBAPool
+	closed int32
+}
+
+func newImageNRGBARef(pix []uint8, img *image.NRGBA, pool *ImageNRGBAPool) *ImageNRGBARef {
+	return &ImageNRGBARef{
+		Img:    img,
+		pix:    pix,
+		pool:   pool,
+		closed: refInit,
+	}
+}
+
+func (b *ImageNRGBARef) Image() *image.NRGBA {
+	return b.Img
+}
+
+func (b *ImageNRGBARef) isClosed() bool {
+	return atomic.LoadInt32(&b.closed) == refClosed
+}
+
+func (b *ImageNRGBARef) setFinalizer() {
+	runtime.SetFinalizer(b, finalizeRef)
+}
+
+func (b *ImageNRGBARef) Release() {
+	if atomic.CompareAndSwapInt32(&b.closed, refInit, refClosed) {
+		runtime.SetFinalizer(b, nil) // clear
+		b.pool.Put(b.pix)
+		b.Img = nil
+		b.pix = nil
+		b.pool = nil
+	}
+}
+
 type ImageYCbCrRef struct {
 	Img    *image.YCbCr
 	pix    []uint8
diff --git a/ref_test.go b/ref_test.go
index df2024c..fddea18 100644
--- a/ref_test.go
+++ b/ref_test.go
@@ -56,6 +56,18 @@ func TestRefValue(t *testing.T) {
 			tt.Errorf("same value")
 		}
 	})
+	t.Run("image.NRGBA", func(tt *testing.T) {
+		img := image.NewNRGBA(image.Rect(0, 0, 10, 10))
+		b := newImageNRGBARef(img.Pix, img, nil)
+
+		d := b.Image()
+		if d.Bounds().Eq(image.Rect(0, 0, 10, 10)) != true {
+			tt.Errorf("same size")
+		}
+		if bytes.Equal(d.Pix, img.Pix) != true {
+			tt.Errorf("same value")
+		}
+	})
 	t.Run("image.YCbCr", func(tt *testing.T) {
 		img := image.NewYCbCr(image.Rect(0, 0, 10, 10), image.YCbCrSubsampleRatio420)
 		c := cap(img.Y) + cap(img.Cb) + cap(img.Cr)
@@ -124,6 +136,12 @@ func TestRefRelease(t *testing.T) {
 		b := newImageRGBARef(img.Pix, img, p)
 		testRelease(tt, b)
 	})
+	t.Run("image.NRGBA", func(tt *testing.T) {
+		img := image.NewNRGBA(image.Rect(0, 0, 10, 10))
+		p := NewImageNRGBAPool(1, image.Rect(0, 0, 10, 10))
+		b := newImageNRGBARef(img.Pix, img, p)
+		testRelease(tt, b)
+	})
 	t.Run("image.YCbCrPool", func(tt *testing.T) {
 		img := image.NewYCbCr(image.Rect(0, 0, 10, 10), image.YCbCrSubsampleRatio420)
 		p := NewImageYCbCrPool(1, image.Rect(0, 0, 10, 10), image.YCbCrSubsampleRatio420)