Support for depth buffers in compositor and viewer node

Support for only alpha images in compositor and viewer node

8 files changed, 117 insertions, 41 deletions

diff --git a/source/blender/compositor/nodes/COM_CompositorNode.cpp b/source/blender/compositor/nodes/COM_CompositorNode.cpp
index 28e466203c4..8a84908f478 100644
--- a/source/blender/compositor/nodes/COM_CompositorNode.cpp
+++ b/source/blender/compositor/nodes/COM_CompositorNode.cpp
@@ -33,13 +33,14 @@ void CompositorNode::convertToOperations(ExecutionSystem *graph, CompositorConte
 {
 	InputSocket *imageSocket = this->getInputSocket(0);
 	InputSocket *alphaSocket = this->getInputSocket(1);
-	if (imageSocket->isConnected()) {
-		CompositorOperation *colorAlphaProg = new CompositorOperation();
-		colorAlphaProg->setRenderData(context->getRenderData());
-		colorAlphaProg->setbNodeTree(context->getbNodeTree());
-		imageSocket->relinkConnections(colorAlphaProg->getInputSocket(0));
-		alphaSocket->relinkConnections(colorAlphaProg->getInputSocket(1));
-		graph->addOperation(colorAlphaProg);
-		addPreviewOperation(graph, colorAlphaProg->getInputSocket(0));
-	}
+	InputSocket *depthSocket = this->getInputSocket(2);
+
+	CompositorOperation *compositorOperation = new CompositorOperation();
+	compositorOperation->setRenderData(context->getRenderData());
+	compositorOperation->setbNodeTree(context->getbNodeTree());
+	imageSocket->relinkConnections(compositorOperation->getInputSocket(0), 0, graph);
+	alphaSocket->relinkConnections(compositorOperation->getInputSocket(1));
+	depthSocket->relinkConnections(compositorOperation->getInputSocket(2));
+	graph->addOperation(compositorOperation);
+	addPreviewOperation(graph, compositorOperation->getInputSocket(0));
 }
diff --git a/source/blender/compositor/nodes/COM_ViewerNode.cpp b/source/blender/compositor/nodes/COM_ViewerNode.cpp
index 1205767cb28..88ce0ff2016 100644
--- a/source/blender/compositor/nodes/COM_ViewerNode.cpp
+++ b/source/blender/compositor/nodes/COM_ViewerNode.cpp
@@ -35,23 +35,32 @@ void ViewerNode::convertToOperations(ExecutionSystem *graph, CompositorContext *
 {
 	InputSocket *imageSocket = this->getInputSocket(0);
 	InputSocket *alphaSocket = this->getInputSocket(1);
+	InputSocket *depthSocket = this->getInputSocket(2);
 	Image *image = (Image *)this->getbNode()->id;
 	ImageUser *imageUser = (ImageUser *) this->getbNode()->storage;
 	bNode *editorNode = this->getbNode();
-	if (imageSocket->isConnected()) {
-		ViewerOperation *viewerOperation = new ViewerOperation();
-		viewerOperation->setColorManagement(context->getRenderData()->color_mgt_flag & R_COLOR_MANAGEMENT);
-		viewerOperation->setColorPredivide(context->getRenderData()->color_mgt_flag & R_COLOR_MANAGEMENT_PREDIVIDE);
-		viewerOperation->setbNodeTree(context->getbNodeTree());
-		viewerOperation->setImage(image);
-		viewerOperation->setImageUser(imageUser);
-		viewerOperation->setActive((editorNode->flag & NODE_DO_OUTPUT) && this->isInActiveGroup());
-		viewerOperation->setChunkOrder((OrderOfChunks)editorNode->custom1);
-		viewerOperation->setCenterX(editorNode->custom3);
-		viewerOperation->setCenterY(editorNode->custom4);
-		imageSocket->relinkConnections(viewerOperation->getInputSocket(0), 0, graph);
-		alphaSocket->relinkConnections(viewerOperation->getInputSocket(1));
-		graph->addOperation(viewerOperation);
-		addPreviewOperation(graph, viewerOperation->getInputSocket(0));
+	ViewerOperation *viewerOperation = new ViewerOperation();
+	viewerOperation->setColorManagement(context->getRenderData()->color_mgt_flag & R_COLOR_MANAGEMENT);
+	viewerOperation->setColorPredivide(context->getRenderData()->color_mgt_flag & R_COLOR_MANAGEMENT_PREDIVIDE);
+	viewerOperation->setbNodeTree(context->getbNodeTree());
+	viewerOperation->setImage(image);
+	viewerOperation->setImageUser(imageUser);
+	viewerOperation->setActive((editorNode->flag & NODE_DO_OUTPUT) && this->isInActiveGroup());
+	viewerOperation->setChunkOrder((OrderOfChunks)editorNode->custom1);
+	viewerOperation->setCenterX(editorNode->custom3);
+	viewerOperation->setCenterY(editorNode->custom4);
+
+	viewerOperation->setResolutionInputSocketIndex(0);
+	if (!imageSocket->isConnected())
+	{
+		if (alphaSocket->isConnected()) {
+			viewerOperation->setResolutionInputSocketIndex(1);
+		}
 	}
+
+	imageSocket->relinkConnections(viewerOperation->getInputSocket(0), 0, graph);
+	alphaSocket->relinkConnections(viewerOperation->getInputSocket(1));
+	depthSocket->relinkConnections(viewerOperation->getInputSocket(2));
+	graph->addOperation(viewerOperation);
+	addPreviewOperation(graph, viewerOperation->getInputSocket(0));
 }
diff --git a/source/blender/compositor/operations/COM_CompositorOperation.cpp b/source/blender/compositor/operations/COM_CompositorOperation.cpp
index 43aad4f19d9..57a4639ecba 100644
--- a/source/blender/compositor/operations/COM_CompositorOperation.cpp
+++ b/source/blender/compositor/operations/COM_CompositorOperation.cpp
@@ -41,11 +41,14 @@ CompositorOperation::CompositorOperation() : NodeOperation()
 {
 	this->addInputSocket(COM_DT_COLOR);
 	this->addInputSocket(COM_DT_VALUE);
+	this->addInputSocket(COM_DT_VALUE);
 
 	this->setRenderData(NULL);
 	this->m_outputBuffer = NULL;
+	this->m_depthBuffer = NULL;
 	this->m_imageInput = NULL;
 	this->m_alphaInput = NULL;
+	this->m_depthInput = NULL;
 }
 
 void CompositorOperation::initExecution()
@@ -53,9 +56,13 @@ void CompositorOperation::initExecution()
 	// When initializing the tree during initial load the width and height can be zero.
 	this->m_imageInput = getInputSocketReader(0);
 	this->m_alphaInput = getInputSocketReader(1);
+	this->m_depthInput = getInputSocketReader(2);
 	if (this->getWidth() * this->getHeight() != 0) {
 		this->m_outputBuffer = (float *) MEM_callocN(this->getWidth() * this->getHeight() * 4 * sizeof(float), "CompositorOperation");
 	}
+	if (this->m_depthInput != NULL) {
+		this->m_depthBuffer = (float *) MEM_callocN(this->getWidth() * this->getHeight() * sizeof(float), "CompositorOperation");
+	}
 }
 
 void CompositorOperation::deinitExecution()
@@ -70,11 +77,18 @@ void CompositorOperation::deinitExecution()
 				MEM_freeN(rr->rectf);
 			}
 			rr->rectf = this->m_outputBuffer;
+			if (rr->rectz != NULL) {
+				MEM_freeN(rr->rectz);
+			}
+			rr->rectz = this->m_depthBuffer;
 		}
 		else {
 			if (this->m_outputBuffer) {
 				MEM_freeN(this->m_outputBuffer);
 			}
+			if (this->m_depthBuffer) {
+				MEM_freeN(this->m_depthBuffer);
+			}
 		}
 
 		BLI_lock_thread(LOCK_DRAW_IMAGE);
