<commit_before>
#include "dwi/shells.h"

#include "file/config.h"

namespace MR
{
  namespace DWI
  {

    using namespace App;

    const OptionGroup ShellOption = OptionGroup ("DW Shell selection options")
      + Option ("shell",
          "specify one or more diffusion-weighted gradient shells to use during processing. "
          "The shell structure is determined automatically from the gradient encoding (either "
          "stored in the image header or provided using the -grad option); use this option to "
          "list the b-value(s) corresponding to the shell(s) that you wish to be used for "
          "processing."
          "Note that some commands are incompatible with multiple shells, and will throw an "
          "error if you provide more than one b-value.")
        + Argument ("list").type_sequence_int();

    // TODO Want to re-name this config option to BZeroThreshold
    const float bzero_threshold = File::Config::get_float ("BValueThreshold", 10.0);

    Shell::Shell (const Math::Matrix<float>& grad, const std::vector<size_t>& indices) :
        volumes (indices),
        mean (0.0),
        stdev (0.0),
        min (std::numeric_limits<float>::max()),
        max (0.0)
    {
      assert (volumes.size());
      for (std::vector<size_t>::const_iterator i = volumes.begin(); i != volumes.end(); i++) {
        const float b = grad (*i, 3);
        mean += b;
        min = std::min (min, b);
        max = std::max (min, b);
      }
      mean /= float(volumes.size());
      for (std::vector<size_t>::const_iterator i = volumes.begin(); i != volumes.end(); i++)
        stdev += Math::pow2 (grad (*i, 3) - mean);
      stdev = Math::sqrt (stdev / (volumes.size() - 1));
    }

    void Shells::select_shells (const bool keep_bzero, const bool force_single_shell)
    {

      // Easiest way to restrict processing to particular shells is to simply erase
      //   the unwanted shells; makes it command independent

      BitSet to_retain (shells.size(), 0);
      if (keep_bzero && smallest().is_bzero())
        to_retain[0] = true;

      App::Options opt = App::get_options ("shell");
      if (opt.size()) {

        std::vector<int> desired_bvalues = opt[0][0];
        if (desired_bvalues.size() > 1 && !(desired_bvalues.front() == 0) && force_single_shell)
          throw Exception ("Command not compatible with multiple non-zero b-shells");

        for (std::vector<int>::const_iterator b = desired_bvalues.begin(); b != desired_bvalues.end(); ++b) {

          if (*b < 0)
            throw Exception ("Cannot select shells corresponding to negative b-values");

          // Automatically select a b=0 shell if the requested b-value is zero
          if (*b <= bzero_threshold) {

            if (smallest().is_bzero()) {
              to_retain[0] = true;
              DEBUG ("User requested b-value " + str(*b) + "; got b=0 shell : " + str(smallest().get_mean()) + " +- " + str(smallest().get_stdev()) + " with " + str(smallest().count()) + " volumes");
            } else {
              throw Exception ("User selected b=0 shell, but no such data exists");
            }

          } else {

            // First, see if the b-value lies within the range of one of the shells
            // If this doesn't occur, need to make a decision based on the shell distributions
            // Few ways this could be done:
            // * Compute number of standard deviations away from each shell mean, see if there's a clear winner
            //   - Won't work if any of the standard deviations are zero
            // * Assume each is a Poisson distribution, see if there's a clear winner
            // Prompt warning if decision is slightly askew, exception if ambiguous
            bool shell_selected = false;
            for (size_t s = 0; s != shells.size(); ++s) {
              if ((*b >= shells[s].get_min()) && (*b <= shells[s].get_max())) {
                to_retain[s] = true;
                shell_selected = true;
                DEBUG ("User requested b-value " + str(*b) + "; got shell " + str(s) + ": " + str(shells[s].get_mean()) + " +- " + str(shells[s].get_stdev()) + " with " + str(shells[s].count()) + " volumes");
              }
            }
            if (!shell_selected) {

              // First, check to see if all non-zero shells have a non-zero standard deviation
              bool zero_stdev = false;
              for (std::vector<Shell>::const_iterator s = shells.begin(); s != shells.end(); ++s) {
                if (!s->is_bzero() && !s->get_stdev()) {
                  zero_stdev = true;
                  break;
                }
              }

              size_t best_shell = 0;
              float best_num_stdevs = std::numeric_limits<float>::max();
              bool ambiguous = false;
              for (size_t s = 0; s != shells.size(); ++s) {
                const float num_stdev = Math::abs ((*b - shells[s].get_mean()) / (zero_stdev ? Math::sqrt (shells[s].get_mean()) : shells[s].get_stdev()));
                const bool within_range = (num_stdev < (zero_stdev ? 1.0 : 5.0));
                if (within_range) {
                  if (!shell_selected) {
                    best_shell = s;
                    best_num_stdevs = num_stdev;
                  } else {
                    // More than one shell plausible; decide whether or not the decision is unambiguous
                    if (num_stdev < 0.1 * best_num_stdevs) {
                      best_shell = s;
                      best_num_stdevs = num_stdev;
                    } else if (num_stdev < 10.0 * best_num_stdevs) {
                      ambiguous = true;
                    } // No need to do anything if the existing selection is significantly better than this shell
                  }
                }
              }

              if (ambiguous) {
                std::string bvalues;
                for (size_t s = 0; s != shells.size(); ++s) {
                  if (bvalues.size())
                    bvalues += ", ";
                  bvalues += str(shells[s].get_mean()) + " +- " + str(shells[s].get_stdev());
                }
                throw Exception ("Unable to robustly select desired shell b=" + str(*b) + " (detected shells are: " + bvalues + ")");
              }

              WARN ("User requested shell b=" + str(*b) + "; have selected shell " + str(shells[best_shell].get_mean()) + " +- " + str(shells[best_shell].get_stdev()));

              to_retain[best_shell] = true;

            }

          } // End checking to see if requested shell is b=0

        } // End looping over list of requested b-value shells

      } else {

        if (force_single_shell && !is_single_shell())
          WARN ("Multiple non-zero b-value shells detected; automatically selecting b=" + str(largest().get_mean()) + " with " + str(largest().count()) + " volumes");
        to_retain[shells.size()-1] = true;

      }

      if (to_retain.full()) {
        DEBUG ("No DW shells to be removed");
        return;
      }

      // Erase the unwanted shells
      std::vector<Shell> new_shells;
      for (size_t s = 0; s != shells.size(); ++s) {
        if (to_retain[s])
          new_shells.push_back (shells[s]);
      }
      shells.swap (new_shells);

    }

    void Shells::initialise (const Math::Matrix<float>& grad)
    {
      BValueList bvals (grad.column (3));
      std::vector<size_t> clusters (bvals.size(), 0);
      const size_t num_shells = clusterBvalues (bvals, clusters);

      INFO ("Diffusion gradient encoding data clustered into " + str(num_shells) + " shells");
      if ((num_shells < 1) || (num_shells > Math::sqrt (float(grad.rows()))))
        throw Exception ("Gradient encoding matrix does not represent a HARDI sequence");

      for (size_t shellIdx = 0; shellIdx <= num_shells; shellIdx++) {

        std::vector<size_t> volumes;
        for (size_t volumeIdx = 0; volumeIdx != clusters.size(); ++volumeIdx) {
          if (clusters[volumeIdx] == shellIdx)
            volumes.push_back (volumeIdx);
        }

        if (shellIdx) {
          shells.push_back (Shell (grad, volumes));
        } else if (volumes.size()) {
          std::string unassigned;
          for (size_t i = 0; i != volumes.size(); ++i) {
            if (unassigned.size())
              unassigned += ", ";
            unassigned += str(volumes[i]) + " (" + str(bvals[volumes[i]]) + ")";
          }
          WARN ("The following image volumes were not successfully assigned to a b-value shell:");
          WARN (unassigned);
        }

      }

      std::sort (shells.begin(), shells.end());
    }

    size_t Shells::clusterBvalues (const BValueList& bvals, std::vector<size_t>& clusters) const
    {
      BitSet visited (bvals.size(), false);
      size_t clusterIdx = 0;

      for (size_t ii = 0; ii != bvals.size(); ii++) {
        if (!visited[ii]) {

          visited[ii] = true;
          const float b = bvals[ii];
          std::vector<size_t> neighborIdx;
          regionQuery (bvals, b, neighborIdx);

          if (b > bzero_threshold && neighborIdx.size() < DWI_SHELLS_MIN_LINKAGE) {

            clusters[ii] = 0;

          } else {

            clusters[ii] = ++clusterIdx;
            for (size_t i = 0; i < neighborIdx.size(); i++) {
              if (!visited[neighborIdx[i]]) {
                visited[neighborIdx[i]] = true;
                std::vector<size_t> neighborIdx2;
                regionQuery (bvals, bvals[neighborIdx[i]], neighborIdx2);
                if (neighborIdx2.size() >= DWI_SHELLS_MIN_LINKAGE)
                  for (size_t j = 0; j != neighborIdx2.size(); j++)
                    neighborIdx.push_back (neighborIdx2[j]);
              }
              if (clusters[neighborIdx[i]] == 0)
                clusters[neighborIdx[i]] = clusterIdx;
            }

          }

        }
      }

      clusterIdx = rejectSmallShells (bvals, clusters, clusterIdx);
      return clusterIdx;
    }

    void Shells::regionQuery (const BValueList& bvals, const float b, std::vector<size_t>& idx) const
    {
      for (size_t i = 0; i < bvals.size(); i++) {
        if (Math::abs (b - bvals[i]) < DWI_SHELLS_EPSILON)
          idx.push_back (i);
      }
    }

    size_t Shells::rejectSmallShells (const BValueList& bvals, std::vector<size_t>& clusters, const size_t max_index) const
    {
      // Count the number of volumes in each shell
      std::vector<size_t> counts (max_index + 1, 0);
      // Also need the mean b-value; don't want to reject b=0 due to having too few volumes
      std::vector<float> means (max_index + 1, 0.0);
      for (size_t vol = 0; vol != bvals.size(); ++vol) {
        const size_t shell = clusters[vol];
        ++counts[shell];
        means[shell] += bvals[vol];
      }
      // Generate a mapping from old shell index to new shell index
      //   (if one shell is rejected, the indices of the shells above it need to decrease)
      std::vector<size_t> index_mapping (max_index + 1, 0);
      size_t new_index = 0;
      for (size_t shell = 1; shell <= max_index; ++shell) {
        means[shell] /= float(counts[shell]);
        if ((means[shell] < bzero_threshold) || (counts[shell] >= DWI_SHELLS_MIN_DIRECTIONS))
          index_mapping[shell] = ++new_index;
      }
      // At least one shell rejected
      if (new_index != max_index) {
        // Re-map the shell indices; volumes belong to a rejected shell will be reset to an index of 0
        for (size_t i = 0; i != clusters.size(); ++i)
          clusters[i] = index_mapping[clusters[i]];
      }
      return max_index;
    }

  }
}

<commit_msg>DWI::Shells: update to option documentation<commit_after>
#include "dwi/shells.h"

#include "file/config.h"

namespace MR
{
  namespace DWI
  {

    using namespace App;

    const OptionGroup ShellOption = OptionGroup ("DW Shell selection options")
      + Option ("shell",
          "specify one or more diffusion-weighted gradient shells to use during "
          "processing, as a comma-separated list of the desired approximate b-values. "
          "Note that some commands are incompatible with multiple shells, and "
          "will throw an error if more than one b-value are provided.")
        + Argument ("list").type_sequence_int();

    // TODO Want to re-name this config option to BZeroThreshold
    const float bzero_threshold = File::Config::get_float ("BValueThreshold", 10.0);

    Shell::Shell (const Math::Matrix<float>& grad, const std::vector<size_t>& indices) :
        volumes (indices),
        mean (0.0),
        stdev (0.0),
        min (std::numeric_limits<float>::max()),
        max (0.0)
    {
      assert (volumes.size());
      for (std::vector<size_t>::const_iterator i = volumes.begin(); i != volumes.end(); i++) {
        const float b = grad (*i, 3);
        mean += b;
        min = std::min (min, b);
        max = std::max (min, b);
      }
      mean /= float(volumes.size());
      for (std::vector<size_t>::const_iterator i = volumes.begin(); i != volumes.end(); i++)
        stdev += Math::pow2 (grad (*i, 3) - mean);
      stdev = Math::sqrt (stdev / (volumes.size() - 1));
    }

