mitsuba/include/mitsuba/core/octree.h

/*
    This file is part of Mitsuba, a physically based rendering system.

    Copyright (c) 2007-2014 by Wenzel Jakob and others.

    Mitsuba is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License Version 3
    as published by the Free Software Foundation.

    Mitsuba is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program. If not, see <http://www.gnu.org/licenses/>.
*/

#pragma once
#if !defined(__MITSUBA_CORE_OCTREE_H_)
#define __MITSUBA_CORE_OCTREE_H_

#include <mitsuba/mitsuba.h>
#include <mitsuba/core/atomic.h>
#include <mitsuba/core/timer.h>
#include <mitsuba/core/aabb.h>

MTS_NAMESPACE_BEGIN

/**
 * \brief Lock-free linked list data structure
 *
 * This class provides a very basic linked list data structure whose primary
 * purpose it is to efficiently service append operations from multiple parallel
 * threads. These are internally realized via atomic compare and exchange
 * operations, meaning that no lock must be acquired.
 *
 * \ingroup libcore
 */
template <typename T> class LockFreeList {
public:
	struct ListItem {
		T value;
		ListItem *next;

		inline ListItem(const T &value) :
			value(value), next(NULL) { }
	};

	inline LockFreeList() : m_head(NULL) {}

	~LockFreeList() {
		ListItem *cur = m_head;
		while (cur) {
			ListItem *next = cur->next;
			delete cur;
			cur = next;
		}
	}

	inline const ListItem *head() const {
		return m_head;
	}

	void append(const T &value) {
		ListItem *item = new ListItem(value);
		ListItem **cur = &m_head;

		while (!atomicCompareAndExchangePtr<ListItem>(cur, item, NULL))
			cur = &((*cur)->next);
	}
private:
	ListItem *m_head;
};

/**
 * \brief Generic single-reference static octree
 *
 * This class is currently used to implement BSSRDF evaluation
 * with irradiance point clouds.
 *
 * The \c Item template parameter must implement a function
 * named <tt>getPosition()</tt> that returns a \ref Point.
 *
 * \ingroup libcore
 */