@@ -90,11 +104,16 @@ void CompositorOperation::deinitExecution()
 		if (this->m_outputBuffer) {
 			MEM_freeN(this->m_outputBuffer);
 		}
+		if (this->m_depthBuffer) {
+			MEM_freeN(this->m_depthBuffer);
+		}
 	}
 
 	this->m_outputBuffer = NULL;
+	this->m_depthBuffer = NULL;
 	this->m_imageInput = NULL;
 	this->m_alphaInput = NULL;
+	this->m_depthInput = NULL;
 }
 
 
@@ -102,13 +121,16 @@ void CompositorOperation::executeRegion(rcti *rect, unsigned int tileNumber)
 {
 	float color[8]; // 7 is enough
 	float *buffer = this->m_outputBuffer;
+	float *zbuffer = this->m_depthBuffer;
 
 	if (!buffer) return;
 	int x1 = rect->xmin;
 	int y1 = rect->ymin;
 	int x2 = rect->xmax;
 	int y2 = rect->ymax;
-	int offset = (y1 * this->getWidth() + x1) * COM_NUMBER_OF_CHANNELS;
+	int offset = (y1 * this->getWidth() + x1);
+	int add = (this->getWidth() - (x2 - x1));
+	int offset4 = offset * COM_NUMBER_OF_CHANNELS;
 	int x;
 	int y;
 	bool breaked = false;
@@ -119,13 +141,20 @@ void CompositorOperation::executeRegion(rcti *rect, unsigned int tileNumber)
 			if (this->m_alphaInput != NULL) {
 				this->m_alphaInput->read(&(color[3]), x, y, COM_PS_NEAREST);
 			}
-			copy_v4_v4(buffer + offset, color);
-			offset += COM_NUMBER_OF_CHANNELS;
+			copy_v4_v4(buffer + offset4, color);
+
+			if (this->m_depthInput != NULL) {
+				this->m_depthInput->read(color, x, y, COM_PS_NEAREST);
+				zbuffer[offset] = color[0];
+			}
+			offset4 += COM_NUMBER_OF_CHANNELS;
+			offset++;
 			if (isBreaked()) {
 				breaked = true;
 			}
 		}