    void Shells::select_shells (const bool keep_bzero, const bool force_single_shell)
    {

      // Easiest way to restrict processing to particular shells is to simply erase
      //   the unwanted shells; makes it command independent

      BitSet to_retain (shells.size(), 0);
      if (keep_bzero && smallest().is_bzero())
        to_retain[0] = true;

      App::Options opt = App::get_options ("shell");
      if (opt.size()) {

        std::vector<int> desired_bvalues = opt[0][0];
        if (desired_bvalues.size() > 1 && !(desired_bvalues.front() == 0) && force_single_shell)
          throw Exception ("Command not compatible with multiple non-zero b-shells");

        for (std::vector<int>::const_iterator b = desired_bvalues.begin(); b != desired_bvalues.end(); ++b) {

          if (*b < 0)
            throw Exception ("Cannot select shells corresponding to negative b-values");

          // Automatically select a b=0 shell if the requested b-value is zero
          if (*b <= bzero_threshold) {

            if (smallest().is_bzero()) {
              to_retain[0] = true;
              DEBUG ("User requested b-value " + str(*b) + "; got b=0 shell : " + str(smallest().get_mean()) + " +- " + str(smallest().get_stdev()) + " with " + str(smallest().count()) + " volumes");
            } else {
              throw Exception ("User selected b=0 shell, but no such data exists");
            }

          } else {

            // First, see if the b-value lies within the range of one of the shells
            // If this doesn't occur, need to make a decision based on the shell distributions
            // Few ways this could be done:
            // * Compute number of standard deviations away from each shell mean, see if there's a clear winner
            //   - Won't work if any of the standard deviations are zero
            // * Assume each is a Poisson distribution, see if there's a clear winner
            // Prompt warning if decision is slightly askew, exception if ambiguous
            bool shell_selected = false;
            for (size_t s = 0; s != shells.size(); ++s) {
              if ((*b >= shells[s].get_min()) && (*b <= shells[s].get_max())) {
                to_retain[s] = true;
                shell_selected = true;
                DEBUG ("User requested b-value " + str(*b) + "; got shell " + str(s) + ": " + str(shells[s].get_mean()) + " +- " + str(shells[s].get_stdev()) + " with " + str(shells[s].count()) + " volumes");
              }
            }
            if (!shell_selected) {

              // First, check to see if all non-zero shells have a non-zero standard deviation
              bool zero_stdev = false;
              for (std::vector<Shell>::const_iterator s = shells.begin(); s != shells.end(); ++s) {
                if (!s->is_bzero() && !s->get_stdev()) {
                  zero_stdev = true;
                  break;
                }
              }

              size_t best_shell = 0;
              float best_num_stdevs = std::numeric_limits<float>::max();
              bool ambiguous = false;
              for (size_t s = 0; s != shells.size(); ++s) {
                const float num_stdev = Math::abs ((*b - shells[s].get_mean()) / (zero_stdev ? Math::sqrt (shells[s].get_mean()) : shells[s].get_stdev()));
                const bool within_range = (num_stdev < (zero_stdev ? 1.0 : 5.0));
                if (within_range) {
                  if (!shell_selected) {
                    best_shell = s;
                    best_num_stdevs = num_stdev;
                  } else {
                    // More than one shell plausible; decide whether or not the decision is unambiguous
                    if (num_stdev < 0.1 * best_num_stdevs) {
                      best_shell = s;
                      best_num_stdevs = num_stdev;
                    } else if (num_stdev < 10.0 * best_num_stdevs) {
                      ambiguous = true;
                    } // No need to do anything if the existing selection is significantly better than this shell
                  }
                }
              }

              if (ambiguous) {
                std::string bvalues;
                for (size_t s = 0; s != shells.size(); ++s) {
                  if (bvalues.size())
                    bvalues += ", ";
                  bvalues += str(shells[s].get_mean()) + " +- " + str(shells[s].get_stdev());
                }
                throw Exception ("Unable to robustly select desired shell b=" + str(*b) + " (detected shells are: " + bvalues + ")");
              }

              WARN ("User requested shell b=" + str(*b) + "; have selected shell " + str(shells[best_shell].get_mean()) + " +- " + str(shells[best_shell].get_stdev()));

              to_retain[best_shell] = true;

            }

          } // End checking to see if requested shell is b=0

        } // End looping over list of requested b-value shells

      } else {

        if (force_single_shell && !is_single_shell())
          WARN ("Multiple non-zero b-value shells detected; automatically selecting b=" + str(largest().get_mean()) + " with " + str(largest().count()) + " volumes");
        to_retain[shells.size()-1] = true;

      }

      if (to_retain.full()) {
        DEBUG ("No DW shells to be removed");
        return;
      }

      // Erase the unwanted shells
      std::vector<Shell> new_shells;
      for (size_t s = 0; s != shells.size(); ++s) {
        if (to_retain[s])
          new_shells.push_back (shells[s]);
      }
      shells.swap (new_shells);

    }

    void Shells::initialise (const Math::Matrix<float>& grad)
    {
      BValueList bvals (grad.column (3));
      std::vector<size_t> clusters (bvals.size(), 0);
      const size_t num_shells = clusterBvalues (bvals, clusters);

      INFO ("Diffusion gradient encoding data clustered into " + str(num_shells) + " shells");
      if ((num_shells < 1) || (num_shells > Math::sqrt (float(grad.rows()))))
        throw Exception ("Gradient encoding matrix does not represent a HARDI sequence");

      for (size_t shellIdx = 0; shellIdx <= num_shells; shellIdx++) {

        std::vector<size_t> volumes;
        for (size_t volumeIdx = 0; volumeIdx != clusters.size(); ++volumeIdx) {
          if (clusters[volumeIdx] == shellIdx)
            volumes.push_back (volumeIdx);
        }

        if (shellIdx) {
          shells.push_back (Shell (grad, volumes));
        } else if (volumes.size()) {
          std::string unassigned;
          for (size_t i = 0; i != volumes.size(); ++i) {
            if (unassigned.size())
              unassigned += ", ";
            unassigned += str(volumes[i]) + " (" + str(bvals[volumes[i]]) + ")";
          }
          WARN ("The following image volumes were not successfully assigned to a b-value shell:");
          WARN (unassigned);
        }

      }

      std::sort (shells.begin(), shells.end());
    }

    size_t Shells::clusterBvalues (const BValueList& bvals, std::vector<size_t>& clusters) const
    {
      BitSet visited (bvals.size(), false);
      size_t clusterIdx = 0;

      for (size_t ii = 0; ii != bvals.size(); ii++) {
        if (!visited[ii]) {

          visited[ii] = true;
          const float b = bvals[ii];
          std::vector<size_t> neighborIdx;
          regionQuery (bvals, b, neighborIdx);

          if (b > bzero_threshold && neighborIdx.size() < DWI_SHELLS_MIN_LINKAGE) {

            clusters[ii] = 0;

          } else {

            clusters[ii] = ++clusterIdx;
            for (size_t i = 0; i < neighborIdx.size(); i++) {
              if (!visited[neighborIdx[i]]) {
                visited[neighborIdx[i]] = true;
                std::vector<size_t> neighborIdx2;
                regionQuery (bvals, bvals[neighborIdx[i]], neighborIdx2);
                if (neighborIdx2.size() >= DWI_SHELLS_MIN_LINKAGE)
                  for (size_t j = 0; j != neighborIdx2.size(); j++)
                    neighborIdx.push_back (neighborIdx2[j]);
              }
              if (clusters[neighborIdx[i]] == 0)
                clusters[neighborIdx[i]] = clusterIdx;
            }

          }

        }
      }

      clusterIdx = rejectSmallShells (bvals, clusters, clusterIdx);
      return clusterIdx;
    }

    void Shells::regionQuery (const BValueList& bvals, const float b, std::vector<size_t>& idx) const
    {
      for (size_t i = 0; i < bvals.size(); i++) {
        if (Math::abs (b - bvals[i]) < DWI_SHELLS_EPSILON)
          idx.push_back (i);
      }
    }

    size_t Shells::rejectSmallShells (const BValueList& bvals, std::vector<size_t>& clusters, const size_t max_index) const
    {
      // Count the number of volumes in each shell
      std::vector<size_t> counts (max_index + 1, 0);
      // Also need the mean b-value; don't want to reject b=0 due to having too few volumes
      std::vector<float> means (max_index + 1, 0.0);
      for (size_t vol = 0; vol != bvals.size(); ++vol) {
        const size_t shell = clusters[vol];
        ++counts[shell];
        means[shell] += bvals[vol];
      }
      // Generate a mapping from old shell index to new shell index
      //   (if one shell is rejected, the indices of the shells above it need to decrease)
      std::vector<size_t> index_mapping (max_index + 1, 0);
      size_t new_index = 0;
      for (size_t shell = 1; shell <= max_index; ++shell) {
        means[shell] /= float(counts[shell]);
        if ((means[shell] < bzero_threshold) || (counts[shell] >= DWI_SHELLS_MIN_DIRECTIONS))
          index_mapping[shell] = ++new_index;
      }
      // At least one shell rejected
      if (new_index != max_index) {
        // Re-map the shell indices; volumes belong to a rejected shell will be reset to an index of 0
        for (size_t i = 0; i != clusters.size(); ++i)
          clusters[i] = index_mapping[clusters[i]];
      }
      return max_index;
    }

  }
}

<|endoftext|>


