RootInputFile.cc

Go to the documentation of this file.

#include "RootIOSvc/RootInputFile.h"
#include "RootIOSvc/RootIOBaseObject.h"
#include "RootIOSvc/RootIOUserData.h"
#include "RootIOSvc/RootIOUserDataProxy.h"

#include "TFile.h"
#include "TTree.h"
#include "TBranchElement.h"
#include "TLeaf.h"
#include "TDirectory.h"
#include "TKey.h"
#include "TIterator.h"

#include <string>
#include <vector>


static void handle_user_data(TTree* tree, const std::string& path)
{
    TObjArray* branches = tree->GetListOfBranches();
    size_t nbranches = branches->GetEntries();

    Dyb::MsgStreamMember log("RootInputFile::handle_user_data");
    log << MSG::DEBUG << "Found Tree (nBranches=" << nbranches 
        << ") at path: " << path << endreq;

    // Check the tree for branches besides ones made from a
    // RootIOBaseObject
    std::vector<TBranch*> known_unknowns, unknown_unknowns;
    for (size_t ind=0; ind<nbranches; ++ind) {

        TBranch* branch = static_cast<TBranch*>(branches->At(ind));
        TLeaf* leaf = static_cast<TLeaf*>(branch->GetListOfLeaves()->At(0));

        log << MSG::DEBUG << "\tCheckingName: " 
            << branch->GetName() 
            << " (" << leaf->GetTypeName() << ")"
            << endreq;

        // Let's do a little dance.  It's called, "find who we know".
        std::string type = leaf->GetTypeName();
        if (type == "Int_t" || type == "Float_t") {
            log << MSG::DEBUG << "\t\tlikely branch (known scalar): \"" 
                << branch->GetName() << "\"" << endreq;
            unknown_unknowns.push_back(branch); // Outlander!
        }
        else {
            TClass c(type.c_str());
            if (c.InheritsFrom("RootIOBaseObject")) {
                known_unknowns.push_back(branch); // I know you!
            }
            else {
                log << MSG::DEBUG << "\t\tlikely branch (something): \"" 
                    << branch->GetName() << "\"" << endreq;
                unknown_unknowns.push_back(branch); // Outlander!
            }
        }
    }

    RootIOUserData ud;          // get proxies for given stream
    RootIOUserData::ProxyCollection& proxies = ud.input(path);

    for (size_t ind=0; ind<unknown_unknowns.size(); ++ind) {
        TBranch* branch = unknown_unknowns[ind];

        std::string className = branch->ClassName(); // eg "TBranch"

        TLeaf* leaf = static_cast<TLeaf*>(branch->GetListOfLeaves()->At(0));

        std::string branchName = branch->GetName(); // eg "count"
        std::string branchType = leaf->GetTypeName(); // eg "Int_t", "vector<int>"

        log << MSG::DEBUG << "\t\t" << ind 
            << ": got className: \"" << className << "\"" 
            << " branchName: \"" << branchName << "\"" 
            << " branchType: \"" << branchType << "\"" 
            << endreq;

        char* addr = 0;
        if (className == "TBranchElement") { // not a simple scalar
            TBranchElement *be  = dynamic_cast<TBranchElement*>(branch);
            addr = be->GetObject();
            if (!addr) {
                be->SetAddress(0); // let ROOT memory manage
                addr = be->GetObject();
            }
        }
        else {                               // simple scalar
            addr = (char*)(leaf->GetValuePointer());
        }
        if (!addr) {
            log << MSG::ERROR
                << "Failed to get branch adddress for variable \"" << branchName
                << "\" at path \"" << path << "\". I will be crashing now, thank you." << endreq;
            assert (0);
        }

        RootIOUserDataProxy* udp = proxies[branchName];
        if (!udp) {
            udp = new RootIOUserDataProxy(branchName,branchType,addr);
            proxies[branchName] = udp;
        }
        else {
            udp->resetAddress(addr);
        }

    } // loop over unknown unknowns

}     // handle_user_data()


