Introduce SubresourceStorage (2/N): Merge

This CL adds the Merge() operation to SubresourceStorage() that allows
modifying the content of a storage with another storage.

Bug: dawn:441

Change-Id: I28e3cd7bc967056eda2c387b2b6e164eb370a241
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/35520
Reviewed-by: Austin Eng <enga@chromium.org>
Reviewed-by: Ben Clayton <bclayton@google.com>
Commit-Queue: Corentin Wallez <cwallez@chromium.org>
diff --git a/src/dawn_native/SubresourceStorage.h b/src/dawn_native/SubresourceStorage.h
index 38f7400..a10aca7 100644
--- a/src/dawn_native/SubresourceStorage.h
+++ b/src/dawn_native/SubresourceStorage.h
@@ -81,9 +81,23 @@
     //
     // T must be a copyable type that supports equality comparison with ==.
     //
-    // TODO(cwallez@chromium.org): Add the Merge() operation.
+    // The implementation of functions in this file can have a lot of control flow and corner cases
+    // so each modification should come with extensive tests and ensure 100% code coverage of the
+    // modified functions. See instructions at
+    // https://chromium.googlesource.com/chromium/src/+/master/docs/testing/code_coverage.md#local-coverage-script
+    // to run the test with code coverage. A command line that worked in the past (with the right
+    // GN args for the out/coverage directory in a Chromium checkout) is:
+    //
+    /*
+       python tools/code_coverage/coverage.py dawn_unittests -b out/coverage -o out/report -c \
+           "out/coverage/dawn_unittests --gtest_filter=SubresourceStorage\*" -f \
+           third_party/dawn/src/dawn_native
+    */
+    //
     // TODO(cwallez@chromium.org): Inline the storage for aspects to avoid allocating when
     // possible.
+    // TODO(cwallez@chromium.org): Make the recompression optional, the calling code should know
+    // if recompression can happen or not in Update() and Merge()
     template <typename T>
     class SubresourceStorage {
       public:
@@ -98,7 +112,7 @@
         // same for multiple subresources.
         const T& Get(Aspect aspect, uint32_t arrayLayer, uint32_t mipLevel) const;
 
-        // Given an iterateFunc that's a function or function-like objet that call be called with
+        // Given an iterateFunc that's a function or function-like objet that can be called with
         // arguments of type (const SubresourceRange& range, const T& data) and returns void,
         // calls it with aggregate ranges if possible, such that each subresource is part of
         // exactly one of the ranges iterateFunc is called with (and obviously data is the value
@@ -110,7 +124,7 @@
         template <typename F>
         void Iterate(F&& iterateFunc) const;
 
-        // Given an updateFunc that's a function or function-like objet that call be called with
+        // Given an updateFunc that's a function or function-like objet that can be called with
         // arguments of type (const SubresourceRange& range, T* data) and returns void,
         // calls it with ranges that in aggregate form `range` and pass for each of the
         // sub-ranges a pointer to modify the value for that sub-range. For example:
@@ -125,10 +139,25 @@
         template <typename F>
         void Update(const SubresourceRange& range, F&& updateFunc);
 
+        // Given a mergeFunc that's a function or a function-like object that can be called with
+        // arguments of type (const SubresourceRange& range, T* data, const U& otherData) and
+        // returns void, calls it with ranges that in aggregate form the full resources and pass
+        // for each of the sub-ranges a pointer to modify the value for that sub-range and the
+        // corresponding value from other for that sub-range. For example:
+        //
+        //   subresources.Merge(otherUsages,
+        //       [](const SubresourceRange&, T* data, const T& otherData) {
+        //          *data |= otherData;
+        //       });
+        //
+        // /!\ WARNING: mergeFunc should never use range to compute the update to data otherwise
+        // your code is likely to break when compression happens. Range should only be used for
+        // side effects like using it to compute a Vulkan pipeline barrier.
+        template <typename U, typename F>
+        void Merge(const SubresourceStorage<U>& other, F&& mergeFunc);
+
         // Other operations to consider:
         //
-        //  - Merge(Range, SubresourceStorage<U>, mergeFunc) that takes the values from the other
-        //    storage and modifies the value of the current storage with it.
         //  - UpdateTo(Range, T) that updates the range to a constant value.
 
         // Methods to query the internal state of SubresourceStorage for testing.
@@ -139,6 +168,9 @@
         bool IsLayerCompressedForTesting(Aspect aspect, uint32_t layer) const;
 
       private:
+        template <typename U>
+        friend class SubresourceStorage;
+
         void DecompressAspect(uint32_t aspectIndex);
         void RecompressAspect(uint32_t aspectIndex);
 
@@ -255,6 +287,63 @@
     }
 
     template <typename T>