<commit_before>#pragma once
/*
* Covariant Script C/C++ Native Interface
* CNI's effect is to forward CovScript function call to C/C++ function and need to cooperate with cs::native_interface to use.
* Please do not bind a CNI directly to CovScript. CovScript can not directly identify a CNI.
* CNI can recognize any form of C/C++ functions, but please note that CNI does not support implicit conversion of parameter types.
* To avoid compatibility issues, use CovScript built-in types as much as possible.
* You can use the reference to modify the value of the CovScript passed parameters, but please note that this may cause an error because some of the variables in CovScript are protected.
* To improve the efficiency of a function, you can use constant references instead of values. Please note that CNI does not currently support forwarding right-value references.
* If you want to receive a variable parameter, just use cov::any as the argument type.
* This program was written by Micheal Lee(mikecovlee@163.com)。
*/
#include "../include/mozart/bind.hpp"
#include "../include/mozart/traits.hpp"
#include "./core.hpp"
#include "./arglist.hpp"
namespace cs {
	template<typename T>struct convert {
		static inline const T& get_val(cov::any& val)
		{
			return val.const_val<T>();
		}
	};
	template<typename T>struct convert<const T&> {
		static inline const T& get_val(cov::any& val)
		{
			return val.const_val<T>();
		}
	};
	template<typename T>struct convert<T&> {
		static inline T& get_val(cov::any& val)
		{
			return val.val<T>(true);
		}
	};
	template<>struct convert<const cov::any&> {
		static inline const cov::any& get_val(const cov::any& val)
		{
			return val;
		}
	};
	template<>struct convert<cov::any&> {
		static inline cov::any& get_val(cov::any& val)
		{
			return val;
		}
	};
	template<typename T,typename X>class cni_helper;
	template<typename X>class cni_helper<void(*)(),X> {
		X mFunc;
	public:
		cni_helper()=delete;
		cni_helper(const cni_helper&)=default;
		cni_helper(const X& func):mFunc(func) {}
		cov::any call(array& args) const
		{
			if(!args.empty())
				throw syntax_error("Wrong size of the arguments.Expected 0");
			mFunc();
			return cov::any::make<number>(0);
		}
	};
	template<typename RetT,typename X>class cni_helper<RetT(*)(),X> {
		X mFunc;
	public:
		cni_helper()=delete;
		cni_helper(const cni_helper&)=default;
		cni_helper(const X& func):mFunc(func) {}
		cov::any call(array& args) const
		{
			if(!args.empty())
				throw syntax_error("Wrong size of the arguments.Expected 0");
			return std::move(mFunc());
		}
	};
	template<typename...ArgsT,typename X>
	class cni_helper<void(*)(ArgsT...),X> {
		using args_t=typename cov::type_list::make<ArgsT...>::result;
		X mFunc;
		template<int...S>void _call(array& args,const cov::sequence<S...>&) const
		{
			mFunc(convert<ArgsT>::get_val(args.at(S))...);
		}
	public:
		cni_helper()=delete;
		cni_helper(const cni_helper&)=default;
		cni_helper(const X& func):mFunc(func) {}
		cov::any call(array& args) const
		{
			arglist::check<ArgsT...>(args);
			_call(args,cov::make_sequence<cov::type_list::get_size<args_t>::result>::result);
			return cov::any::make<number>(0);
		}
	};
	template<typename RetT,typename...ArgsT,typename X>
	class cni_helper<RetT(*)(ArgsT...),X> {
		using args_t=typename cov::type_list::make<ArgsT...>::result;
		X mFunc;
		template<int...S>RetT _call(array& args,const cov::sequence<S...>&) const
		{
			return std::move(mFunc(convert<ArgsT>::get_val(args.at(S))...));
		}
	public:
		cni_helper()=delete;
		cni_helper(const cni_helper&)=default;
		cni_helper(const X& func):mFunc(func) {}
		cov::any call(array& args) const
		{
			arglist::check<ArgsT...>(args);
			return std::move(_call(args,cov::make_sequence<cov::type_list::get_size<args_t>::result>::result));
		}
	};
	template<typename T>struct cni_modify {
		using type=T;
	};
	template<typename RetT,typename...ArgsT>struct cni_modify<RetT(ArgsT...)> {
		using type=RetT(*)(ArgsT...);
	};
	class cni final {
		class cni_base {
		public:
			cni_base()=default;
			cni_base(const cni_base&)=default;
			virtual ~cni_base()=default;
			virtual cni_base* clone()=0;
			virtual cov::any call(array&) const=0;
		};
		template<typename T>class cni_holder final:public cni_base {
			cni_helper<typename cov::function_parser<T>::type::common_type,T> mCni;
		public:
			cni_holder()=delete;
			cni_holder(const cni_holder&)=default;
			cni_holder(const T& func):mCni(func) {}
			virtual ~cni_holder()=default;
			virtual cni_base* clone() override
			{
				return new cni_holder(*this);
			}
			virtual cov::any call(array& args) const override
			{
				return mCni.call(args);
			}
		};
		template<typename T>struct construct_helper {
			template<typename X>static cni_base* construct(X&& val)
			{
				return new cni_holder<T>(std::forward<X>(val));
			}
		};
		cni_base* mCni=nullptr;
	public:
		cni()=delete;
		cni(const cni& c):mCni(c.mCni->clone()) {}
		template<typename T>cni(T&& val):mCni(construct_helper<typename cni_modify<typename cov::remove_reference<T>::type>::type>::construct(std::forward<T>(val))) {}
		~cni()
		{
			delete mCni;
		}
		cov::any operator()(array& args) const
		{
			return mCni->call(args);
		}
	};
	template<>struct cni::construct_helper<const cni> {
		static cni_base* construct(const cni& c)
		{
			return c.mCni->clone();
		}
	};
	template<>struct cni::construct_helper<cni> {
		static cni_base* construct(const cni& c)
		{
			return c.mCni->clone();
		}
	};
}
<commit_msg>重新改回std::function<commit_after>#pragma once
/*
* Covariant Script C/C++ Native Interface
* CNI's effect is to forward CovScript function call to C/C++ function and need to cooperate with cs::native_interface to use.
* Please do not bind a CNI directly to CovScript. CovScript can not directly identify a CNI.
* CNI can recognize any form of C/C++ functions, but please note that CNI does not support implicit conversion of parameter types.
* To avoid compatibility issues, use CovScript built-in types as much as possible.
* You can use the reference to modify the value of the CovScript passed parameters, but please note that this may cause an error because some of the variables in CovScript are protected.
* To improve the efficiency of a function, you can use constant references instead of values. Please note that CNI does not currently support forwarding right-value references.
* If you want to receive a variable parameter, just use cov::any as the argument type.
* This program was written by Micheal Lee(mikecovlee@163.com)。
*/
#include "../include/mozart/bind.hpp"
#include "../include/mozart/traits.hpp"
#include "./core.hpp"
#include "./arglist.hpp"
#include <functional>
namespace cs {
	template<typename T>struct convert {
		static inline const T& get_val(cov::any& val)
		{
			return val.const_val<T>();
		}
	};
	template<typename T>struct convert<const T&> {
		static inline const T& get_val(cov::any& val)
		{
			return val.const_val<T>();
		}
	};
	template<typename T>struct convert<T&> {
		static inline T& get_val(cov::any& val)
		{
			return val.val<T>(true);
		}
	};
	template<>struct convert<const cov::any&> {
		static inline const cov::any& get_val(const cov::any& val)
		{
			return val;
		}
	};
	template<>struct convert<cov::any&> {
		static inline cov::any& get_val(cov::any& val)
		{
			return val;
		}
	};
	template<typename T>class cni_helper;
	template<>class cni_helper<void(*)()> {
		std::function<void()> mFunc;
	public:
		cni_helper()=delete;
		cni_helper(const cni_helper&)=default;
		cni_helper(const std::function<void()>& func):mFunc(func) {}
		cov::any call(array& args) const
		{
			if(!args.empty())
				throw syntax_error("Wrong size of the arguments.Expected 0");
			mFunc();
			return cov::any::make<number>(0);
		}
	};
	template<typename RetT>class cni_helper<RetT(*)()> {
		std::function<RetT()> mFunc;
	public:
		cni_helper()=delete;
		cni_helper(const cni_helper&)=default;
		cni_helper(const std::function<RetT()>& func):mFunc(func) {}
		cov::any call(array& args) const
		{
			if(!args.empty())
				throw syntax_error("Wrong size of the arguments.Expected 0");
			return std::move(mFunc());
		}
	};
	template<typename...ArgsT>
	class cni_helper<void(*)(ArgsT...)> {
		using args_t=typename cov::type_list::make<ArgsT...>::result;
		std::function<void(ArgsT...)> mFunc;
		template<int...S>void _call(array& args,const cov::sequence<S...>&) const
		{
			mFunc(convert<ArgsT>::get_val(args.at(S))...);
		}
	public:
		cni_helper()=delete;
		cni_helper(const cni_helper&)=default;
		cni_helper(const std::function<void(ArgsT...)>& func):mFunc(func) {}
		cov::any call(array& args) const
		{
			arglist::check<ArgsT...>(args);
			_call(args,cov::make_sequence<cov::type_list::get_size<args_t>::result>::result);
			return cov::any::make<number>(0);
		}
	};
	template<typename RetT,typename...ArgsT>
	class cni_helper<RetT(*)(ArgsT...)> {
		using args_t=typename cov::type_list::make<ArgsT...>::result;
		std::function<RetT(ArgsT...)> mFunc;
		template<int...S>RetT _call(array& args,const cov::sequence<S...>&) const
		{
			return std::move(mFunc(convert<ArgsT>::get_val(args.at(S))...));
		}
	public:
		cni_helper()=delete;
		cni_helper(const cni_helper&)=default;
		cni_helper(const std::function<RetT(ArgsT...)>& func):mFunc(func) {}
		cov::any call(array& args) const
		{
			arglist::check<ArgsT...>(args);
			return std::move(_call(args,cov::make_sequence<cov::type_list::get_size<args_t>::result>::result));
		}
	};
	template<typename T>struct cni_modify {
		using type=T;
	};
	template<typename RetT,typename...ArgsT>struct cni_modify<RetT(ArgsT...)> {
		using type=RetT(*)(ArgsT...);
	};
	class cni final {
		class cni_base {
		public:
			cni_base()=default;
			cni_base(const cni_base&)=default;
			virtual ~cni_base()=default;
			virtual cni_base* clone()=0;
			virtual cov::any call(array&) const=0;
		};
		template<typename T>class cni_holder final:public cni_base {
			cni_helper<typename cov::function_parser<T>::type::common_type> mCni;
		public:
			cni_holder()=delete;
			cni_holder(const cni_holder&)=default;
			cni_holder(const T& func):mCni(func) {}
			virtual ~cni_holder()=default;
			virtual cni_base* clone() override
			{
				return new cni_holder(*this);
			}
			virtual cov::any call(array& args) const override
			{
				return mCni.call(args);
			}
		};
		template<typename T>struct construct_helper {
			template<typename X>static cni_base* construct(X&& val)
			{
				return new cni_holder<T>(std::forward<X>(val));
			}
		};
		cni_base* mCni=nullptr;
	public:
		cni()=delete;
		cni(const cni& c):mCni(c.mCni->clone()) {}
		template<typename T>cni(T&& val):mCni(construct_helper<typename cni_modify<typename cov::remove_reference<T>::type>::type>::construct(std::forward<T>(val))) {}
		~cni()
		{
			delete mCni;
		}
		cov::any operator()(array& args) const
		{
			return mCni->call(args);
		}
	};
	template<>struct cni::construct_helper<const cni> {
		static cni_base* construct(const cni& c)
		{
			return c.mCni->clone();
		}
	};
	template<>struct cni::construct_helper<cni> {
		static cni_base* construct(const cni& c)
		{
			return c.mCni->clone();
		}
	};
}
<|endoftext|>


<commit_before>/*
 * ---------------------------------------------------------------------
 * Numenta Platform for Intelligent Computing (NuPIC)
 * Copyright (C) 2013, Numenta, Inc.  Unless you have purchased from
 * Numenta, Inc. a separate commercial license for this software code, the
 * following terms and conditions apply:
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 3 as
 * published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 * See the GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see http://www.gnu.org/licenses.
 *
 * http://numenta.org/licenses/
 * ---------------------------------------------------------------------
 */

/** @file
 * Templates implementation for class Grouper
 */

#ifndef NTA_ZETA1_GROUPER_T_HPP
#define NTA_ZETA1_GROUPER_T_HPP

#ifdef NUPIC2
#error "Zeta1Grouper is used only by Zeta1Node, which is not part of NuPIC 2"
#endif

//----------------------------------------------------------------------
namespace nta {

  //--------------------------------------------------------------------------------
  /**
   * begin2 is not used: we do not output anything in learning mode.
   */
  template <typename InIter, typename OutIter>
  void Zeta1Grouper::learn(InIter begin1, OutIter, size_type baby_idx)
  {
    size_type winnerIndex = size_type(*begin1);
    tam_.learn(winnerIndex, baby_idx);
  }