static void process_directory(TDirectory* dir, std::vector<std::string>& ret)
{
    Dyb::MsgStreamMember log("RootInputFile::TreePaths");

    TList* keys = dir->GetListOfKeys();
    TIter root_sucks(keys);
    TKey* key=0;
    while ((key = (TKey*)root_sucks())) { // infinite loop? ROOT does suck a lot!
        std::string dirpath = dir->GetPath();
        dirpath = dirpath.substr(dirpath.rfind(":")+1); 

        TObject* obj = key->ReadObj();

        TDirectory* isDir = dynamic_cast<TDirectory*>(obj);
        if (isDir) {
            process_directory(isDir,ret);
            continue;
        }
        TTree* isTree = dynamic_cast<TTree*>(obj);
        if (isTree) {
            std::string path = "/";
            // Catch trees in root directory
            if(dirpath != "/")
                path = dirpath + "/";
            path += isTree->GetName();
            ret.push_back(path);
            continue;
        }

        log << MSG::WARNING
            << "unexpected object " << obj->GetName()
            << " of type " << obj->ClassName() 
            << " found in directory " << dirpath << endreq;
    }
}

std::vector<std::string> RootInputFile::TreePaths(const std::string& filename)
{
    std::vector<std::string> ret;


    TFile* f = TFile::Open(filename.c_str(),"READ");
    if (!f->IsOpen()) {
        std::cerr << "File " << filename << " can not be opened for reading\n";
        return ret;
    }
    //std::cerr << "File " << filename << " openeded\n";
    
    process_directory(f,ret);
    return ret;
}


//#include <iostream>
//using namespace std;

static int find_uniq(TBranch* branch)
{
    static std::string name("RootIOBaseObject");

    //cerr << "Checking branch " << branch->GetName() << std::endl;

    if (branch->GetName() == name) {
        TBranch* brID = branch->FindBranch("clID");
        RootIOBaseObject riobo;
        brID->SetAddress(&riobo);
        brID->GetEntry(0);
        return riobo.clID;
    }

    TObjArray* branches = branch->GetListOfBranches();
    TIter root_sucks(branches);
    branch=0;
    while ((branch = (TBranch*)root_sucks())) {
        int ret = find_uniq(branch);
        if (ret) return ret;
    }
    return 0;
}

int RootInputFile::TestForObject(const std::string& filename,
                                 const std::string& treepath,
                                 const std::string& branchname,
                                 bool ignore_missing)
{
    Dyb::MsgStreamMember log("RootInputFile::TestForTObject");


    TFile *f = TFile::Open(filename.c_str(),"READ");
    if (!f->IsOpen()) {
        log << MSG::ERROR 
            << "File " << filename 
            << " can not be opened for reading" 
            << endreq;
        return 0;
    }
    log << MSG::DEBUG  << "File " << filename
        << " opened. will now get obj at treepath is " << treepath 
        << endreq;

    TObject* obj = 0;
    if (treepath.rfind('/') == 0) {
        // Catch root-level trees
        log << MSG::VERBOSE << "Try to get root-level tree. get obj at treepath " 
            << treepath.substr(1).c_str() << endreq;
        obj = f->Get(treepath.substr(1).c_str());
    }
    else {
        log << MSG::VERBOSE << "Get obj at treepath " << treepath.c_str() << endreq;
        obj = f->Get(treepath.c_str());
    }
    if (!obj) {
        if (!ignore_missing) 
            log << MSG::ERROR
                << "No object " << treepath
                << " in file " << filename << endreq;
        return 0;
    }
    log << MSG::DEBUG << "Got object at " << treepath << endreq;
    
    TTree* tree = dynamic_cast<TTree*>(obj);
    if (!tree) {
        log << MSG::ERROR
            << "Object at " << treepath
            << " in file " << filename 
            << " is not a TTree" << endreq;
        return 0;
    }
    //cerr << "Object is a TTree\n";

    int nentries = tree->GetEntries();
    if (!nentries) {
        log << MSG::ERROR
            << "Tree " << treepath 
            << " in filename " << filename 
            << " found, but no entries" << endreq;
        return 0;
    }
    //cerr << "Tree has " << nentries << " entries\n";

    TBranch* b = tree->GetBranch(branchname.c_str());
    if (!b) {
        log << MSG::ERROR
            << "Failed to get branch " << branchname << endreq;
        return 0;
    }
    //cerr << "search for uniq id" << std::endl;

    int ret = find_uniq(b);
    f->Close();
    delete f;
    return ret;
}


