Compositor: Add utils for accessing pixel data

source/blender/compositor/realtime_compositor/COM_result.hh

@@ -353,10 +353,34 @@ class Result {
   /* Returns a reference to the domain of the result. See the Domain class. */
   const Domain &domain() const;
 
+  /* Returns a reference to the allocate float data. */
+  float *float_texture();
+
+  /* Loads the float pixel at the given texel coordinates and returns it in a float4. If the number
+   * of channels in the result are less than 4, then the rest of the returned float4 will have its
+   * vales initialized as follows: float4(0, 0, 0, 1). This is similar to how the texelFetch
+   * function in GLSL works. If the result is a single value result, then that single value is
+   * returned for all texel coordinates. */
+  float4 load_pixel(const int2 &texel) const;
+
+  /* Stores the given pixel value in the float pixel at the given texel coordinates. While a float4
+   * is given, only the number of channels of the result will be written, while the rest of the
+   * float4 will be ignored. This is similar to how the imageStore function in GLSL works. */
+  void store_pixel(const int2 &texel, const float4 &pixel_value);
+
  private:
   /* Allocates the texture data for the given size, either on the GPU or CPU based on the result's
    * context. See the allocate_texture method for information about the from_pool argument. */
   void allocate_data(int2 size, bool from_pool);
+
+  /* Computes the number of channels of the result based on its type. */
+  int64_t channels_count() const;
+
+  /* Get a pointer to the float pixel at the given texel position. */
+  float *get_float_pixel(const int2 &texel) const;
+
+  /* Copy the float pixel from the source pointer to the target pointer. */
+  void copy_pixel(float *target, const float *source) const;
 };
 
 }  // namespace blender::realtime_compositor

source/blender/compositor/realtime_compositor/intern/result.cc

@@ -682,6 +682,29 @@ const Domain &Result::domain() const
   return domain_;
 }
 
+float *Result::float_texture()
+{
+  BLI_assert(storage_type_ == ResultStorageType::FloatCPU);
+  return float_texture_;
+}
+
+float4 Result::load_pixel(const int2 &texel) const
+{
+  float4 pixel_value = float4(0.0f, 0.0f, 0.0f, 1.0f);
+  if (is_single_value_) {
+    this->copy_pixel(pixel_value, float_texture_);
+  }
+  else {
+    this->copy_pixel(pixel_value, this->get_float_pixel(texel));
+  }
+  return pixel_value;
+}
+
+void Result::store_pixel(const int2 &texel, const float4 &pixel_value)
+{
+  this->copy_pixel(this->get_float_pixel(texel), pixel_value);
+}
+
 void Result::allocate_data(int2 size, bool from_pool)
 {
   if (context_->use_gpu()) {
@@ -702,37 +725,63 @@ void Result::allocate_data(int2 size, bool from_pool)
   else {
     switch (type_) {
       case ResultType::Float:
-        float_texture_ = static_cast<float *>(
-            MEM_malloc_arrayN(size.x * size.y, sizeof(float), __func__));
-        storage_type_ = ResultStorageType::FloatCPU;
-        break;
-
       case ResultType::Vector:
       case ResultType::Color:
-        float_texture_ = static_cast<float *>(
-            MEM_malloc_arrayN(size.x * size.y, sizeof(float4), __func__));
-        storage_type_ = ResultStorageType::FloatCPU;
-        break;
-
       case ResultType::Float2:
-        float_texture_ = static_cast<float *>(
-            MEM_malloc_arrayN(size.x * size.y, sizeof(float2), __func__));
-        storage_type_ = ResultStorageType::FloatCPU;
-        break;
-
       case ResultType::Float3:
-        float_texture_ = static_cast<float *>(
-            MEM_malloc_arrayN(size.x * size.y, sizeof(float3), __func__));
+        float_texture_ = static_cast<float *>(MEM_malloc_arrayN(
+            int64_t(size.x) * int64_t(size.y), this->channels_count() * sizeof(float), __func__));
         storage_type_ = ResultStorageType::FloatCPU;
         break;
-
       case ResultType::Int2:
-        integer_texture_ = static_cast<int *>(
-            MEM_malloc_arrayN(size.x * size.y, sizeof(int2), __func__));
+        integer_texture_ = static_cast<int *>(MEM_malloc_arrayN(
+            int64_t(size.x) * int64_t(size.y), this->channels_count() * sizeof(int), __func__));
         storage_type_ = ResultStorageType::IntegerCPU;
         break;
     }
   }
 }
 
+int64_t Result::channels_count() const
+{
+  switch (type_) {
+    case ResultType::Float:
+      return 1;
+    case ResultType::Float2:
+    case ResultType::Int2:
+      return 2;
+    case ResultType::Float3:
+      return 3;
+    case ResultType::Vector:
+    case ResultType::Color:
+      return 4;
+  }
+  return 4;
+}
+
+float *Result::get_float_pixel(const int2 &texel) const
+{
+  return float_texture_ + (texel.y * domain_.size.x + texel.x) * this->channels_count();
+}
+
+void Result::copy_pixel(float *target, const float *source) const
+{
+  switch (type_) {
+    case ResultType::Float:
+      *target = *source;
+      break;
+    case ResultType::Float2:
+    case ResultType::Int2:
+      copy_v2_v2(target, source);
+      break;
+    case ResultType::Float3:
+      copy_v3_v3(target, source);
+      break;
+    case ResultType::Vector:
+    case ResultType::Color:
+      copy_v4_v4(target, source);
+      break;
+  }
+}
+
 }  // namespace blender::realtime_compositor