  //--------------------------------------------------------------------------------
  /**
   * x is the output of the coincidence detector
   *
   * If using Time Based Inference ('tbi' set in the constructor), then the inference output
   * is computed by treating the Time Adjacency Matrix (tam) as a set of cell weights between
   * "cells" in each group. The intent here is to have the inference output for each group
   * increase in certainty as we see successive coincidence inputs that are in that group.
   *
   * For the TBI computation, each group is assigned 1 cell per coincidence in the
   * group. Each cell's output is updated after each time step based on the following
   * equation:
   *                             N   /		        	 \
   * cellOut (t) = bottomUp * ( SUM | cellWeight  * cellOut (t-1) |  + A0  )
   *        j              j    i=0  \         ij          i     /
   *
   * The net inference output for each group is then the max of all the cell outputs for
   * that group.
   *
   * Each group has it's own cellWeight matrix, which is produced by extracting the entries
   * from the TAM corresponding to the coincidences in that group, and then normalizing down
   * the columns so that the sum of the weights in each column is 1.
   *
   * The cellOuts for each group are kept unique from each other - when we have overlapping
   * groups for example, cell 0's output in group A will not necessarily be the same value
   * as cell 0's output in group B - this is because we only consider the contribution
   * from other cells *in the same group* when we peform the above cellOut computation.
   *
   * The A0 contribution can be considered as the likelihood that this cell is a start of
   * the group, or rather, it is the sum contribution from all the cells in the other
   * groups. Without this factor of course, none of the cell outputs would ever be non-zero.
   * In the end, the exact value chosen for A0 is immaterial since we are only looking at
   * the relative output strengths of each group.
   *
   * cellOut is a joint pdf over groups and coincidences.
   */
  template <typename InIter, typename OutIter>
  void Zeta1Grouper::tbiInfer(InIter x, InIter x_end,
			 OutIter y, TBICellOutputsVec::iterator cellOuts)
  {
    using namespace std;

    { // Pre-conditions
      NTA_ASSERT(!tbiCellWeights_.empty())
	<< "Grouper::tbiInfer: Cell weights not initialized";
    } // End pre-conditions

    const value_type A0 = (value_type) 0.1;

    // Compute TBI output.
    OutIter y_begin = y;
    TBICellWeightsVec::iterator w = tbiCellWeights_.begin();
    Groups::const_iterator g = groups_.begin(), g_end = groups_.end();
    size_type g_idx = 0;

    for (; g != g_end; ++g, ++w, ++cellOuts, ++y, ++g_idx)
      {
	// Compute the product of the cellWeights and the current
	// cell outputs.
	// w has size g->size X g->size
	// *cellOuts has size g->size (vector)
	// tbiBuffer_ is sized to be the size of the largest group
	// (only the first g->size positions are used here)
	w->rightVecProd(*cellOuts, tbiBuffer_);

	// Add A0 to each cell output and multiply by the bottom up input
	AGroup::const_iterator s = g->begin(), s_end = g->end();
	TBICellOutputs::iterator cell = cellOuts->begin();
	TBICellOutputs::const_iterator b = tbiBuffer_.begin();
	value_type maxCellOut = 0.0;

	// In case HOT is used, we need to convert the HOT state
	// index to its original coincidence. If HOT is not used,
	// getHOTCoincidence reduces to identity.
	for (; s != s_end; ++s, ++b, ++cell) {
	  size_type c = tam_.getHOTCoincidence(*s);
	  *cell = (*b + A0) * (*(x + c));
	  maxCellOut = max(*cell, maxCellOut);
	}

	*y = maxCellOut;
      }

    if(rescaleTBI_) {
      // Scale the group outputs so that the max is the same as the max of the inputs.
      // This preserves the relative strength between the group output and blank score
      // computed by the spatial pooler.
      value_type maxInValue = * std::max_element(x, x_end);
      if (maxInValue > (value_type) 0)
        normalize_max(y_begin, y_begin + groups_.size(), maxInValue);
    }
  }

  //--------------------------------------------------------------------------------
  template <typename InIter, typename OutIter>
  void Zeta1Grouper::infer(InIter x, InIter x_end, OutIter y, size_type tbi_idx)
  {
    switch (mode_) {

    case maxProp:
      // For each row, find the max corresponding to a non-zero
      weights_.vecMaxAtNZ(x, y);
      break;

    case sumProp:
      weights_.rightVecProd(x, y);
      break;

    case tbi:
      {
	if (tbiCellWeights_.empty())
	  tbi_create_();

	tbiInfer(x, x_end, y, tbiCellOutputs_[tbi_idx].begin());
      }
      break;
    } // switch mode
  }

  //--------------------------------------------------------------------------------

} // end namespace nta

#endif // NTA_ZETA1_GROUPER_T_HPP
<commit_msg>remove Zeta1Grouper header, not used in NUPIC2<commit_after><|endoftext|>


<commit_before>
int main(int argc, char* argv[])
{

}<commit_msg>Add basic command line flag handling using cppassist/cmdline<commit_after>#include <iostream>
#include <string>

#include <cppassist/cmdline/ArgumentParser.h>

int main(int argc, char* argv[])
{
  cppassist::ArgumentParser argParser;
  argParser.parse(argc, argv);

  auto inFileName = argParser.value("--input");
  auto outFileName = argParser.value("--output");

  auto outFileType = outFileName.substr(outFileName.find_last_of('.') + 1);

  auto fromGeneratedKernel = argParser.isSet("-k");

  if (fromGeneratedKernel)
  {
    std::cout << "Converting kernel \"" << inFileName << "\" to output file \"" << outFileName << "\" (type: " << outFileType << ")" <<std::endl;
  }
  else
  {
    std::cout << "Using kernel description \"" << inFileName << "\" to generate kernel \"" << outFileName << "\" (type: " << outFileType << ")" << std::endl;
  }

}<|endoftext|>


<commit_before>#define __STDC_FORMAT_MACROS

#include <stdio.h>
#include <iostream>
#include <time.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <algorithm>

#include <sinsp.h>
#include "lua_parser.h"
#include "digwatch.h"
#include "utils.h"

#include <unistd.h>
#include <getopt.h>

lua_parser* g_lua_parser;

static void usage();

//
// Program help
//
static void usage()
{
    printf(
	   "Usage: digwatch [options] [-p <output_format>] [filter]\n\n"
	   "Options:\n"
	   " -h, --help         Print this page\n"
	   " -M <num_seconds>   Stop collecting after <num_seconds> reached.\n"
	   " -N                 Don't convert port numbers to names.\n"
	   " -n <num>, --numevents=<num>\n"
	   "                    Stop capturing after <num> events\n"
	   " -u <filename>, --user-parser <filename>\n"
	   "                    Name of lua file containing parser\n"
	   " -r <readfile>, --read=<readfile>\n"
	   "                    Read the events from <readfile>.\n"
	   " --unbuffered       Turn off output buffering. This causes every single line\n"
	   "                    emitted by digwatch to be flushed, which generates higher CPU\n"
	   "                    usage but is useful when piping digwatch's output into another\n"
	   "                    process or into a script.\n"
	   "\n"
    );
}

//
// Event processing loop
//
captureinfo do_inspect(sinsp* inspector,
		       uint64_t cnt,
		       int duration_to_tot,
		       sinsp_evt_formatter* formatter)
{
	captureinfo retval;
	int32_t res;
	sinsp_evt* ev;
	string line;
        int duration_start = 0;

	//
	// Loop through the events
	//
	duration_start = ((double)clock()) / CLOCKS_PER_SEC;
	while(1)
	{
		if(duration_to_tot > 0)
		{
			int duration_tot = ((double)clock()) / CLOCKS_PER_SEC - duration_start;
			if(duration_tot >= duration_to_tot)
			{
				break;
			}
		}
		if(retval.m_nevts == cnt)
		{
			//
			// End of capture, either because the user stopped it, or because
			// we reached the event count specified with -n.
			//
			break;
		}

		res = inspector->next(&ev);

		if(res == SCAP_TIMEOUT)
		{
			continue;
		}
		else if(res == SCAP_EOF)
		{
			break;
		}
		else if(res != SCAP_SUCCESS)
		{
			//
			// Event read error.
			// Notify the chisels that we're exiting, and then die with an error.
			//
			cerr << "res = " << res << endl;
			throw sinsp_exception(inspector->getlasterr().c_str());
		}

		retval.m_nevts++;

		if(!inspector->is_debug_enabled() &&
			ev->get_category() & EC_INTERNAL)
		{
			continue;
		}
		if(formatter->tostring(ev, &line))
		{
			cout << line;
			cout << endl;
		}

	}

	return retval;
}

//
// ARGUMENT PARSING AND PROGRAM SETUP
//
digwatch_init_res digwatch_init(int argc, char **argv)
{
	digwatch_init_res res;
	sinsp* inspector = NULL;
	int op;
	uint64_t cnt = -1;
	sinsp_evt::param_fmt event_buffer_format = sinsp_evt::PF_NORMAL;
	int duration_to_tot = 0;
	captureinfo cinfo;
	string output_format;
	int long_index = 0;
	string user_parser;

	static struct option long_options[] =
	{
		{"help", no_argument, 0, 'h' },
		{"numevents", required_argument, 0, 'n' },
		{"user-parser", required_argument, 0, 'u' },
		{"readfile", required_argument, 0, 'r' },
		{"unbuffered", no_argument, 0, 0 },
		{0, 0, 0, 0}
	};

	output_format = "*%evt.num %evt.outputtime %evt.cpu %proc.name (%thread.tid) %evt.dir %evt.type %evt.info";

	try
	{
		inspector = new sinsp();

		//
		// Parse the args
		//
		while((op = getopt_long(argc, argv,
                                        "hM:Nn:r:u:",
                                        long_options, &long_index)) != -1)
		{
			switch(op)
			{
			case 'h':
				usage();
				delete inspector;
				return digwatch_init_res(EXIT_SUCCESS);
			case 'M':
				duration_to_tot = atoi(optarg);
				if(duration_to_tot <= 0)
				{
					throw sinsp_exception(string("invalid duration") + optarg);
					res.m_res = EXIT_FAILURE;
					goto exit;
				}
				break;
			case 'N':
				inspector->set_hostname_and_port_resolution_mode(false);
				break;
			case 'n':
				try
				{
					cnt = sinsp_numparser::parseu64(optarg);
				}
				catch(...)
				{
					throw sinsp_exception("can't parse the -n argument, make sure it's a number");
				}

				if(cnt <= 0)
				{
					throw sinsp_exception(string("invalid event count ") + optarg);
					res.m_res = EXIT_FAILURE;
					goto exit;
				}
				break;
			case 'u':
				user_parser = optarg;
				break;
			case '?':
				delete inspector;
				return digwatch_init_res(EXIT_FAILURE);
				break;
			default:
				break;
			}

		}

		inspector->set_buffer_format(event_buffer_format);

		string filter;

		//
		// the filter is at the end of the command line
		//
		if(optind < argc)
		{
#ifdef HAS_FILTERING
			for(int32_t j = optind ; j < argc; j++)
			{
				filter += argv[j];
				if(j < argc - 1)
				{
					filter += " ";
				}
			}

#else
			fprintf(stderr, "filtering not compiled.\n");
			res.m_res = EXIT_FAILURE;
			goto exit;
#endif
		}

		//
		// Create the event formatter
		//
		sinsp_evt_formatter formatter(inspector, output_format);

		g_lua_parser = new lua_parser(inspector, user_parser);

		if(filter.size())
		{
			inspector->set_filter(filter);
		}
		inspector->open("");

		cinfo = do_inspect(inspector,
				   cnt,
				   duration_to_tot,
				   &formatter);

		inspector->close();
	}
	catch(sinsp_exception& e)
	{
		cerr << e.what() << endl;
		res.m_res = EXIT_FAILURE;
	}
	catch(...)
	{
		printf("Exeception\n");
		res.m_res = EXIT_FAILURE;
	}

exit:

	if(inspector)
	{
		delete inspector;
	}

	return res;
}

//
// MAIN
//
int main(int argc, char **argv)
{
	digwatch_init_res res;

	res = digwatch_init(argc, argv);

	return res.m_res;
}
<commit_msg>Pass filter object into sinsp::set_filter()<commit_after>#define __STDC_FORMAT_MACROS

#include <stdio.h>
#include <iostream>
#include <time.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <algorithm>

#include <sinsp.h>
#include "lua_parser.h"
#include "digwatch.h"
#include "utils.h"

#include <unistd.h>
#include <getopt.h>

lua_parser* g_lua_parser;

static void usage();

//
// Program help
//
static void usage()
{
    printf(
	   "Usage: digwatch [options] [-p <output_format>] [filter]\n\n"
	   "Options:\n"
	   " -h, --help         Print this page\n"
	   " -M <num_seconds>   Stop collecting after <num_seconds> reached.\n"
	   " -N                 Don't convert port numbers to names.\n"
	   " -n <num>, --numevents=<num>\n"
	   "                    Stop capturing after <num> events\n"
	   " -u <filename>, --user-parser <filename>\n"
	   "                    Name of lua file containing parser\n"
	   " -r <readfile>, --read=<readfile>\n"
	   "                    Read the events from <readfile>.\n"
	   " --unbuffered       Turn off output buffering. This causes every single line\n"
	   "                    emitted by digwatch to be flushed, which generates higher CPU\n"
	   "                    usage but is useful when piping digwatch's output into another\n"
	   "                    process or into a script.\n"
	   "\n"
    );
}

