tdf_less_templates.cpp

/*
Proof of concept of a TDataFrame with a lot less templates.

Issues:
- circular dependency nightmare in which all the different nodes have to know about each other and call methods on
  each other. Could be solved introducing a common parent class, but that would probably mean virtual dispatch?
*/

#include "TTreeReader.h"
#include "TTreeReaderValue.h"
#include "TTree.h"
#include "TFile.h"
#include <cassert>
#include <memory>
#include <string>
#include <vector>

class TDFValueBase {
public:
   virtual ~TDFValueBase() {}
   virtual void MakeProxy(TTreeReader &r) = 0;
};

template <typename BranchType>
class TDFValue : public TDFValueBase {
   std::unique_ptr<TTreeReaderValue<BranchType>> fReaderValue;
   std::string fBranchName;

public:
   TDFValue(const std::string &b) : fBranchName(b) {}
   void MakeProxy(TTreeReader &r) { fReaderValue.reset(new TTreeReaderValue<BranchType>(r, fBranchName.c_str())); }
   BranchType GetValue() const { return **fReaderValue; }
};

class TDFImpl; // forward decl for TDFFilter
class TDFFilter {
   std::function<bool(void)> fFunctor;
   TDFImpl &fImpl;
   std::unique_ptr<TDFValueBase> fValue; // TODO make this a vector of values

public:
   template <typename Functor>
   TDFFilter(Functor f, const std::string &branch, TDFImpl &impl)
      : fImpl(impl), fValue(new TDFValue<int>(branch)) // TODO deduce branch type from Functor signature
   {
      auto vRawPtr = static_cast<TDFValue<int> *>(fValue.get());
      fFunctor = [f, vRawPtr]() { return f(vRawPtr->GetValue()); };
   }

   TDFImpl &GetImpl() const { return fImpl; }
   void InitValues(TTreeReader &r) { fValue->MakeProxy(r); }
   bool Check() { return fFunctor(); }
};

class TDFImpl {
   TTree &fTree;
   std::vector<std::shared_ptr<TDFFilter>> fFilters;

public:
   TDFImpl(TTree &t) : fTree(t) {}
   void Book(const std::shared_ptr<TDFFilter> &filter) { fFilters.emplace_back(filter); }
   TDFImpl &GetImpl() { return *this; }
   void Run()
   {
      TTreeReader r(&fTree);
      for (auto &f : fFilters) f->InitValues(r);
     
      while (r.Next()) {
         for (auto &f : fFilters)
            f.Run();
      }
   }
};

template <typename Node>
class TDFInterface {
   using NodePtr = std::shared_ptr<Node>;
   NodePtr fNodePtr;

public:
   TDFInterface(const NodePtr &n) : fNodePtr(n) {}

   template <typename Functor>
   TDFInterface<TDFFilter> Filter(Functor functor, const std::string &branch)
   {
      auto &impl = fNodePtr->GetImpl();
      auto filter = std::make_shared<TDFFilter>(functor, branch, impl);
      impl.Book(filter);
      TDFInterface<TDFFilter> tdf(filter);
      return tdf;
   }

   void Run() { fNodePtr->GetImpl().Run(); }
};

class TDataFrame : public TDFInterface<TDFImpl> {
public:
   TDataFrame(TTree &t) : TDFInterface<TDFImpl>(std::make_shared<TDFImpl>(t)) {}
};

int main()
{
   // create tree
   TTree t("t", "t");
   int a = 21;
   t.Branch("a", &a);
   t.Fill();

   // process tree
   auto isTwentyone = [](int a) { return a == 21; };
   TDataFrame d(t);
   auto d2 = d.Filter(isTwentyone, "a");
   d2.Run();

   // auto byTwo = [](int a) { return a * 2; };
   // auto max = d.Filter(IsTwentyone, "a").Define("b", byTwo, "a").Range(1).Max("b");
   // assert(max == 42, "");
   return 0;
}