diff --git a/services/surfaceflinger/DisplayHardware/HWC2.cpp b/services/surfaceflinger/DisplayHardware/HWC2.cpp
index ae6ba98..87a0e9a 100644
--- a/services/surfaceflinger/DisplayHardware/HWC2.cpp
+++ b/services/surfaceflinger/DisplayHardware/HWC2.cpp
@@ -317,9 +317,12 @@
             "Capability size has changed");
     uint32_t numCapabilities = 0;
     mHwcDevice->getCapabilities(mHwcDevice, &numCapabilities, nullptr);
-    mCapabilities.resize(numCapabilities);
-    auto asInt = reinterpret_cast<int32_t*>(mCapabilities.data());
+    std::vector<Capability> capabilities(numCapabilities);
+    auto asInt = reinterpret_cast<int32_t*>(capabilities.data());
     mHwcDevice->getCapabilities(mHwcDevice, &numCapabilities, asInt);
+    for (auto capability : capabilities) {
+        mCapabilities.emplace(capability);
+    }
 }
 
 bool Device::hasCapability(HWC2::Capability capability) const
diff --git a/services/surfaceflinger/DisplayHardware/HWC2.h b/services/surfaceflinger/DisplayHardware/HWC2.h
index 68781da..beb7bc6 100644
--- a/services/surfaceflinger/DisplayHardware/HWC2.h
+++ b/services/surfaceflinger/DisplayHardware/HWC2.h
@@ -33,6 +33,7 @@
 #include <functional>
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 namespace android {
@@ -66,7 +67,7 @@
 
     std::string dump() const;
 
-    const std::vector<Capability>& getCapabilities() const {
+    const std::unordered_set<Capability>& getCapabilities() const {
         return mCapabilities;
     };
 
@@ -181,7 +182,7 @@
     HWC2_PFN_SET_LAYER_VISIBLE_REGION mSetLayerVisibleRegion;
     HWC2_PFN_SET_LAYER_Z_ORDER mSetLayerZOrder;
 
-    std::vector<Capability> mCapabilities;
+    std::unordered_set<Capability> mCapabilities;
     std::unordered_map<hwc2_display_t, std::weak_ptr<Display>> mDisplays;
 
     HotplugCallback mHotplug;
diff --git a/services/surfaceflinger/DisplayHardware/HWComposer.cpp b/services/surfaceflinger/DisplayHardware/HWComposer.cpp
index 133e5f1..199848f 100644
--- a/services/surfaceflinger/DisplayHardware/HWComposer.cpp
+++ b/services/surfaceflinger/DisplayHardware/HWComposer.cpp
@@ -144,6 +144,11 @@
     mRemainingHwcVirtualDisplays = mHwcDevice->getMaxVirtualDisplayCount();
 }
 
+bool HWComposer::hasCapability(HWC2::Capability capability) const
+{
+    return mHwcDevice->getCapabilities().count(capability) > 0;
+}
+
 bool HWComposer::isValidDisplay(int32_t displayId) const {
     return static_cast<size_t>(displayId) < mDisplayData.size() &&
             mDisplayData[displayId].hwcDisplay;
@@ -692,6 +697,28 @@
     return NO_ERROR;
 }
 
+status_t HWComposer::setColorTransform(int32_t displayId,
+        const mat4& transform) {
+    if (!isValidDisplay(displayId)) {
+        ALOGE("setColorTransform: Display %d is not valid", displayId);
+        return BAD_INDEX;
+    }
+
+    auto& displayData = mDisplayData[displayId];
+    bool isIdentity = transform == mat4();
+    auto error = displayData.hwcDisplay->setColorTransform(transform,
+            isIdentity ? HAL_COLOR_TRANSFORM_IDENTITY :
+            HAL_COLOR_TRANSFORM_ARBITRARY_MATRIX);
+    if (error != HWC2::Error::None) {
+        ALOGE("setColorTransform: Failed to set transform on display %d: "
+                "%s (%d)", displayId, to_string(error).c_str(),
+                static_cast<int32_t>(error));
+        return UNKNOWN_ERROR;
+    }
+
+    return NO_ERROR;
+}
+
 void HWComposer::disconnectDisplay(int displayId) {
     LOG_ALWAYS_FATAL_IF(displayId < 0);
     auto& displayData = mDisplayData[displayId];
diff --git a/services/surfaceflinger/DisplayHardware/HWComposer.h b/services/surfaceflinger/DisplayHardware/HWComposer.h
index aa233df..17676ae 100644
--- a/services/surfaceflinger/DisplayHardware/HWComposer.h
+++ b/services/surfaceflinger/DisplayHardware/HWComposer.h
@@ -81,6 +81,8 @@
 
     void setEventHandler(EventHandler* handler);
 
+    bool hasCapability(HWC2::Capability capability) const;
+
     // Attempts to allocate a virtual display. If the virtual display is created
     // on the HWC device, outId will contain its HWC ID.
     status_t allocateVirtualDisplay(uint32_t width, uint32_t height,
@@ -104,6 +106,9 @@
     // set active config
     status_t setActiveConfig(int32_t displayId, size_t configId);
 
+    // Sets a color transform to be applied to the result of composition
+    status_t setColorTransform(int32_t displayId, const mat4& transform);
+
     // reset state when an external, non-virtual display is disconnected
     void disconnectDisplay(int32_t displayId);