//
// Event processing loop
//
captureinfo do_inspect(sinsp* inspector,
		       uint64_t cnt,
		       int duration_to_tot,
		       sinsp_evt_formatter* formatter)
{
	captureinfo retval;
	int32_t res;
	sinsp_evt* ev;
	string line;
        int duration_start = 0;

	//
	// Loop through the events
	//
	duration_start = ((double)clock()) / CLOCKS_PER_SEC;
	while(1)
	{
		if(duration_to_tot > 0)
		{
			int duration_tot = ((double)clock()) / CLOCKS_PER_SEC - duration_start;
			if(duration_tot >= duration_to_tot)
			{
				break;
			}
		}
		if(retval.m_nevts == cnt)
		{
			//
			// End of capture, either because the user stopped it, or because
			// we reached the event count specified with -n.
			//
			break;
		}

		res = inspector->next(&ev);

		if(res == SCAP_TIMEOUT)
		{
			continue;
		}
		else if(res == SCAP_EOF)
		{
			break;
		}
		else if(res != SCAP_SUCCESS)
		{
			//
			// Event read error.
			// Notify the chisels that we're exiting, and then die with an error.
			//
			cerr << "res = " << res << endl;
			throw sinsp_exception(inspector->getlasterr().c_str());
		}

		retval.m_nevts++;

		if(!inspector->is_debug_enabled() &&
			ev->get_category() & EC_INTERNAL)
		{
			continue;
		}
		if(formatter->tostring(ev, &line))
		{
			cout << line;
			cout << endl;
		}

	}

	return retval;
}

//
// ARGUMENT PARSING AND PROGRAM SETUP
//
digwatch_init_res digwatch_init(int argc, char **argv)
{
	digwatch_init_res res;
	sinsp* inspector = NULL;
	int op;
	uint64_t cnt = -1;
	sinsp_evt::param_fmt event_buffer_format = sinsp_evt::PF_NORMAL;
	int duration_to_tot = 0;
	captureinfo cinfo;
	string output_format;
	int long_index = 0;
	string user_parser;

	static struct option long_options[] =
	{
		{"help", no_argument, 0, 'h' },
		{"numevents", required_argument, 0, 'n' },
		{"user-parser", required_argument, 0, 'u' },
		{"readfile", required_argument, 0, 'r' },
		{"unbuffered", no_argument, 0, 0 },
		{0, 0, 0, 0}
	};

	output_format = "*%evt.num %evt.outputtime %evt.cpu %proc.name (%thread.tid) %evt.dir %evt.type %evt.info";

	try
	{
		inspector = new sinsp();

		//
		// Parse the args
		//
		while((op = getopt_long(argc, argv,
                                        "hM:Nn:r:u:",
                                        long_options, &long_index)) != -1)
		{
			switch(op)
			{
			case 'h':
				usage();
				delete inspector;
				return digwatch_init_res(EXIT_SUCCESS);
			case 'M':
				duration_to_tot = atoi(optarg);
				if(duration_to_tot <= 0)
				{
					throw sinsp_exception(string("invalid duration") + optarg);
					res.m_res = EXIT_FAILURE;
					goto exit;
				}
				break;
			case 'N':
				inspector->set_hostname_and_port_resolution_mode(false);
				break;
			case 'n':
				try
				{
					cnt = sinsp_numparser::parseu64(optarg);
				}
				catch(...)
				{
					throw sinsp_exception("can't parse the -n argument, make sure it's a number");
				}

				if(cnt <= 0)
				{
					throw sinsp_exception(string("invalid event count ") + optarg);
					res.m_res = EXIT_FAILURE;
					goto exit;
				}
				break;
			case 'u':
				user_parser = optarg;
				break;
			case '?':
				delete inspector;
				return digwatch_init_res(EXIT_FAILURE);
				break;
			default:
				break;
			}

		}

		inspector->set_buffer_format(event_buffer_format);

		string filter;

		//
		// the filter is at the end of the command line
		//
		if(optind < argc)
		{
#ifdef HAS_FILTERING
			for(int32_t j = optind ; j < argc; j++)
			{
				filter += argv[j];
				if(j < argc - 1)
				{
					filter += " ";
				}
			}

#else
			fprintf(stderr, "filtering not compiled.\n");
			res.m_res = EXIT_FAILURE;
			goto exit;
#endif
		}

		//
		// Create the event formatter
		//
		sinsp_evt_formatter formatter(inspector, output_format);

		g_lua_parser = new lua_parser(inspector, user_parser);

		inspector->set_filter(g_lua_parser->m_filter);
		inspector->open("");

		cinfo = do_inspect(inspector,
				   cnt,
				   duration_to_tot,
				   &formatter);

		inspector->close();
	}
	catch(sinsp_exception& e)
	{
		cerr << e.what() << endl;
		res.m_res = EXIT_FAILURE;
	}
	catch(...)
	{
		printf("Exeception\n");
		res.m_res = EXIT_FAILURE;
	}

exit:

	if(inspector)
	{
		delete inspector;
	}

	return res;
}

//
// MAIN
//
int main(int argc, char **argv)
{
	digwatch_init_res res;

	res = digwatch_init(argc, argv);

	return res.m_res;
}
<|endoftext|>


<commit_before>/*
 * Launch and manage FastCGI child processes.
 *
 * author: Max Kellermann <mk@cm4all.com>
 */

#include "Stock.hxx"
#include "Quark.hxx"
#include "stock/MapStock.hxx"
#include "stock/Class.hxx"
#include "stock/Item.hxx"
#include "child_stock.hxx"
#include "child_manager.hxx"
#include "spawn/Prepared.hxx"
#include "spawn/ChildOptions.hxx"
#include "spawn/JailConfig.hxx"
#include "gerrno.h"
#include "pool.hxx"
#include "event/Event.hxx"
#include "event/Callback.hxx"
#include "util/ConstBuffer.hxx"

#include <daemon/log.h>

#include <glib.h>

#include <assert.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <string.h>

#ifdef __linux
#include <sched.h>
#endif

struct FcgiStock {
    StockMap *hstock;
    StockMap *child_stock;

    void FadeAll() {
        hstock_fade_all(*hstock);
        hstock_fade_all(*child_stock);
    }
};

struct FcgiChildParams {
    const char *executable_path;

    ConstBuffer<const char *> args;

    const ChildOptions *options;

    FcgiChildParams(const char *_executable_path,
                    ConstBuffer<const char *> _args,
                    const ChildOptions &_options)
        :executable_path(_executable_path), args(_args),
         options(&_options) {}

    const char *GetStockKey(struct pool &pool) const;
};

struct FcgiConnection final : PoolStockItem {
    JailParams jail_params;

    struct jail_config jail_config;

    StockItem *child = nullptr;

    int fd = -1;
    Event event;

    /**
     * Is this a fresh connection to the FastCGI child process?
     */
    bool fresh;

    /**
     * Shall the FastCGI child process be killed?
     */
    bool kill;

    /**
     * Was the current request aborted by the fcgi_client caller?
     */
    bool aborted;

    explicit FcgiConnection(struct pool &_pool, CreateStockItem c)
        :PoolStockItem(_pool, c) {}

    gcc_pure
    const char *GetStockKey() const {
        return child_stock_item_key(child);
    }

    void EventCallback(evutil_socket_t fd, short events);

    /* virtual methods from class StockItem */
    bool Borrow(gcc_unused void *ctx) override;
    bool Release(gcc_unused void *ctx) override;
    void Destroy(void *ctx) override;
};

const char *
FcgiChildParams::GetStockKey(struct pool &pool) const
{
    const char *key = executable_path;

    for (auto i : args)
        key = p_strcat(&pool, key, " ", i, nullptr);

    for (auto i : options->env)
        key = p_strcat(&pool, key, "$", i, nullptr);

    char options_buffer[4096];
    *options->MakeId(options_buffer) = 0;
    if (*options_buffer != 0)
        key = p_strcat(&pool, key, options_buffer, nullptr);

    return key;
}

/*
 * libevent callback
 *
 */

inline void
FcgiConnection::EventCallback(evutil_socket_t _fd, short events)
{
    assert(_fd == fd);

    if ((events & EV_TIMEOUT) == 0) {
        char buffer;
        ssize_t nbytes = recv(_fd, &buffer, sizeof(buffer), MSG_DONTWAIT);
        if (nbytes < 0)
            daemon_log(2, "error on idle FastCGI connection '%s': %s\n",
                       GetStockKey(), strerror(errno));
        else if (nbytes > 0)
            daemon_log(2, "unexpected data from idle FastCGI connection '%s'\n",
                       GetStockKey());
    }

    InvokeIdleDisconnect();
    pool_commit();
}

/*
 * child_stock class
 *
 */

static bool
fcgi_child_stock_prepare(void *info, int fd,
                         PreparedChildProcess &p, GError **error_r)
{
    const auto &params = *(const FcgiChildParams *)info;
    const ChildOptions &options = *params.options;

    p.stdin_fd = fd;

    /* the FastCGI protocol defines a channel for stderr, so we could
       close its "real" stderr here, but many FastCGI applications
       don't use the FastCGI protocol to send error messages, so we
       just keep it open */

    int null_fd = open("/dev/null", O_WRONLY|O_CLOEXEC|O_NOCTTY);
    if (null_fd >= 0)
        p.stdout_fd = null_fd;

    p.Append(params.executable_path);
    for (auto i : params.args)
        p.Append(i);

    return options.CopyTo(p, true, nullptr, error_r);
}

static const ChildStockClass fcgi_child_stock_class = {
    .shutdown_signal = SIGUSR1,
    .socket_type = nullptr,
    .prepare = fcgi_child_stock_prepare,
};

/*
 * stock class
 *
 */

static void
fcgi_stock_create(void *ctx, struct pool &parent_pool, CreateStockItem c,
                  const char *key, void *info,
                  struct pool &caller_pool,
                  gcc_unused struct async_operation_ref &async_ref)
{
    FcgiStock *fcgi_stock = (FcgiStock *)ctx;
    FcgiChildParams *params = (FcgiChildParams *)info;

    assert(key != nullptr);
    assert(params != nullptr);
    assert(params->executable_path != nullptr);

    auto &pool = *pool_new_linear(&parent_pool, "fcgi_connection", 2048);
    auto *connection = NewFromPool<FcgiConnection>(pool, pool, c);

    const ChildOptions *const options = params->options;
    if (options->jail.enabled) {
        connection->jail_params.CopyFrom(pool, options->jail);

        if (!jail_config_load(&connection->jail_config,
                              "/etc/cm4all/jailcgi/jail.conf", &pool)) {
            GError *error = g_error_new(fcgi_quark(), 0,
                                        "Failed to load /etc/cm4all/jailcgi/jail.conf");
            connection->InvokeCreateError(error);
            return;
        }
    } else
        connection->jail_params.enabled = false;

    GError *error = nullptr;
    connection->child = hstock_get_now(*fcgi_stock->child_stock, caller_pool,
                                       key, params, &error);
    if (connection->child == nullptr) {
        g_prefix_error(&error, "failed to start to FastCGI server '%s': ",
                       key);
        connection->InvokeCreateError(error);
        return;
    }

    connection->fd = child_stock_item_connect(connection->child, &error);
    if (connection->fd < 0) {
        g_prefix_error(&error, "failed to connect to FastCGI server '%s': ",
                       key);

        connection->kill = true;
        connection->InvokeCreateError(error);
        return;
    }

    connection->fresh = true;
    connection->kill = false;
    connection->aborted = false;

    connection->event.Set(connection->fd, EV_READ|EV_TIMEOUT,
                          MakeEventCallback(FcgiConnection, EventCallback),
                          connection);

    connection->InvokeCreateSuccess();
}