-		offset += (this->getWidth() - (x2 - x1)) * COM_NUMBER_OF_CHANNELS;
+		offset += add;
+		offset4 += add * COM_NUMBER_OF_CHANNELS;
 	}
 }
 
diff --git a/source/blender/compositor/operations/COM_CompositorOperation.h b/source/blender/compositor/operations/COM_CompositorOperation.h
index 23d34abbfff..491fe3eb4e4 100644
--- a/source/blender/compositor/operations/COM_CompositorOperation.h
+++ b/source/blender/compositor/operations/COM_CompositorOperation.h
@@ -42,6 +42,11 @@ private:
 	float *m_outputBuffer;
 
 	/**
+	 * @brief reference to the output depth float buffer
+	 */
+	float *m_depthBuffer;
+
+	/**
 	 * @brief local reference to the input image operation
 	 */
 	SocketReader *m_imageInput;
@@ -50,6 +55,11 @@ private:
 	 * @brief local reference to the input alpha operation
 	 */
 	SocketReader *m_alphaInput;
+
+	/**
+	 * @brief local reference to the depth operation
+	 */
+	SocketReader *m_depthInput;
 public:
 	CompositorOperation();
 	void executeRegion(rcti *rect, unsigned int tileNumber);
diff --git a/source/blender/compositor/operations/COM_ViewerBaseOperation.cpp b/source/blender/compositor/operations/COM_ViewerBaseOperation.cpp
index 2470b239987..a5060f42e3a 100644
--- a/source/blender/compositor/operations/COM_ViewerBaseOperation.cpp
+++ b/source/blender/compositor/operations/COM_ViewerBaseOperation.cpp
@@ -43,9 +43,11 @@ ViewerBaseOperation::ViewerBaseOperation() : NodeOperation()
 	this->setImage(NULL);
 	this->setImageUser(NULL);
 	this->m_outputBuffer = NULL;
