Node UI: Add square and diamond socked draw styles

Currently not used by Blender's node trees

D2814 by @charlie

5 files changed, 102 insertions, 35 deletions

diff --git a/source/blender/editors/space_node/drawnode.c b/source/blender/editors/space_node/drawnode.c
index f58b4050834..6185f4afbe2 100644
--- a/source/blender/editors/space_node/drawnode.c
+++ b/source/blender/editors/space_node/drawnode.c
@@ -638,7 +638,7 @@ static void node_draw_reroute(const bContext *C, ARegion *ar, SpaceNode *UNUSED(
 	 * highlight also if node itself is selected, since we don't display the node body separately!
 	 */
 	for (sock = node->inputs.first; sock; sock = sock->next) {
-		node_socket_circle_draw(C, ntree, node, sock, socket_size, (sock->flag & SELECT) || (node->flag & SELECT));
+		node_socket_draw(C, ntree, node, sock, socket_size, (sock->flag & SELECT) || (node->flag & SELECT));
 	}
 
 	UI_block_end(C, node->block);
diff --git a/source/blender/editors/space_node/node_draw.c b/source/blender/editors/space_node/node_draw.c
index 6e25c87d274..213e326b1a6 100644
--- a/source/blender/editors/space_node/node_draw.c
+++ b/source/blender/editors/space_node/node_draw.c
@@ -617,34 +617,22 @@ static void node_draw_mute_line(View2D *v2d, SpaceNode *snode, bNode *node)
 	glDisable(GL_LINE_SMOOTH);
 }
 
-/* this might have some more generic use */
-static void node_circle_draw(float x, float y, float size, const float col[4], int highlight)
+static void node_socket_shape_draw(
+        float x, float y, float size, const float col[4], bool highlight,
+        const float coords[][2], int coords_len)
 {
-	/* 16 values of sin function */
-	static const float si[16] = {
-		0.00000000f, 0.39435585f, 0.72479278f, 0.93775213f,
-		0.99871650f, 0.89780453f, 0.65137248f, 0.29936312f,
-		-0.10116832f, -0.48530196f, -0.79077573f, -0.96807711f,
-		-0.98846832f, -0.84864425f, -0.57126821f, -0.20129852f
-	};
-	/* 16 values of cos function */
-	static const float co[16] = {
-		1.00000000f, 0.91895781f, 0.68896691f, 0.34730525f,
-		-0.05064916f, -0.44039415f, -0.75875812f, -0.95413925f,
-		-0.99486932f, -0.87434661f, -0.61210598f, -0.25065253f,
-		0.15142777f, 0.52896401f, 0.82076344f, 0.97952994f,
-	};
 	int a;
-	
+
 	glColor4fv(col);
-	
+
 	glEnable(GL_BLEND);
 	glBegin(GL_POLYGON);
-	for (a = 0; a < 16; a++)
-		glVertex2f(x + size * si[a], y + size * co[a]);
+	for (a = 0; a < coords_len; a++) {
+		glVertex2f(x + size * coords[a][0], y + size * coords[a][1]);
+	}
 	glEnd();
 	glDisable(GL_BLEND);
-	
+
 	if (highlight) {
 		UI_ThemeColor(TH_TEXT_HI);
 		glLineWidth(1.5f);
@@ -655,14 +643,16 @@ static void node_circle_draw(float x, float y, float size, const float col[4], i
 	glEnable(GL_BLEND);
 	glEnable(GL_LINE_SMOOTH);
 	glBegin(GL_LINE_LOOP);
-	for (a = 0; a < 16; a++)
-		glVertex2f(x + size * si[a], y + size * co[a]);
+	for (a = 0; a < coords_len; a++) {
+		glVertex2f(x + size * coords[a][0], y + size * coords[a][1]);
+	}
 	glEnd();
 	glDisable(GL_LINE_SMOOTH);
 	glDisable(GL_BLEND);
 }
 
-void node_socket_circle_draw(const bContext *C, bNodeTree *ntree, bNode *node, bNodeSocket *sock, float size, int highlight)
+
+void node_socket_draw(const bContext *C, bNodeTree *ntree, bNode *node, bNodeSocket *sock, float size, bool highlight)
 {
 	PointerRNA ptr, node_ptr;
 	float color[4];
@@ -670,7 +660,60 @@ void node_socket_circle_draw(const bContext *C, bNodeTree *ntree, bNode *node, b
 	RNA_pointer_create((ID *)ntree, &RNA_NodeSocket, sock, &ptr);
 	RNA_pointer_create((ID *)ntree, &RNA_Node, node, &node_ptr);
 	sock->typeinfo->draw_color((bContext *)C, &ptr, &node_ptr, color);
-	node_circle_draw(sock->locx, sock->locy, size, color, highlight);
+
+	/* 16 values of {sin, cos} function */
+	const float shape_circle[16][2] = {
+		{0.00000000f, 1.00000000f},
+		{0.39435585f, 0.91895781f},
+		{0.72479278f, 0.68896691f},
+		{0.93775213f, 0.34730525f},
+		{0.99871650f, -0.05064916f},
+		{0.89780453f, -0.44039415f},
+		{0.65137248f, -0.75875812f},
+		{0.29936312f, -0.95413925f},
+		{-0.10116832f, -0.99486932f},
+		{-0.48530196f, -0.87434661f},
+		{-0.79077573f, -0.61210598f},
+		{-0.96807711f, -0.25065253f},
+		{-0.98846832f, 0.15142777f},
+		{-0.84864425f, 0.52896401f},
+		{-0.57126821f, 0.82076344f},
+		{-0.20129852f, 0.97952994f }
+	};
+
+	const float shape_diamond[4][2] = {
+		{0.0f, 1.2f},
+		{1.2f, 0.0f},
+		{0.0f, -1.2f},
+		{-1.2f, 0.0f},
+	};
+
+	const float shape_square[4][2] = {
+		{-0.9f, 0.9f},
+		{0.9f, 0.9f},
+		{0.9f, -0.9f},
+		{-0.9f, -0.9f},
+	};
+
+	const float (*shape)[2];
+	int shape_len;
+	switch (sock->draw_shape) {
+		default:
+		case SOCK_DRAW_SHAPE_CIRCLE:
+			shape = shape_circle;
+			shape_len = ARRAY_SIZE(shape_circle);
+			break;
+		case SOCK_DRAW_SHAPE_DIAMOND:
+			shape = shape_diamond;
+			shape_len = ARRAY_SIZE(shape_diamond);
+			break;
+		case SOCK_DRAW_SHAPE_SQUARE:
+			shape = shape_square;
+			shape_len = ARRAY_SIZE(shape_square);
+			break;
+	}
+
+	node_socket_shape_draw(sock->locx, sock->locy, size, color, highlight, shape, shape_len);
 }
 
 /* **************  Socket callbacks *********** */
@@ -935,7 +978,7 @@ static void node_draw_basis(const bContext *C, ARegion *ar, SpaceNode *snode, bN
 		if (nodeSocketIsHidden(sock))
 			continue;
 		
-		node_socket_circle_draw(C, ntree, node, sock, NODE_SOCKSIZE, sock->flag & SELECT);
+		node_socket_draw(C, ntree, node, sock, NODE_SOCKSIZE, sock->flag & SELECT);
 	}
 	
 	/* socket outputs */
@@ -943,7 +986,7 @@ static void node_draw_basis(const bContext *C, ARegion *ar, SpaceNode *snode, bN
 		if (nodeSocketIsHidden(sock))
 			continue;
 		
-		node_socket_circle_draw(C, ntree, node, sock, NODE_SOCKSIZE, sock->flag & SELECT);
+		node_socket_draw(C, ntree, node, sock, NODE_SOCKSIZE, sock->flag & SELECT);
 	}
 	
 	/* preview */
@@ -1066,13 +1109,15 @@ static void node_draw_hidden(const bContext *C, ARegion *ar, SpaceNode *snode, b
 
 	/* sockets */
 	for (sock = node->inputs.first; sock; sock = sock->next) {
-		if (!nodeSocketIsHidden(sock))
-			node_socket_circle_draw(C, ntree, node, sock, socket_size, sock->flag & SELECT);
+		if (!nodeSocketIsHidden(sock)) {
+			node_socket_draw(C, ntree, node, sock, socket_size, sock->flag & SELECT);
+		}
 	}
 	
 	for (sock = node->outputs.first; sock; sock = sock->next) {
-		if (!nodeSocketIsHidden(sock))
-			node_socket_circle_draw(C, ntree, node, sock, socket_size, sock->flag & SELECT);
+		if (!nodeSocketIsHidden(sock)) {
+			node_socket_draw(C, ntree, node, sock, socket_size, sock->flag & SELECT);
+		}
 	}
 	
 	UI_block_end(C, node->block);
diff --git a/source/blender/editors/space_node/node_intern.h b/source/blender/editors/space_node/node_intern.h
index 6b8fa0b88fe..352f9e51012 100644
--- a/source/blender/editors/space_node/node_intern.h
+++ b/source/blender/editors/space_node/node_intern.h
@@ -67,8 +67,9 @@ void snode_group_offset(struct SpaceNode *snode, float *x, float *y);	/* transfo
 
 /* node_draw.c */
 int node_get_colorid(struct bNode *node);
-void node_socket_circle_draw(const struct bContext *C, struct bNodeTree *ntree, struct bNode *node,
-                             struct bNodeSocket *sock, float size, int highlight);
+void node_socket_draw(
+        const struct bContext *C, struct bNodeTree *ntree, struct bNode *node,
+        struct bNodeSocket *sock, float size, bool highlight);
 int node_get_resize_cursor(int directions);
 void node_draw_shadow(struct SpaceNode *snode, struct bNode *node, float radius, float alpha);
 void node_draw_default(const struct bContext *C, struct ARegion *ar, struct SpaceNode *snode,
diff --git a/source/blender/makesdna/DNA_node_types.h b/source/blender/makesdna/DNA_node_types.h
index b922ac072b0..e6bc315b728 100644
--- a/source/blender/makesdna/DNA_node_types.h
+++ b/source/blender/makesdna/DNA_node_types.h
@@ -111,7 +111,7 @@ typedef struct bNodeSocket {
 	short stack_index;			/* local stack index */
 	/* XXX deprecated, kept for forward compatibility */
 	short stack_type  DNA_DEPRECATED;
-	int pad;
+	char draw_shape, pad[3];
 	
 	void *cache;				/* cached data from execution */
 	
@@ -143,6 +143,13 @@ typedef enum eNodeSocketDatatype {
 	SOCK_STRING			= 7
 } eNodeSocketDatatype;
 
+/* socket shape */
+typedef enum eNodeSocketDrawShape {
+	SOCK_DRAW_SHAPE_CIRCLE = 0,
+	SOCK_DRAW_SHAPE_SQUARE = 1,
+	SOCK_DRAW_SHAPE_DIAMOND = 2
+} eNodeSocketDrawShape;
+
 /* socket side (input/output) */
 typedef enum eNodeSocketInOut {
 	SOCK_IN = 1,
diff --git a/source/blender/makesrna/intern/rna_nodetree.c b/source/blender/makesrna/intern/rna_nodetree.c
index 4de91a1c57c..ee4608b550f 100644
--- a/source/blender/makesrna/intern/rna_nodetree.c
+++ b/source/blender/makesrna/intern/rna_nodetree.c
@@ -73,6 +73,13 @@ EnumPropertyItem rna_enum_node_socket_in_out_items[] = {
 };
 
 #ifndef RNA_RUNTIME
+static EnumPropertyItem rna_enum_node_socket_draw_shape_items[] = {
+	{SOCK_DRAW_SHAPE_CIRCLE, "CIRCLE", 0, "Circle", ""},
+	{SOCK_DRAW_SHAPE_SQUARE, "SQUARE", 0, "Square", ""},
+	{SOCK_DRAW_SHAPE_DIAMOND, "DIAMOND", 0, "Diamond", ""},
+	{0, NULL, 0, NULL, NULL }
+};
+
 static EnumPropertyItem node_socket_type_items[] = {
 	{SOCK_CUSTOM,  "CUSTOM",    0,    "Custom",    ""},
 	{SOCK_FLOAT,   "VALUE",     0,    "Value",     ""},
@@ -6987,6 +6994,13 @@ static void rna_def_node_socket(BlenderRNA *brna)
 	RNA_def_property_ui_text(prop, "Type", "Data type");
 	RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_NodeSocket_update");
 
+	prop = RNA_def_property(srna, "draw_shape", PROP_ENUM, PROP_NONE);
+	RNA_def_property_enum_sdna(prop, NULL, "draw_shape");
+	RNA_def_property_enum_items(prop, rna_enum_node_socket_draw_shape_items);
+	RNA_def_property_enum_default(prop, SOCK_DRAW_SHAPE_CIRCLE);
+	RNA_def_property_ui_text(prop, "Shape", "Socket shape");
+	RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_NodeSocket_update");
+
 	/* registration */
 	prop = RNA_def_property(srna, "bl_idname", PROP_STRING, PROP_NONE);
 	RNA_def_property_string_sdna(prop, NULL, "typeinfo->idname");