Playing with multidrawindirect in osg

osg_indirect.cpp

#include <osg/GL2Extensions>
#include <osg/Notify>
#include <osg/ref_ptr>
#include <osg/Geode>
#include <osg/Geometry>
#include <osg/Point>
#include <osg/Vec3>
#include <osg/Vec4>
#include <osg/Program>
#include <osg/Shader>
#include <osg/BlendFunc>
#include <osgEarth/Random>

#include <iostream>
#include <osgViewer/Viewer>
#include <osgViewer/ViewerEventHandlers>
#include <osgGA/TrackballManipulator>

#include <osg/PrimitiveSetIndirect>
#include <osg/BufferIndexBinding>

osg::Vec4 randomColor()
{
    float r = (float)rand() / (float)RAND_MAX;
    float g = (float)rand() / (float)RAND_MAX;
    float b = (float)rand() / (float)RAND_MAX;
    return osg::Vec4(r, g, b, 1.0f);
}

osg::Vec3 randomVert()
{
    float x = (float)rand() / (float)RAND_MAX;
    float y = (float)rand() / (float)RAND_MAX;
    float z = (float)rand() / (float)RAND_MAX;
    return osg::Vec3(x, y, z);
}

float random()
{
    return (float)rand() / (float)RAND_MAX;
}

struct Point
{
public:
    Point(const osg::Vec3f& pos, const osg::Vec4f& col) :
        position(pos),
        color(col)
    {
    }

    osg::Vec3f position;
    osg::Vec4 color;
};

typedef std::vector< Point > PointList;

class PointChunk : public osg::Node
{
public:
    PointChunk();
    osg::BoundingSphere computeBound() const;

    void traverse(osg::NodeVisitor& nv);
    PointList points;
    unsigned int index;
    osg::DrawArraysIndirectCommand cmd;
};

typedef std::vector< osg::ref_ptr< PointChunk > > PointChunkList;

struct TraversedChunks : public osg::Referenced
{
    PointChunkList chunks;
};

PointChunk::PointChunk() :
    osg::Node()
{
}

osg::BoundingSphere PointChunk::computeBound() const
{
    osg::BoundingSphere bs;
    for (PointList::const_iterator itr = points.begin(); itr != points.end(); ++itr)
    {
        bs.expandBy(itr->position);
    }
    //OSG_NOTICE << "PointChunk::computeBound count= " << points.size() << " " << bs.center().x() << ", " << bs.center().y() << ", " << bs.center().z() << " radius=" << bs.radius() << std::endl;
    return bs;
}

void PointChunk::traverse(osg::NodeVisitor& nv)
{
    if (nv.getVisitorType() == osg::NodeVisitor::CULL_VISITOR)
    {
        //osgUtil::CullVisitor* cv = nv.asCullVisitor();
        //if (!cv->isCulled(this->getBound()))
        {
            TraversedChunks* chunks = dynamic_cast<TraversedChunks*>(nv.getUserData());
            if (chunks)
            {
                chunks->chunks.push_back(this);
            }
        }
        osg::Node::traverse(nv);
    }
    else
    {
        osg::Node::traverse(nv);
    }
}




class PointManager : public osg::Group
{
public:
    PointManager()
    {
        osg::Geometry* geom = new osg::Geometry;
        geom->setUseVertexBufferObjects(true);

        osg::MultiDrawArraysIndirect* drawArrays = new osg::MultiDrawArraysIndirect(osg::PrimitiveSet::POINTS);
        _commands = new osg::DefaultIndirectCommandDrawArrays();
        drawArrays->setIndirectCommandArray(_commands);

        _verts = new osg::Vec3Array();
        geom->setVertexArray(_verts);

        _colors = new osg::Vec4Array;
        geom->setColorArray(_colors, osg::Array::BIND_PER_VERTEX);

        geom->addPrimitiveSet(drawArrays);
        addChild(geom);
    }

    virtual osg::BoundingBox computeBoundingBox() const
    {
        osg::BoundingBox bb;
        for (unsigned int i = 0; i < _verts->size(); i++)
        {
            bb.expandBy((*_verts)[i]);
        }
        return bb;
    }

    osg::BoundingSphere computeBound() const
    {
        osg::BoundingSphere bs;
        for (unsigned int i = 0; i < _verts->size(); i++)
        {
            bs.expandBy((*_verts)[i]);
        }
        return bs;
    }

    void addPoints(PointChunk* chunk)
    {
        ///create indirect command
        osg::DrawArraysIndirectCommand cmd;
        cmd.count = chunk->points.size();
        cmd.instanceCount = 1;
        cmd.first = _verts->size();
        for (PointList::const_iterator itr = chunk->points.begin(); itr != chunk->points.end(); ++itr)
        {
            _verts->push_back(itr->position);
            _colors->push_back(itr->color);
        }
        // Clear the points, we don't need them anymore in the chunk.
        //chunk->points.clear();
        chunk->cmd = cmd;
        chunk->index = _commands->size();
        _chunks.push_back(chunk);
        _commands->push_back(cmd);
        _verts->dirty();
        _colors->dirty();
        _commands->dirty();
        dirtyBound();
    }