+	this->m_depthBuffer = NULL;
 	this->m_outputBufferDisplay = NULL;
 	this->m_active = false;
 	this->m_doColorManagement = true;
+	this->m_doDepthBuffer = false;
 }
 
 void ViewerBaseOperation::initExecution()
@@ -61,8 +63,8 @@ void ViewerBaseOperation::initImage()
 	ImBuf *ibuf = BKE_image_acquire_ibuf(anImage, this->m_imageUser, &this->m_lock);
 	
 	if (!ibuf) return;
+	BLI_lock_thread(LOCK_DRAW_IMAGE);
 	if (ibuf->x != (int)getWidth() || ibuf->y != (int)getHeight()) {
-		BLI_lock_thread(LOCK_DRAW_IMAGE);
 
 		imb_freerectImBuf(ibuf);
 		imb_freerectfloatImBuf(ibuf);
@@ -73,12 +75,21 @@ void ViewerBaseOperation::initImage()
 		imb_addrectfloatImBuf(ibuf);
 		anImage->ok = IMA_OK_LOADED;
 
-		BLI_unlock_thread(LOCK_DRAW_IMAGE);
 	}
+	if (m_doDepthBuffer) 
+	{
+		addzbuffloatImBuf(ibuf);
+	}
+	BLI_unlock_thread(LOCK_DRAW_IMAGE);
+	
 	
 	/* now we combine the input with ibuf */
 	this->m_outputBuffer = ibuf->rect_float;
 	this->m_outputBufferDisplay = (unsigned char *)ibuf->rect;
+	if (m_doDepthBuffer)
+	{
+		this->m_depthBuffer = ibuf->zbuf_float;
+	}
 	
 	BKE_image_release_ibuf(this->m_image, this->m_lock);
 }
diff --git a/source/blender/compositor/operations/COM_ViewerBaseOperation.h b/source/blender/compositor/operations/COM_ViewerBaseOperation.h
index f3fd1e9c9df..d90eb343f6c 100644
--- a/source/blender/compositor/operations/COM_ViewerBaseOperation.h
+++ b/source/blender/compositor/operations/COM_ViewerBaseOperation.h
@@ -29,6 +29,7 @@
 class ViewerBaseOperation : public NodeOperation {
 protected:
 	float *m_outputBuffer;
+	float *m_depthBuffer;
 	unsigned char *m_outputBufferDisplay;
 	Image *m_image;
 	ImageUser *m_imageUser;
@@ -39,6 +40,7 @@ protected:
 	OrderOfChunks m_chunkOrder;
 	bool m_doColorManagement;
 	bool m_doColorPredivide;
+	bool m_doDepthBuffer;
 
 public:
 	bool isOutputOperation(bool rendering) const { return isActiveViewerOutput(); }
diff --git a/source/blender/compositor/operations/COM_ViewerOperation.cpp b/source/blender/compositor/operations/COM_ViewerOperation.cpp
index f7c2ff93b3e..fac90ba2a9e 100644
--- a/source/blender/compositor/operations/COM_ViewerOperation.cpp
+++ b/source/blender/compositor/operations/COM_ViewerOperation.cpp
@@ -43,9 +43,11 @@ ViewerOperation::ViewerOperation() : ViewerBaseOperation()
 {
 	this->addInputSocket(COM_DT_COLOR);
 	this->addInputSocket(COM_DT_VALUE);
+	this->addInputSocket(COM_DT_VALUE);
 
 	this->m_imageInput = NULL;
 	this->m_alphaInput = NULL;
+	this->m_depthInput = NULL;
 }
 
 void ViewerOperation::initExecution()
@@ -53,6 +55,8 @@ void ViewerOperation::initExecution()
 	// When initializing the tree during initial load the width and height can be zero.
 	this->m_imageInput = getInputSocketReader(0);
 	this->m_alphaInput = getInputSocketReader(1);
+	this->m_depthInput = getInputSocketReader(2);
+	this->m_doDepthBuffer = (this->m_depthInput != NULL);
 	ViewerBaseOperation::initExecution();
 }
 
