mybvh.cpp

#include <algorithm>
#include <cassert>
#include <vector>

#include "common\intersection.h"
#include "common\reflectcuts.h"
#include "common\accel.h"

#include "math\ray.h"
#include "math\math.h"
#include "math\aabb.h"

// written in same style as qbvh
class alignas(32) BvhNode
{
public:
	Aabb bbox[2]; // 4 * 3 * 2 = 24 bytes
	// child[0] is always current index + 1
	int32_t child[2]; // 8 bytes
}; // 32 bytes


class Bvh : public Accel
{
public:
	Bvh() {};
	bool intersect(Intersection * isectPtr, const Ray & r) const override;
	bool intersectP(const Ray & r) const override;
	void build(const std::vector<shared_ptr<Shape>> & hitablePtrs) override;
	Aabb computeBbox() const override;

	friend class Qbvh4;

private:
	bool intersectBranch(Intersection * isectPtr, const Ray & r, const size_t index) const;
	bool intersectBranchSimple(Intersection * isectPtr, const Ray & r, const size_t index) const;
	bool intersectBranchP(const Ray & r, const size_t index) const;
	int buildBranch(std::vector<std::pair<shared_ptr<Shape>, Aabb>> * hitablePairsPtr, size_t begin, size_t end, const Aabb & bound);

private:

	Aabb bound;
	std::vector<shared_ptr<Shape>> _mShapePtrs;
	std::vector<BvhNode> _mBvhNodes;
};

bool Bvh::intersect(Intersection * isectPtr, const Ray & r) const
{
	isectPtr->t = r.tmax;
	if (!bound.intersectP(r))
	{
		return false;
	}
	return this->intersectBranchSimple(isectPtr, r, 0);
}

bool Bvh::intersectP(const Ray & r) const
{
	if (!bound.intersectP(r))
	{
		return false;
	}
	return this->intersectBranchP(r, 0);
}

void Bvh::build(const std::vector<shared_ptr<Shape>>& shapePtrs)
{
	// refine all shape
	std::vector<shared_ptr<Shape>> refinedShapePtrs;
	for (size_t i = 0;i < shapePtrs.size();i++)
	{
		if (shapePtrs[i]->canIntersect())
		{
			refinedShapePtrs.push_back(shapePtrs[i]);
		}
		else
		{
			shapePtrs[i]->refine(&refinedShapePtrs);
		}
	}

	// construct shape pairs
	size_t numShapes = refinedShapePtrs.size();
	std::vector<std::pair<shared_ptr<Shape>, Aabb>> shapePairs;
	Aabb sumBound;

	for (size_t i = 0;i < numShapes;i++)
	{
		Aabb bound = refinedShapePtrs[i]->computeBbox();
		sumBound = Aabb::Union(bound, sumBound);
		shapePairs.push_back(std::pair<shared_ptr<Shape>, Aabb>(refinedShapePtrs[i], bound));
	}
	this->buildBranch(&shapePairs, 0, refinedShapePtrs.size(), sumBound);

	this->bound = sumBound;
	// populate data in _mShapePtrs
	for (size_t i = 0;i < numShapes;i++)
	{
		this->_mShapePtrs.push_back(shapePairs[i].first);
	}
}

Aabb Bvh::computeBbox() const
{
	return this->bound;
}

bool Bvh::intersectBranch(Intersection * isectPtr, const Ray & r, const size_t index) const
{
	const BvhNode & currentNode = this->_mBvhNodes[index];

	bool isIntersect = false;

	Float tNear[2];
	bool isHit[2];
	size_t childIndices[2];
	childIndices[0] = 0;
	childIndices[1] = 1;

	for (size_t i = 0;i < 2;i++)
	{
		Float tFar;
		if (currentNode.bbox[i].intersect(&tNear[i], &tFar, r))
		{
			tNear[i] = r.tmax;
			isHit[i] = true;
		}
		else
		{
			isHit[i] = false;
		}
	}

	if (tNear[0] > tNear[1])
	{
		std::swap(isHit[0], isHit[1]);
		std::swap(childIndices[0], childIndices[1]);
	}

	for (size_t i = 0;i < 2;i++)
	{
		size_t childIndex = childIndices[i];

		if (isHit[i])
		{
			// check if leaf
			if (currentNode.child[childIndex] < 0)
			{
				int32_t shapeIndex = -currentNode.child[childIndex] - 1;
				bool isHit = this->_mShapePtrs[shapeIndex]->intersect(isectPtr, Ray(r.origin, r.direction, r.tmin, isectPtr->t));
				isIntersect = isHit || isIntersect;
			}
			else if (currentNode.child[childIndex] > 0)
			{
				bool isHit = intersectBranch(isectPtr, Ray(r.origin, r.direction, r.tmin, isectPtr->t), currentNode.child[childIndex]);
				isIntersect = isHit || isIntersect;
			}
		}
	}

	return isIntersect;
}

bool Bvh::intersectBranchSimple(Intersection * isectPtr, const Ray & r, const size_t index) const
{
	const BvhNode & currentNode = this->_mBvhNodes[index];
	
	bool isIntersect = false;

	for (size_t i = 0;i < 2;i++)
	{
		size_t childIndex = i;

		if (currentNode.bbox[i].intersectP(r))
		{
			// check if leaf
			if (currentNode.child[childIndex] < 0)
			{
				int32_t shapeIndex = -currentNode.child[childIndex] - 1;
				bool isHit = this->_mShapePtrs[shapeIndex]->intersect(isectPtr, Ray(r.origin, r.direction, r.tmin, isectPtr->t));
				isIntersect = isHit || isIntersect;
			}
			else if (currentNode.child[childIndex] > 0)
			{
				bool isHit = intersectBranch(isectPtr, Ray(r.origin, r.direction, r.tmin, isectPtr->t), currentNode.child[childIndex]);
				isIntersect = isHit || isIntersect;
			}
		}
	}

	return isIntersect;
}

