Remove SkImage_Base::peekProxy.

This is the only place we provide a proxy from a SkImage other than
asView() and is only used in one place in the library. Add a utility
function to replace uses in tests based on asView().

Bug: skia:11208
Change-Id: Ica6eae6e5f153f208523331566acc0360925a20e
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/373617
Commit-Queue: Brian Salomon <bsalomon@google.com>
Reviewed-by: Robert Phillips <robertphillips@google.com>
diff --git a/gm/bc1_transparency.cpp b/gm/bc1_transparency.cpp
index 0808496..d52be82 100644
--- a/gm/bc1_transparency.cpp
+++ b/gm/bc1_transparency.cpp
@@ -15,6 +15,7 @@
 #include "src/gpu/GrImageContextPriv.h"
 #include "src/image/SkImage_Base.h"
 #include "src/image/SkImage_GpuBase.h"
+#include "tools/gpu/ProxyUtils.h"
 
 constexpr int kImgWidth  = 16;
 constexpr int kImgHeight = 8;
@@ -121,8 +122,8 @@
     bool isCompressed = false;
     if (image && image->isTextureBacked()) {
         const GrCaps* caps = as_IB(image)->context()->priv().caps();
-
-        GrTextureProxy* proxy = as_IB(image)->peekProxy();
+        GrTextureProxy* proxy = sk_gpu_test::GetTextureImageProxy(image.get(),
+                                                                  canvas->recordingContext());
         isCompressed = caps->isFormatCompressed(proxy->backendFormat());
     }
 
diff --git a/gm/compressed_textures.cpp b/gm/compressed_textures.cpp
index b0c5c36..63c7606 100644
--- a/gm/compressed_textures.cpp
+++ b/gm/compressed_textures.cpp
@@ -31,6 +31,7 @@
 #include "src/image/SkImage_Base.h"
 #include "src/image/SkImage_GpuBase.h"
 #include "third_party/etc1/etc1.h"
+#include "tools/gpu/ProxyUtils.h"
 
 class GrSurfaceDrawContext;
 
@@ -275,8 +276,8 @@
         bool isCompressed = false;
         if (image->isTextureBacked()) {
             const GrCaps* caps = as_IB(image)->context()->priv().caps();
-
-            GrTextureProxy* proxy = as_IB(image)->peekProxy();
+            GrTextureProxy* proxy = sk_gpu_test::GetTextureImageProxy(image,
+                                                                      canvas->recordingContext());
             isCompressed = caps->isFormatCompressed(proxy->backendFormat());
         }
 
diff --git a/gm/exoticformats.cpp b/gm/exoticformats.cpp
index ed1a392..3bd0fad 100644
--- a/gm/exoticformats.cpp
+++ b/gm/exoticformats.cpp
@@ -18,6 +18,7 @@
 #include "src/gpu/gl/GrGLDefines.h"
 #include "src/image/SkImage_Base.h"
 #include "src/image/SkImage_GpuBase.h"
+#include "tools/gpu/ProxyUtils.h"
 
 #include "tools/Resources.h"
 
@@ -382,8 +383,8 @@
         bool isCompressed = false;
         if (image->isTextureBacked()) {
             const GrCaps* caps = as_IB(image)->context()->priv().caps();
-
-            GrTextureProxy* proxy = as_IB(image)->peekProxy();
+            GrTextureProxy* proxy = sk_gpu_test::GetTextureImageProxy(image,
+                                                                      canvas->recordingContext());
             isCompressed = caps->isFormatCompressed(proxy->backendFormat());
         }
 
diff --git a/src/image/SkImage_Base.h b/src/image/SkImage_Base.h
index 0f58f9f..c2093e4 100644
--- a/src/image/SkImage_Base.h
+++ b/src/image/SkImage_Base.h
@@ -88,11 +88,6 @@
         return GrSemaphoresSubmitted::kNo;
     }
 