bool
FcgiConnection::Borrow(gcc_unused void *ctx)
{
    /* check the connection status before using it, just in case the
       FastCGI server has decided to close the connection before
       fcgi_connection_event_callback() got invoked */
    char buffer;
    ssize_t nbytes = recv(fd, &buffer, sizeof(buffer), MSG_DONTWAIT);
    if (nbytes > 0) {
        daemon_log(2, "unexpected data from idle FastCGI connection '%s'\n",
                   GetStockKey());
        return false;
    } else if (nbytes == 0) {
        /* connection closed (not worth a log message) */
        return false;
    } else if (errno != EAGAIN) {
        daemon_log(2, "error on idle FastCGI connection '%s': %s\n",
                   GetStockKey(), strerror(errno));
        return false;
    }

    event.Delete();
    aborted = false;
    return true;
}

bool
FcgiConnection::Release(gcc_unused void *ctx)
{
    static constexpr struct timeval tv = {
        .tv_sec = 300,
        .tv_usec = 0,
    };

    fresh = false;
    event.Add(tv);
    return true;
}

void
FcgiConnection::Destroy(void *ctx)
{
    FcgiStock *fcgi_stock = (FcgiStock *)ctx;

    if (fd >= 0) {
        event.Delete();
        close(fd);
    }

    if (child != nullptr)
        child_stock_put(fcgi_stock->child_stock, child, kill);

    PoolStockItem::Destroy(ctx);
}

static constexpr StockClass fcgi_stock_class = {
    .create = fcgi_stock_create,
};

/*
 * interface
 *
 */

FcgiStock *
fcgi_stock_new(struct pool *pool, unsigned limit, unsigned max_idle)
{
    auto fcgi_stock = NewFromPool<FcgiStock>(*pool);
    fcgi_stock->child_stock = child_stock_new(pool, limit, max_idle,
                                              &fcgi_child_stock_class);
    fcgi_stock->hstock = hstock_new(*pool, fcgi_stock_class, fcgi_stock,
                                    limit, max_idle);

    return fcgi_stock;
}

void
fcgi_stock_free(FcgiStock *fcgi_stock)
{
    hstock_free(fcgi_stock->hstock);
    hstock_free(fcgi_stock->child_stock);
}

void
fcgi_stock_fade_all(FcgiStock &fs)
{
    fs.FadeAll();
}

StockItem *
fcgi_stock_get(FcgiStock *fcgi_stock, struct pool *pool,
               const ChildOptions &options,
               const char *executable_path,
               ConstBuffer<const char *> args,
               GError **error_r)
{
    auto params = NewFromPool<FcgiChildParams>(*pool, executable_path,
                                               args, options);

    return hstock_get_now(*fcgi_stock->hstock, *pool,
                          params->GetStockKey(*pool), params,
                          error_r);
}

int
fcgi_stock_item_get_domain(gcc_unused const StockItem &item)
{
    return AF_UNIX;
}

int
fcgi_stock_item_get(const StockItem &item)
{
    const auto *connection = (const FcgiConnection *)&item;

    assert(connection->fd >= 0);

    return connection->fd;
}

const char *
fcgi_stock_translate_path(const StockItem &item,
                          const char *path, struct pool *pool)
{
    const auto *connection = (const FcgiConnection *)&item;

    if (!connection->jail_params.enabled)
        /* no JailCGI - application's namespace is the same as ours,
           no translation needed */
        return path;

    const char *jailed = jail_translate_path(&connection->jail_config, path,
                                             connection->jail_params.home_directory,
                                             pool);
    return jailed != nullptr ? jailed : path;
}

void
fcgi_stock_put(FcgiStock *fcgi_stock, StockItem &item,
               bool destroy)
{
    auto *connection = (FcgiConnection *)&item;

    if (connection->fresh && connection->aborted && destroy)
        /* the fcgi_client caller has aborted the request before the
           first response on a fresh connection was received: better
           kill the child process, it may be failing on us
           completely */
        connection->kill = true;

    hstock_put(*fcgi_stock->hstock, connection->GetStockKey(),
               item, destroy);
}

void
fcgi_stock_aborted(StockItem &item)
{
    auto *connection = (FcgiConnection *)&item;

    connection->aborted = true;
}
<commit_msg>fcgi/Stock: convert pointer to reference<commit_after>/*
 * Launch and manage FastCGI child processes.
 *
 * author: Max Kellermann <mk@cm4all.com>
 */

#include "Stock.hxx"
#include "Quark.hxx"
#include "stock/MapStock.hxx"
#include "stock/Class.hxx"
#include "stock/Item.hxx"
#include "child_stock.hxx"
#include "child_manager.hxx"
#include "spawn/Prepared.hxx"
#include "spawn/ChildOptions.hxx"
#include "spawn/JailConfig.hxx"
#include "gerrno.h"
#include "pool.hxx"
#include "event/Event.hxx"
#include "event/Callback.hxx"
#include "util/ConstBuffer.hxx"

#include <daemon/log.h>

#include <glib.h>

#include <assert.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <string.h>

#ifdef __linux
#include <sched.h>
#endif

struct FcgiStock {
    StockMap *hstock;
    StockMap *child_stock;

    void FadeAll() {
        hstock_fade_all(*hstock);
        hstock_fade_all(*child_stock);
    }
};

struct FcgiChildParams {
    const char *executable_path;

    ConstBuffer<const char *> args;

    const ChildOptions &options;

    FcgiChildParams(const char *_executable_path,
                    ConstBuffer<const char *> _args,
                    const ChildOptions &_options)
        :executable_path(_executable_path), args(_args),
         options(_options) {}

    const char *GetStockKey(struct pool &pool) const;
};

struct FcgiConnection final : PoolStockItem {
    JailParams jail_params;

    struct jail_config jail_config;

    StockItem *child = nullptr;

    int fd = -1;
    Event event;

    /**
     * Is this a fresh connection to the FastCGI child process?
     */
    bool fresh;

    /**
     * Shall the FastCGI child process be killed?
     */
    bool kill;

    /**
     * Was the current request aborted by the fcgi_client caller?
     */
    bool aborted;

    explicit FcgiConnection(struct pool &_pool, CreateStockItem c)
        :PoolStockItem(_pool, c) {}

    gcc_pure
    const char *GetStockKey() const {
        return child_stock_item_key(child);
    }

    void EventCallback(evutil_socket_t fd, short events);

    /* virtual methods from class StockItem */
    bool Borrow(gcc_unused void *ctx) override;
    bool Release(gcc_unused void *ctx) override;
    void Destroy(void *ctx) override;
};

const char *
FcgiChildParams::GetStockKey(struct pool &pool) const
{
    const char *key = executable_path;

    for (auto i : args)
        key = p_strcat(&pool, key, " ", i, nullptr);

    for (auto i : options.env)
        key = p_strcat(&pool, key, "$", i, nullptr);

    char options_buffer[4096];
    *options.MakeId(options_buffer) = 0;
    if (*options_buffer != 0)
        key = p_strcat(&pool, key, options_buffer, nullptr);

    return key;
}

/*
 * libevent callback
 *
 */

inline void
FcgiConnection::EventCallback(evutil_socket_t _fd, short events)
{
    assert(_fd == fd);

    if ((events & EV_TIMEOUT) == 0) {
        char buffer;
        ssize_t nbytes = recv(_fd, &buffer, sizeof(buffer), MSG_DONTWAIT);
        if (nbytes < 0)
            daemon_log(2, "error on idle FastCGI connection '%s': %s\n",
                       GetStockKey(), strerror(errno));
        else if (nbytes > 0)
            daemon_log(2, "unexpected data from idle FastCGI connection '%s'\n",
                       GetStockKey());
    }

    InvokeIdleDisconnect();
    pool_commit();
}

/*
 * child_stock class
 *
 */

static bool
fcgi_child_stock_prepare(void *info, int fd,
                         PreparedChildProcess &p, GError **error_r)
{
    const auto &params = *(const FcgiChildParams *)info;
    const ChildOptions &options = params.options;

    p.stdin_fd = fd;

    /* the FastCGI protocol defines a channel for stderr, so we could
       close its "real" stderr here, but many FastCGI applications
       don't use the FastCGI protocol to send error messages, so we
       just keep it open */

    int null_fd = open("/dev/null", O_WRONLY|O_CLOEXEC|O_NOCTTY);
    if (null_fd >= 0)
        p.stdout_fd = null_fd;

    p.Append(params.executable_path);
    for (auto i : params.args)
        p.Append(i);

    return options.CopyTo(p, true, nullptr, error_r);
}

static const ChildStockClass fcgi_child_stock_class = {
    .shutdown_signal = SIGUSR1,
    .socket_type = nullptr,
    .prepare = fcgi_child_stock_prepare,
};

/*
 * stock class
 *
 */

static void
fcgi_stock_create(void *ctx, struct pool &parent_pool, CreateStockItem c,
                  const char *key, void *info,
                  struct pool &caller_pool,
                  gcc_unused struct async_operation_ref &async_ref)
{
    FcgiStock *fcgi_stock = (FcgiStock *)ctx;
    FcgiChildParams *params = (FcgiChildParams *)info;

    assert(key != nullptr);
    assert(params != nullptr);
    assert(params->executable_path != nullptr);

    auto &pool = *pool_new_linear(&parent_pool, "fcgi_connection", 2048);
    auto *connection = NewFromPool<FcgiConnection>(pool, pool, c);

    const ChildOptions &options = params->options;
    if (options.jail.enabled) {
        connection->jail_params.CopyFrom(pool, options.jail);

        if (!jail_config_load(&connection->jail_config,
                              "/etc/cm4all/jailcgi/jail.conf", &pool)) {
            GError *error = g_error_new(fcgi_quark(), 0,
                                        "Failed to load /etc/cm4all/jailcgi/jail.conf");
            connection->InvokeCreateError(error);
            return;
        }
    } else
        connection->jail_params.enabled = false;

    GError *error = nullptr;
    connection->child = hstock_get_now(*fcgi_stock->child_stock, caller_pool,
                                       key, params, &error);
    if (connection->child == nullptr) {
        g_prefix_error(&error, "failed to start to FastCGI server '%s': ",
                       key);
        connection->InvokeCreateError(error);
        return;
    }

    connection->fd = child_stock_item_connect(connection->child, &error);
    if (connection->fd < 0) {
        g_prefix_error(&error, "failed to connect to FastCGI server '%s': ",
                       key);

        connection->kill = true;
        connection->InvokeCreateError(error);
        return;
    }

    connection->fresh = true;
    connection->kill = false;
    connection->aborted = false;

    connection->event.Set(connection->fd, EV_READ|EV_TIMEOUT,
                          MakeEventCallback(FcgiConnection, EventCallback),
                          connection);

    connection->InvokeCreateSuccess();
}

bool
FcgiConnection::Borrow(gcc_unused void *ctx)
{
    /* check the connection status before using it, just in case the
       FastCGI server has decided to close the connection before
       fcgi_connection_event_callback() got invoked */
    char buffer;
    ssize_t nbytes = recv(fd, &buffer, sizeof(buffer), MSG_DONTWAIT);
    if (nbytes > 0) {
        daemon_log(2, "unexpected data from idle FastCGI connection '%s'\n",
                   GetStockKey());
        return false;
    } else if (nbytes == 0) {
        /* connection closed (not worth a log message) */
        return false;
    } else if (errno != EAGAIN) {
        daemon_log(2, "error on idle FastCGI connection '%s': %s\n",
                   GetStockKey(), strerror(errno));
        return false;
    }

    event.Delete();
    aborted = false;
    return true;
}

bool
FcgiConnection::Release(gcc_unused void *ctx)
{
    static constexpr struct timeval tv = {
        .tv_sec = 300,
        .tv_usec = 0,
    };

    fresh = false;
    event.Add(tv);
    return true;
}

void
FcgiConnection::Destroy(void *ctx)
{
    FcgiStock *fcgi_stock = (FcgiStock *)ctx;

    if (fd >= 0) {
        event.Delete();
        close(fd);
    }

    if (child != nullptr)
        child_stock_put(fcgi_stock->child_stock, child, kill);

    PoolStockItem::Destroy(ctx);
}

static constexpr StockClass fcgi_stock_class = {
    .create = fcgi_stock_create,
};

/*
 * interface
 *
 */

