Making Simple Deform Stretch Not Terrible

I’m currently working on making the Stretch mode in the Simple Deform modifier not terrible.

So far it seems to be okay, but one thing I can’t figure out is why the Z component doesn’t normalize properly based on the size of the object. Even when removing the code that scales the Z component, objects with a radius other than 1 don’t work properly in this mode, as if the code that normalizes the Z component doesn’t execute for some reason.

What I mean by “normalization” is remapping the Z component to the range [-1,1] based on the height of the object. This seems to work fine for twist/bend/taper modes, but not for stretch mode (even when removing the Z scale code like mentioned before). I dunno why, and I’d like some insight as to where in the code I can look.

[EDIT]

Here’s a diff of what I have so far if anyone wants to take a look themselves:

diff --git a/source/blender/modifiers/intern/MOD_simpledeform.c b/source/blender/modifiers/intern/MOD_simpledeform.c
index ec89176f97e..a0eea8e0aa6 100644
--- a/source/blender/modifiers/intern/MOD_simpledeform.c
+++ b/source/blender/modifiers/intern/MOD_simpledeform.c
@@ -118,13 +118,16 @@ static void simpleDeform_stretch(const float factor,
                                  float r_co[3])
 {
   float x = r_co[0], y = r_co[1], z = r_co[2];
-  float scale;
+  float ascale, wscale;
 
-  scale = (z * z * factor - factor + 1.0f);
+  /* Scale along main axis */
+  ascale = factor+1.0f;
+  /* Scale width-wise */
+  wscale = (sqrtf(1.0f-(z*z))*(sqrtf(ascale)-1.0f))+1.0f;
 
-  r_co[0] = x * scale;
-  r_co[1] = y * scale;
-  r_co[2] = z * (1.0f + factor);
+  r_co[0] = x / wscale;
+  r_co[1] = y / wscale;
+  r_co[2] = z * ascale;
 
   add_v3_v3(r_co, dcut);
 }