+    template <typename U, typename F>
+    void SubresourceStorage<T>::Merge(const SubresourceStorage<U>& other, F&& mergeFunc) {
+        ASSERT(mAspects == other.mAspects);
+        ASSERT(mArrayLayerCount == other.mArrayLayerCount);
+        ASSERT(mMipLevelCount == other.mMipLevelCount);
+
+        for (Aspect aspect : IterateEnumMask(mAspects)) {
+            uint32_t aspectIndex = GetAspectIndex(aspect);
+
+            // If the other storage's aspect is compressed we don't need to decompress anything
+            // in `this` and can just iterate through it, merging with `other`'s constant value for
+            // the aspect. For code simplicity this can be done with a call to Update().
+            if (other.mAspectCompressed[aspectIndex]) {
+                const U& otherData = other.Data(aspectIndex);
+                Update(SubresourceRange::MakeFull(aspect, mArrayLayerCount, mMipLevelCount),
+                       [&](const SubresourceRange& subrange, T* data) {
+                           mergeFunc(subrange, data, otherData);
+                       });
+                continue;
+            }
+
+            // Other doesn't have the aspect compressed so we must do at least per-layer merging.
+            if (mAspectCompressed[aspectIndex]) {
+                DecompressAspect(aspectIndex);
+            }
+
+            for (uint32_t layer = 0; layer < mArrayLayerCount; layer++) {
+                // Similarly to above, use a fast path if other's layer is compressed.
+                if (other.LayerCompressed(aspectIndex, layer)) {
+                    const U& otherData = other.Data(aspectIndex, layer);
+                    Update(GetFullLayerRange(aspect, layer),
+                           [&](const SubresourceRange& subrange, T* data) {
+                               mergeFunc(subrange, data, otherData);
+                           });
+                    continue;
+                }
+
+                // Sad case, other is decompressed for this layer, do per-level merging.
+                if (LayerCompressed(aspectIndex, layer)) {
+                    DecompressLayer(aspectIndex, layer);
+                }
+
+                for (uint32_t level = 0; level < mMipLevelCount; level++) {
+                    SubresourceRange updateRange =
+                        SubresourceRange::MakeSingle(aspect, layer, level);
+                    mergeFunc(updateRange, &Data(aspectIndex, layer, level),
+                              other.Data(aspectIndex, layer, level));
+                }
+
+                RecompressLayer(aspectIndex, layer);
+            }
+
+            RecompressAspect(aspectIndex);
+        }
+    }
+
+    template <typename T>
     template <typename F>
     void SubresourceStorage<T>::Iterate(F&& iterateFunc) const {
         for (Aspect aspect : IterateEnumMask(mAspects)) {
diff --git a/src/tests/unittests/SubresourceStorageTests.cpp b/src/tests/unittests/SubresourceStorageTests.cpp
index 2f2ef0b..c128bed 100644
--- a/src/tests/unittests/SubresourceStorageTests.cpp
+++ b/src/tests/unittests/SubresourceStorageTests.cpp
@@ -49,6 +49,19 @@
         }
     }
 
+    template <typename U, typename F>
+    void Merge(const SubresourceStorage<U>& other, F&& mergeFunc) {
+        for (Aspect aspect : IterateEnumMask(mAspects)) {
+            for (uint32_t layer = 0; layer < mArrayLayerCount; layer++) {
+                for (uint32_t level = 0; level < mMipLevelCount; level++) {
+                    SubresourceRange range = SubresourceRange::MakeSingle(aspect, layer, level);
+                    mergeFunc(range, &mData[GetDataIndex(aspect, layer, level)],
+                              other.Get(aspect, layer, level));
+                }
+            }
+        }
+    }
+
     const T& Get(Aspect aspect, uint32_t arrayLayer, uint32_t mipLevel) const {
         return mData[GetDataIndex(aspect, arrayLayer, mipLevel)];
     }
@@ -352,15 +365,15 @@
     CheckLayerCompressed(s, Aspect::Stencil, 3, true);
 
     {
-        SubresourceRange range(Aspect::Depth | Aspect::Stencil, {0, kLayers}, {5, 2});
+        SubresourceRange range(Aspect::Depth, {0, kLayers}, {5, 2});
         CallUpdateOnBoth(&s, &f, range, [](const SubresourceRange&, int* data) { *data *= 3; });
     }
 
-    // The layers had to be decompressed.
+    // The layers had to be decompressed in depth
     CheckLayerCompressed(s, Aspect::Depth, 2, false);
     CheckLayerCompressed(s, Aspect::Depth, 3, false);