@@ -60,6 +64,7 @@ void ViewerOperation::deinitExecution()
 {
 	this->m_imageInput = NULL;
 	this->m_alphaInput = NULL;
+	this->m_depthInput = NULL;
 	ViewerBaseOperation::deinitExecution();
 }
 
@@ -67,47 +72,55 @@ void ViewerOperation::deinitExecution()
 void ViewerOperation::executeRegion(rcti *rect, unsigned int tileNumber)
 {
 	float *buffer = this->m_outputBuffer;
+	float *depthbuffer = this->m_depthBuffer;
 	unsigned char *bufferDisplay = this->m_outputBufferDisplay;
 	if (!buffer) return;
 	const int x1 = rect->xmin;
 	const int y1 = rect->ymin;
 	const int x2 = rect->xmax;
 	const int y2 = rect->ymax;
-	const int offsetadd = (this->getWidth() - (x2 - x1)) * 4;
-	int offset = (y1 * this->getWidth() + x1) * 4;
-	float alpha[4], srgb[4];
+	const int offsetadd = (this->getWidth() - (x2 - x1));
+	const int offsetadd4 = offsetadd * 4;
+	int offset = (y1 * this->getWidth() + x1);
+	int offset4 = offset * 4;
+	float alpha[4], srgb[4], depth[4];
 	int x;
 	int y;
 	bool breaked = false;
 
 	for (y = y1; y < y2 && (!breaked); y++) {
 		for (x = x1; x < x2; x++) {
-			this->m_imageInput->read(&(buffer[offset]), x, y, COM_PS_NEAREST);
+			this->m_imageInput->read(&(buffer[offset4]), x, y, COM_PS_NEAREST);
 			if (this->m_alphaInput != NULL) {
 				this->m_alphaInput->read(alpha, x, y, COM_PS_NEAREST);
-				buffer[offset + 3] = alpha[0];
+				buffer[offset4 + 3] = alpha[0];
 			}
+			if (m_depthInput) {
+				this->m_depthInput->read(depth, x, y, COM_PS_NEAREST);
+				depthbuffer[offset] = depth[0];
+			} 
 			if (this->m_doColorManagement) {
 				if (this->m_doColorPredivide) {
-					linearrgb_to_srgb_predivide_v4(srgb, buffer + offset);
+					linearrgb_to_srgb_predivide_v4(srgb, buffer + offset4);
 				}
 				else {
-					linearrgb_to_srgb_v4(srgb, buffer + offset);
+					linearrgb_to_srgb_v4(srgb, buffer + offset4);
 				}
 			}
 			else {
-				copy_v4_v4(srgb, buffer + offset);
+				copy_v4_v4(srgb, buffer + offset4);
 			}
 
-			rgba_float_to_uchar(bufferDisplay + offset, srgb);
+			rgba_float_to_uchar(bufferDisplay + offset4, srgb);
 
-			offset += 4;
+			offset ++;
+			offset4 += 4;
 		}
 		if (isBreaked()) {
 			breaked = true;
 		}
-
 		offset += offsetadd;
+		offset4 += offsetadd4;
 	}
 	updateImage(rect);
 }
diff --git a/source/blender/compositor/operations/COM_ViewerOperation.h b/source/blender/compositor/operations/COM_ViewerOperation.h
index d900d8db408..262efd87dba 100644
--- a/source/blender/compositor/operations/COM_ViewerOperation.h
+++ b/source/blender/compositor/operations/COM_ViewerOperation.h
@@ -31,6 +31,7 @@ class ViewerOperation : public ViewerBaseOperation {
 private:
 	SocketReader *m_imageInput;
 	SocketReader *m_alphaInput;
+	SocketReader *m_depthInput;
 
 public:
 	ViewerOperation();