template <typename Item, typename NodeData> class StaticOctree {
public:
	struct OctreeNode {
		bool leaf : 1;
		NodeData data;

		union {
			struct {
				OctreeNode *children[8];
			};

			struct {
				uint32_t offset;
				uint32_t count;
			};
		};

		~OctreeNode() {
			if (!leaf) {
				for (int i=0; i<8; ++i) {
					if (children[i])
						delete children[i];
				}
			}
		}
	};

	/**
	 * \brief Create a new octree
	 * \param maxDepth
	 *     Maximum tree depth (24 by default)
	 * \param maxItems
	 *     Maximum items per interior node (8 by default)
	 *
	 * By default, the maximum tree depth is set to 16
	 */
	inline StaticOctree(const AABB &aabb, uint32_t maxDepth = 24, uint32_t maxItems = 8) :
		m_aabb(aabb), m_maxDepth(maxDepth), m_maxItems(maxItems), m_root(NULL) { }

	/// Release all memory
	~StaticOctree() {
		if (m_root)
			delete m_root;
	}

	void build() {
		SLog(EDebug, "Building an octree over " SIZE_T_FMT " data points (%s)..",
			m_items.size(), memString(m_items.size() * sizeof(Item)).c_str());

		ref<Timer> timer = new Timer();
		std::vector<uint32_t> perm(m_items.size()), temp(m_items.size());

		for (uint32_t i=0; i<m_items.size(); ++i)
			perm[i] = i;

		/* Build the kd-tree and compute a suitable permutation of the elements */
		m_root = build(m_aabb, 0, &perm[0], &temp[0], &perm[0], &perm[0] + m_items.size());

		/* Apply the permutation */
		permute_inplace(&m_items[0], perm);

		SLog(EDebug, "Done (took %i ms)" , timer->getMilliseconds());
	}

protected:
	struct LabelOrdering : public std::binary_function<uint32_t, uint32_t, bool> {
		LabelOrdering(const std::vector<Item> &items) : m_items(items) { }

		inline bool operator()(uint32_t a, uint32_t b) const {
			return m_items[a].label < m_items[b].label;
		}

		const std::vector<Item> &m_items;
	};

	/// Return the AABB for a child of the specified index
	inline AABB childBounds(int child, const AABB &nodeAABB, const Point &center) const {
		AABB childAABB;
		childAABB.min.x = (child & 4) ? center.x : nodeAABB.min.x;
		childAABB.max.x = (child & 4) ? nodeAABB.max.x : center.x;
		childAABB.min.y = (child & 2) ? center.y : nodeAABB.min.y;
		childAABB.max.y = (child & 2) ? nodeAABB.max.y : center.y;
		childAABB.min.z = (child & 1) ? center.z : nodeAABB.min.z;
		childAABB.max.z = (child & 1) ? nodeAABB.max.z : center.z;
		return childAABB;
	}

	OctreeNode *build(const AABB &aabb, uint32_t depth, uint32_t *base,
			uint32_t *temp, uint32_t *start, uint32_t *end) {
		if (start == end) {
			return NULL;
		} else if ((uint32_t) (end-start) < m_maxItems || depth > m_maxDepth) {
			OctreeNode *result = new OctreeNode();
			result->count = (uint32_t) (end-start);
			result->offset = (uint32_t) (start-base);
			result->leaf = true;
			return result;
		}

		Point center = aabb.getCenter();
		uint32_t nestedCounts[8];
		memset(nestedCounts, 0, sizeof(uint32_t)*8);

		/* Label all items */
		for (uint32_t *it = start; it != end; ++it) {
			Item &item = m_items[*it];
			const Point &p = item.getPosition();

			uint8_t label = 0;
			if (p.x > center.x) label |= 4;
			if (p.y > center.y) label |= 2;
			if (p.z > center.z) label |= 1;

			AABB bounds = childBounds(label, aabb, center);
			SAssert(bounds.contains(p));

			item.label = label;
			nestedCounts[label]++;
		}

		uint32_t nestedOffsets[9];
		nestedOffsets[0] = 0;
		for (int i=1; i<=8; ++i)
			nestedOffsets[i] = nestedOffsets[i-1] + nestedCounts[i-1];

		/* Sort by label */
		for (uint32_t *it = start; it != end; ++it) {
			int offset = nestedOffsets[m_items[*it].label]++;
			temp[offset] = *it;
		}
		memcpy(start, temp, (end-start) * sizeof(uint32_t));

		/* Recurse */
		OctreeNode *result = new OctreeNode();
		for (int i=0; i<8; i++) {
			AABB bounds = childBounds(i, aabb, center);

			uint32_t *it = start + nestedCounts[i];
			result->children[i] = build(bounds, depth+1, base, temp, start, it);
			start = it;
		}

		result->leaf = false;

		return result;
	}

	inline StaticOctree() : m_root(NULL) { }
protected:
	AABB m_aabb;
	std::vector<Item> m_items;
	uint32_t m_maxDepth;
	uint32_t m_maxItems;
	OctreeNode *m_root;
};

/**
 * \brief Generic multiple-reference octree with support for parallel dynamic updates
 *
 * Based on the excellent implementation in PBRT. Modifications are
 * the addition of a bounding sphere query and support for multithreading.
 *
 * This class is currently used to implement irradiance caching.
 *
 * \ingroup libcore
 */