    void removePoints(PointChunk* chunk)
    {
        _commands->erase(_commands->begin() + chunk->index);
        _chunks.erase(_chunks.begin() + chunk->index);
        _commands->dirty();
    }

    void PointManager::traverse(osg::NodeVisitor& nv)
    {
        if (nv.getVisitorType() == osg::NodeVisitor::CULL_VISITOR)
        {
            osgUtil::CullVisitor* cv = nv.asCullVisitor();
            osg::ref_ptr< TraversedChunks > chunks = new TraversedChunks;
            cv->setUserData(chunks.get());
            osg::Group::traverse(nv);
            //OSG_NOTICE << "Traversed " << chunks->chunks.size() << " chunks" << std::endl;
            _commands->clear();
            for (unsigned int i = 0; i < chunks->chunks.size(); i++)
            {
                _commands->push_back(chunks->chunks[i].get()->cmd);
            }
            //OSG_NOTICE << "Got " << _commands->size() << " commands" << std::endl;
            _commands->dirty();
        }
        else
        {
            osg::Group::traverse(nv);
        }
    }

    osg::DefaultIndirectCommandDrawArrays* _commands;
    osg::Vec3Array* _verts;
    osg::Vec4Array* _colors;
    PointChunkList _chunks;
};




int main(int argc, char**argv)
{

    osg::ArgumentParser arguments(&argc, argv);
    arguments.getApplicationUsage()->setDescription(arguments.getApplicationName() + " is the example which demonstrates Multi Indirect Draw with basevertex");
    arguments.getApplicationUsage()->setCommandLineUsage(arguments.getApplicationName() + " [options] ");
    arguments.getApplicationUsage()->addCommandLineOption("-h or --help", "Display this information");
    arguments.getApplicationUsage()->addCommandLineOption("--numX", "square count on X");
    arguments.getApplicationUsage()->addCommandLineOption("--numY", "square count on Y");
    arguments.getApplicationUsage()->addCommandLineOption("--classic", "disable MDI and use classic DrawElements");

    if (arguments.read("-h") || arguments.read("--help"))
    {
        arguments.getApplicationUsage()->write(std::cout);
        return 1;
    }

    int MAXX = 200;
    int MAXY = 200;
    arguments.read("--numX", MAXX);
    arguments.read("--numY", MAXY);

    bool indirect = true;
    if (arguments.read("--classic"))
    {
        indirect = false;
    }


    osg::Group* root = new osg::Group;

    if (indirect)
    {
        OSG_NOTICE << "Indirect" << std::endl;
        PointManager* manager = new PointManager;
        root->addChild(manager);
        for (unsigned int c = 0; c < MAXX; c++)
        {
            for (unsigned int r = 0; r < MAXY; r++)
            {
                PointChunk *chunk = new PointChunk;
                osg::Vec4 color = randomColor();
                for (unsigned int j = 0; j < 100; j++)
                {
                    float x = c + random();
                    float y = r + random();
                    float z = 0;
                    chunk->points.push_back(Point(osg::Vec3(x, y, z), color));
                }
                chunk->dirtyBound();
                manager->addPoints(chunk);
                manager->addChild(chunk);
                if (c % 2 == 0) chunk->setNodeMask(0);
            }
        }
    }
    else
    {
        OSG_NOTICE << "Classic" << std::endl;
        for (unsigned int c = 0; c < MAXX; c++)
        {
            for (unsigned int r = 0; r < MAXY; r++)
            {
                osg::Geometry* geom = new osg::Geometry;
                osg::Vec3Array* verts = new osg::Vec3Array;
                geom->setVertexArray(verts);
                osg::Vec4Array* colors = new osg::Vec4Array;
                geom->setColorArray(colors, osg::Array::BIND_PER_VERTEX);

                osg::Vec4 color = randomColor();
                for (unsigned int j = 0; j < 100; j++)
                {
                    float x = c + random();
                    float y = r + random();
                    float z = 0;
                    verts->push_back(osg::Vec3(x, y, z));
                    colors->push_back(color);
                }
                geom->addPrimitiveSet(new osg::DrawArrays(GL_POINTS, 0, verts->size()));
                if (c % 2 == 0) geom->setNodeMask(0);
                root->addChild(geom);
            }
        }
    }

    root->getOrCreateStateSet()->setMode(GL_LIGHTING, osg::StateAttribute::OFF);

    osgViewer::Viewer viewer;
    viewer.addEventHandler(new osgViewer::StatsHandler);
    viewer.setCameraManipulator(new osgGA::TrackballManipulator());
    viewer.setSceneData(root);

    viewer.realize();

    while (!viewer.done())
    {
        viewer.frame();
    }
    return 0;
}