Add RuntimeShader#setWorkingColorSpace

Bug: 299670828
Test: android.uirendering.cts.testclasses.RuntimeShaderTests#testWorkingColorSpace
Flag: com.android.graphics.hwui.flags.shader_color_space
Change-Id: I0aed0e414dc25a1f30ca6e8d8c182bd53ff8c186

4 files changed, 54 insertions, 4 deletions

diff --git a/core/api/current.txt b/core/api/current.txt
index b7f7a7f9e779..9ec3e0f35657 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -17603,6 +17603,7 @@ package android.graphics {
     method public void setIntUniform(@NonNull String, int, int, int);
     method public void setIntUniform(@NonNull String, int, int, int, int);
     method public void setIntUniform(@NonNull String, @NonNull int[]);
+    method @FlaggedApi("com.android.graphics.hwui.flags.shader_color_space") public void setWorkingColorSpace(@Nullable android.graphics.ColorSpace);
   }
 
   @FlaggedApi("com.android.graphics.hwui.flags.runtime_color_filters_blenders") public class RuntimeXfermode extends android.graphics.Xfermode {
diff --git a/graphics/java/android/graphics/RuntimeShader.java b/graphics/java/android/graphics/RuntimeShader.java
index 3543e991924e..9016724b765e 100644
--- a/graphics/java/android/graphics/RuntimeShader.java
+++ b/graphics/java/android/graphics/RuntimeShader.java
@@ -20,6 +20,7 @@ import android.annotation.ColorInt;
 import android.annotation.ColorLong;
 import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.util.ArrayMap;
 import android.view.Window;
 
@@ -76,6 +77,7 @@ import libcore.util.NativeAllocationRegistry;
  * Additionally, if the shader is invoked by another using {@link #setInputShader(String, Shader)},
  * then that parent shader may modify the input coordinates arbitrarily.</p>
  *
+ * <a id="agsl-and-color-spaces"/>
  * <h3>AGSL and Color Spaces</h3>
  * <p>Android Graphics and by extension {@link RuntimeShader} are color managed.  The working
  * {@link ColorSpace} for an AGSL shader is defined to be the color space of the destination, which
@@ -267,6 +269,8 @@ public class RuntimeShader extends Shader {
     private ArrayMap<String, ColorFilter> mColorFilterUniforms = new ArrayMap<>();
     private ArrayMap<String, RuntimeXfermode> mXfermodeUniforms = new ArrayMap<>();
 
+    private ColorSpace mWorkingColorSpace = null;
+
 
     /**
      * Creates a new RuntimeShader.
@@ -286,6 +290,35 @@ public class RuntimeShader extends Shader {
     }
 
     /**
+     * Sets the working color space for this shader. That is, the shader will be evaluated
+     * in the given colorspace before being converted to the output destination's colorspace.
+     *
+     * <p>By default the RuntimeShader is evaluated in the context of the
+     * <a href="#agsl-and-color-spaces">destination colorspace</a>. By calling this method
+     * that can be overridden to force the shader to be evaluated in the given colorspace first
+     * before then being color converted to the destination colorspace.</p>
+     *
+     * @param colorSpace The ColorSpace to evaluate in. Must be an {@link ColorSpace#getModel() RGB}
+     *                   ColorSpace. Passing null restores default behavior of working in the
+     *                   destination colorspace.
+     * @throws IllegalArgumentException If the colorspace is not RGB
+     */
+    @FlaggedApi(Flags.FLAG_SHADER_COLOR_SPACE)
+    public void setWorkingColorSpace(@Nullable ColorSpace colorSpace) {
+        if (colorSpace != null && colorSpace.getModel() != ColorSpace.Model.RGB) {
+            throw new IllegalArgumentException("ColorSpace must be RGB, given " + colorSpace);
+        }
+        if (mWorkingColorSpace != colorSpace) {
+            mWorkingColorSpace = colorSpace;
+            if (mWorkingColorSpace != null) {
+                // Just to enforce this can be resolved instead of erroring out later
+                mWorkingColorSpace.getNativeInstance();
+            }
+            discardNativeInstance();
+        }
+    }
+
+    /**
      * Sets the uniform color value corresponding to this shader.  If the shader does not have a
      * uniform with that name or if the uniform is declared with a type other than vec3 or vec4 and
      * corresponding layout(color) annotation then an IllegalArgumentException is thrown.
@@ -578,7 +611,8 @@ public class RuntimeShader extends Shader {
     /** @hide */
     @Override
     protected long createNativeInstance(long nativeMatrix, boolean filterFromPaint) {
-        return nativeCreateShader(mNativeInstanceRuntimeShaderBuilder, nativeMatrix);
+        return nativeCreateShader(mNativeInstanceRuntimeShaderBuilder, nativeMatrix,
+                mWorkingColorSpace != null ? mWorkingColorSpace.getNativeInstance() : 0);
     }
 
     /** @hide */
@@ -588,7 +622,8 @@ public class RuntimeShader extends Shader {
 
     private static native long nativeGetFinalizer();
     private static native long nativeCreateBuilder(String agsl);
-    private static native long nativeCreateShader(long shaderBuilder, long matrix);
+    private static native long nativeCreateShader(long shaderBuilder, long matrix,
+            long colorSpacePtr);
     private static native void nativeUpdateUniforms(
             long shaderBuilder, String uniformName, float[] uniforms, boolean isColor);
     private static native void nativeUpdateUniforms(
diff --git a/libs/hwui/aconfig/hwui_flags.aconfig b/libs/hwui/aconfig/hwui_flags.aconfig
index 5e71d3360f39..2cddf24f6c72 100644
--- a/libs/hwui/aconfig/hwui_flags.aconfig
+++ b/libs/hwui/aconfig/hwui_flags.aconfig
@@ -127,6 +127,14 @@ flag {
 }
 
 flag {
+  name: "shader_color_space"
+  is_exported: true
+  namespace: "core_graphics"
+  description: "API to set the working colorspace of a Shader or ColorFilter"
+  bug: "299670828"
+}
+
+flag {
   name: "query_global_priority"
   namespace: "core_graphics"
   description: "Attempt to query whether the vulkan driver supports the requested global priority before queue creation."
diff --git a/libs/hwui/jni/Shader.cpp b/libs/hwui/jni/Shader.cpp
index eadb9dea566f..c02508977eba 100644
--- a/libs/hwui/jni/Shader.cpp
+++ b/libs/hwui/jni/Shader.cpp
@@ -266,11 +266,17 @@ static jlong RuntimeShader_getNativeFinalizer(JNIEnv*, jobject) {
     return static_cast<jlong>(reinterpret_cast<uintptr_t>(&SkRuntimeShaderBuilder_delete));
 }
 
-static jlong RuntimeShader_create(JNIEnv* env, jobject, jlong shaderBuilder, jlong matrixPtr) {
+static jlong RuntimeShader_create(JNIEnv* env, jobject, jlong shaderBuilder, jlong matrixPtr,
+                                  jlong colorSpacePtr) {
     SkRuntimeShaderBuilder* builder = reinterpret_cast<SkRuntimeShaderBuilder*>(shaderBuilder);
     const SkMatrix* matrix = reinterpret_cast<const SkMatrix*>(matrixPtr);
+    auto colorSpace = GraphicsJNI::getNativeColorSpace(colorSpacePtr);
     sk_sp<SkShader> shader = builder->makeShader(matrix);
     ThrowIAE_IfNull(env, shader);
+    if (colorSpace) {
+        shader = shader->makeWithWorkingColorSpace(colorSpace);
+        ThrowIAE_IfNull(env, shader);
+    }
     return reinterpret_cast<jlong>(shader.release());
 }
 
@@ -379,7 +385,7 @@ static const JNINativeMethod gComposeShaderMethods[] = {
 
 static const JNINativeMethod gRuntimeShaderMethods[] = {
         {"nativeGetFinalizer", "()J", (void*)RuntimeShader_getNativeFinalizer},
-        {"nativeCreateShader", "(JJ)J", (void*)RuntimeShader_create},
+        {"nativeCreateShader", "(JJJ)J", (void*)RuntimeShader_create},
         {"nativeCreateBuilder", "(Ljava/lang/String;)J", (void*)RuntimeShader_createShaderBuilder},
         {"nativeUpdateUniforms", "(JLjava/lang/String;[FZ)V",
          (void*)RuntimeShader_updateFloatArrayUniforms},