bool Bvh::intersectBranchP(const Ray & r, const size_t index) const
{
	const BvhNode & currentNode = this->_mBvhNodes[index];
	for (size_t i = 0;i < 2;i++)
	{
		if (currentNode.bbox[i].intersectP(r))
		{
			if (currentNode.child[i] < 0)
			{
				int32_t shapeIndex = -currentNode.child[i] - 1;
				if (this->_mShapePtrs[shapeIndex]->intersectP(r))
				{
					return true;
				}
			}
			else if (currentNode.child[i] > 0)
			{
				if (intersectBranchP(r, currentNode.child[i]))
				{
					return true;
				}
			}
		}
	}

	return false;
}

int Bvh::buildBranch(std::vector<std::pair<shared_ptr<Shape>, Aabb>> * shapePairsPtr, size_t begin, size_t end, const Aabb & bound)
{
	const size_t numShapes = end - begin;
	std::vector<std::pair<shared_ptr<Shape>, Aabb>> & shapePairs = *shapePairsPtr;
	BvhNode result;

	size_t maxExtentDim = bound.maxExtent();

	// handle a very rare special case where end - begin = 1
	if (numShapes == 0)
	{
		result.child[0] = 0;
		result.child[1] = 0;

		int32_t currentNodeIndex = static_cast<int32_t>(this->_mBvhNodes.size());
		this->_mBvhNodes.push_back(result);

		return currentNodeIndex + 1;
	}
	else if (numShapes == 1)
	{
		Aabb bbox = shapePairs[begin].second;
		result.bbox[0] = bbox;
		result.child[0] = -(int32_t(begin) + 1);
		result.child[1] = 0; // means it doesn't have child

		int32_t currentNodeIndex = static_cast<int32_t>(this->_mBvhNodes.size());
		this->_mBvhNodes.push_back(result);

		return currentNodeIndex + 1;
	}

	// sort bboxes in limited range
	std::sort(shapePairs.begin() + begin, shapePairs.begin() + end,
		[&](const std::pair<shared_ptr<Shape>, Aabb> & a, const std::pair<shared_ptr<Shape>, Aabb> & b) {
			return a.second.pMin[maxExtentDim] < b.second.pMin[maxExtentDim];
	});

	std::vector<Float> leftArea = std::vector<Float>(numShapes),
		rightArea = std::vector<Float>(numShapes);
	{
		// populate leftArea vector
		Aabb leftBBoxTmp;
		for (size_t i = 0;i < numShapes;i++)
		{
			leftBBoxTmp = Aabb::Union(leftBBoxTmp, shapePairs[i + begin].second);
			leftArea[i] = leftBBoxTmp.surfaceArea();
		}

		// populate rightArea vector
		Aabb rightBBoxTmp;
		for (int i = static_cast<int>(numShapes - 1);i >= 0;i--)
		{
			rightBBoxTmp = Aabb::Union(rightBBoxTmp, shapePairs[i + begin].second);
			rightArea[i] = rightBBoxTmp.surfaceArea();
		}
	}

	// find index where it has minSahCost
	Float minSahCost = std::numeric_limits<Float>::infinity();
	size_t minSahIndex = begin;
	for (size_t i = 0;i < numShapes - 1;i++)
	{
		Float sahCost = (i + 1) * leftArea[i] + (numShapes - i - 1) * rightArea[i + 1];
		//std::cout << "sahCost : " << sahCost << std::endl; _getch();
		if (sahCost < minSahCost)
		{
			minSahCost = sahCost;
			minSahIndex = i + begin;
		}
	}

	// compute bbox for left child
	Aabb leftBBox, rightBBox;
	for (size_t i = begin;i < minSahIndex + 1;i++)
	{
		leftBBox = Aabb::Union(leftBBox, shapePairs[i].second);
	}
	result.bbox[0] = leftBBox;

	// compute bbox for right child
	for (size_t i = minSahIndex + 1;i < end;i++)
	{
		rightBBox = Aabb::Union(rightBBox, shapePairs[i].second);
	}
	result.bbox[1] = rightBBox;

	size_t currentNodeIndex = this->_mBvhNodes.size();
	this->_mBvhNodes.push_back(result);
	int32_t resultIndex = static_cast<int32_t>(currentNodeIndex + 1);

	if (minSahIndex == begin)
	{
		this->_mBvhNodes[currentNodeIndex].child[0] = -static_cast<int32_t>(begin + 1);
	}
	else
	{
		this->_mBvhNodes[currentNodeIndex].child[0] = static_cast<int32_t>(resultIndex);
		resultIndex = this->buildBranch(shapePairsPtr, begin, minSahIndex + 1, leftBBox);
	}

	if (minSahIndex == end - 2)
	{
		this->_mBvhNodes[currentNodeIndex].child[1] = -static_cast<int32_t>(end - 1 + 1);
	}
	else
	{
		this->_mBvhNodes[currentNodeIndex].child[1] = static_cast<int32_t>(resultIndex);
		resultIndex = this->buildBranch(shapePairsPtr, minSahIndex + 1, end, rightBBox);
	}

	return static_cast<int32_t>(resultIndex);
}