RootInputFile::RootInputFile(const std::string& filename, 
                             const std::string& treepath,
                             const std::string& branchname)
    : m_filename(filename)
    , m_treepath(treepath)
    , m_branchname(branchname)
    , m_pruneSimHits(0)
    , m_file(0), m_tree(0)
    , m_entry(-1), m_entries(-1)
    , m_addr(0)
    , log("RootInputFile")
{
  //Prune SimHitHeader when reading
  char* pruneSimHitHeader = getenv("NUWA_PRUNESIMHITHEADER");
  if(0 != pruneSimHitHeader) m_pruneSimHits = 1;
}
RootInputFile::~RootInputFile()
{
}

bool RootInputFile::open()
{
    // is this a dummy entry in the input stream?
    if (m_filename == "") {
        m_entries = 0;
        return true;
    }

    if (!m_file) {
        log << MSG::DEBUG
            << "openning " << m_filename << " for " << m_treepath 
            << endreq;

        m_file = TFile::Open(m_filename.c_str(),"READ");
    }
    if (!m_tree) {
        std::string treepath = m_treepath;
        // Catch top-level objects
        if(treepath.rfind("/") == 0) treepath = treepath.substr(1);
        TObject* obj = m_file->Get(treepath.c_str());
        if (!obj) {
            log << MSG::ERROR << "open(): No such object at " << m_treepath
                << " in " << m_filename << endreq;
            return false;
        }
        m_tree = dynamic_cast<TTree*>(obj);
        if (!m_tree) {
            log << MSG::ERROR
                << "open(): Object at " << m_treepath
                << " in " << m_filename 
                << " not a TTree\n";
            return false;
        }
        if (m_entries < 0) {
            m_entries = m_tree->GetEntries();
        }
    }
    if (!m_tree) return false;

    handle_user_data(m_tree,m_treepath);

    return true;
}

bool RootInputFile::leave()
{
    log << MSG::DEBUG
        << "leaving " << m_filename << " for " << m_treepath
        << endreq;

    // The opened file must be closed to avoid memory leak
    if(m_file) {
        m_file->Close();
        delete m_file;
        m_file = 0;
    }
    m_tree = 0;
    m_addr = 0;
    m_entry = -1;

    return true;
}

int RootInputFile::entry()
{
    return m_entry;
}

bool RootInputFile::setEntry(int entry)
{
    if (!this->open()) return false;
    if (entry >= m_entries) {
        log << MSG::ERROR << "setEntry("<<entry<<"): Entry too large >= " << m_entries
            << endreq;
        return false;
    }
    m_entry = entry;
    return true;
}


bool RootInputFile::setAddr(void* addr)
{
    if (addr == m_addr) return true;
    TBranch* b = m_tree->GetBranch(m_branchname.c_str());
    if (!b) {
        log << MSG::ERROR << "setAddr(): no such branch named " 
            << m_branchname << endreq;
        return false;
    }
    
    log << MSG::DEBUG 
        << "setAddr("<<(void*)addr<<") for " << m_branchname << " in " << m_filename
        << endreq;


    if(1 == m_pruneSimHits && m_branchname == "Sim_SimHeader")
    {
      TBranch* hits = b->FindBranch("hits");
      if(hits)
      {
        hits->DeleteBaskets("all");
        m_pruneSimHits = 2;
      }
    }           
    b->SetAddress(addr);
    m_addr = addr;
    return true;
}

bool RootInputFile::read(void* addr)
{
    int nbytes = 0;
    return this->read(addr,nbytes);
}

