Draw individual face's material and shading correctly in the PBVH.

Previously, the shading and material was set once per PBVHNode when
drawing. This is still the case, but PBVHNodes are now built to
contain only one material and shading mode.

This is done with an extra partitioning step; once the number of
primitives in the node falls below the PBVH leaf limit, it's
primitives are checked for matching materials. If more than one
material or shading mode is present in the node, it is split and
partitioned (partitioned by material rather than 3D location.)

Given a sufficiently 'annoying' input, like a dense mesh with
thousands of materials randomly scattered across it, this could
greatly increase PBVH build time (since nodes might end up containing
a single primitive), but in general this shouldn't come up.

In order to support materials for grids, the CCGDM is building another
grid array (of DMFaceMat structs). This could be used to replace
CCGDM.faceFlag for some small memory savings (TODO).

8 files changed, 201 insertions, 75 deletions

diff --git a/source/blender/blenkernel/BKE_DerivedMesh.h b/source/blender/blenkernel/BKE_DerivedMesh.h
index b9e2fddb895..8c03624dbf7 100644
--- a/source/blender/blenkernel/BKE_DerivedMesh.h
+++ b/source/blender/blenkernel/BKE_DerivedMesh.h
@@ -112,6 +112,12 @@ typedef struct DMGridAdjacency {
 	int rotation[4];
 } DMGridAdjacency;
 