FcgiStock *
fcgi_stock_new(struct pool *pool, unsigned limit, unsigned max_idle)
{
    auto fcgi_stock = NewFromPool<FcgiStock>(*pool);
    fcgi_stock->child_stock = child_stock_new(pool, limit, max_idle,
                                              &fcgi_child_stock_class);
    fcgi_stock->hstock = hstock_new(*pool, fcgi_stock_class, fcgi_stock,
                                    limit, max_idle);

    return fcgi_stock;
}

void
fcgi_stock_free(FcgiStock *fcgi_stock)
{
    hstock_free(fcgi_stock->hstock);
    hstock_free(fcgi_stock->child_stock);
}

void
fcgi_stock_fade_all(FcgiStock &fs)
{
    fs.FadeAll();
}

StockItem *
fcgi_stock_get(FcgiStock *fcgi_stock, struct pool *pool,
               const ChildOptions &options,
               const char *executable_path,
               ConstBuffer<const char *> args,
               GError **error_r)
{
    auto params = NewFromPool<FcgiChildParams>(*pool, executable_path,
                                               args, options);

    return hstock_get_now(*fcgi_stock->hstock, *pool,
                          params->GetStockKey(*pool), params,
                          error_r);
}

int
fcgi_stock_item_get_domain(gcc_unused const StockItem &item)
{
    return AF_UNIX;
}

int
fcgi_stock_item_get(const StockItem &item)
{
    const auto *connection = (const FcgiConnection *)&item;

    assert(connection->fd >= 0);

    return connection->fd;
}

const char *
fcgi_stock_translate_path(const StockItem &item,
                          const char *path, struct pool *pool)
{
    const auto *connection = (const FcgiConnection *)&item;

    if (!connection->jail_params.enabled)
        /* no JailCGI - application's namespace is the same as ours,
           no translation needed */
        return path;

    const char *jailed = jail_translate_path(&connection->jail_config, path,
                                             connection->jail_params.home_directory,
                                             pool);
    return jailed != nullptr ? jailed : path;
}

void
fcgi_stock_put(FcgiStock *fcgi_stock, StockItem &item,
               bool destroy)
{
    auto *connection = (FcgiConnection *)&item;

    if (connection->fresh && connection->aborted && destroy)
        /* the fcgi_client caller has aborted the request before the
           first response on a fresh connection was received: better
           kill the child process, it may be failing on us
           completely */
        connection->kill = true;

    hstock_put(*fcgi_stock->hstock, connection->GetStockKey(),
               item, destroy);
}

void
fcgi_stock_aborted(StockItem &item)
{
    auto *connection = (FcgiConnection *)&item;

    connection->aborted = true;
}
<|endoftext|>


<commit_before>/* -*- Mode: C; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */

/**
 * @copyright 2013 Couchbase, Inc.
 *
 * @author Filipe Manana  <filipe@couchbase.com>
 * @author Aliaksey Kandratsenka <alk@tut.by> (small optimization)
 *
 * 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.
 **/

#include "file_merger.h"
#include <stdlib.h>
#include <assert.h>
#include <string.h>

typedef struct {
    void      *data;
    unsigned  file;
} record_t;

#define FREE_RECORD(ctx, rec)                                \
    do {                                                     \
        (*(ctx)->free_record)((rec)->data, (ctx)->user_ctx); \
        free((rec));                                         \
    } while (0)

struct file_merger_ctx_t;

typedef struct {
    struct file_merger_ctx_t  *ctx;
    record_t                  **data;
    unsigned                  count;
} sorted_vector_t;

typedef struct file_merger_ctx_t {
    unsigned                            num_files;
    FILE                                **files;
    FILE                                *dest_file;
    file_merger_read_record_t           read_record;
    file_merger_write_record_t          write_record;
    file_merger_record_free_t           free_record;
    file_merger_compare_records_t       compare_records;
    file_merger_deduplicate_records_t   dedup_records;
    file_merger_feed_record_t           feed_record;
    void                                *user_ctx;
    sorted_vector_t                     sorted_vector;
} file_merger_ctx_t;

static int  init_sorted_vector(sorted_vector_t *sorted_vector, unsigned max_elements, file_merger_ctx_t *ctx);
static void sorted_vector_destroy(sorted_vector_t *sorted_vector);
static void sorted_vector_pop(sorted_vector_t *sorted_vector,
                              record_t ***records,
                              size_t *n);
static int  sorted_vector_add(sorted_vector_t *sorted_vector, record_t *record);

static file_merger_error_t do_merge_files(file_merger_ctx_t *ctx);

file_merger_error_t merge_files(const char *source_files[],
                                unsigned num_files,
                                const char *dest_file,
                                file_merger_read_record_t read_record,
                                file_merger_write_record_t write_record,
                                file_merger_feed_record_t feed_record,
                                file_merger_compare_records_t compare_records,
                                file_merger_deduplicate_records_t dedup_records,
                                file_merger_record_free_t free_record,
                                int skip_writeback,
                                void *user_ctx)
{
    file_merger_ctx_t ctx;
    file_merger_error_t ret;
    unsigned i, j;

    if (num_files == 0) {
        return FILE_MERGER_ERROR_BAD_ARG;
    }

    ctx.num_files = num_files;
    ctx.read_record = read_record;
    ctx.write_record = write_record;
    ctx.free_record = free_record;
    ctx.compare_records = compare_records;
    ctx.user_ctx = user_ctx;
    ctx.feed_record = feed_record;
    ctx.dedup_records = dedup_records;

    if (feed_record && skip_writeback) {
        ctx.dest_file = NULL;
    } else {
        ctx.dest_file = fopen(dest_file, "ab");
    }

    if (!init_sorted_vector(&ctx.sorted_vector, num_files, &ctx)) {
        return FILE_MERGER_ERROR_ALLOC;
    }

    if (feed_record == NULL && ctx.dest_file == NULL) {
        sorted_vector_destroy(&ctx.sorted_vector);
        return FILE_MERGER_ERROR_OPEN_FILE;
    }

    ctx.files = (FILE **) malloc(sizeof(FILE *) * num_files);

    if (ctx.files == NULL) {
        sorted_vector_destroy(&ctx.sorted_vector);
        fclose(ctx.dest_file);
        return FILE_MERGER_ERROR_ALLOC;
    }

    for (i = 0; i < num_files; ++i) {
        ctx.files[i] = fopen(source_files[i], "rb");

        if (ctx.files[i] == NULL) {
            for (j = 0; j < i; ++j) {
                fclose(ctx.files[j]);
            }
            free(ctx.files);
            fclose(ctx.dest_file);
            sorted_vector_destroy(&ctx.sorted_vector);

            return FILE_MERGER_ERROR_OPEN_FILE;
        }
    }

    ret = do_merge_files(&ctx);

    for (i = 0; i < ctx.num_files; ++i) {
        if (ctx.files[i] != NULL) {
            fclose(ctx.files[i]);
        }
    }
    free(ctx.files);
    sorted_vector_destroy(&ctx.sorted_vector);
    if (ctx.dest_file) {
        fclose(ctx.dest_file);
    }

    return ret;
}

static file_merger_error_t do_merge_files(file_merger_ctx_t *ctx)
{
    unsigned i;

    for (i = 0; i < ctx->num_files; ++i) {
        FILE *f = ctx->files[i];
        int record_len;
        void *record_data;
        record_t *record;

        record_len = (*ctx->read_record)(f, &record_data, ctx->user_ctx);

        if (record_len == 0) {
            fclose(f);
            ctx->files[i] = NULL;
        } else if (record_len < 0) {
            return (file_merger_error_t) record_len;
        } else {
            int rv;
            record = (record_t *) malloc(sizeof(*record));
            if (record == NULL) {
                return FILE_MERGER_ERROR_ALLOC;
            }
            record->data = record_data;
            record->file = i;
            rv = sorted_vector_add(&ctx->sorted_vector, record);
            if (!rv) {
                return FILE_MERGER_SORT_ERROR;
            }
        }
    }

    while (ctx->sorted_vector.count != 0) {
        record_t **records;
        size_t n;
        size_t i;
        void *record_data;
        int record_len;
        file_merger_error_t ret;

        /* The head of the list is the required item which needs to be written
         * to the output destination records file.
         * For each duplicated item that is eliminated (elements of linked
         * list), we need to read one record from the same file as the
         * duplicated record came from and add them to the sort vector.
         */
        sorted_vector_pop(&ctx->sorted_vector, &records, &n);
        assert(records != NULL);
        assert(n != 0);

        if (ctx->feed_record) {
            ret = (*ctx->feed_record)(records[0]->data, ctx->user_ctx);
            if (ret != FILE_MERGER_SUCCESS) {
                FREE_RECORD(ctx, records[0]);
                return ret;
            }
        } else {
            assert(ctx->dest_file != NULL);
        }

        if (ctx->dest_file) {
            ret = (*ctx->write_record)(ctx->dest_file, records[0]->data, ctx->user_ctx);
            if (ret != FILE_MERGER_SUCCESS) {
                FREE_RECORD(ctx, records[0]);
                return ret;
            }
        }

        for (i = 0; i < n; i++) {
            record_len = (*ctx->read_record)(ctx->files[records[i]->file],
                                             &record_data,
                                             ctx->user_ctx);
            if (record_len == 0) {
                fclose(ctx->files[records[i]->file]);
                ctx->files[records[i]->file] = NULL;
                FREE_RECORD(ctx, records[i]);

            } else if (record_len < 0) {
                FREE_RECORD(ctx, records[i]);

                return (file_merger_error_t) record_len;
            } else {
                int rv;
                (*ctx->free_record)(records[i]->data, ctx->user_ctx);
                records[i]->data = record_data;
                rv = sorted_vector_add(&ctx->sorted_vector, records[i]);
                if (!rv) {
                    return FILE_MERGER_SORT_ERROR;
                }
            }
        }

        free(records);
    }

    return FILE_MERGER_SUCCESS;
}

static int init_sorted_vector(sorted_vector_t *sorted_vector,
                              unsigned max_elements,
                              file_merger_ctx_t *ctx)
{
    sorted_vector->data = (record_t **) malloc(sizeof(record_t *) * max_elements);
    if (sorted_vector->data == NULL) {
        return 0;
    }

    sorted_vector->count = 0;
    sorted_vector->ctx = ctx;

    return 1;
}

static void sorted_vector_destroy(sorted_vector_t *sorted_vector)
{
    unsigned i;

    for (i = 0; i < sorted_vector->count; ++i) {
        FREE_RECORD(sorted_vector->ctx, sorted_vector->data[i]);
    }

    free(sorted_vector->data);
}

#define SORTED_VECTOR_CMP(h, a, b)  \
    (*(h)->ctx->compare_records)((a)->data, (b)->data, (h)->ctx->user_ctx)

static void sorted_vector_pop(sorted_vector_t *sorted_vector,
                              record_t ***records,
                              size_t *n)
{
    record_t *head;
    file_merger_record_t **duplicates;
    size_t i, j;

    if (sorted_vector->count == 0) {
        *records = NULL;
        *n = 0;
        return;
    }

    /* For deduplication, return the list of records whose keys are equal.
     * Hence they can be eliminated from the sort vector and least element can
     * be picked for writing out to the output file.
     * */
    head = sorted_vector->data[0];

    for (i = 1; sorted_vector->ctx->dedup_records != NULL && i < sorted_vector->count; i++) {
        if (SORTED_VECTOR_CMP(sorted_vector, head, sorted_vector->data[i]) != 0) {
            break;
        }
    }

    *records = (record_t **) malloc(sizeof(record_t *) * i);
    memcpy(*records, sorted_vector->data, sizeof(record_t *) * i);
    *n = i;

    if (i > 1) {
        record_t *tmp;
        duplicates = (file_merger_record_t **) *records;
        j = sorted_vector->ctx->dedup_records(duplicates, i, sorted_vector->ctx->user_ctx);
        tmp = (*records)[0];
        (*records)[0] = (*records)[j];
        (*records)[j] = tmp;
    }

    sorted_vector->count -= i;
    memmove(sorted_vector->data + 0, sorted_vector->data + i,
            sizeof(sorted_vector->data[0]) * (sorted_vector->count));
}