-    // Return the proxy if this image is backed by a single proxy. For YUVA images, this
-    // will return nullptr unless the YUVA planes have been converted to RGBA in which case
-    // that single backing proxy will be returned.
-    virtual GrTextureProxy* peekProxy() const { return nullptr; }
-
     // Returns a GrSurfaceProxyView representation of the image, if possible. This also returns
     // a color type. This may be different than the image's color type when the image is not
     // texture-backed and the capabilities of the GPU require a data type conversion to put
diff --git a/src/image/SkImage_Gpu.cpp b/src/image/SkImage_Gpu.cpp
index a99f927..09d9f03 100644
--- a/src/image/SkImage_Gpu.cpp
+++ b/src/image/SkImage_Gpu.cpp
@@ -63,6 +63,10 @@
 
 SkImage_Gpu::~SkImage_Gpu() {}
 
+bool SkImage_Gpu::surfaceMustCopyOnWrite(GrSurfaceProxy* surfaceProxy) const {
+    return surfaceProxy->underlyingUniqueID() == fView.proxy()->underlyingUniqueID();
+}
+
 GrSemaphoresSubmitted SkImage_Gpu::onFlush(GrDirectContext* dContext, const GrFlushInfo& info) {
     if (!fContext->priv().matches(dContext) || dContext->abandoned()) {
         if (info.fSubmittedProc) {
diff --git a/src/image/SkImage_Gpu.h b/src/image/SkImage_Gpu.h
index 8896c59..1421dcc 100644
--- a/src/image/SkImage_Gpu.h
+++ b/src/image/SkImage_Gpu.h
@@ -37,14 +37,17 @@
 
     ~SkImage_Gpu() override;
 
+    // If this is image is a cached SkSurface snapshot then this method is called by the SkSurface
+    // before a write to check if the surface must make a copy to avoid modifying the image's
+    // contents.
+    bool surfaceMustCopyOnWrite(GrSurfaceProxy* surfaceProxy) const;
+
     bool onHasMipmaps() const override {
         return fView.asTextureProxy()->mipmapped() == GrMipmapped::kYes;
     }
 
     GrSemaphoresSubmitted onFlush(GrDirectContext*, const GrFlushInfo&) override;
 
-    GrTextureProxy* peekProxy() const override { return fView.asTextureProxy(); }
-
     GrBackendTexture onGetBackendTexture(bool flushPendingGrContextIO,
                                          GrSurfaceOrigin* origin) const final;
 
diff --git a/src/image/SkImage_GpuYUVA.cpp b/src/image/SkImage_GpuYUVA.cpp
index e561283..38867f2 100644
--- a/src/image/SkImage_GpuYUVA.cpp
+++ b/src/image/SkImage_GpuYUVA.cpp
@@ -141,8 +141,6 @@
     return fYUVAProxies.mipmapped() == GrMipmapped::kYes;
 }
 
-GrTextureProxy* SkImage_GpuYUVA::peekProxy() const { return fRGBView.asTextureProxy(); }
-
 size_t SkImage_GpuYUVA::onTextureSize() const {
     if (fRGBView) {
         return fRGBView.asTextureProxy()->gpuMemorySize();
diff --git a/src/image/SkImage_GpuYUVA.h b/src/image/SkImage_GpuYUVA.h
index 261f2f9..d0741c7 100644
--- a/src/image/SkImage_GpuYUVA.h
+++ b/src/image/SkImage_GpuYUVA.h
@@ -34,10 +34,6 @@
 
     GrSemaphoresSubmitted onFlush(GrDirectContext*, const GrFlushInfo&) override;
 
-    // This returns the single backing proxy if the YUV channels have already been flattened but
-    // nullptr if they have not.
-    GrTextureProxy* peekProxy() const override;
-
     bool onIsTextureBacked() const override {
         // We should have YUVA proxies or a RGBA proxy,but not both.
         SkASSERT(fYUVAProxies.isValid() != SkToBool(fRGBView));
diff --git a/src/image/SkSurface_Gpu.cpp b/src/image/SkSurface_Gpu.cpp
index ccc7ab7..a1d20e7 100644
--- a/src/image/SkSurface_Gpu.cpp
+++ b/src/image/SkSurface_Gpu.cpp
@@ -188,13 +188,10 @@
 
     // are we sharing our backing proxy with the image? Note this call should never create a new
     // image because onCopyOnWrite is only called when there is a cached image.
-    sk_sp<SkImage> image(this->refCachedImage());
+    sk_sp<SkImage> image = this->refCachedImage();
     SkASSERT(image);
 
-    GrSurfaceProxy* imageProxy = ((SkImage_Base*) image.get())->peekProxy();
-    SkASSERT(imageProxy);
-
-    if (sdc->asSurfaceProxy()->underlyingUniqueID() == imageProxy->underlyingUniqueID()) {
+    if (static_cast<SkImage_Gpu*>(image.get())->surfaceMustCopyOnWrite(sdc->asSurfaceProxy())) {
         fDevice->replaceSurfaceDrawContext(mode);
     } else if (kDiscard_ContentChangeMode == mode) {
         this->SkSurface_Gpu::onDiscard();
diff --git a/tests/BackendAllocationTest.cpp b/tests/BackendAllocationTest.cpp
index f34d124..577496d 100644
--- a/tests/BackendAllocationTest.cpp
+++ b/tests/BackendAllocationTest.cpp
@@ -20,6 +20,7 @@
 #include "tests/TestUtils.h"
 #include "tools/ToolUtils.h"
 #include "tools/gpu/ManagedBackendTexture.h"
+#include "tools/gpu/ProxyUtils.h"
 
 #ifdef SK_GL
 #include "src/gpu/gl/GrGLCaps.h"
@@ -92,9 +93,7 @@
             ERRORF(reporter, "Couldn't make SkImage from backendTexture for %s\n",
                    ToolUtils::colortype_name(skColorType));
         } else {
-            SkImage_Base* ib = as_IB(img);
-
-            GrTextureProxy* proxy = ib->peekProxy();
+            GrTextureProxy* proxy = sk_gpu_test::GetTextureImageProxy(img.get(), dContext);
             REPORTER_ASSERT(reporter, proxy);
 
             REPORTER_ASSERT(reporter, mipMapped == proxy->proxyMipmapped());
diff --git a/tests/BackendSurfaceMutableStateTest.cpp b/tests/BackendSurfaceMutableStateTest.cpp
index ef2de28..0eed3c3 100644
--- a/tests/BackendSurfaceMutableStateTest.cpp
+++ b/tests/BackendSurfaceMutableStateTest.cpp
@@ -14,6 +14,7 @@
 #include "src/gpu/GrTextureProxy.h"
 #include "src/image/SkImage_Base.h"
 #include "tests/Test.h"
+#include "tools/gpu/ProxyUtils.h"
 
 #ifdef SK_VULKAN
 #include "src/gpu/vk/GrVkGpu.h"
@@ -61,7 +62,8 @@
                                                            kRGBA_8888_SkColorType,
                                                            kPremul_SkAlphaType, nullptr);
 
-    GrSurfaceProxy* proxy = as_IB(wrappedImage)->peekProxy();
+    GrSurfaceProxy* proxy = sk_gpu_test::GetTextureImageProxy(wrappedImage.get(), dContext);
+
     REPORTER_ASSERT(reporter, proxy);
     REPORTER_ASSERT(reporter, proxy->isInstantiated());
     GrTexture* texture = proxy->peekTexture();
diff --git a/tests/DeferredDisplayListTest.cpp b/tests/DeferredDisplayListTest.cpp
index 0928dd0..a12e1a7 100644
--- a/tests/DeferredDisplayListTest.cpp
+++ b/tests/DeferredDisplayListTest.cpp
@@ -45,6 +45,7 @@
 #include "tools/gpu/BackendSurfaceFactory.h"
 #include "tools/gpu/GrContextFactory.h"
 #include "tools/gpu/ManagedBackendTexture.h"
+#include "tools/gpu/ProxyUtils.h"
 
 #include <initializer_list>
 #include <memory>
@@ -1219,7 +1220,7 @@
             }
             REPORTER_ASSERT(reporter, image);
 
-            GrTextureProxy* backingProxy = ((SkImage_GpuBase*) image.get())->peekProxy();
+            GrTextureProxy* backingProxy = sk_gpu_test::GetTextureImageProxy(image.get(), context);
 
             REPORTER_ASSERT(reporter, backingProxy->mipmapped() == mipMapped);
             if (GR_GL_TEXTURE_2D == target) {
diff --git a/tests/GLBackendSurfaceTest.cpp b/tests/GLBackendSurfaceTest.cpp
index b028f36..5387e8e 100644
--- a/tests/GLBackendSurfaceTest.cpp
+++ b/tests/GLBackendSurfaceTest.cpp
@@ -23,6 +23,7 @@
 #include "src/gpu/gl/GrGLCaps.h"
 #include "src/gpu/gl/GrGLTexture.h"
 #include "src/image/SkImage_Base.h"
+#include "tools/gpu/ProxyUtils.h"
 
 static bool sampler_params_invalid(const GrGLTextureParameters& parameters) {
     return SkScalarIsNaN(parameters.samplerOverriddenState().fMaxLOD);
@@ -69,7 +70,7 @@
                                      kRGBA_8888_SkColorType, kPremul_SkAlphaType, nullptr);
     REPORTER_ASSERT(reporter, wrappedImage);
 
-    GrSurfaceProxy* proxy = as_IB(wrappedImage)->peekProxy();
+    GrSurfaceProxy* proxy = sk_gpu_test::GetTextureImageProxy(wrappedImage.get(), dContext);
     REPORTER_ASSERT(reporter, proxy);
     REPORTER_ASSERT(reporter, proxy->isInstantiated());
     auto texture = static_cast<GrGLTexture*>(proxy->peekTexture());
diff --git a/tests/GrMipMappedTest.cpp b/tests/GrMipMappedTest.cpp
index 08eebaf..6a81d68 100644
--- a/tests/GrMipMappedTest.cpp
+++ b/tests/GrMipMappedTest.cpp
@@ -31,6 +31,7 @@
 #include "tools/gpu/BackendSurfaceFactory.h"
 #include "tools/gpu/BackendTextureImageFactory.h"
 #include "tools/gpu/ManagedBackendTexture.h"
+#include "tools/gpu/ProxyUtils.h"
 
 static constexpr int kSize = 8;
 
@@ -86,12 +87,7 @@
                                                  sk_gpu_test::ManagedBackendTexture::ReleaseProc,
                                                  mbet->releaseContext());
                 REPORTER_ASSERT(reporter, (mipmapped == GrMipmapped::kYes) == image->hasMipmaps());
-                GrSurfaceProxy* surfaceProxy = as_IB(image)->peekProxy();
-                REPORTER_ASSERT(reporter, surfaceProxy);
-                if (!proxy) {
-                    continue;
-                }
-                proxy = sk_ref_sp(surfaceProxy->asTextureProxy());
+                proxy = sk_ref_sp(sk_gpu_test::GetTextureImageProxy(image.get(), dContext));
             }
             REPORTER_ASSERT(reporter, proxy);
             if (!proxy) {
@@ -136,7 +132,7 @@
                     dContext, ii, SkColors::kTransparent, betMipmapped);
             REPORTER_ASSERT(reporter, (betMipmapped == GrMipmapped::kYes) == image->hasMipmaps());
 
-            GrTextureProxy* proxy = as_IB(image)->peekProxy();
+            GrTextureProxy* proxy = sk_gpu_test::GetTextureImageProxy(image.get(), dContext);
             REPORTER_ASSERT(reporter, proxy);
             if (!proxy) {
                 return;
@@ -304,7 +300,7 @@
             sk_sp<SkImage> image = surface->makeImageSnapshot();
             REPORTER_ASSERT(reporter, willUseMips == image->hasMipmaps());
             REPORTER_ASSERT(reporter, image);
-            texProxy = as_IB(image)->peekProxy();
+            texProxy = sk_gpu_test::GetTextureImageProxy(image.get(), dContext);
             REPORTER_ASSERT(reporter, mipmapped == texProxy->mipmapped());
 
             texProxy->instantiate(resourceProvider);
diff --git a/tests/GrTextureMipMapInvalidationTest.cpp b/tests/GrTextureMipMapInvalidationTest.cpp
index c751074..2f7d1f3 100644
--- a/tests/GrTextureMipMapInvalidationTest.cpp
+++ b/tests/GrTextureMipMapInvalidationTest.cpp
@@ -13,6 +13,7 @@
 #include "src/image/SkImage_Base.h"
 #include "src/image/SkImage_GpuBase.h"
 #include "tests/Test.h"
+#include "tools/gpu/ProxyUtils.h"
 
 // Tests that MIP maps are created and invalidated as expected when drawing to and from GrTextures.
 DEF_GPUTEST_FOR_RENDERING_CONTEXTS(GrTextureMipMapInvalidationTest, reporter, ctxInfo) {
@@ -22,15 +23,19 @@
     }
 
     auto isMipped = [reporter](SkSurface* surf) {
-        SkImage_GpuBase* image = static_cast<SkImage_GpuBase*>(as_IB(surf->makeImageSnapshot()));
-        bool proxyIsMipmapped = image->peekProxy()->mipmapped() == GrMipmapped::kYes;
+        sk_sp<SkImage> image = surf->makeImageSnapshot();
+        GrTextureProxy* proxy = sk_gpu_test::GetTextureImageProxy(image.get(),
+                                                                  surf->recordingContext());
+        bool proxyIsMipmapped = proxy->mipmapped() == GrMipmapped::kYes;
         REPORTER_ASSERT(reporter, proxyIsMipmapped == image->hasMipmaps());
         return image->hasMipmaps();
     };
 
     auto mipsAreDirty = [](SkSurface* surf) {
-        SkImage_GpuBase* image = static_cast<SkImage_GpuBase*>(as_IB(surf->makeImageSnapshot()));
-        return image->peekProxy()->peekTexture()->mipmapsAreDirty();
+        sk_sp<SkImage> image = surf->makeImageSnapshot();
+        GrTextureProxy* proxy = sk_gpu_test::GetTextureImageProxy(image.get(),
+                                                                  surf->recordingContext());
+        return proxy->peekTexture()->mipmapsAreDirty();
     };
 
     auto info = SkImageInfo::MakeN32Premul(256, 256);
diff --git a/tests/ImageTest.cpp b/tests/ImageTest.cpp
index dac50a7..6133f4d 100644
--- a/tests/ImageTest.cpp
+++ b/tests/ImageTest.cpp
@@ -39,6 +39,7 @@
 #include "tools/Resources.h"
 #include "tools/ToolUtils.h"
 #include "tools/gpu/ManagedBackendTexture.h"
+#include "tools/gpu/ProxyUtils.h"
 
 using namespace sk_gpu_test;
 
@@ -404,9 +405,7 @@
             GrTextureProxy* origProxy = nullptr;
             bool origIsMippedTexture = false;
 
-            if (auto sp = as_IB(image)->peekProxy()) {
-                origProxy = sp->asTextureProxy();
-                SkASSERT(origProxy);
+            if ((origProxy = sk_gpu_test::GetTextureImageProxy(image.get(), dContext))) {
                 REPORTER_ASSERT(reporter, (origProxy->mipmapped() == GrMipmapped::kYes) ==
                                           image->hasMipmaps());
                 origIsMippedTexture = image->hasMipmaps();
@@ -425,7 +424,9 @@
                     ERRORF(reporter, "makeTextureImage returned non-texture image.");
                     continue;
                 }
-                GrTextureProxy* copyProxy = as_IB(texImage)->peekProxy()->asTextureProxy();
+
+                GrTextureProxy* copyProxy = sk_gpu_test::GetTextureImageProxy(texImage.get(),
+                                                                              dContext);
                 SkASSERT(copyProxy);
                 // Did we ask for MIPs on a context that supports them?
                 bool validRequestForMips = (mipmapped == GrMipmapped::kYes &&
@@ -443,7 +444,8 @@
                         ERRORF(reporter, "makeTextureImage made unnecessary texture copy.");
                     }
                 } else {
-                    auto* texProxy = as_IB(texImage)->peekProxy()->asTextureProxy();
+                    GrTextureProxy* texProxy = sk_gpu_test::GetTextureImageProxy(texImage.get(),
+                                                                                 dContext);
                     REPORTER_ASSERT(reporter, !texProxy->getUniqueKey().isValid());
                     REPORTER_ASSERT(reporter, texProxy->isBudgeted() == budgeted);
                 }
diff --git a/tests/SurfaceTest.cpp b/tests/SurfaceTest.cpp
index 7756bd0..cbe8370 100644
--- a/tests/SurfaceTest.cpp
+++ b/tests/SurfaceTest.cpp
@@ -33,6 +33,7 @@
 #include "tools/ToolUtils.h"
 #include "tools/gpu/BackendSurfaceFactory.h"
 #include "tools/gpu/ManagedBackendTexture.h"
+#include "tools/gpu/ProxyUtils.h"
 
 #include <functional>
 #include <initializer_list>
@@ -481,18 +482,19 @@
     canvas2->clear(5);
     sk_sp<SkImage> image4(surface2->makeImageSnapshot());
 
-    SkImage_GpuBase* gpuImage1 = static_cast<SkImage_GpuBase*>(as_IB(image1));
-    SkImage_GpuBase* gpuImage2 = static_cast<SkImage_GpuBase*>(as_IB(image2));
-    SkImage_GpuBase* gpuImage3 = static_cast<SkImage_GpuBase*>(as_IB(image3));
-    SkImage_GpuBase* gpuImage4 = static_cast<SkImage_GpuBase*>(as_IB(image4));
+    auto imageProxy = [ctx = surface1->recordingContext()](SkImage* img) {
+        GrTextureProxy* proxy = sk_gpu_test::GetTextureImageProxy(img, ctx);
+        SkASSERT(proxy);
+        return proxy;
+    };
 
-    REPORTER_ASSERT(reporter, gpuImage4->peekProxy() != gpuImage3->peekProxy());
+    REPORTER_ASSERT(reporter, imageProxy(image4.get()) != imageProxy(image3.get()));
     // The following assertion checks crbug.com/263329
-    REPORTER_ASSERT(reporter, gpuImage4->peekProxy() != gpuImage2->peekProxy());
-    REPORTER_ASSERT(reporter, gpuImage4->peekProxy() != gpuImage1->peekProxy());
-    REPORTER_ASSERT(reporter, gpuImage3->peekProxy() != gpuImage2->peekProxy());
-    REPORTER_ASSERT(reporter, gpuImage3->peekProxy() != gpuImage1->peekProxy());
-    REPORTER_ASSERT(reporter, gpuImage2->peekProxy() != gpuImage1->peekProxy());
+    REPORTER_ASSERT(reporter, imageProxy(image4.get()) != imageProxy(image2.get()));
+    REPORTER_ASSERT(reporter, imageProxy(image4.get()) != imageProxy(image1.get()));
+    REPORTER_ASSERT(reporter, imageProxy(image3.get()) != imageProxy(image2.get()));
+    REPORTER_ASSERT(reporter, imageProxy(image3.get()) != imageProxy(image1.get()));
+    REPORTER_ASSERT(reporter, imageProxy(image2.get()) != imageProxy(image1.get()));
 }
 DEF_GPUTEST_FOR_RENDERING_CONTEXTS(SurfaceCRBug263329_Gpu, reporter, ctxInfo) {
     for (auto& surface_func : { &create_gpu_surface, &create_gpu_scratch_surface }) {
@@ -531,18 +533,15 @@
     return proxy->isBudgeted();
 }
 
-static SkBudgeted is_budgeted(SkImage* image) {
-    return ((SkImage_Gpu*)image)->peekProxy()->isBudgeted();
-}
-
-static SkBudgeted is_budgeted(const sk_sp<SkImage> image) {
-    return is_budgeted(image.get());
+static SkBudgeted is_budgeted(SkImage* image, GrRecordingContext* rc) {
+    return sk_gpu_test::GetTextureImageProxy(image, rc)->isBudgeted();
 }
 
 DEF_GPUTEST_FOR_RENDERING_CONTEXTS(SurfaceBudget, reporter, ctxInfo) {
     SkImageInfo info = SkImageInfo::MakeN32Premul(8,8);
+    GrDirectContext* dContext = ctxInfo.directContext();
     for (auto budgeted : { SkBudgeted::kNo, SkBudgeted::kYes }) {
-        auto surface(SkSurface::MakeRenderTarget(ctxInfo.directContext(), budgeted, info));
+        auto surface(SkSurface::MakeRenderTarget(dContext, budgeted, info));
         SkASSERT(surface);
         REPORTER_ASSERT(reporter, budgeted == is_budgeted(surface));
 
@@ -551,14 +550,14 @@
         // Initially the image shares a texture with the surface, and the
         // the budgets should always match.
         REPORTER_ASSERT(reporter, budgeted == is_budgeted(surface));
-        REPORTER_ASSERT(reporter, budgeted == is_budgeted(image));
+        REPORTER_ASSERT(reporter, budgeted == is_budgeted(image.get(), dContext));
 
         // Now trigger copy-on-write
         surface->getCanvas()->clear(SK_ColorBLUE);
 
         // They don't share a texture anymore but the budgets should still match.
         REPORTER_ASSERT(reporter, budgeted == is_budgeted(surface));
-        REPORTER_ASSERT(reporter, budgeted == is_budgeted(image));
+        REPORTER_ASSERT(reporter, budgeted == is_budgeted(image.get(), dContext));
     }
 }
 
diff --git a/tests/VkBackendSurfaceTest.cpp b/tests/VkBackendSurfaceTest.cpp
index 11fec2a..8a6e08b 100644
--- a/tests/VkBackendSurfaceTest.cpp
+++ b/tests/VkBackendSurfaceTest.cpp
@@ -28,6 +28,7 @@
 #include "src/image/SkSurface_Gpu.h"
 #include "tests/Test.h"
 #include "tools/gpu/ManagedBackendTexture.h"
+#include "tools/gpu/ProxyUtils.h"
 
 DEF_GPUTEST_FOR_VULKAN_CONTEXT(VkImageLayoutTest, reporter, ctxInfo) {
     auto dContext = ctxInfo.directContext();
@@ -72,7 +73,7 @@
             mbet->releaseContext());
     REPORTER_ASSERT(reporter, wrappedImage.get());
 
-    GrSurfaceProxy* proxy = as_IB(wrappedImage)->peekProxy();
+    GrSurfaceProxy* proxy = sk_gpu_test::GetTextureImageProxy(wrappedImage.get(), dContext);
     REPORTER_ASSERT(reporter, proxy);
     REPORTER_ASSERT(reporter, proxy->isInstantiated());
     GrTexture* texture = proxy->peekTexture();
diff --git a/tools/DDLPromiseImageHelper.cpp b/tools/DDLPromiseImageHelper.cpp
index 762ca2c..66af4cb 100644
--- a/tools/DDLPromiseImageHelper.cpp
+++ b/tools/DDLPromiseImageHelper.cpp
@@ -345,15 +345,6 @@
             curImage.callbackContext(i)->wasAddedToImage();
         }
 
-#ifdef SK_DEBUG
-        {
-            // By the peekProxy contract this image should not have a single backing proxy so
-            // should return null. The call should also not trigger the conversion to RGBA.
-            SkImage_GpuYUVA* yuva = reinterpret_cast<SkImage_GpuYUVA*>(image.get());
-            SkASSERT(!yuva->peekProxy());
-            SkASSERT(!yuva->peekProxy());  // the first call didn't force a conversion to RGBA
-        }
-#endif
     } else {
         const GrBackendFormat& backendFormat = curImage.backendFormat(0);
         SkASSERT(backendFormat.isValid());
diff --git a/tools/gpu/ProxyUtils.cpp b/tools/gpu/ProxyUtils.cpp
index 7c65774..bf1d7b3 100644
--- a/tools/gpu/ProxyUtils.cpp
+++ b/tools/gpu/ProxyUtils.cpp
@@ -10,18 +10,45 @@
 #include "include/core/SkColor.h"
 #include "include/gpu/GrBackendSurface.h"
 #include "include/gpu/GrDirectContext.h"
+#include "include/private/GrImageContext.h"
 #include "src/gpu/GrDirectContextPriv.h"
 #include "src/gpu/GrDrawingManager.h"
 #include "src/gpu/GrGpu.h"
+#include "src/gpu/GrImageContextPriv.h"
 #include "src/gpu/GrPixmap.h"
 #include "src/gpu/GrProgramInfo.h"
 #include "src/gpu/GrProxyProvider.h"
 #include "src/gpu/GrSurfaceContext.h"
 #include "src/gpu/SkGr.h"
 #include "src/gpu/ops/GrSimpleMeshDrawOpHelper.h"
+#include "src/image/SkImage_Base.h"
 
 namespace sk_gpu_test {
 
+GrTextureProxy* GetTextureImageProxy(SkImage* image, GrRecordingContext* rContext) {
+    if (!image->isTextureBacked() || as_IB(image)->isYUVA()) {
+        return nullptr;
+    }
+    if (!rContext) {
+        // If the image is backed by a recording context we'll use that.
+        GrImageContext* iContext = as_IB(image)->context();
+        SkASSERT(iContext);
+        rContext = iContext->priv().asRecordingContext();
+        if (!rContext) {
+            return nullptr;
+        }
+    }
+    auto [view, ct] = as_IB(image)->asView(rContext, GrMipmapped::kNo);
+    if (!view) {
+        // With the above checks we expect this to succeed unless there is a context mismatch.
+        SkASSERT(!image->isValid(rContext));
+        return nullptr;
+    }
+    GrSurfaceProxy* proxy = view.proxy();
+    SkASSERT(proxy->asTextureProxy());
+    return proxy->asTextureProxy();
+}
+
 GrSurfaceProxyView MakeTextureProxyViewFromData(GrDirectContext* dContext,
                                                 GrRenderable renderable,
                                                 GrSurfaceOrigin origin,
diff --git a/tools/gpu/ProxyUtils.h b/tools/gpu/ProxyUtils.h
index 4533394..7ce0122 100644
--- a/tools/gpu/ProxyUtils.h
+++ b/tools/gpu/ProxyUtils.h
@@ -19,6 +19,9 @@
 
 namespace sk_gpu_test {
 
+/** Returns the proxy backing an image if it is texture backed, otherwise nullptr. */
+GrTextureProxy* GetTextureImageProxy(SkImage*, GrRecordingContext*);
+
 /** Makes a texture proxy containing the passed in color data. */
 GrSurfaceProxyView MakeTextureProxyViewFromData(GrDirectContext*,
                                                 GrRenderable,