SumRegister.h

/* Copyright (C) 2020 IBM Corp.
 * This program is Licensed under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *   http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License. See accompanying LICENSE file.
 */
 
#ifndef HELIB_SUMREGISTER_H
#define HELIB_SUMREGISTER_H
 
#include <iostream>
#include <vector>
#include <cmath>
#include <memory>
 
namespace helib {
 
template <typename T>
class SumRegister
{
private:
  std::vector<std::unique_ptr<T>> intermediateResults;
 
  unsigned int maxNumOfInputs;
  unsigned int remainingInputs;
  bool flushRequiredFlag = false;
  bool resultFlag = false;
  size_t depth = 0;
 
public:
  SumRegister(unsigned int _maxNumOfInputs) :
      maxNumOfInputs(_maxNumOfInputs), remainingInputs(_maxNumOfInputs)
  {
    // log2 returns -inf if arg is zero!
    if (_maxNumOfInputs != 0)
      this->depth = std::ceil(std::log2(_maxNumOfInputs));
 
    // Set if NOT power of 2
    if ((_maxNumOfInputs & (_maxNumOfInputs - 1)) || (_maxNumOfInputs == 0)) {
      this->flushRequiredFlag = true;
    }
 
    // Default size will be 1.
    this->intermediateResults = std::vector<std::unique_ptr<T>>(depth + 1);
  }
 
  void add(std::unique_ptr<T>& t)
  {
    if (this->remainingInputs == 0) {
      return;
    }
    this->remainingInputs--;
 
    if (this->intermediateResults.at(0) != nullptr) {
      *this->intermediateResults.at(0) += *t;
      if (depth > 0) {
        for (size_t i = 1, j = 0; i < this->intermediateResults.size();
             i++, j++) {
          if (intermediateResults[i] != nullptr) {
            *this->intermediateResults[i] += *this->intermediateResults[j];
            this->intermediateResults[j].reset();
          } else {
            this->intermediateResults[i] =
                std::move(this->intermediateResults[j]);
            break;
          }
        }
      }
 
      if (intermediateResults.at(depth) != nullptr)
        resultFlag = true;
 
    } else {
      intermediateResults.at(0) = std::move(t);
    }
 
    if ((this->remainingInputs == 0) && this->flushRequiredFlag) {
      this->flush();
    }
  }
 
  std::unique_ptr<T> getResult()
  {
    return std::move(intermediateResults.at(depth));
  }
 
  bool hasResult() const { return resultFlag; }
 
  size_t getDepth() const { return depth; }
 
  void flush()
  {
    // Flushing with a result should do nothing.
    if (this->hasResult()) {
      return;
    }
 
    for (size_t i = 0, j = 1; i < depth; i++, j++) {
      if (this->intermediateResults[i] == nullptr)
        continue;
 
      if (this->intermediateResults[j] == nullptr) {
        this->intermediateResults[j] = std::move(this->intermediateResults[i]);
      } else {
        *this->intermediateResults[j] += *this->intermediateResults[i];
        this->intermediateResults[i].reset();
      }
    }
 
    if (intermediateResults.at(depth) != nullptr)
      resultFlag = true;
  }
 
  void print()
  {
    std::cout << "Current values\n";
    for (size_t i = 0; i < this->intermediateResults.size(); i++) {
      std::cout << "[" << i << "]: "
                << (this->intermediateResults[i] != nullptr
                        ? *this->intermediateResults[i]
                        : 0)
                << " (" << this->intermediateResults[i] << ")" << '\n';
    }
  }
 
  std::vector<std::unique_ptr<T>>& getIntermediates()
  {
    return intermediateResults;
  };
};
 
} // namespace helib
#endif // HELIB_SUMREGISTER_H