static int sorted_vector_add(sorted_vector_t *sorted_vector, record_t *record)
{
    unsigned l, r;

    if (sorted_vector->count == sorted_vector->ctx->num_files) {
        /* sorted_vector full */
        return 0;
    }

    l = 0;
    r = sorted_vector->count;
    while (r - l > 1) {
        unsigned pos = (l + r) / 2;

        if (SORTED_VECTOR_CMP(sorted_vector, record, sorted_vector->data[pos]) < 0) {
            r = pos;
        } else {
            l = pos;
        }
    }

    if (l == 0 && r != 0) {
        if (SORTED_VECTOR_CMP(sorted_vector, record, sorted_vector->data[0]) < 0) {
            r = 0;
        }
    }

    if (r < sorted_vector->count) {
        memmove(sorted_vector->data + r + 1,
                sorted_vector->data + r,
                sizeof(sorted_vector->data[0]) * (sorted_vector->count - r));
    }

    sorted_vector->count += 1;

    sorted_vector->data[r] = record;

    return 1;
}
<commit_msg>MB-17154: Free records in file merger correctly<commit_after>/* -*- Mode: C; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */

/**
 * @copyright 2013 Couchbase, Inc.
 *
 * @author Filipe Manana  <filipe@couchbase.com>
 * @author Aliaksey Kandratsenka <alk@tut.by> (small optimization)
 *
 * 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.
 **/

#include "file_merger.h"
#include <stdlib.h>
#include <assert.h>
#include <string.h>

typedef struct {
    void      *data;
    unsigned  file;
} record_t;

#define FREE_RECORD(ctx, rec)                                \
    do {                                                     \
        (*(ctx)->free_record)((rec)->data, (ctx)->user_ctx); \
        free((rec));                                         \
    } while (0)

struct file_merger_ctx_t;

typedef struct {
    struct file_merger_ctx_t  *ctx;
    record_t                  **data;
    unsigned                  count;
} sorted_vector_t;

typedef struct file_merger_ctx_t {
    unsigned                            num_files;
    FILE                                **files;
    FILE                                *dest_file;
    file_merger_read_record_t           read_record;
    file_merger_write_record_t          write_record;
    file_merger_record_free_t           free_record;
    file_merger_compare_records_t       compare_records;
    file_merger_deduplicate_records_t   dedup_records;
    file_merger_feed_record_t           feed_record;
    void                                *user_ctx;
    sorted_vector_t                     sorted_vector;
} file_merger_ctx_t;

static int  init_sorted_vector(sorted_vector_t *sorted_vector, unsigned max_elements, file_merger_ctx_t *ctx);
static void sorted_vector_destroy(sorted_vector_t *sorted_vector);
static void sorted_vector_pop(sorted_vector_t *sorted_vector,
                              record_t ***records,
                              size_t *n);
static int  sorted_vector_add(sorted_vector_t *sorted_vector, record_t *record);

static file_merger_error_t do_merge_files(file_merger_ctx_t *ctx);
static void free_all_records(file_merger_ctx_t *ctx, record_t **records,
                             int offset, int num);

file_merger_error_t merge_files(const char *source_files[],
                                unsigned num_files,
                                const char *dest_file,
                                file_merger_read_record_t read_record,
                                file_merger_write_record_t write_record,
                                file_merger_feed_record_t feed_record,
                                file_merger_compare_records_t compare_records,
                                file_merger_deduplicate_records_t dedup_records,
                                file_merger_record_free_t free_record,
                                int skip_writeback,
                                void *user_ctx)
{
    file_merger_ctx_t ctx;
    file_merger_error_t ret;
    unsigned i, j;

    if (num_files == 0) {
        return FILE_MERGER_ERROR_BAD_ARG;
    }

    ctx.num_files = num_files;
    ctx.read_record = read_record;
    ctx.write_record = write_record;
    ctx.free_record = free_record;
    ctx.compare_records = compare_records;
    ctx.user_ctx = user_ctx;
    ctx.feed_record = feed_record;
    ctx.dedup_records = dedup_records;

    if (feed_record && skip_writeback) {
        ctx.dest_file = NULL;
    } else {
        ctx.dest_file = fopen(dest_file, "ab");
    }

    if (!init_sorted_vector(&ctx.sorted_vector, num_files, &ctx)) {
        return FILE_MERGER_ERROR_ALLOC;
    }

    if (feed_record == NULL && ctx.dest_file == NULL) {
        sorted_vector_destroy(&ctx.sorted_vector);
        return FILE_MERGER_ERROR_OPEN_FILE;
    }

    ctx.files = (FILE **) malloc(sizeof(FILE *) * num_files);

    if (ctx.files == NULL) {
        sorted_vector_destroy(&ctx.sorted_vector);
        fclose(ctx.dest_file);
        return FILE_MERGER_ERROR_ALLOC;
    }

    for (i = 0; i < num_files; ++i) {
        ctx.files[i] = fopen(source_files[i], "rb");

        if (ctx.files[i] == NULL) {
            for (j = 0; j < i; ++j) {
                fclose(ctx.files[j]);
            }
            free(ctx.files);
            fclose(ctx.dest_file);
            sorted_vector_destroy(&ctx.sorted_vector);

            return FILE_MERGER_ERROR_OPEN_FILE;
        }
    }

    ret = do_merge_files(&ctx);

    for (i = 0; i < ctx.num_files; ++i) {
        if (ctx.files[i] != NULL) {
            fclose(ctx.files[i]);
        }
    }
    free(ctx.files);
    sorted_vector_destroy(&ctx.sorted_vector);
    if (ctx.dest_file) {
        fclose(ctx.dest_file);
    }

    return ret;
}

static file_merger_error_t do_merge_files(file_merger_ctx_t *ctx)
{
    unsigned i;

    for (i = 0; i < ctx->num_files; ++i) {
        FILE *f = ctx->files[i];
        int record_len;
        void *record_data;
        record_t *record;

        record_len = (*ctx->read_record)(f, &record_data, ctx->user_ctx);

        if (record_len == 0) {
            fclose(f);
            ctx->files[i] = NULL;
        } else if (record_len < 0) {
            return (file_merger_error_t) record_len;
        } else {
            int rv;
            record = (record_t *) malloc(sizeof(*record));
            if (record == NULL) {
                return FILE_MERGER_ERROR_ALLOC;
            }
            record->data = record_data;
            record->file = i;
            rv = sorted_vector_add(&ctx->sorted_vector, record);
            if (!rv) {
                FREE_RECORD(ctx, record);
                return FILE_MERGER_SORT_ERROR;
            }
        }
    }

    while (ctx->sorted_vector.count != 0) {
        record_t **records;
        size_t n;
        size_t i;
        void *record_data;
        int record_len;
        file_merger_error_t ret;

        /* The head of the list is the required item which needs to be written
         * to the output destination records file.
         * For each duplicated item that is eliminated (elements of linked
         * list), we need to read one record from the same file as the
         * duplicated record came from and add them to the sort vector.
         */
        sorted_vector_pop(&ctx->sorted_vector, &records, &n);
        assert(records != NULL);
        assert(n != 0);

        if (ctx->feed_record) {
            ret = (*ctx->feed_record)(records[0]->data, ctx->user_ctx);
            if (ret != FILE_MERGER_SUCCESS) {
                free_all_records(ctx, records, 0, n);
                return ret;
            }
        } else {
            assert(ctx->dest_file != NULL);
        }

        if (ctx->dest_file) {
            ret = (*ctx->write_record)(ctx->dest_file, records[0]->data, ctx->user_ctx);
            if (ret != FILE_MERGER_SUCCESS) {
                free_all_records(ctx, records, 0, n);
                return ret;
            }
        }

        for (i = 0; i < n; i++) {
            record_len = (*ctx->read_record)(ctx->files[records[i]->file],
                                             &record_data,
                                             ctx->user_ctx);
            if (record_len == 0) {
                fclose(ctx->files[records[i]->file]);
                ctx->files[records[i]->file] = NULL;
                FREE_RECORD(ctx, records[i]);

            } else if (record_len < 0) {
                free_all_records(ctx, records, i, n);
                return (file_merger_error_t) record_len;
            } else {
                int rv;
                (*ctx->free_record)(records[i]->data, ctx->user_ctx);
                records[i]->data = record_data;
                rv = sorted_vector_add(&ctx->sorted_vector, records[i]);
                if (!rv) {
                    free_all_records(ctx, records, i, n);
                    return FILE_MERGER_SORT_ERROR;
                }
            }
        }

        free(records);
    }

    return FILE_MERGER_SUCCESS;
}

static void free_all_records(file_merger_ctx_t *ctx, record_t **records,
                             int offset, int num) {
    for (; offset < num; offset++) {
        FREE_RECORD(ctx, records[offset]);
    }
    free(records);
}

static int init_sorted_vector(sorted_vector_t *sorted_vector,
                              unsigned max_elements,
                              file_merger_ctx_t *ctx)
{
    sorted_vector->data = (record_t **) malloc(sizeof(record_t *) * max_elements);
    if (sorted_vector->data == NULL) {
        return 0;
    }

    sorted_vector->count = 0;
    sorted_vector->ctx = ctx;

    return 1;
}

static void sorted_vector_destroy(sorted_vector_t *sorted_vector)
{
    unsigned i;

    for (i = 0; i < sorted_vector->count; ++i) {
        FREE_RECORD(sorted_vector->ctx, sorted_vector->data[i]);
    }

    free(sorted_vector->data);
}

#define SORTED_VECTOR_CMP(h, a, b)  \
    (*(h)->ctx->compare_records)((a)->data, (b)->data, (h)->ctx->user_ctx)

static void sorted_vector_pop(sorted_vector_t *sorted_vector,
                              record_t ***records,
                              size_t *n)
{
    record_t *head;
    file_merger_record_t **duplicates;
    size_t i, j;

    if (sorted_vector->count == 0) {
        *records = NULL;
        *n = 0;
        return;
    }

    /* For deduplication, return the list of records whose keys are equal.
     * Hence they can be eliminated from the sort vector and least element can
     * be picked for writing out to the output file.
     * */
    head = sorted_vector->data[0];

    for (i = 1; sorted_vector->ctx->dedup_records != NULL && i < sorted_vector->count; i++) {
        if (SORTED_VECTOR_CMP(sorted_vector, head, sorted_vector->data[i]) != 0) {
            break;
        }
    }

    *records = (record_t **) malloc(sizeof(record_t *) * i);
    memcpy(*records, sorted_vector->data, sizeof(record_t *) * i);
    *n = i;

    if (i > 1) {
        record_t *tmp;
        duplicates = (file_merger_record_t **) *records;
        j = sorted_vector->ctx->dedup_records(duplicates, i, sorted_vector->ctx->user_ctx);
        tmp = (*records)[0];
        (*records)[0] = (*records)[j];
        (*records)[j] = tmp;
    }

    sorted_vector->count -= i;
    memmove(sorted_vector->data + 0, sorted_vector->data + i,
            sizeof(sorted_vector->data[0]) * (sorted_vector->count));
}

static int sorted_vector_add(sorted_vector_t *sorted_vector, record_t *record)
{
    unsigned l, r;

    if (sorted_vector->count == sorted_vector->ctx->num_files) {
        /* sorted_vector full */
        return 0;
    }

    l = 0;
    r = sorted_vector->count;
    while (r - l > 1) {
        unsigned pos = (l + r) / 2;

        if (SORTED_VECTOR_CMP(sorted_vector, record, sorted_vector->data[pos]) < 0) {
            r = pos;
        } else {
            l = pos;
        }
    }

    if (l == 0 && r != 0) {
        if (SORTED_VECTOR_CMP(sorted_vector, record, sorted_vector->data[0]) < 0) {
            r = 0;
        }
    }

    if (r < sorted_vector->count) {
        memmove(sorted_vector->data + r + 1,
                sorted_vector->data + r,
                sizeof(sorted_vector->data[0]) * (sorted_vector->count - r));
    }

    sorted_vector->count += 1;

    sorted_vector->data[r] = record;

    return 1;
}
<|endoftext|>