+/* keep in sync with MFace/MPoly types */
+typedef struct DMFlagMat {
+	short mat_nr;
+	char flag;
+} DMFlagMat;
+
 typedef enum DerivedMeshType {
 	DM_TYPE_CDDM,
 	DM_TYPE_EDITBMESH,
@@ -218,6 +224,8 @@ struct DerivedMesh {
 	DMGridData **(*getGridData)(DerivedMesh *dm);
 	DMGridAdjacency *(*getGridAdjacency)(DerivedMesh *dm);
 	int *(*getGridOffset)(DerivedMesh *dm);
+	DMFlagMat *(*getGridFlagMats)(DerivedMesh *dm);
+	
 
 	/* Iterate over each mapped vertex in the derived mesh, calling the
 	 * given function with the original vert and the mapped vert's new
diff --git a/source/blender/blenkernel/BKE_subsurf.h b/source/blender/blenkernel/BKE_subsurf.h
index a85ef88473b..f3cdd858e7d 100644
--- a/source/blender/blenkernel/BKE_subsurf.h
+++ b/source/blender/blenkernel/BKE_subsurf.h
@@ -31,7 +31,7 @@
  *  \ingroup bke
  */
 
-struct DMFaceFlags;
+struct DMFlagMat;
 struct DMGridAdjacency;
 struct DMGridData;
 struct DerivedMesh;
@@ -76,7 +76,7 @@ typedef struct CCGDerivedMesh {
 			int startFace; struct CCGFace *face;} *faceMap;
 
 	short *edgeFlags;
-	struct DMFaceFlags *faceFlags;
+	struct DMFlagMat *faceFlags;
 
 	int *reverseFaceMap;
 
@@ -91,6 +91,7 @@ typedef struct CCGDerivedMesh {
 	struct DMGridAdjacency *gridAdjacency;
 	int *gridOffset;
 	struct CCGFace **gridFaces;
+	struct DMFlagMat *gridFlagMats;
 
 	struct {
 		struct MultiresModifierData *mmd;
diff --git a/source/blender/blenkernel/intern/cdderivedmesh.c b/source/blender/blenkernel/intern/cdderivedmesh.c
index 04d438ebf0c..b868d6fbcea 100644
--- a/source/blender/blenkernel/intern/cdderivedmesh.c
+++ b/source/blender/blenkernel/intern/cdderivedmesh.c
@@ -529,8 +529,7 @@ static void cdDM_drawFacesSolid(DerivedMesh *dm,
 			if(!setMaterial(mface->mat_nr+1, NULL))
 				return;
 
-			glShadeModel((mface->flag & ME_SMOOTH)? GL_SMOOTH: GL_FLAT);
-			BLI_pbvh_draw(cddm->pbvh, partial_redraw_planes, face_nors, (mface->flag & ME_SMOOTH));
+			BLI_pbvh_draw(cddm->pbvh, partial_redraw_planes, face_nors, setMaterial);
 			glShadeModel(GL_FLAT);
 		}
 
diff --git a/source/blender/blenkernel/intern/subsurf_ccg.c b/source/blender/blenkernel/intern/subsurf_ccg.c
index 765684e36f1..56f998cf2aa 100644
--- a/source/blender/blenkernel/intern/subsurf_ccg.c
+++ b/source/blender/blenkernel/intern/subsurf_ccg.c
@@ -71,12 +71,6 @@
 
 #include "CCGSubSurf.h"
 
-/* keep in sync with MPoly types */
-typedef struct DMFaceFlags {
-	short mat_nr;
-	char flag;
-} DMFaceFlags;
-
 extern GLubyte stipple_quarttone[128]; /* glutil.c, bad level data */
 
 static CCGDerivedMesh *getCCGDerivedMesh(CCGSubSurf *ss,
@@ -926,7 +920,7 @@ static void ccgDM_getFinalFace(DerivedMesh *dm, int faceNum, MFace *mf)
 	int grid;
 	int x, y;
 	/*int lastface = ccgSubSurf_getNumFaces(ss) - 1;*/ /*UNUSED*/
-	DMFaceFlags *faceFlags = ccgdm->faceFlags;
+	DMFlagMat *faceFlags = ccgdm->faceFlags;
 
 	memset(mf, 0, sizeof(*mf));
 	if (faceNum >= ccgdm->dm.numTessFaceData)
@@ -1113,7 +1107,7 @@ static void ccgDM_copyFinalFaceArray(DerivedMesh *dm, MFace *mface)
 	int gridSize = ccgSubSurf_getGridSize(ss);
 	int edgeSize = ccgSubSurf_getEdgeSize(ss);
 	int i = 0;
-	DMFaceFlags *faceFlags = ccgdm->faceFlags;
+	DMFlagMat *faceFlags = ccgdm->faceFlags;
 
 	totface = ccgSubSurf_getNumFaces(ss);
 	for(index = 0; index < totface; index++) {
@@ -1155,7 +1149,7 @@ static void ccgDM_copyFinalLoopArray(DerivedMesh *dm, MLoop *mloop)
 	int edgeSize = ccgSubSurf_getEdgeSize(ss);
 	int i = 0;
 	MLoop *mv;
-	/* DMFaceFlags *faceFlags = ccgdm->faceFlags; */ /* UNUSED */
+	/* DMFlagMat *faceFlags = ccgdm->faceFlags; */ /* UNUSED */
 
 	if (!ccgdm->ehash) {
 		MEdge *medge;
@@ -1221,7 +1215,7 @@ static void ccgDM_copyFinalPolyArray(DerivedMesh *dm, MPoly *mface)
 	int gridSize = ccgSubSurf_getGridSize(ss);
 	/* int edgeSize = ccgSubSurf_getEdgeSize(ss); */ /* UNUSED */
 	int i = 0, k = 0;
-	DMFaceFlags *faceFlags = ccgdm->faceFlags;
+	DMFlagMat *faceFlags = ccgdm->faceFlags;
 
 	totface = ccgSubSurf_getNumFaces(ss);
 	for(index = 0; index < totface; index++) {
@@ -1545,7 +1539,7 @@ static void ccgDM_drawFacesSolid(DerivedMesh *dm, float (*partial_redraw_planes)
 	CCGDerivedMesh *ccgdm = (CCGDerivedMesh*) dm;
 	CCGSubSurf *ss = ccgdm->ss;
 	int gridSize = ccgSubSurf_getGridSize(ss);
-	DMFaceFlags *faceFlags = ccgdm->faceFlags;
+	DMFlagMat *faceFlags = ccgdm->faceFlags;
 	int step = (fast)? gridSize-1: 1;
 	int i, totface = ccgSubSurf_getNumFaces(ss);
 	int drawcurrent = 0, matnr = -1, shademodel = -1;
@@ -1558,8 +1552,7 @@ static void ccgDM_drawFacesSolid(DerivedMesh *dm, float (*partial_redraw_planes)
 			if(!setMaterial(faceFlags[0].mat_nr+1, NULL))
 				return;
 
-			glShadeModel((faceFlags[0].flag & ME_SMOOTH)? GL_SMOOTH: GL_FLAT);
-			BLI_pbvh_draw(ccgdm->pbvh, partial_redraw_planes, NULL, (faceFlags[0].flag & ME_SMOOTH));
+			BLI_pbvh_draw(ccgdm->pbvh, partial_redraw_planes, NULL, setMaterial);
 			glShadeModel(GL_FLAT);
 		}
 
@@ -1647,7 +1640,7 @@ static void ccgDM_drawMappedFacesGLSL(DerivedMesh *dm,
 	int gridSize = ccgSubSurf_getGridSize(ss);
 	int gridFaces = gridSize - 1;
 	int edgeSize = ccgSubSurf_getEdgeSize(ss);
-	DMFaceFlags *faceFlags = ccgdm->faceFlags;
+	DMFlagMat *faceFlags = ccgdm->faceFlags;
 	int a, b, i, doDraw, numVerts, matnr, new_matnr, totface;
 
 	ccgdm_pbvh_update(ccgdm);
@@ -1790,7 +1783,7 @@ static void ccgDM_drawMappedFacesMat(DerivedMesh *dm, void (*setMaterial)(void *
 	int gridSize = ccgSubSurf_getGridSize(ss);
 	int gridFaces = gridSize - 1;
 	int edgeSize = ccgSubSurf_getEdgeSize(ss);
-	DMFaceFlags *faceFlags = ccgdm->faceFlags;
+	DMFlagMat *faceFlags = ccgdm->faceFlags;
 	int a, b, i, numVerts, matnr, new_matnr, totface;
 
 	ccgdm_pbvh_update(ccgdm);
@@ -1936,7 +1929,7 @@ static void ccgDM_drawFacesTex_common(DerivedMesh *dm,
 	CCGSubSurf *ss = ccgdm->ss;
 	MCol *mcol = dm->getTessFaceDataArray(dm, CD_WEIGHT_MCOL);
 	MTFace *tf = DM_get_tessface_data_layer(dm, CD_MTFACE);
-	DMFaceFlags *faceFlags = ccgdm->faceFlags;
+	DMFlagMat *faceFlags = ccgdm->faceFlags;
 	int i, totface, flag, gridSize = ccgSubSurf_getGridSize(ss);
 	int gridFaces = gridSize - 1;
 
@@ -2131,7 +2124,7 @@ static void ccgDM_drawMappedFaces(DerivedMesh *dm,
 	CCGSubSurf *ss = ccgdm->ss;
 	MCol *mcol= NULL;
 	int i, gridSize = ccgSubSurf_getGridSize(ss);
-	DMFaceFlags *faceFlags = ccgdm->faceFlags;
+	DMFlagMat *faceFlags = ccgdm->faceFlags;
 	int gridFaces = gridSize - 1, totface;
 
 	/* currently unused -- each original face is handled separately */
@@ -2356,6 +2349,7 @@ static void ccgDM_release(DerivedMesh *dm)
 		if(ccgdm->gridData) MEM_freeN(ccgdm->gridData);
 		if(ccgdm->gridAdjacency) MEM_freeN(ccgdm->gridAdjacency);
 		if(ccgdm->gridOffset) MEM_freeN(ccgdm->gridOffset);
+		if(ccgdm->gridFlagMats) MEM_freeN(ccgdm->gridFlagMats);
 		if(ccgdm->freeSS) ccgSubSurf_free(ccgdm->ss);
 		if(ccgdm->fmap) MEM_freeN(ccgdm->fmap);
 		if(ccgdm->fmap_mem) MEM_freeN(ccgdm->fmap_mem);
@@ -2633,6 +2627,7 @@ static void ccgdm_create_grids(DerivedMesh *dm)
 	CCGSubSurf *ss= ccgdm->ss;
 	DMGridData **gridData;
 	DMGridAdjacency *gridAdjacency, *adj;
+	DMFlagMat *gridFlagMats;
 	CCGFace **gridFaces;
 	int *gridOffset;
 	int index, numFaces, numGrids, S, gIndex /*, gridSize*/;
@@ -2659,6 +2654,7 @@ static void ccgdm_create_grids(DerivedMesh *dm)
 	gridData = MEM_mallocN(sizeof(DMGridData*)*numGrids, "ccgdm.gridData");
 	gridAdjacency = MEM_mallocN(sizeof(DMGridAdjacency)*numGrids, "ccgdm.gridAdjacency");
 	gridFaces = MEM_mallocN(sizeof(CCGFace*)*numGrids, "ccgdm.gridFaces");
+	gridFlagMats = MEM_mallocN(sizeof(DMFlagMat)*numGrids, "ccgdm.gridFlagMats");
 
 	for(gIndex = 0, index = 0; index < numFaces; index++) {
 		CCGFace *f = ccgdm->faceMap[index].face;
@@ -2670,6 +2666,7 @@ static void ccgdm_create_grids(DerivedMesh *dm)
 
 			gridData[gIndex] = ccgSubSurf_getFaceGridDataArray(ss, f, S);
 			gridFaces[gIndex] = f;
+			gridFlagMats[gIndex] = ccgdm->faceFlags[index];
 
 			adj = &gridAdjacency[gIndex];
 
@@ -2688,6 +2685,7 @@ static void ccgdm_create_grids(DerivedMesh *dm)
 	ccgdm->gridFaces = gridFaces;
 	ccgdm->gridAdjacency = gridAdjacency;
 	ccgdm->gridOffset = gridOffset;
+	ccgdm->gridFlagMats = gridFlagMats;
 }
 
 static DMGridData **ccgDM_getGridData(DerivedMesh *dm)
@@ -2714,6 +2712,14 @@ static int *ccgDM_getGridOffset(DerivedMesh *dm)
 	return ccgdm->gridOffset;
 }
 
+static DMFlagMat *ccgDM_getGridFlagMats(DerivedMesh *dm)
+{
+	CCGDerivedMesh *ccgdm= (CCGDerivedMesh*)dm;
+	
+	ccgdm_create_grids(dm);
+	return ccgdm->gridFlagMats;
+}
+
 static ListBase *ccgDM_getPolyMap(Object *ob, DerivedMesh *dm)
 {
 	CCGDerivedMesh *ccgdm= (CCGDerivedMesh*)dm;
@@ -2783,7 +2789,7 @@ static struct PBVH *ccgDM_getPBVH(Object *ob, DerivedMesh *dm)
 
 		ob->sculpt->pbvh= ccgdm->pbvh = BLI_pbvh_new();
 		BLI_pbvh_build_grids(ccgdm->pbvh, ccgdm->gridData, ccgdm->gridAdjacency,
-			numGrids, gridSize, (void**)ccgdm->gridFaces);
+			numGrids, gridSize, (void**)ccgdm->gridFaces, ccgdm->gridFlagMats);
 	} else if(ob->type == OB_MESH) {
 		Mesh *me= ob->data;
 		ob->sculpt->pbvh= ccgdm->pbvh = BLI_pbvh_new();
@@ -2819,7 +2825,7 @@ static CCGDerivedMesh *getCCGDerivedMesh(CCGSubSurf *ss,
 	int vertNum, edgeNum, faceNum;
 	int *vertOrigIndex, *faceOrigIndex, *polyOrigIndex, *base_polyOrigIndex; /* *edgeOrigIndex - as yet, unused  */
 	short *edgeFlags;
-	DMFaceFlags *faceFlags;
+	DMFlagMat *faceFlags;
 	int *loopidx = NULL, *vertidx = NULL;
 	BLI_array_declare(loopidx);
 	BLI_array_declare(vertidx);
@@ -2895,6 +2901,7 @@ static CCGDerivedMesh *getCCGDerivedMesh(CCGSubSurf *ss,
 	ccgdm->dm.getGridData = ccgDM_getGridData;
 	ccgdm->dm.getGridAdjacency = ccgDM_getGridAdjacency;
 	ccgdm->dm.getGridOffset = ccgDM_getGridOffset;
+	ccgdm->dm.getGridFlagMats = ccgDM_getGridFlagMats;
 	ccgdm->dm.getPolyMap = ccgDM_getPolyMap;
 	ccgdm->dm.getPBVH = ccgDM_getPBVH;
 
@@ -2984,7 +2991,7 @@ static CCGDerivedMesh *getCCGDerivedMesh(CCGSubSurf *ss,
 	
 	/*CDDM hack*/
 	edgeFlags = ccgdm->edgeFlags = MEM_callocN(sizeof(short)*totedge, "edgeFlags");
-	faceFlags = ccgdm->faceFlags = MEM_callocN(sizeof(DMFaceFlags)*totface, "faceFlags");
+	faceFlags = ccgdm->faceFlags = MEM_callocN(sizeof(DMFlagMat)*totface, "faceFlags");
 
 	vertOrigIndex = DM_get_vert_data_layer(&ccgdm->dm, CD_ORIGINDEX);
 	/*edgeOrigIndex = DM_get_edge_data_layer(&ccgdm->dm, CD_ORIGINDEX);*/
diff --git a/source/blender/blenlib/BLI_pbvh.h b/source/blender/blenlib/BLI_pbvh.h
index 85dd1c1a197..f0d56a3cf52 100644
--- a/source/blender/blenlib/BLI_pbvh.h
+++ b/source/blender/blenlib/BLI_pbvh.h
@@ -26,13 +26,14 @@
  *  \brief A BVH for high poly meshes.
  */
 
-struct MFace;
-struct MVert;
+struct DMFlagMat;
 struct DMGridAdjacency;
 struct DMGridData;
+struct ListBase;
+struct MFace;
+struct MVert;
 struct PBVH;
 struct PBVHNode;
-struct ListBase;
 
 typedef struct PBVH PBVH;
 typedef struct PBVHNode PBVHNode;
@@ -56,7 +57,7 @@ void BLI_pbvh_build_mesh(PBVH *bvh, struct MFace *faces, struct MVert *verts,
 			int totface, int totvert);
 void BLI_pbvh_build_grids(PBVH *bvh, struct DMGridData **grids,
 	struct DMGridAdjacency *gridadj, int totgrid,
-	int gridsize, void **gridfaces);
+	int gridsize, void **gridfaces, struct DMFlagMat *flagmats);
 void BLI_pbvh_free(PBVH *bvh);
 
 /* Hierarchical Search in the BVH, two methods:
@@ -85,7 +86,8 @@ int BLI_pbvh_node_raycast(PBVH *bvh, PBVHNode *node, float (*origco)[3],
 
 void BLI_pbvh_node_draw(PBVHNode *node, void *data);
 int BLI_pbvh_node_planes_contain_AABB(PBVHNode *node, void *data);
-void BLI_pbvh_draw(PBVH *bvh, float (*planes)[4], float (*face_nors)[3], int smooth);
+void BLI_pbvh_draw(PBVH *bvh, float (*planes)[4], float (*face_nors)[3],
+				   int (*setMaterial)(int, void *attribs));
 
 /* Node Access */
 
diff --git a/source/blender/blenlib/intern/pbvh.c b/source/blender/blenlib/intern/pbvh.c
index f8562df07a8..46d38e10585 100644
--- a/source/blender/blenlib/intern/pbvh.c
+++ b/source/blender/blenlib/intern/pbvh.c
@@ -136,6 +136,7 @@ struct PBVH {
 	DMGridData **grids;
 	DMGridAdjacency *gridadj;
 	void **gridfaces;
+	const DMFlagMat *grid_flag_mats;
 	int totgrid;
 	int gridsize;
 
@@ -261,6 +262,17 @@ static void update_node_vb(PBVH *bvh, PBVHNode *node)
 //	BB_expand(&node->vb, co);
 //}
 
+static int face_materials_match(const MFace *f1, const MFace *f2)
+{
+	return ((f1->flag & ME_SMOOTH) == (f2->flag & ME_SMOOTH) &&
+			(f1->mat_nr == f2->mat_nr));
+}
+
+static int grid_materials_match(const DMFlagMat *f1, const DMFlagMat *f2)
+{
+	return ((f1->flag & ME_SMOOTH) == (f2->flag & ME_SMOOTH) &&
+			(f1->mat_nr == f2->mat_nr));
+}
 
 /* Adapted from BLI_kdopbvh.c */
 /* Returns the index of the first element on the right of the partition */
@@ -280,6 +292,38 @@ static int partition_indices(int *prim_indices, int lo, int hi, int axis,
 	}
 }
 
+/* Returns the index of the first element on the right of the partition */
+static int partition_indices_material(PBVH *bvh, int lo, int hi)
+{
+	const MFace *faces = bvh->faces;
+	const DMFlagMat *flagmats = bvh->grid_flag_mats;
+	const int *indices = bvh->prim_indices;
+	const void *first;
+	int i=lo, j=hi;
+
+	if(bvh->faces)
+		first = &faces[bvh->prim_indices[lo]];
+	else
+		first = &flagmats[bvh->prim_indices[lo]];
+
+	for(;;) {
+		if(bvh->faces) {
+			for(; face_materials_match(first, &faces[indices[i]]); i++);
+			for(; !face_materials_match(first, &faces[indices[j]]); j--);
+		}
+		else {
+			for(; grid_materials_match(first, &flagmats[indices[i]]); i++);
+			for(; !grid_materials_match(first, &flagmats[indices[j]]); j--);
+		}
+		
+		if(!(i < j))
+			return i;
+
+		SWAP(int, bvh->prim_indices[i], bvh->prim_indices[j]);
+		i++;
+	}
+}
+
 static void grow_nodes(PBVH *bvh, int totnode)
 {
 	if(totnode > bvh->node_mem_count) {
@@ -431,6 +475,38 @@ static void build_leaf(PBVH *bvh, int node_index, BBC *prim_bbc,
 		build_grids_leaf_node(bvh, bvh->nodes + node_index);
 }
 
+/* Return zero if all primitives in the node can be drawn with the
+   same material (including flat/smooth shading), non-zerootherwise */
+int leaf_needs_material_split(PBVH *bvh, int offset, int count)
+{
+	int i, prim;
+
+	if(count <= 1)
+		return 0;
+
+	if(bvh->faces) {
+		const MFace *first = &bvh->faces[bvh->prim_indices[offset]];
+
+		for(i = offset + count - 1; i > offset; --i) {
+			prim = bvh->prim_indices[i];
+			if(!face_materials_match(first, &bvh->faces[prim]))
+				return 1;
+		}
+	}
+	else {
+		const DMFlagMat *first = &bvh->grid_flag_mats[bvh->prim_indices[offset]];
+
+		for(i = offset + count - 1; i > offset; --i) {
+			prim = bvh->prim_indices[i];
+			if(!grid_materials_match(first, &bvh->grid_flag_mats[prim]))
+				return 1;
+		}
+	}
+
+	return 0;
+}
+
+
 /* Recursively build a node in the tree
  *
  * vb is the voxel box around all of the primitives contained in
@@ -443,15 +519,18 @@ static void build_leaf(PBVH *bvh, int node_index, BBC *prim_bbc,
  */
 
 static void build_sub(PBVH *bvh, int node_index, BB *cb, BBC *prim_bbc,
-		   int offset, int count)
+					  int offset, int count)
 {
-	int i, axis, end;
+	int i, axis, end, below_leaf_limit;
 	BB cb_backing;
 
 	/* Decide whether this is a leaf or not */
-	if(count <= bvh->leaf_limit) {
-		build_leaf(bvh, node_index, prim_bbc, offset, count);
-		return;
+	below_leaf_limit = count <= bvh->leaf_limit;
+	if(below_leaf_limit) {
+		if(!leaf_needs_material_split(bvh, offset, count)) {
+			build_leaf(bvh, node_index, prim_bbc, offset, count);
+			return;
+		}
 	}
 
 	/* Add two child nodes */
@@ -461,26 +540,33 @@ static void build_sub(PBVH *bvh, int node_index, BB *cb, BBC *prim_bbc,
 	/* Update parent node bounding box */
 	update_vb(bvh, &bvh->nodes[node_index], prim_bbc, offset, count);
 
-	/* Find axis with widest range of primitive centroids */
-	if(!cb) {
-		cb = &cb_backing;
-		BB_reset(cb);
-		for(i = offset + count - 1; i >= offset; --i)
-			BB_expand(cb, prim_bbc[bvh->prim_indices[i]].bcentroid);
+	if(!below_leaf_limit) {
+		/* Find axis with widest range of primitive centroids */
+		if(!cb) {
+			cb = &cb_backing;
+			BB_reset(cb);
+			for(i = offset + count - 1; i >= offset; --i)
+				BB_expand(cb, prim_bbc[bvh->prim_indices[i]].bcentroid);
+		}
+		axis = BB_widest_axis(cb);
+
+		/* Partition primitives along that axis */
+		end = partition_indices(bvh->prim_indices,
+								offset, offset + count - 1,
+								axis,
+								(cb->bmax[axis] + cb->bmin[axis]) * 0.5f,
+								prim_bbc);
+	}
+	else {
+		/* Partition primitives by material */
+		end = partition_indices_material(bvh, offset, offset + count - 1);
 	}
-	axis = BB_widest_axis(cb);
-
-	/* Partition primitives along that axis */
-	end = partition_indices(bvh->prim_indices, offset, offset + count - 1,
-	                        axis,
-	                        (cb->bmax[axis] + cb->bmin[axis]) * 0.5f,
-	                        prim_bbc);
 
 	/* Build children */
 	build_sub(bvh, bvh->nodes[node_index].children_offset, NULL,
-		  prim_bbc, offset, end - offset);
+			  prim_bbc, offset, end - offset);
 	build_sub(bvh, bvh->nodes[node_index].children_offset + 1, NULL,
-		  prim_bbc, end, offset + count - end);
+			  prim_bbc, end, offset + count - end);
 }
 
 static void pbvh_build(PBVH *bvh, BB *cb, BBC *prim_bbc, int totprim)
@@ -550,7 +636,7 @@ void BLI_pbvh_build_mesh(PBVH *bvh, MFace *faces, MVert *verts, int totface, int
 
 /* Do a full rebuild with on Grids data structure */
 void BLI_pbvh_build_grids(PBVH *bvh, DMGridData **grids, DMGridAdjacency *gridadj,
-	int totgrid, int gridsize, void **gridfaces)
+	int totgrid, int gridsize, void **gridfaces, DMFlagMat *flagmats)
 {
 	BBC *prim_bbc = NULL;
 	BB cb;
@@ -559,6 +645,7 @@ void BLI_pbvh_build_grids(PBVH *bvh, DMGridData **grids, DMGridAdjacency *gridad
 	bvh->grids= grids;
 	bvh->gridadj= gridadj;
 	bvh->gridfaces= gridfaces;
+	bvh->grid_flag_mats= flagmats;
 	bvh->totgrid= totgrid;
 	bvh->gridsize= gridsize;
 	bvh->leaf_limit = MAX2(LEAF_LIMIT/((gridsize-1)*(gridsize-1)), 1);
@@ -1027,7 +1114,7 @@ static void pbvh_update_BB_redraw(PBVH *bvh, PBVHNode **nodes,
 	}
 }
 
-static void pbvh_update_draw_buffers(PBVH *bvh, PBVHNode **nodes, int totnode, int smooth)
+static void pbvh_update_draw_buffers(PBVH *bvh, PBVHNode **nodes, int totnode)
 {
 	PBVHNode *node;
 	int n;
@@ -1040,18 +1127,17 @@ static void pbvh_update_draw_buffers(PBVH *bvh, PBVHNode **nodes, int totnode, i
 			if(bvh->grids) {
 				GPU_update_grid_buffers(node->draw_buffers,
 						   bvh->grids,
+						   bvh->grid_flag_mats,
 						   node->prim_indices,
 						   node->totprim,
-						   bvh->gridsize,
-						   smooth);
+						   bvh->gridsize);
 			}
 			else {
 				GPU_update_mesh_buffers(node->draw_buffers,
 						   bvh->verts,
 						   node->vert_indices,
 						   node->uniq_verts +
-						   node->face_verts,
-		                   smooth);
+						   node->face_verts);
 			}
 
 			node->flag &= ~PBVH_UpdateDrawBuffers;
@@ -1420,7 +1506,7 @@ int BLI_pbvh_node_raycast(PBVH *bvh, PBVHNode *node, float (*origco)[3],
 
 //#include <GL/glew.h>
 
-void BLI_pbvh_node_draw(PBVHNode *node, void *UNUSED(data))
+void BLI_pbvh_node_draw(PBVHNode *node, void *setMaterial)
 {
 #if 0
 	/* XXX: Just some quick code to show leaf nodes in different colors */
@@ -1438,7 +1524,7 @@ void BLI_pbvh_node_draw(PBVHNode *node, void *UNUSED(data))
 
 	glColor3f(1, 0, 0);
 #endif
-	GPU_draw_buffers(node->draw_buffers);
+	GPU_draw_buffers(node->draw_buffers, setMaterial);
 }
 
 /* Adapted from:
@@ -1473,7 +1559,8 @@ int BLI_pbvh_node_planes_contain_AABB(PBVHNode *node, void *data)
 	return 1;
 }
 
-void BLI_pbvh_draw(PBVH *bvh, float (*planes)[4], float (*face_nors)[3], int smooth)
+void BLI_pbvh_draw(PBVH *bvh, float (*planes)[4], float (*face_nors)[3],
+				   int (*setMaterial)(int, void *attribs))
 {
 	PBVHNode **nodes;
 	int totnode;
@@ -1482,7 +1569,7 @@ void BLI_pbvh_draw(PBVH *bvh, float (*planes)[4], float (*face_nors)[3], int smo
 		&nodes, &totnode);
 
 	pbvh_update_normals(bvh, nodes, totnode, face_nors);
-	pbvh_update_draw_buffers(bvh, nodes, totnode, smooth);
+	pbvh_update_draw_buffers(bvh, nodes, totnode);
 
 	if(nodes) MEM_freeN(nodes);
 
@@ -1491,7 +1578,7 @@ void BLI_pbvh_draw(PBVH *bvh, float (*planes)[4], float (*face_nors)[3], int smo
 				planes, BLI_pbvh_node_draw, NULL);
 	}
 	else {
-		BLI_pbvh_search_callback(bvh, NULL, NULL, BLI_pbvh_node_draw, NULL);
+		BLI_pbvh_search_callback(bvh, NULL, NULL, BLI_pbvh_node_draw, setMaterial);
 	}
 }
 
diff --git a/source/blender/gpu/GPU_buffers.h b/source/blender/gpu/GPU_buffers.h
index b89219c2d56..2e260ab8a2e 100644
--- a/source/blender/gpu/GPU_buffers.h
+++ b/source/blender/gpu/GPU_buffers.h
@@ -40,6 +40,7 @@
 #endif
 
 struct DerivedMesh;
+struct DMFlagMat;
 struct DMGridData;
 struct GHash;
 struct DMGridData;
@@ -161,14 +162,15 @@ GPU_Buffers *GPU_build_mesh_buffers(int (*face_vert_indices)[4],
 			struct MFace *mface, int *face_indices, int totface);
 
 void GPU_update_mesh_buffers(GPU_Buffers *buffers, struct MVert *mvert,
-			int *vert_indices, int totvert, int smooth);
+			int *vert_indices, int totvert);
 
 GPU_Buffers *GPU_build_grid_buffers(int totgrid, int gridsize);
 
 void GPU_update_grid_buffers(GPU_Buffers *buffers, struct DMGridData **grids,
-	int *grid_indices, int totgrid, int gridsize, int smooth);
+							 const struct DMFlagMat *grid_flag_mats,
+							 int *grid_indices, int totgrid, int gridsize);
 
-void GPU_draw_buffers(GPU_Buffers *buffers);
+void GPU_draw_buffers(GPU_Buffers *buffers, int (*setMaterial)(int, void *attribs));
 
 void GPU_free_buffers(GPU_Buffers *buffers);
 
diff --git a/source/blender/gpu/intern/gpu_buffers.c b/source/blender/gpu/intern/gpu_buffers.c
index 6e7ec98e65b..cd931450497 100644
--- a/source/blender/gpu/intern/gpu_buffers.c
+++ b/source/blender/gpu/intern/gpu_buffers.c
@@ -1277,16 +1277,16 @@ struct GPU_Buffers {
 
 	/* grid pointers */
 	DMGridData **grids;
+	const DMFlagMat *grid_flag_mats;
 	int *grid_indices;
 	int totgrid;
 	int gridsize;
 
 	unsigned int tot_tri, tot_quad;
-	int smooth;
 };
 
 void GPU_update_mesh_buffers(GPU_Buffers *buffers, MVert *mvert,
-			int *vert_indices, int totvert, int smooth)
+			int *vert_indices, int totvert)
 {
 	VertexBufferFormat *vert_data;
 	int i;
@@ -1319,7 +1319,6 @@ void GPU_update_mesh_buffers(GPU_Buffers *buffers, MVert *mvert,
 	}
 
 	buffers->mvert = mvert;
-	buffers->smooth = smooth;
 }
 
 GPU_Buffers *GPU_build_mesh_buffers(int (*face_vert_indices)[4],
@@ -1390,7 +1389,7 @@ GPU_Buffers *GPU_build_mesh_buffers(int (*face_vert_indices)[4],
 }
 
 void GPU_update_grid_buffers(GPU_Buffers *buffers, DMGridData **grids,
-	int *grid_indices, int totgrid, int gridsize, int smooth)
+	const DMFlagMat *grid_flag_mats, int *grid_indices, int totgrid, int gridsize)
 {
 	DMGridData *vert_data;
 	int i, j, k, totvert;
@@ -1399,6 +1398,8 @@ void GPU_update_grid_buffers(GPU_Buffers *buffers, DMGridData **grids,
 
 	/* Build VBO */
 	if(buffers->vert_buf) {
+		int smooth = grid_flag_mats[grid_indices[0]].flag & ME_SMOOTH;
+
 		glBindBufferARB(GL_ARRAY_BUFFER_ARB, buffers->vert_buf);
 		glBufferDataARB(GL_ARRAY_BUFFER_ARB,
 		                sizeof(DMGridData) * totvert,
@@ -1442,7 +1443,7 @@ void GPU_update_grid_buffers(GPU_Buffers *buffers, DMGridData **grids,
 	buffers->grid_indices = grid_indices;
 	buffers->totgrid = totgrid;
 	buffers->gridsize = gridsize;
-	buffers->smooth = smooth;
+	buffers->grid_flag_mats = grid_flag_mats;
 
 	//printf("node updated %p\n", buffers);
 }
@@ -1524,7 +1525,7 @@ GPU_Buffers *GPU_build_grid_buffers(int totgrid, int gridsize)
 
 #undef FILL_QUAD_BUFFER
 
-static void gpu_draw_buffers_legacy_mesh(GPU_Buffers *buffers)
+static void gpu_draw_buffers_legacy_mesh(GPU_Buffers *buffers, int smooth)
 {
 	const MVert *mvert = buffers->mvert;
 	int i, j;
@@ -1536,7 +1537,7 @@ static void gpu_draw_buffers_legacy_mesh(GPU_Buffers *buffers)
 
 		glBegin((f->v4)? GL_QUADS: GL_TRIANGLES);
 
-		if(buffers->smooth) {
+		if(smooth) {
 			for(j = 0; j < S; j++) {
 				glNormal3sv(mvert[fv[j]].no);
 				glVertex3fv(mvert[fv[j]].co);
@@ -1562,11 +1563,11 @@ static void gpu_draw_buffers_legacy_mesh(GPU_Buffers *buffers)
 	}
 }
 
-static void gpu_draw_buffers_legacy_grids(GPU_Buffers *buffers)
+static void gpu_draw_buffers_legacy_grids(GPU_Buffers *buffers, int smooth)
 {
 	int i, x, y, gridsize = buffers->gridsize;
 
-	if(buffers->smooth) {
+	if(smooth) {
 		for(i = 0; i < buffers->totgrid; ++i) {
 			DMGridData *grid = buffers->grids[buffers->grid_indices[i]];
 
@@ -1612,8 +1613,27 @@ static void gpu_draw_buffers_legacy_grids(GPU_Buffers *buffers)
 	}
 }
 
-void GPU_draw_buffers(GPU_Buffers *buffers)
+void GPU_draw_buffers(GPU_Buffers *buffers, int (*setMaterial)(int, void *attribs))
 {
+	int smooth = 0;
+
+	if(buffers->totface) {
+		const MFace *f = &buffers->mface[buffers->face_indices[0]];
+		if(!setMaterial(f->mat_nr+1, NULL))
+			return;
+
+		smooth = f->flag & ME_SMOOTH;
+		glShadeModel(smooth ? GL_SMOOTH: GL_FLAT);
+	}
+	else if(buffers->totgrid) {
+		const DMFlagMat *f = &buffers->grid_flag_mats[buffers->grid_indices[0]];
+		if(!setMaterial(f->mat_nr+1, NULL))
+			return;
+
+		smooth = f->flag & ME_SMOOTH;
+		glShadeModel(smooth ? GL_SMOOTH: GL_FLAT);
+	}
+
 	if(buffers->vert_buf && buffers->index_buf) {
 		glEnableClientState(GL_VERTEX_ARRAY);
 		glEnableClientState(GL_NORMAL_ARRAY);
@@ -1642,10 +1662,10 @@ void GPU_draw_buffers(GPU_Buffers *buffers)
 	}
 	/* fallbacks if we are out of memory or VBO is disabled */
 	else if(buffers->totface) {
-		gpu_draw_buffers_legacy_mesh(buffers);
+		gpu_draw_buffers_legacy_mesh(buffers, smooth);
 	}
 	else if(buffers->totgrid) {
-		gpu_draw_buffers_legacy_grids(buffers);
+		gpu_draw_buffers_legacy_grids(buffers, smooth);
 	}
 }