From 2a7cc5a8db735c1799f93e814bfdcf113628c638 Mon Sep 17 00:00:00 2001
From: Nico Burns <nico@nicoburns.com>
Date: Sun, 22 Oct 2023 22:36:32 +0100
Subject: [PATCH] Add ability to set measure function

---
 ctaffy/src/tree.rs  | 68 ++++++++++++++++++++++++++++++++++++++++++---
 ctaffy/src/value.rs | 29 +++++++++++++++++++
 2 files changed, 93 insertions(+), 4 deletions(-)

diff --git a/ctaffy/src/tree.rs b/ctaffy/src/tree.rs
index cc655263a..8ab0c2495 100644
--- a/ctaffy/src/tree.rs
+++ b/ctaffy/src/tree.rs
@@ -1,12 +1,28 @@
 use super::{
-    bail, bail_if_null, ok, try_or, TaffyFFIDefault, TaffyFFIResult, TaffyLayout, TaffyResult, TaffyReturnCode,
-    TaffyStyleMutRef,
+    bail, bail_if_null, ok, try_or, TaffyFFIDefault, TaffyFFIResult, TaffyLayout, TaffyMeasureMode, TaffyResult,
+    TaffyReturnCode, TaffySize, TaffyStyleMutRef,
 };
+use ::core::ffi::c_void;
 use taffy::prelude as core;
+use taffy::style::AvailableSpace;
 use taffy::Taffy as CoreTaffy;
 
+pub type TaffyMeasureFunction = extern "C" fn(
+    width_measure_mode: TaffyMeasureMode,
+    width: f32,
+    height_measure_mode: TaffyMeasureMode,
+    height: f32,
+    context: *mut c_void,
+) -> TaffySize;
+
+#[allow(dead_code)] // false positive
+struct NodeContext {
+    context: *mut c_void,
+    measure_function: TaffyMeasureFunction,
+}
+
 pub struct TaffyTree {
-    inner: CoreTaffy,
+    inner: CoreTaffy<NodeContext>,
 }
 pub type TaffyTreeOwnedRef = *mut TaffyTree;
 pub type TaffyTreeMutRef = *mut TaffyTree;
@@ -94,7 +110,33 @@ pub unsafe extern "C" fn TaffyTree_ComputeLayout(
             width: available_space_from_f32(available_width),
             height: available_space_from_f32(available_height),
         };
-        try_or!(InvalidNodeId, tree.inner.compute_layout(node_id.into(), available_space));
+        try_or!(
+            InvalidNodeId,
+            tree.inner.compute_layout_with_measure(
+                node_id.into(),
+                available_space,
+                |known_dimensions, available_space, _node_id, node_context| {
+                    let (width, width_measure_mode) = match (known_dimensions.width, available_space.width) {
+                        (Some(width), _) => (width, TaffyMeasureMode::Exact),
+                        (None, AvailableSpace::Definite(width)) => (width, TaffyMeasureMode::FitContent),
+                        (None, AvailableSpace::MaxContent) => (f32::INFINITY, TaffyMeasureMode::MaxContent),
+                        (None, AvailableSpace::MinContent) => (f32::INFINITY, TaffyMeasureMode::MinContent),
+                    };
+                    let (height, height_measure_mode) = match (known_dimensions.height, available_space.height) {
+                        (Some(height), _) => (height, TaffyMeasureMode::Exact),
+                        (None, AvailableSpace::Definite(height)) => (height, TaffyMeasureMode::FitContent),
+                        (None, AvailableSpace::MaxContent) => (f32::INFINITY, TaffyMeasureMode::MaxContent),
+                        (None, AvailableSpace::MinContent) => (f32::INFINITY, TaffyMeasureMode::MinContent),
+                    };
+                    match node_context {
+                        Some(NodeContext { measure_function, context }) => {
+                            measure_function(width_measure_mode, width, height_measure_mode, height, *context).into()
+                        }
+                        _ => core::Size::ZERO,
+                    }
+                }
+            )
+        );
         TaffyReturnCode::Ok
     })
 }
@@ -165,6 +207,24 @@ pub unsafe extern "C" fn TaffyTree_GetStyleMut(
     })
 }
 
+/// Create a new Node in the TaffyTree. Returns a NodeId handle to the node.
+#[no_mangle]
+#[allow(clippy::missing_safety_doc)]
+pub unsafe extern "C" fn TaffyTree_SetNodeContext(
+    raw_tree: TaffyTreeMutRef,
+    node_id: TaffyNodeId,
+    measure_function: TaffyMeasureFunction,
+    context: *mut c_void,
+) -> TaffyReturnCode {
+    with_tree_mut!(raw_tree, tree, {
+        try_or!(
+            InvalidNodeId,
+            tree.inner.set_node_context(node_id.into(), Some(NodeContext { measure_function, context }))
+        );
+        ok!(TaffyReturnCode::Ok);
+    })
+}
+
 /// Create a new Node in the TaffyTree. Returns a NodeId handle to the node.
 #[no_mangle]
 #[allow(clippy::missing_safety_doc)]
diff --git a/ctaffy/src/value.rs b/ctaffy/src/value.rs
index 3250a65cc..543125204 100644
--- a/ctaffy/src/value.rs
+++ b/ctaffy/src/value.rs
@@ -46,6 +46,35 @@ pub enum TaffyUnit {
     Fr,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[repr(C)]
+pub enum TaffyMeasureMode {
+    /// A none value (used to unset optional fields)
+    Exact,
+    /// Fixed Length (pixel) value
+    FitContent,
+    /// Percentage value
+    MinContent,
+    /// Min-content size
+    MaxContent,
+}
+
+#[derive(Debug, Clone, Copy)]
+#[repr(C)]
+pub struct TaffySize {
+    width: f32,
+    height: f32,
+}
+impl From<TaffySize> for core::Size<f32> {
+    #[inline(always)]
+    fn from(value: TaffySize) -> Self {
+        core::Size {
+            width: value.width,
+            height: value.height,
+        }
+    }
+}
+
 #[repr(C)]
 pub struct TaffyLayout {
     pub x: f32,