template <typename Item> class DynamicOctree {
public:
	/**
	 * \brief Create a new octree
	 *
	 * By default, the maximum tree depth is set to 24
	 */
	inline DynamicOctree(const AABB &aabb, uint32_t maxDepth = 24)
	 : m_aabb(aabb), m_maxDepth(maxDepth) {
	}

	/// Insert an item with the specified cell coverage
	inline void insert(const Item &value, const AABB &coverage) {
		insert(&m_root, m_aabb, value, coverage,
			coverage.getExtents().lengthSquared(), 0);
	}

	/// Execute <tt>functor.operator()</tt> on all records, which potentially overlap \c p
	template <typename Functor> inline void lookup(const Point &p, Functor &functor) const {
		if (!m_aabb.contains(p))
			return;
		lookup(&m_root, m_aabb, p, functor);
	}

	/// Execute <tt>functor.operator()</tt> on all records, which potentially overlap \c bsphere
	template <typename Functor> inline void searchSphere(const BSphere &sphere, Functor &functor) {
		if (!m_aabb.overlaps(sphere))
			return;
		searchSphere(&m_root, m_aabb, sphere, functor);
	}

	inline const AABB &getAABB() const { return m_aabb; }
private:
	struct OctreeNode {
	public:
		OctreeNode() {
			for (int i=0; i<8; ++i)
				children[i] = NULL;
		}

		~OctreeNode() {
			for (int i=0; i<8; ++i) {
				if (children[i])
					delete children[i];
			}
		}

		OctreeNode *children[8];
		LockFreeList<Item> data;
	};

	/// Return the AABB for a child of the specified index
	inline AABB childBounds(int child, const AABB &nodeAABB, const Point &center) const {
		AABB childAABB;
		childAABB.min.x = (child & 4) ? center.x : nodeAABB.min.x;
		childAABB.max.x = (child & 4) ? nodeAABB.max.x : center.x;
		childAABB.min.y = (child & 2) ? center.y : nodeAABB.min.y;
		childAABB.max.y = (child & 2) ? nodeAABB.max.y : center.y;
		childAABB.min.z = (child & 1) ? center.z : nodeAABB.min.z;
		childAABB.max.z = (child & 1) ? nodeAABB.max.z : center.z;
		return childAABB;
	}

	void insert(OctreeNode *node, const AABB &nodeAABB, const Item &value,
			const AABB &coverage, Float diag2, uint32_t depth) {
		/* Add the data item to the current octree node if the max. tree
		   depth is reached or the data item's coverage area is smaller
		   than the current node size */
		if (depth == m_maxDepth ||
			(nodeAABB.getExtents().lengthSquared() < diag2)) {
			node->data.append(value);
			return;
		}

		/* Otherwise: test for overlap */
		const Point center = nodeAABB.getCenter();

		/* Otherwise: test for overlap */
		bool x[2] = { coverage.min.x <= center.x, coverage.max.x > center.x };
		bool y[2] = { coverage.min.y <= center.y, coverage.max.y > center.y };
		bool z[2] = { coverage.min.z <= center.z, coverage.max.z > center.z };
		bool over[8] = { x[0] && y[0] && z[0], x[0] && y[0] && z[1],
						 x[0] && y[1] && z[0], x[0] && y[1] && z[1],
						 x[1] && y[0] && z[0], x[1] && y[0] && z[1],
						 x[1] && y[1] && z[0], x[1] && y[1] && z[1] };

		/* Recurse */
		for (int child=0; child<8; ++child) {
			if (!over[child])
				continue;
			if (!node->children[child]) {
				OctreeNode *newNode = new OctreeNode();
				if (!atomicCompareAndExchangePtr<OctreeNode>(&node->children[child], newNode, NULL))
					delete newNode;
			}
			const AABB childAABB(childBounds(child, nodeAABB, center));
			insert(node->children[child], childAABB,
				value, coverage, diag2, depth+1);
		}
	}

	/// Internal lookup procedure - const version
	template <typename Functor> inline void lookup(const OctreeNode *node,
			const AABB &nodeAABB, const Point &p, Functor &functor) const {
		const Point center = nodeAABB.getCenter();

		const typename LockFreeList<Item>::ListItem *item = node->data.head();
		while (item) {
			functor(item->value);
			item = item->next;
		}

		int child = (p.x > center.x ? 4 : 0)
				+ (p.y > center.y ? 2 : 0)
				+ (p.z > center.z ? 1 : 0);

		OctreeNode *childNode = node->children[child];

		if (childNode) {
			const AABB childAABB(childBounds(child, nodeAABB, center));
			lookup(node->children[child], childAABB, p, functor);
		}
	}

	template <typename Functor> inline void searchSphere(OctreeNode *node,
			const AABB &nodeAABB, const BSphere &sphere,
			Functor &functor) {
		const Point center = nodeAABB.getCenter();

		const typename LockFreeList<Item>::ListItem *item = node->data.head();
		while (item) {
			functor(item->value);
			item = item->next;
		}

		// Potential for much optimization..
		for (int child=0; child<8; ++child) {
			if (node->children[child]) {
				const AABB childAABB(childBounds(child, nodeAABB, center));
				if (childAABB.overlaps(sphere))
					searchSphere(node->children[child], childAABB, sphere, functor);
			}
		}
	}
private:
	OctreeNode m_root;
	AABB m_aabb;
	uint32_t m_maxDepth;
};

MTS_NAMESPACE_END

#endif /* __MITSUBA_CORE_OCTREE_H_ */