Don't leak.

Change-Id: I756d54feeef4ff73f17c80d4b4405de56a5787a6

1 files changed, 61 insertions, 23 deletions

diff --git a/core/jni/android_view_Surface.cpp b/core/jni/android_view_Surface.cpp
index 9a85edc..206e320 100644
--- a/core/jni/android_view_Surface.cpp
+++ b/core/jni/android_view_Surface.cpp
@@ -32,6 +32,7 @@
 #include <SkCanvas.h>
 #include <SkBitmap.h>
 #include <SkRegion.h>
+#include <SkPixelRef.h>
 
 #include "jni.h"
 #include <android_runtime/AndroidRuntime.h>
@@ -437,9 +438,15 @@ static void Surface_unfreezeDisplay(
     }
 }
 
-class ScreenshotBitmap : public SkBitmap {
+class ScreenshotPixelRef : public SkPixelRef {
 public:
-    ScreenshotBitmap() {
+    ScreenshotPixelRef(SkColorTable* ctable) {
+        fCTable = ctable;
+        ctable->safeRef();
+        setImmutable();
+    }
+    virtual ~ScreenshotPixelRef() {
+        SkSafeUnref(fCTable);
     }
 
     status_t update(int width, int height) {
@@ -450,40 +457,71 @@ public:
             return res;
         }
 
-        void const* base = mScreenshot.getPixels();
-        uint32_t w = mScreenshot.getWidth();
-        uint32_t h = mScreenshot.getHeight();
-        uint32_t s = mScreenshot.getStride();
-        uint32_t f = mScreenshot.getFormat();
+        return NO_ERROR;
+    }
 
-        ssize_t bpr = s * android::bytesPerPixel(f);
-        setConfig(convertPixelFormat(f), w, h, bpr);
-        if (f == PIXEL_FORMAT_RGBX_8888) {
-            setIsOpaque(true);
-        }
-        if (w > 0 && h > 0) {
-            setPixels((void*)base);
-        } else {
-            // be safe with an empty bitmap.
-            setPixels(NULL);
-        }
+    uint32_t getWidth() const {
+        return mScreenshot.getWidth();
+    }
 
-        return NO_ERROR;
+    uint32_t getHeight() const {
+        return mScreenshot.getHeight();
+    }
+
+    uint32_t getStride() const {
+        return mScreenshot.getStride();
+    }
+
+    uint32_t getFormat() const {
+        return mScreenshot.getFormat();
+    }
+
+protected:
+    // overrides from SkPixelRef
+    virtual void* onLockPixels(SkColorTable** ct) {
+        *ct = fCTable;
+        return (void*)mScreenshot.getPixels();
+    }
+
+    virtual void onUnlockPixels() {
     }
 
 private:
     ScreenshotClient mScreenshot;
+    SkColorTable*    fCTable;
+
+    typedef SkPixelRef INHERITED;
 };
 
 static jobject Surface_screenshot(JNIEnv* env, jobject clazz, jint width, jint height)
 {
-    ScreenshotBitmap* bitmap = new ScreenshotBitmap();
-
-    if (bitmap->update(width, height) != NO_ERROR) {
-        delete bitmap;
+    ScreenshotPixelRef* pixels = new ScreenshotPixelRef(NULL);
+    if (pixels->update(width, height) != NO_ERROR) {
+        delete pixels;
         return 0;
     }
 
+    uint32_t w = pixels->getWidth();
+    uint32_t h = pixels->getHeight();
+    uint32_t s = pixels->getStride();
+    uint32_t f = pixels->getFormat();
+    ssize_t bpr = s * android::bytesPerPixel(f);
+
+    SkBitmap* bitmap = new SkBitmap();
+    bitmap->setConfig(convertPixelFormat(f), w, h, bpr);
+    if (f == PIXEL_FORMAT_RGBX_8888) {
+        bitmap->setIsOpaque(true);
+    }
+
+    if (w > 0 && h > 0) {
+        bitmap->setPixelRef(pixels)->unref();
+        bitmap->lockPixels();
+    } else {
+        // be safe with an empty bitmap.
+        delete pixels;
+        bitmap->setPixels(NULL);
+    }
+
     return GraphicsJNI::createBitmap(env, bitmap, false, NULL);
 }