-    CheckLayerCompressed(s, Aspect::Stencil, 2, false);
-    CheckLayerCompressed(s, Aspect::Stencil, 3, false);
+    CheckLayerCompressed(s, Aspect::Stencil, 2, true);
+    CheckLayerCompressed(s, Aspect::Stencil, 3, true);
 
     // Update completely. Without a single value recompression shouldn't happen.
     {
@@ -445,6 +458,172 @@
     CheckLayerCompressed(s, Aspect::Color, 1, false);
 }
 
+// The tests for Merge() all follow the same as the Update() tests except that they use Update()
+// to set up the test storages.
+
+// Similar to CallUpdateOnBoth but for Merge
+template <typename T, typename U, typename F>
+void CallMergeOnBoth(SubresourceStorage<T>* s,
+                     FakeStorage<T>* f,
+                     const SubresourceStorage<U>& other,
+                     F&& mergeFunc) {
+    RangeTracker tracker(*s);
+
+    s->Merge(other, [&](const SubresourceRange& range, T* data, const U& otherData) {
+        tracker.Track(range);
+        mergeFunc(range, data, otherData);
+    });
+    f->Merge(other, mergeFunc);
+
+    tracker.CheckTrackedExactly(
+        SubresourceRange::MakeFull(f->mAspects, f->mArrayLayerCount, f->mMipLevelCount));
+    f->CheckSameAs(*s);
+}
+
+// Test merging two fully compressed single-aspect resources.
+TEST(SubresourceStorageTest, MergeFullWithFullSingleAspect) {
+    SubresourceStorage<int> s(Aspect::Color, 4, 6);
+    FakeStorage<int> f(Aspect::Color, 4, 6);
+
+    // Merge the whole resource in a single call.
+    SubresourceStorage<bool> other(Aspect::Color, 4, 6, true);
+    CallMergeOnBoth(&s, &f, other, [](const SubresourceRange&, int* data, bool other) {
+        if (other) {
+            *data = 13;
+        }
+    });
+
+    CheckAspectCompressed(s, Aspect::Color, true);
+}
+
+// Test merging two fully compressed multi-aspect resources.
+TEST(SubresourceStorageTest, MergeFullWithFullMultiAspect) {
+    SubresourceStorage<int> s(Aspect::Depth | Aspect::Stencil, 6, 7);
+    FakeStorage<int> f(Aspect::Depth | Aspect::Stencil, 6, 7);
+
+    // Merge the whole resource in a single call.
+    SubresourceStorage<bool> other(Aspect::Depth | Aspect::Stencil, 6, 7, true);
+    CallMergeOnBoth(&s, &f, other, [](const SubresourceRange&, int* data, bool other) {
+        if (other) {
+            *data = 13;
+        }
+    });
+
+    CheckAspectCompressed(s, Aspect::Depth, true);
+    CheckAspectCompressed(s, Aspect::Stencil, true);
+}
+
+// Test merging a fully compressed resource in a resource with the "cross band" pattern.
+//  - The first band is full layers [2, 3] on both aspects
+//  - The second band is full mips [5, 6] on one aspect.
+// This provides coverage of using a single piece of data from `other` to update all of `s`
+TEST(SubresourceStorageTest, MergeFullInTwoBand) {
+    const uint32_t kLayers = 5;
+    const uint32_t kLevels = 9;
+    SubresourceStorage<int> s(Aspect::Depth | Aspect::Stencil, kLayers, kLevels);
+    FakeStorage<int> f(Aspect::Depth | Aspect::Stencil, kLayers, kLevels);
+
+    // Update the two bands
+    {
+        SubresourceRange range(Aspect::Depth | Aspect::Stencil, {2, 2}, {0, kLevels});
+        CallUpdateOnBoth(&s, &f, range, [](const SubresourceRange&, int* data) { *data += 3; });
+    }
+    {
+        SubresourceRange range(Aspect::Depth, {0, kLayers}, {5, 2});
+        CallUpdateOnBoth(&s, &f, range, [](const SubresourceRange&, int* data) { *data += 5; });
+    }
+
+    // Merge the fully compressed resource.
+    SubresourceStorage<int> other(Aspect::Depth | Aspect::Stencil, kLayers, kLevels, 17);
+    CallMergeOnBoth(&s, &f, other,
+                    [](const SubresourceRange&, int* data, int other) { *data += other; });
+
+    // The layers traversed by the mip band are still uncompressed.
+    CheckLayerCompressed(s, Aspect::Depth, 1, false);
+    CheckLayerCompressed(s, Aspect::Depth, 2, false);
+    CheckLayerCompressed(s, Aspect::Depth, 3, false);
+    CheckLayerCompressed(s, Aspect::Depth, 4, false);
+
+    // Stencil is decompressed but all its layers are still compressed because there wasn't the mip
+    // band.
+    CheckAspectCompressed(s, Aspect::Stencil, false);
+    CheckLayerCompressed(s, Aspect::Stencil, 1, true);
+    CheckLayerCompressed(s, Aspect::Stencil, 2, true);
+    CheckLayerCompressed(s, Aspect::Stencil, 3, true);
+    CheckLayerCompressed(s, Aspect::Stencil, 4, true);
+}
+// Test the reverse, mergign two-bands in a full resource. This provides coverage for decompressing
+// aspects / and partilly layers to match the compression of `other`
+TEST(SubresourceStorageTest, MergeTwoBandInFull) {
+    const uint32_t kLayers = 5;
+    const uint32_t kLevels = 9;
+    SubresourceStorage<int> s(Aspect::Depth | Aspect::Stencil, kLayers, kLevels, 75);
+    FakeStorage<int> f(Aspect::Depth | Aspect::Stencil, kLayers, kLevels, 75);
+
+    // Update the two bands
+    SubresourceStorage<int> other(Aspect::Depth | Aspect::Stencil, kLayers, kLevels);
+    {
+        SubresourceRange range(Aspect::Depth | Aspect::Stencil, {2, 2}, {0, kLevels});
+        other.Update(range, [](const SubresourceRange&, int* data) { *data += 3; });
+    }
+    {
+        SubresourceRange range(Aspect::Depth, {0, kLayers}, {5, 2});
+        other.Update(range, [](const SubresourceRange&, int* data) { *data += 5; });
+    }
+
+    // Merge the fully compressed resource.
+    CallMergeOnBoth(&s, &f, other,
+                    [](const SubresourceRange&, int* data, int other) { *data += other; });
+
+    // The layers traversed by the mip band are still uncompressed.
+    CheckLayerCompressed(s, Aspect::Depth, 1, false);
+    CheckLayerCompressed(s, Aspect::Depth, 2, false);
+    CheckLayerCompressed(s, Aspect::Depth, 3, false);
+    CheckLayerCompressed(s, Aspect::Depth, 4, false);
+
+    // Stencil is decompressed but all its layers are still compressed because there wasn't the mip
+    // band.
+    CheckAspectCompressed(s, Aspect::Stencil, false);
+    CheckLayerCompressed(s, Aspect::Stencil, 1, true);
+    CheckLayerCompressed(s, Aspect::Stencil, 2, true);
+    CheckLayerCompressed(s, Aspect::Stencil, 3, true);
+    CheckLayerCompressed(s, Aspect::Stencil, 4, true);
+}
+
+// Test merging storage with a layer band in a stipple patterned storage. This provide coverage
+// for the code path that uses the same layer data for other multiple times.
+TEST(SubresourceStorageTest, MergeLayerBandInStipple) {
+    const uint32_t kLayers = 3;
+    const uint32_t kLevels = 5;
+
+    SubresourceStorage<int> s(Aspect::Color, kLayers, kLevels);
+    FakeStorage<int> f(Aspect::Color, kLayers, kLevels);
+    SubresourceStorage<int> other(Aspect::Color, kLayers, kLevels);
+
+    for (uint32_t layer = 0; layer < kLayers; layer++) {
+        for (uint32_t level = 0; level < kLevels; level++) {
+            if ((layer + level) % 2 == 0) {
+                SubresourceRange range = SubresourceRange::MakeSingle(Aspect::Color, layer, level);
+                CallUpdateOnBoth(&s, &f, range,
+                                 [](const SubresourceRange&, int* data) { *data += 17; });
+            }
+        }
+        if (layer % 2 == 0) {
+            other.Update({Aspect::Color, {layer, 1}, {0, kLevels}},
+                         [](const SubresourceRange&, int* data) { *data += 8; });
+        }
+    }
+
+    // Merge the band in the stipple.
+    CallMergeOnBoth(&s, &f, other,
+                    [](const SubresourceRange&, int* data, int other) { *data += other; });
+
+    // None of the resulting layers are compressed.
+    CheckLayerCompressed(s, Aspect::Color, 0, false);
+    CheckLayerCompressed(s, Aspect::Color, 1, false);
+    CheckLayerCompressed(s, Aspect::Color, 2, false);
+}
+
 // Bugs found while testing:
 //  - mLayersCompressed not initialized to true.
 //  - DecompressLayer setting Compressed to true instead of false.