Armature: Fix/Change bone axes display.

Now the axes are displayed correctly at the tip of the bone and with the
axes names.

I've made some modifications though:
- Axes are colored. (should not be in object mode but that's TODO)
- Axes ends are not flat arrows anymore. Replaced with a small diamond.
- Axes names are now scale by their respective axes instead of being
  affected by other axes.
- Changed axes names "font" to be a bit more sexy.

6 files changed, 211 insertions, 4 deletions

diff --git a/source/blender/draw/CMakeLists.txt b/source/blender/draw/CMakeLists.txt
index 1550f48fec5..53d46575a62 100644
--- a/source/blender/draw/CMakeLists.txt
+++ b/source/blender/draw/CMakeLists.txt
@@ -224,6 +224,7 @@ data_to_c_simple(modes/shaders/common_globals_lib.glsl SRC)
 data_to_c_simple(modes/shaders/common_view_lib.glsl SRC)
 data_to_c_simple(modes/shaders/common_fxaa_lib.glsl SRC)
 data_to_c_simple(modes/shaders/common_fullscreen_vert.glsl SRC)
+data_to_c_simple(modes/shaders/armature_axes_vert.glsl SRC)
 data_to_c_simple(modes/shaders/armature_sphere_solid_vert.glsl SRC)
 data_to_c_simple(modes/shaders/armature_sphere_solid_frag.glsl SRC)
 data_to_c_simple(modes/shaders/armature_sphere_outline_vert.glsl SRC)
diff --git a/source/blender/draw/intern/draw_armature.c b/source/blender/draw/intern/draw_armature.c
index d2a36f2794c..e4b9ddc7631 100644
--- a/source/blender/draw/intern/draw_armature.c
+++ b/source/blender/draw/intern/draw_armature.c
@@ -315,8 +315,7 @@ static void drw_shgroup_bone_point(
 static void drw_shgroup_bone_axes(const float (*bone_mat)[4], const float color[4])
 {
 	if (g_data.bone_axes == NULL) {
-		struct Gwn_Batch *geom = DRW_cache_bone_arrows_get();
-		g_data.bone_axes = shgroup_instance_wire(g_data.pass_bone_wire, geom);
+		g_data.bone_axes = shgroup_instance_bone_axes(g_data.pass_bone_wire);
 	}
 	float final_bonemat[4][4];
 	mul_m4_m4m4(final_bonemat, g_data.ob->obmat, bone_mat);
diff --git a/source/blender/draw/intern/draw_cache.c b/source/blender/draw/intern/draw_cache.c
index 909e577bd4f..ebcad9cb7d5 100644
--- a/source/blender/draw/intern/draw_cache.c
+++ b/source/blender/draw/intern/draw_cache.c
@@ -2104,10 +2104,164 @@ Gwn_Batch *DRW_cache_bone_point_wire_outline_get(void)
 	return SHC.drw_bone_point_wire;
 }
 
+static void set_bone_axis_vert(
+        Gwn_VertBuf *vbo, uint axis, uint pos, uint col,
+        unsigned int *v, const float *a, const float *p, const float *c)
+{
+	GWN_vertbuf_attr_set(vbo, axis, *v, a);
+	GWN_vertbuf_attr_set(vbo, pos,  *v, p);
+	GWN_vertbuf_attr_set(vbo, col,  *v, c);
+	*v += 1;
+}
+
+#define S_X 0.0215f
+#define S_Y 0.025f
+static float x_axis_name[4][2] = {
+	{ 0.9f * S_X,  1.0f * S_Y}, {-1.0f * S_X, -1.0f * S_Y},
+	{-0.9f * S_X,  1.0f * S_Y}, { 1.0f * S_X, -1.0f * S_Y}
+};
+#define X_LEN (sizeof(x_axis_name) / (sizeof(float) * 2))
+#undef S_X
+#undef S_Y
+
+#define S_X 0.0175f
+#define S_Y 0.025f
+static float y_axis_name[6][2] = {
+	{-1.0f * S_X,  1.0f * S_Y}, { 0.0f * S_X, -0.1f * S_Y},
+	{ 1.0f * S_X,  1.0f * S_Y}, { 0.0f * S_X, -0.1f * S_Y},
+	{ 0.0f * S_X, -0.1f * S_Y}, { 0.0f * S_X, -1.0f * S_Y}
+};
+#define Y_LEN (sizeof(y_axis_name) / (sizeof(float) * 2))
+#undef S_X
+#undef S_Y
+
+#define S_X 0.02f
+#define S_Y 0.025f
+static float z_axis_name[10][2] = {
+	{-0.95f * S_X,  1.00f * S_Y}, { 0.95f * S_X,  1.00f * S_Y},
+	{ 0.95f * S_X,  1.00f * S_Y}, { 0.95f * S_X,  0.90f * S_Y},
+	{ 0.95f * S_X,  0.90f * S_Y}, {-1.00f * S_X, -0.90f * S_Y},
+	{-1.00f * S_X, -0.90f * S_Y}, {-1.00f * S_X, -1.00f * S_Y},
+	{-1.00f * S_X, -1.00f * S_Y}, { 1.00f * S_X, -1.00f * S_Y}
+};
+#define Z_LEN (sizeof(z_axis_name) / (sizeof(float) * 2))
+#undef S_X
+#undef S_Y
+
+#define S_X 0.007f
+#define S_Y 0.007f
+static float axis_marker[8][2] = {
+#if 0 /* square */
+	{-1.0f * S_X,  1.0f * S_Y}, { 1.0f * S_X,  1.0f * S_Y},
+	{ 1.0f * S_X,  1.0f * S_Y}, { 1.0f * S_X, -1.0f * S_Y},
+	{ 1.0f * S_X, -1.0f * S_Y}, {-1.0f * S_X, -1.0f * S_Y},
+	{-1.0f * S_X, -1.0f * S_Y}, {-1.0f * S_X,  1.0f * S_Y}
+#else /* diamond */
+	{-S_X,  0.f}, { 0.f,  S_Y},
+	{ 0.f,  S_Y}, { S_X,  0.f},
+	{ S_X,  0.f}, { 0.f, -S_Y},
+	{ 0.f, -S_Y}, {-S_X,  0.f}
+#endif
+};
+#define MARKER_LEN (sizeof(axis_marker) / (sizeof(float) * 2))
+#define MARKER_FILL_LAYER 6
+#undef S_X
+#undef S_Y
+
+#define S_X 0.0007f
+#define S_Y 0.0007f
+#define O_X  0.001f
+#define O_Y -0.001f
+static float axis_name_shadow[8][2] = {
+	{-S_X + O_X,  S_Y + O_Y}, { S_X + O_X,  S_Y + O_Y},
+	{ S_X + O_X,  S_Y + O_Y}, { S_X + O_X, -S_Y + O_Y},
+	{ S_X + O_X, -S_Y + O_Y}, {-S_X + O_X, -S_Y + O_Y},
+	{-S_X + O_X, -S_Y + O_Y}, {-S_X + O_X,  S_Y + O_Y}
+};
+// #define SHADOW_RES (sizeof(axis_name_shadow) / (sizeof(float) * 2))
+#define SHADOW_RES 0
+#undef O_X
+#undef O_Y
+#undef S_X
+#undef S_Y
+
 Gwn_Batch *DRW_cache_bone_arrows_get(void)
 {
 	if (!SHC.drw_bone_arrows) {
-		Gwn_VertBuf *vbo = fill_arrows_vbo(0.25f);
+		/* Position Only 3D format */
+		static Gwn_VertFormat format = { 0 };
+		static struct { uint axis, pos, col; } attr_id;
+		if (format.attrib_ct == 0) {
+			attr_id.axis = GWN_vertformat_attr_add(&format, "axis", GWN_COMP_F32, 1, GWN_FETCH_FLOAT);
+			attr_id.pos = GWN_vertformat_attr_add(&format, "screenPos", GWN_COMP_F32, 2, GWN_FETCH_FLOAT);
+			attr_id.col = GWN_vertformat_attr_add(&format, "colorAxis", GWN_COMP_F32, 3, GWN_FETCH_FLOAT);
+		}
+
+		/* Line */
+		Gwn_VertBuf *vbo = GWN_vertbuf_create_with_format(&format);
+		GWN_vertbuf_data_alloc(vbo, (2 + MARKER_LEN * MARKER_FILL_LAYER) * 3 +
+		                            (X_LEN + Y_LEN + Z_LEN) * (1 + SHADOW_RES));
+
+		unsigned int v = 0;
+
+		for (int axis = 0; axis < 3; axis++) {
+			float pos[2] = {0.0f, 0.0f};
+			float c[3] = {0.0f, 0.0f, 0.0f};
+			float a = 0.0f;
+			/* center to axis line */
+			set_bone_axis_vert(vbo, attr_id.axis, attr_id.pos, attr_id.col, &v, &a, pos, c);
+			c[axis] = 0.5f;
+			a = axis + 0.25f;
+			set_bone_axis_vert(vbo, attr_id.axis, attr_id.pos, attr_id.col, &v, &a, pos, c);
+
+			/* Axis end marker */
+			for (int j = 1; j < MARKER_FILL_LAYER + 1; ++j) {
+				for (int i = 0; i < MARKER_LEN; ++i) {
+					float tmp[2];
+					mul_v2_v2fl(tmp, axis_marker[i], j / (float)MARKER_FILL_LAYER);
+					set_bone_axis_vert(vbo, attr_id.axis, attr_id.pos, attr_id.col,
+					                   &v, &a, tmp, c);
+				}
+			}
+
+			a = axis + 0.31f;
+			/* Axis name */
+			int axis_v_ct;
+			float (*axis_verts)[2];
+			if (axis == 0) {
+				axis_verts = x_axis_name;
+				axis_v_ct = X_LEN;
+			}
+			else if (axis == 1) {
+				axis_verts = y_axis_name;
+				axis_v_ct = Y_LEN;
+			}
+			else {
+				axis_verts = z_axis_name;
+				axis_v_ct = Z_LEN;
+			}
+
+			/* Axis name shadows */
+			copy_v3_fl(c, 0.0f);
+			c[axis] = 0.3f;
+			for (int j = 0; j < SHADOW_RES; ++j) {
+				for (int i = 0; i < axis_v_ct; ++i) {
+					float tmp[2];
+					add_v2_v2v2(tmp, axis_verts[i], axis_name_shadow[j]);
+					set_bone_axis_vert(vbo, attr_id.axis, attr_id.pos, attr_id.col,
+					                   &v, &a, tmp, c);
+				}
+			}
+
+			/* Axis name */
+			copy_v3_fl(c, 0.1f);
+			c[axis] = 1.0f;
+			for (int i = 0; i < axis_v_ct; ++i) {
+				set_bone_axis_vert(vbo, attr_id.axis, attr_id.pos, attr_id.col,
+				                   &v, &a, axis_verts[i], c);
+			}
+		}
+
 		SHC.drw_bone_arrows = GWN_batch_create_ex(GWN_PRIM_LINES, vbo, NULL, GWN_BATCH_OWNS_VBO);
 	}
 	return SHC.drw_bone_arrows;
diff --git a/source/blender/draw/intern/draw_common.c b/source/blender/draw/intern/draw_common.c
index ac5d6b87f4b..b2b25410ce4 100644
--- a/source/blender/draw/intern/draw_common.c
+++ b/source/blender/draw/intern/draw_common.c
@@ -153,6 +153,7 @@ void DRW_globals_update(void)
 
 /* ********************************* SHGROUP ************************************* */
 
+extern char datatoc_armature_axes_vert_glsl[];
 extern char datatoc_armature_sphere_solid_vert_glsl[];
 extern char datatoc_armature_sphere_solid_frag_glsl[];
 extern char datatoc_armature_sphere_outline_vert_glsl[];
@@ -171,6 +172,7 @@ extern char datatoc_object_mball_handles_vert_glsl[];
 static struct {
 	struct GPUShader *shape_outline;
 	struct GPUShader *shape_solid;
+	struct GPUShader *bone_axes;
 	struct GPUShader *bone_envelope;
 	struct GPUShader *bone_envelope_distance;
 	struct GPUShader *bone_envelope_outline;
@@ -435,6 +437,27 @@ DRWShadingGroup *shgroup_spot_instance(DRWPass *pass, struct Gwn_Batch *geom)
 	return grp;
 }
 
+DRWShadingGroup *shgroup_instance_bone_axes(DRWPass *pass)
+{
+	if (g_shaders.bone_axes == NULL) {
+		g_shaders.bone_axes = DRW_shader_create(
+		            datatoc_armature_axes_vert_glsl, NULL,
+		            datatoc_gpu_shader_flat_color_frag_glsl, NULL);
+	}
+
+	DRW_shgroup_instance_format(g_formats.instance_color, {
+		{"InstanceModelMatrix", DRW_ATTRIB_FLOAT, 16},
+		{"color"              , DRW_ATTRIB_FLOAT, 4}
+	});
+
+	DRWShadingGroup *grp = DRW_shgroup_instance_create(g_shaders.bone_axes,
+	                                                   pass, DRW_cache_bone_arrows_get(),
+	                                                   g_formats.instance_bone_outline);
+	DRW_shgroup_uniform_vec3(grp, "screenVecs[0]", DRW_viewport_screenvecs_get(), 2);
+
+	return grp;
+}
+
 DRWShadingGroup *shgroup_instance_bone_envelope_outline(DRWPass *pass)
 {
 	if (g_shaders.bone_envelope_outline == NULL) {
@@ -610,7 +633,6 @@ DRWShadingGroup *shgroup_instance_bone_sphere_outline(DRWPass *pass)
 }
 
 
-
 /* ******************************************** COLOR UTILS *********************************************** */
 
 /* TODO FINISH */
diff --git a/source/blender/draw/intern/draw_common.h b/source/blender/draw/intern/draw_common.h
index 0acd68dcd5d..a81daa9949d 100644
--- a/source/blender/draw/intern/draw_common.h
+++ b/source/blender/draw/intern/draw_common.h
@@ -116,6 +116,7 @@ struct DRWShadingGroup *shgroup_camera_instance(struct DRWPass *pass, struct Gwn
 struct DRWShadingGroup *shgroup_distance_lines_instance(struct DRWPass *pass, struct Gwn_Batch *geom);
 struct DRWShadingGroup *shgroup_spot_instance(struct DRWPass *pass, struct Gwn_Batch *geom);
 struct DRWShadingGroup *shgroup_instance_mball_handles(struct DRWPass *pass);
+struct DRWShadingGroup *shgroup_instance_bone_axes(struct DRWPass *pass);
 struct DRWShadingGroup *shgroup_instance_bone_envelope_distance(struct DRWPass *pass);
 struct DRWShadingGroup *shgroup_instance_bone_envelope_outline(struct DRWPass *pass);
 struct DRWShadingGroup *shgroup_instance_bone_envelope_solid(struct DRWPass *pass);
diff --git a/source/blender/draw/modes/shaders/armature_axes_vert.glsl b/source/blender/draw/modes/shaders/armature_axes_vert.glsl
new file mode 100644
index 00000000000..3a70c5d8515
--- /dev/null
+++ b/source/blender/draw/modes/shaders/armature_axes_vert.glsl
@@ -0,0 +1,30 @@
+
+uniform mat4 ViewProjectionMatrix;
+uniform vec3 screenVecs[3];
+
+/* ---- Instanciated Attribs ---- */
+in float axis; /* position on the axis. [0.0-1.0] is X axis, [1.0-2.0] is Y, etc... */
+in vec2 screenPos;
+in vec3 colorAxis;
+
+/* ---- Per instance Attribs ---- */
+in mat4 InstanceModelMatrix;
+in vec3 color;
+
+flat out vec4 finalColor;
+
+void main()
+{
+	vec3 chosen_axis = InstanceModelMatrix[int(axis)].xyz;
+	vec3 y_axis = InstanceModelMatrix[1].xyz;
+	vec3 bone_loc = InstanceModelMatrix[3].xyz;
+	vec3 wpos = bone_loc + y_axis + chosen_axis * fract(axis);
+	vec3 spos = screenVecs[0].xyz * screenPos.x + screenVecs[1].xyz * screenPos.y;
+	/* Scale uniformly by axis length */
+	spos *= length(chosen_axis);
+	
+	gl_Position = ViewProjectionMatrix * vec4(wpos + spos, 1.0);
+
+	finalColor.rgb = mix(color, colorAxis, 0.9);
+	finalColor.a = 1.0;
+}