RootOutputStream.cc

Go to the documentation of this file.

#include "RootIOSvc/RootOutputStream.h"
#include "RootIOSvc/RootIOUserData.h"
#include "RootIOSvc/RootIOUserDataProxy.h"
#include "RootIOSvc/RootIOBaseObject.h"
#include "RootOutputFileManager.h"

#include "TFile.h"
#include "TTree.h"

#include "TClass.h"             // debug

#include <map>

using namespace std;

RootOutputStream::RootOutputStream(void* addr, 
                                   const string& classname,
                                   const string& treepath,
                                   const string& branchname)
    : RootIOStream(addr)
    , m_file(0)
    , m_tree(0)
    , m_dir(0)
    , m_classname(classname)
    , m_treepath(treepath)
    , m_branchname(branchname)
    , m_entries(0)
    , log("RootOutputStream")
    , m_majorIndex("")
    , m_minorIndex("")
{
    TClass* cl = TClass::GetClass(classname.c_str());
    if (!cl) {
        log << MSG::ERROR << "Failed to get class for " 
            << classname << endreq;
    }

}


RootOutputStream::~RootOutputStream()
{
    this->close();
}

void RootOutputStream::setIndexExpression(const std::string& majorIndex, const std::string& minorIndex)
{
    m_majorIndex = majorIndex;
    m_minorIndex = minorIndex;
}

string RootOutputStream::path()
{
    return m_treepath;
}


bool RootOutputStream::write()
{
    if (!m_tree) {
        log << MSG::ERROR << "write(): no tree" 
            << endreq;
        return false;
    }

    int nbytes = m_tree->Fill();
    log << MSG::DEBUG << "Wrote " << nbytes
        << " bytes to entry " << m_entries
        << " of tree " << m_treepath
        << endreq;

    // Debugging #445.
    //TObject* dir = m_file->Get("Event/Sim");
    //if (dir) dir->ls();

    ++m_entries;
    return nbytes > 0;    
}

int RootOutputStream::entries()
{
    return m_entries;
}

int RootOutputStream::fileEntries()
{
    return m_tree->GetEntries();
}

bool RootOutputStream::close()
{
    if (!m_file) return true;
    if (!m_tree) return true;
    
    log << MSG::DEBUG << "Closing file " << m_file->GetName() << endreq;

    if (m_majorIndex.size() || m_minorIndex.size()) {
        log << MSG::DEBUG
            << "Building index using major=\"" << m_majorIndex
            << "\", minor=\"" << m_minorIndex << "\"" << endreq;
        m_tree->BuildIndex(m_majorIndex.c_str(),m_minorIndex.c_str());
    }

    m_dir->cd();
    m_tree->Write(NULL,TObject::kOverwrite); // for kOverwrite see #445
    RootOutputFileManager::get().close_file(m_file);
    m_file = 0;
    m_tree = 0;
    m_dir = 0;
    return true;
}


bool RootOutputStream::newFile(const string& filename,
                               const string& treepath,
                               const string& branchname)
{
    // Only start a new file if filename differes
    if (m_file && filename == m_file->GetName()) {
        return true;
    }

    if ("" != treepath) m_treepath = treepath;
    if ("" != branchname) m_branchname = branchname;

    if ("" == m_treepath || "" == m_branchname) {
        log << MSG::ERROR 
            << "newFile(): tree path or branch name are empty"
            << endreq;
        return false;
    }

    this->close();

    m_file = RootOutputFileManager::get().get_file(filename);

    // Make the directories up to but not including last one which is
    // the tree name.
    m_dir = m_file;
    string::size_type last = 0, slash = m_treepath.find('/');
    for (; slash != string::npos; slash = m_treepath.find('/',last)) {
        if (!slash) {
            last = slash + 1;
            continue;   // skip initial '/'
        }
        string subdir = m_treepath.substr(last,slash-last);
        TDirectory* dir = m_dir->GetDirectory(subdir.c_str());
        if (dir) m_dir = dir;
        else m_dir = m_dir->mkdir(subdir.c_str());
        last = slash + 1;
    }
    if (last) {
        string subdir = m_treepath.substr(0,last-1);
        m_file->cd(subdir.c_str());
    }
        
    string title = "Tree at " + m_treepath + " holding " + m_branchname;
    string treename = m_treepath.substr(last);

    m_tree = new TTree(treename.c_str(),title.c_str());
    m_tree->Branch(m_branchname.c_str(),m_classname.c_str(),m_addr);

    bool ok = handle_user_data();
    if (!ok) return false;

    log << MSG::DEBUG
        << "newFile(): " << title << " in file " << filename
        << endreq;

    return true;
}

bool RootOutputStream::handle_user_data()
{
    RootIOUserData ud;
    RootIOUserData::ProxyCollection& proxies = ud.output(m_treepath);
    RootIOUserData::ProxyCollection::iterator it, done = proxies.end();

    log << MSG::DEBUG << "newFile(): handling user data with " 
        << proxies.size() << " proxies at treepath=" << m_treepath
        << " (one of " << ud.outputMap().size() << ")"
        << endreq;

    bool ret = true;
    
    for (it = proxies.begin(); it != done; ++it) {
        RootIOUserDataProxy* udp = it->second;
        bool ok = udp->branch(m_tree);
        if (ok) {
            log << MSG::DEBUG << "newFile(): initiated user data type: "
                << udp->cppType() << " name: " << udp->varName() << endreq;
        }
        else {
            log << MSG::ERROR << "newFile(): failed to initiate user data type: "
                << udp->cppType() << " name: " << udp->varName() << endreq;
            ret = false;
        }
    }
    return ret;
}