1 #ifndef KDTREE_H

     2 #define KDTREE_H

     3 

     4 #include <vector>

     5 #include <iostream>

     6 #include <fstream>

     7 

     8 #include "scene.h"

     9 

    10 using namespace std;

    11 

    12 class ShapeList: public vector<Shape*>

    13 {

    14 };

    15 

    16 class Container

    17 {

    18 protected:

    19 	BBox bbox;

    20 public:

    21 	ShapeList shapes;

    22 	Container(): bbox(), shapes() {};

    23 	virtual ~Container() {};

    24 	virtual void addShape(Shape* aShape);

    25 	//void addShapeNoExtend(Shape* aShape) { shapes.push_back(aShape); };

    26 	virtual Shape *nearest_intersection(const Shape *origin_shape, const Ray &ray,

    27 		Float &nearest_distance);

    28 	virtual void optimize() {};

    29 };

    30 

    31 class KdNode

    32 {

    33 	Float split;

    34 	short axis; /* 0,1,2 => x,y,z; 3 => leaf */

    35 public:

    36 	union {

    37 		KdNode *children;

    38 		ShapeList *shapes;

    39 	};

    40 

    41 	KdNode() : axis(3) { shapes = new ShapeList(); };

    42 	~KdNode();

    43 

    44 	void setAxis(short aAxis) { axis = aAxis; };

    45 	short getAxis() { return axis; };

    46 

    47 	void setSplit(Float aSplit) { split = aSplit; };

    48 	Float getSplit() { return split; };

    49 

    50 	void setLeaf() { axis = 3; };

    51 	bool isLeaf() { return axis == 3; };

    52 

    53 	KdNode *getLeftChild() { return children; };

    54 	KdNode *getRightChild() { return children+1; };

    55 

    56 	void addShape(Shape* aShape) { shapes->push_back(aShape); };

    57 

    58 	void subdivide(BBox bbox, int depth);

    59 };

    60 

    61 class KdTree: public Container

    62 {

    63 	KdNode *root;

    64 	bool built;

    65 	int max_depth;

    66 public:

    67 	KdTree() : Container(), root(NULL), built(false), max_depth(32) {};

    68 	~KdTree() { if (root) delete root; };

    69 	void addShape(Shape* aShape) { Container::addShape(aShape); built = false; };

    70 	Shape *nearest_intersection(const Shape *origin_shape, const Ray &ray,

    71 		Float &nearest_distance);

    72 	void optimize() { build(); };

    73 	void build();

    74 	void save(ostream &str, KdNode *node = NULL);

    75 	void load(istream &str, KdNode *node = NULL);

    76 	void setMaxDepth(int md) { max_depth = md; };

    77 };

    78 

    79 #endif