bool RootInputFile::read(void* addr, int &nbytes)
{
    nbytes = 0;
    if (!this->open()) return false;
    if (m_entry < 0 || m_entry >= m_entries) {
        log << MSG::ERROR << "read(void*): Bad entry: " << m_entry << endreq;
        return false;
    }

    if (!this->setAddr(addr)) return false;

    log << MSG::DEBUG
        << "read("<<(void*)addr<<") entry=" << m_entry
        << " from file " << m_filename
        << " and branch " << m_branchname
        << endreq;
    nbytes = m_tree->GetEntry(m_entry);
    return nbytes > 0;
}

int RootInputFile::entries()
{
    if (m_entries < 0) {        // never opened this file
        if (!this->open()) return 0;
        this->leave();          // may not want to leave open long term, leave it as found it
    }
    return m_entries;
}

bool RootInputFile::next(int steps)
{
    return this->setEntry(m_entry+steps);
}

bool RootInputFile::beginning()
{
    if (!this->open()) return false;
    m_entry = 0;
    return true;
}
bool RootInputFile::ending()
{
    if (!this->open()) return false;
    m_entry = m_entries-1;
    return true;
}


bool RootInputFile::prev(int steps)
{
    return this->setEntry(m_entry-steps);
}


// ----- RootInputFileList ----- //

RootInputFileList::RootInputFileList()
    : std::vector<RootInputFile*>()
    , m_index(-1)
    , log("RootInputFileList")
{
}

RootInputFileList::~RootInputFileList()
{
}

int RootInputFileList::index()
{
    return m_index;
}

RootInputFile* RootInputFileList::current()
{
    if (m_index < 0 || m_index >= (int)size()) return 0;
    return (*this)[m_index];
}

bool RootInputFileList::next()
{
    if (m_index + 1 >= (int)size()) {
        log << MSG::WARNING << "next(): already at last file"
            << endreq;
        return false;
    }
    RootInputFile* rif = current();
    if (rif) rif->leave();
    ++m_index;
    rif = current();
    if (!rif) {
        log << MSG::ERROR
            << "next(): failed to make next file curren"
            << endreq;
        return false;
    }
    return true;
}

bool RootInputFileList::last()
{
    if (m_index+1 == (int)size()) return true;
    RootInputFile* rif = current();
    if (rif) rif->leave();
    m_index = (int)size() - 1;
    rif = current();
    if (!rif) {
        log << MSG::ERROR
            << "last(): falied to make last file current"
            << endreq;
        return false;
    }
    return true;
}

bool RootInputFileList::prev()
{
    if (m_index <= 0) {
        log << MSG::ERROR
            << "prev(): already at first file"
            << endreq;
        return false;
    }
    RootInputFile* rif = current();
    if (rif) rif->leave();
    --m_index;
    rif = current();
    if (!rif) {
        log << MSG::ERROR
            << "prev(): failed to make prev file curren"
            << endreq;
        return false;
    }
    return true;
}
    


bool RootInputFileList::first()
{
    if (!m_index) return true;
    RootInputFile* rif = current();
    if (rif) rif->leave();
    m_index = 0;
    rif = current();
    if (!rif) {
        log << MSG::ERROR
            << "first(): failed to make first file current"
            << endreq;
        return false;
    }
    return true;
}

bool RootInputFileList::jump(int index)
{
    if (m_index == index) return true;
    RootInputFile* rif = current();
    if (rif) rif->leave();
    m_index = index;
    rif = current();
    if (!rif) {
        log << MSG::ERROR
            << "goto("<<index<<") failed"
            << endreq;
        return false;
    }
    return true;
}

int RootInputFileList::entriesBefore(int index)
{
    if (index < 0 || index >= (int)size()) return -1;
    // cound up entries up to but not including current index
    int totEntries = 0;
    for (int ind = 0; ind < index; ++ind) {
        int entries = (*this)[ind]->entries();
        if (entries < 0) return -1;
        totEntries += entries;
    }
    return totEntries;
}