<commit_before>#ifndef INCLUDE_AL_GRAPHICS_SHADER_HPP
#define INCLUDE_AL_GRAPHICS_SHADER_HPP

/*	Allocore --
	Multimedia / virtual environment application class library

	Copyright (C) 2009. AlloSphere Research Group, Media Arts & Technology, UCSB.
	Copyright (C) 2012. The Regents of the University of California.
	All rights reserved.

	Redistribution and use in source and binary forms, with or without
	modification, are permitted provided that the following conditions are met:

		Redistributions of source code must retain the above copyright notice,
		this list of conditions and the following disclaimer.

		Redistributions in binary form must reproduce the above copyright
		notice, this list of conditions and the following disclaimer in the
		documentation and/or other materials provided with the distribution.

		Neither the name of the University of California nor the names of its
		contributors may be used to endorse or promote products derived from
		this software without specific prior written permission.

	THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
	AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
	IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
	ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
	LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
	CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
	SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
	INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
	CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
	ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
	POSSIBILITY OF SUCH DAMAGE.

	File description:
	Wrappers to OpenGL GLSL shaders

	File author(s):
	Lance Putnam, 2010, putnam.lance@gmail.com
	Graham Wakefield, 2010, grrrwaaa@gmail.com
	Wesley Smith, 2010, wesley.hoke@gmail.com
*/

#include <string>
#include <unordered_map>
#include "allocore/graphics/al_GPUObject.hpp"
#include "allocore/math/al_Mat.hpp"
#include "allocore/math/al_Quat.hpp"
#include "allocore/math/al_Vec.hpp"

#define AL_SHADER_MAX_LOG_SIZE	4096

namespace al{

/// Shader abstract base class
/// @ingroup allocore
class ShaderBase : public GPUObject{
public:

	virtual ~ShaderBase(){}

	/// Returns info log or 0 if none
	const char * log() const;

	/// Prints info log, if any
	void printLog() const;

protected:
	virtual void get(int pname, void * params) const = 0;
	virtual void getLog(char * buf) const = 0;
};

/// Shader object

/// A shader object represents your source code. You are able to pass your
/// source code to a shader object and compile the shader object.
/// @ingroup allocore
class Shader : public ShaderBase{
public:

	enum Type {
		VERTEX,
		FRAGMENT,
		GEOMETRY
	};

	Shader(const std::string& source="", Shader::Type type=FRAGMENT);

	/// This will automatically delete the shader object when it is no longer
	/// attached to any program object.
	virtual ~Shader(){ destroy(); }

	Shader& source(const std::string& v);
	Shader& source(const std::string& v, Shader::Type type);
	Shader& compile();
	bool compiled() const;

	Shader::Type type() const { return mType; }

private:
	std::string mSource;
	Shader::Type mType;
	void sendSource();

	virtual void get(int pname, void * params) const;
	virtual void getLog(char * buf) const;

	virtual void onCreate();
	virtual void onDestroy();
};

/// Shader program object

/// A program object represents a useable part of render pipeline.
/// It links one or more shader units into a single program object.
/// @ingroup allocore
class ShaderProgram : public ShaderBase{
public:

	ShaderProgram();

	/// Any attached shaders will automatically be detached, but not deleted.
	virtual ~ShaderProgram();

	/// Attach shader to program

	/// The input shader will be compiled if necessary.
	///
	ShaderProgram& attach(Shader& s);

	/// Detach shader from program
	const ShaderProgram& detach(const Shader& s) const;

	/// Link attached shaders

	/// @param[in] doValidate	validate program after linking;
	///		You might not want to do this if you need to set uniforms before
	///		validating, e.g., when using different texture sampler types in the
	///		same shader.
	const ShaderProgram& link(bool doValidate=true) const;

	/// Compile and link shader sources

	/// If called outside of a graphics context, compiling and linking will
	/// occur once the first time the shader is bound.
	bool compile(
		const std::string& vertSource,
		const std::string& fragSource,
		const std::string& geomSource=""
	);

	const ShaderProgram& use();

	/// Get whether program is active
	bool active() const { return mActive; }

	/// Set whether program is active
	ShaderProgram& active(bool v){ mActive=v; return *this; }

	/// Toggle active state
	ShaderProgram& toggleActive(){ mActive^=true; return *this; }

	/// Begin use of shader program
	bool begin();

	/// End use of shader program
	void end() const;

	/// Returns whether program linked successfully
	bool linked() const;

	/// Returns whether linked program can execute in current graphics state
	bool validateProgram(bool printLog=false) const;

	/// Set parameters for geometry shader

	/// @param[in] inPrim	Input primitive to geometry shader
	/// @param[in] outPrim	Output primitive from geometry shader
	/// @param[in] outVert	Number of vertices output from geometry shader
	///						(note the number of input vertices is always 1)
	///
	/// This must be called before attaching the geometry shader.
	/// For GLSL, the geometry shader must include the line:
	///		#extension GL_EXT_geometry_shader4 : enable
	ShaderProgram& setGeometry(int inPrim, int outPrim, unsigned outVert){
		mInPrim=inPrim; mOutPrim=outPrim; mOutVertices=outVert; return *this;
	}

	// These parameters must be set before attaching geometry shaders
	ShaderProgram& setGeometryInputPrimitive(int prim){ mInPrim = prim; return *this; }
	ShaderProgram& setGeometryOutputPrimitive(int prim){ mOutPrim = prim; return *this; }
	ShaderProgram& setGeometryOutputVertices(unsigned i){ mOutVertices = i; return *this; }

	std::vector<std::string>& transformFeedbackVaryings(){ return mTFVaryings; }
	const std::vector<std::string>& transformFeedbackVaryings() const { return mTFVaryings; }

	/// Print out all the input parameters to the shader
	void listParams() const;

	/// Get location of uniform
	int uniform(const char * name) const;

	/// Get location of attribute
	int attribute(const char * name) const;

	const ShaderProgram& uniform(const char * name, Color v) const {
	  return uniform(name, v[0], v[1], v[2], v[3]);
	}

	const ShaderProgram& uniform(int loc, int v) const;
	const ShaderProgram& uniform(int loc, float v) const;
	const ShaderProgram& uniform(int loc, double v) const { return uniform(loc, float(v)); }
	const ShaderProgram& uniform(int loc, float v0, float v1) const;
	const ShaderProgram& uniform(int loc, float v0, float v1, float v2) const;
	const ShaderProgram& uniform(int loc, float v0, float v1, float v2, float v3) const;
	template <typename T>
	const ShaderProgram& uniform(int loc, const Vec<2,T>& v) const {
		return uniform(loc, v.x, v.y);
	}
	template <typename T>
	const ShaderProgram& uniform(int loc, const Vec<3,T>& v) const {
		return uniform(loc, v.x, v.y, v.z);
	}
	template <typename T>
	const ShaderProgram& uniform(int loc, const Vec<4,T>& v) const {
		return uniform(loc, v.x, v.y, v.z, v.w);
	}
	const ShaderProgram& uniformMatrix3(int loc, const float * v, bool transpose=false) const;
	const ShaderProgram& uniformMatrix4(int loc, const float * v, bool transpose=false) const;
	const ShaderProgram& uniform(int loc, const Mat<4,float>& m) const{
		return uniformMatrix4(loc, m.elems());
	}
	template<typename T>
	const ShaderProgram& uniform(int loc, const Mat<4,T>& m) const{
		return uniform(loc, Mat4f(m));
	}

	const ShaderProgram& uniform(const char * name, int v) const;
	const ShaderProgram& uniform(const char * name, float v) const;
	const ShaderProgram& uniform(const char * name, double v) const { return uniform(name, float(v)); }
	const ShaderProgram& uniform(const char * name, float v0, float v1) const;
	const ShaderProgram& uniform(const char * name, float v0, float v1, float v2) const;
	const ShaderProgram& uniform(const char * name, float v0, float v1, float v2, float v3) const;

	template <typename T>
	const ShaderProgram& uniform(const char * name, const Vec<2,T>& v) const {
		return uniform(name, v.x, v.y);
	}

	template <typename T>
	const ShaderProgram& uniform(const char * name, const Vec<3,T>& v) const {
		return uniform(name, v.x, v.y, v.z);
	}

	template <typename T>
	const ShaderProgram& uniform(const char * name, const Vec<4,T>& v) const {
		return uniform(name, v.x, v.y, v.z, v.w);
	}

	const ShaderProgram& uniform(const char * name, const Mat<4,float>& m) const{
		return uniformMatrix4(name, m.elems());
	}

	template<typename T>
	const ShaderProgram& uniform(const char * name, const Mat<4,T>& m) const{
		return uniform(name, Mat4f(m));
	}

	template <typename T>
	const ShaderProgram& uniform(const char * name, const Quat<T>& q) const {
		// note wxyz => xyzw for GLSL vec4:
		return uniform(name, q.x, q.y, q.z, q.w);
	}

	const ShaderProgram& uniform1(const char * name, const float * v, int count=1) const;
	const ShaderProgram& uniform2(const char * name, const float * v, int count=1) const;
	const ShaderProgram& uniform3(const char * name, const float * v, int count=1) const;
	const ShaderProgram& uniform4(const char * name, const float * v, int count=1) const;

	const ShaderProgram& uniformMatrix3(const char * name, const float * v, bool transpose=false) const;
	const ShaderProgram& uniformMatrix4(const char * name, const float * v, bool transpose=false) const;

	const ShaderProgram& attribute(int loc, float v) const;
	const ShaderProgram& attribute(int loc, float v0, float v1) const;
	const ShaderProgram& attribute(int loc, float v0, float v1, float v2) const;
	const ShaderProgram& attribute(int loc, float v0, float v1, float v2, float v3) const;

	const ShaderProgram& attribute(const char * name, float v) const;
	const ShaderProgram& attribute(const char * name, float v0, float v1) const;
	const ShaderProgram& attribute(const char * name, float v0, float v1, float v2) const;
	const ShaderProgram& attribute(const char * name, float v0, float v1, float v2, float v3) const;

	const ShaderProgram& attribute1(const char * name, const float * v) const;
	const ShaderProgram& attribute2(const char * name, const float * v) const;
	const ShaderProgram& attribute3(const char * name, const float * v) const;
	const ShaderProgram& attribute4(const char * name, const float * v) const;
	const ShaderProgram& attribute1(int loc, const double * v) const;
	const ShaderProgram& attribute2(int loc, const double * v) const;
	const ShaderProgram& attribute3(int loc, const double * v) const;
	const ShaderProgram& attribute4(int loc, const double * v) const;

	template<typename T>
	const ShaderProgram& attribute(int loc, const Vec<2,T>& v) const {
		return attribute(loc, v.x, v.y);
	}
	template<typename T>
	const ShaderProgram& attribute(int loc, const Vec<3,T>& v) const {
		return attribute(loc, v.x, v.y, v.z);
	}
	template<typename T>
	const ShaderProgram& attribute(int loc, const Vec<4,T>& v) const {
		return attribute(loc, v.x, v.y, v.z, v.w);
	}
	template<typename T>
	const ShaderProgram& attribute(int loc, const Quat<T>& q) const {
		// note wxyz => xyzw for GLSL vec4:
		return attribute(loc, q.x, q.y, q.z, q.w);
	}

	static void use(unsigned programID);

protected:
	int mInPrim, mOutPrim;	// IO primitives for geometry shaders
	unsigned int mOutVertices;
	std::string mVertSource, mFragSource, mGeomSource;
	mutable std::unordered_map<std::string, int> mUniformLocs, mAttribLocs;
	std::vector<std::string> mTFVaryings;
	bool mActive;

	virtual void get(int pname, void * params) const;
	virtual void getLog(char * buf) const;

	virtual void onCreate();
	virtual void onDestroy();
};

} // ::al

#endif
<commit_msg>fix compile error: al_Shader.hpp:218:50: error: unknown type name \'Color\'; thanks to Sihwa Park<commit_after>#ifndef INCLUDE_AL_GRAPHICS_SHADER_HPP
#define INCLUDE_AL_GRAPHICS_SHADER_HPP

/*	Allocore --
        Multimedia / virtual environment application class library

        Copyright (C) 2009. AlloSphere Research Group, Media Arts & Technology,
   UCSB. Copyright (C) 2012. The Regents of the University of California. All
   rights reserved.

        Redistribution and use in source and binary forms, with or without
        modification, are permitted provided that the following conditions are
   met:

                Redistributions of source code must retain the above copyright
   notice, this list of conditions and the following disclaimer.

                Redistributions in binary form must reproduce the above
   copyright notice, this list of conditions and the following disclaimer in the
                documentation and/or other materials provided with the
   distribution.

                Neither the name of the University of California nor the names
   of its contributors may be used to endorse or promote products derived from
                this software without specific prior written permission.

        THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
   IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
        IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
   PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
   CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
   EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
   PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
   OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
   WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
   OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
   ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

        File description:
        Wrappers to OpenGL GLSL shaders

        File author(s):
        Lance Putnam, 2010, putnam.lance@gmail.com
        Graham Wakefield, 2010, grrrwaaa@gmail.com
        Wesley Smith, 2010, wesley.hoke@gmail.com
*/

#include <string>
#include <unordered_map>
#include "allocore/graphics/al_GPUObject.hpp"
#include "allocore/math/al_Mat.hpp"
#include "allocore/math/al_Quat.hpp"
#include "allocore/math/al_Vec.hpp"
#include "allocore/types/al_Color.hpp"

#define AL_SHADER_MAX_LOG_SIZE 4096

namespace al {

/// Shader abstract base class
/// @ingroup allocore
class ShaderBase : public GPUObject {
 public:
  virtual ~ShaderBase() {}

  /// Returns info log or 0 if none
  const char* log() const;

  /// Prints info log, if any
  void printLog() const;

 protected:
  virtual void get(int pname, void* params) const = 0;
  virtual void getLog(char* buf) const = 0;
};

/// Shader object

/// A shader object represents your source code. You are able to pass your
/// source code to a shader object and compile the shader object.
/// @ingroup allocore
class Shader : public ShaderBase {
 public:
  enum Type { VERTEX, FRAGMENT, GEOMETRY };

  Shader(const std::string& source = "", Shader::Type type = FRAGMENT);

  /// This will automatically delete the shader object when it is no longer
  /// attached to any program object.
  virtual ~Shader() { destroy(); }

  Shader& source(const std::string& v);
  Shader& source(const std::string& v, Shader::Type type);
  Shader& compile();
  bool compiled() const;

  Shader::Type type() const { return mType; }

 private:
  std::string mSource;
  Shader::Type mType;
  void sendSource();

  virtual void get(int pname, void* params) const;
  virtual void getLog(char* buf) const;

  virtual void onCreate();
  virtual void onDestroy();
};

/// Shader program object

/// A program object represents a useable part of render pipeline.
/// It links one or more shader units into a single program object.
/// @ingroup allocore
class ShaderProgram : public ShaderBase {
 public:
  ShaderProgram();

  /// Any attached shaders will automatically be detached, but not deleted.
  virtual ~ShaderProgram();

  /// Attach shader to program

  /// The input shader will be compiled if necessary.
  ///
  ShaderProgram& attach(Shader& s);

  /// Detach shader from program
  const ShaderProgram& detach(const Shader& s) const;

  /// Link attached shaders

  /// @param[in] doValidate	validate program after linking;
  ///		You might not want to do this if you need to set uniforms before
  ///		validating, e.g., when using different texture sampler types in
  ///the 		same shader.
  const ShaderProgram& link(bool doValidate = true) const;

  /// Compile and link shader sources

  /// If called outside of a graphics context, compiling and linking will
  /// occur once the first time the shader is bound.
  bool compile(const std::string& vertSource, const std::string& fragSource,
               const std::string& geomSource = "");

  const ShaderProgram& use();

  /// Get whether program is active
  bool active() const { return mActive; }

  /// Set whether program is active
  ShaderProgram& active(bool v) {
    mActive = v;
    return *this;
  }

  /// Toggle active state
  ShaderProgram& toggleActive() {
    mActive ^= true;
    return *this;
  }

  /// Begin use of shader program
  bool begin();

  /// End use of shader program
  void end() const;

  /// Returns whether program linked successfully
  bool linked() const;

  /// Returns whether linked program can execute in current graphics state
  bool validateProgram(bool printLog = false) const;

  /// Set parameters for geometry shader

  /// @param[in] inPrim	Input primitive to geometry shader
  /// @param[in] outPrim	Output primitive from geometry shader
  /// @param[in] outVert	Number of vertices output from geometry shader
  ///						(note the number of input vertices is always
  ///1)
  ///
  /// This must be called before attaching the geometry shader.
  /// For GLSL, the geometry shader must include the line:
  ///		#extension GL_EXT_geometry_shader4 : enable
  ShaderProgram& setGeometry(int inPrim, int outPrim, unsigned outVert) {
    mInPrim = inPrim;
    mOutPrim = outPrim;
    mOutVertices = outVert;
    return *this;
  }

  // These parameters must be set before attaching geometry shaders
  ShaderProgram& setGeometryInputPrimitive(int prim) {
    mInPrim = prim;
    return *this;
  }
  ShaderProgram& setGeometryOutputPrimitive(int prim) {
    mOutPrim = prim;
    return *this;
  }
  ShaderProgram& setGeometryOutputVertices(unsigned i) {
    mOutVertices = i;
    return *this;
  }

  std::vector<std::string>& transformFeedbackVaryings() { return mTFVaryings; }
  const std::vector<std::string>& transformFeedbackVaryings() const {
    return mTFVaryings;
  }

  /// Print out all the input parameters to the shader
  void listParams() const;

  /// Get location of uniform
  int uniform(const char* name) const;

  /// Get location of attribute
  int attribute(const char* name) const;

  const ShaderProgram& uniform(const char* name, Color v) const {
    return uniform(name, v[0], v[1], v[2], v[3]);
  }

  const ShaderProgram& uniform(int loc, int v) const;
  const ShaderProgram& uniform(int loc, float v) const;
  const ShaderProgram& uniform(int loc, double v) const {
    return uniform(loc, float(v));
  }
  const ShaderProgram& uniform(int loc, float v0, float v1) const;
  const ShaderProgram& uniform(int loc, float v0, float v1, float v2) const;
  const ShaderProgram& uniform(int loc, float v0, float v1, float v2,
                               float v3) const;
  template <typename T>
  const ShaderProgram& uniform(int loc, const Vec<2, T>& v) const {
    return uniform(loc, v.x, v.y);
  }
  template <typename T>
  const ShaderProgram& uniform(int loc, const Vec<3, T>& v) const {
    return uniform(loc, v.x, v.y, v.z);
  }
  template <typename T>
  const ShaderProgram& uniform(int loc, const Vec<4, T>& v) const {
    return uniform(loc, v.x, v.y, v.z, v.w);
  }
  const ShaderProgram& uniformMatrix3(int loc, const float* v,
                                      bool transpose = false) const;
  const ShaderProgram& uniformMatrix4(int loc, const float* v,
                                      bool transpose = false) const;
  const ShaderProgram& uniform(int loc, const Mat<4, float>& m) const {
    return uniformMatrix4(loc, m.elems());
  }
  template <typename T>
  const ShaderProgram& uniform(int loc, const Mat<4, T>& m) const {
    return uniform(loc, Mat4f(m));
  }

  const ShaderProgram& uniform(const char* name, int v) const;
  const ShaderProgram& uniform(const char* name, float v) const;
  const ShaderProgram& uniform(const char* name, double v) const {
    return uniform(name, float(v));
  }
  const ShaderProgram& uniform(const char* name, float v0, float v1) const;
  const ShaderProgram& uniform(const char* name, float v0, float v1,
                               float v2) const;
  const ShaderProgram& uniform(const char* name, float v0, float v1, float v2,
                               float v3) const;

  template <typename T>
  const ShaderProgram& uniform(const char* name, const Vec<2, T>& v) const {
    return uniform(name, v.x, v.y);
  }

  template <typename T>
  const ShaderProgram& uniform(const char* name, const Vec<3, T>& v) const {
    return uniform(name, v.x, v.y, v.z);
  }

  template <typename T>
  const ShaderProgram& uniform(const char* name, const Vec<4, T>& v) const {
    return uniform(name, v.x, v.y, v.z, v.w);
  }

  const ShaderProgram& uniform(const char* name, const Mat<4, float>& m) const {
    return uniformMatrix4(name, m.elems());
  }

  template <typename T>
  const ShaderProgram& uniform(const char* name, const Mat<4, T>& m) const {
    return uniform(name, Mat4f(m));
  }

  template <typename T>
  const ShaderProgram& uniform(const char* name, const Quat<T>& q) const {
    // note wxyz => xyzw for GLSL vec4:
    return uniform(name, q.x, q.y, q.z, q.w);
  }

  const ShaderProgram& uniform1(const char* name, const float* v,
                                int count = 1) const;
  const ShaderProgram& uniform2(const char* name, const float* v,
                                int count = 1) const;
  const ShaderProgram& uniform3(const char* name, const float* v,
                                int count = 1) const;
  const ShaderProgram& uniform4(const char* name, const float* v,
                                int count = 1) const;

  const ShaderProgram& uniformMatrix3(const char* name, const float* v,
                                      bool transpose = false) const;
  const ShaderProgram& uniformMatrix4(const char* name, const float* v,
                                      bool transpose = false) const;

  const ShaderProgram& attribute(int loc, float v) const;
  const ShaderProgram& attribute(int loc, float v0, float v1) const;
  const ShaderProgram& attribute(int loc, float v0, float v1, float v2) const;
  const ShaderProgram& attribute(int loc, float v0, float v1, float v2,
                                 float v3) const;

  const ShaderProgram& attribute(const char* name, float v) const;
  const ShaderProgram& attribute(const char* name, float v0, float v1) const;
  const ShaderProgram& attribute(const char* name, float v0, float v1,
                                 float v2) const;
  const ShaderProgram& attribute(const char* name, float v0, float v1, float v2,
                                 float v3) const;

  const ShaderProgram& attribute1(const char* name, const float* v) const;
  const ShaderProgram& attribute2(const char* name, const float* v) const;
  const ShaderProgram& attribute3(const char* name, const float* v) const;
  const ShaderProgram& attribute4(const char* name, const float* v) const;
  const ShaderProgram& attribute1(int loc, const double* v) const;
  const ShaderProgram& attribute2(int loc, const double* v) const;
  const ShaderProgram& attribute3(int loc, const double* v) const;
  const ShaderProgram& attribute4(int loc, const double* v) const;

  template <typename T>
  const ShaderProgram& attribute(int loc, const Vec<2, T>& v) const {
    return attribute(loc, v.x, v.y);
  }
  template <typename T>
  const ShaderProgram& attribute(int loc, const Vec<3, T>& v) const {
    return attribute(loc, v.x, v.y, v.z);
  }
  template <typename T>
  const ShaderProgram& attribute(int loc, const Vec<4, T>& v) const {
    return attribute(loc, v.x, v.y, v.z, v.w);
  }
  template <typename T>
  const ShaderProgram& attribute(int loc, const Quat<T>& q) const {
    // note wxyz => xyzw for GLSL vec4:
    return attribute(loc, q.x, q.y, q.z, q.w);
  }

  static void use(unsigned programID);

 protected:
  int mInPrim, mOutPrim;  // IO primitives for geometry shaders
  unsigned int mOutVertices;
  std::string mVertSource, mFragSource, mGeomSource;
  mutable std::unordered_map<std::string, int> mUniformLocs, mAttribLocs;
  std::vector<std::string> mTFVaryings;
  bool mActive;

  virtual void get(int pname, void* params) const;
  virtual void getLog(char* buf) const;

  virtual void onCreate();
  virtual void onDestroy();
};

}  // namespace al

#endif
<|endoftext|>


<commit_before>#include <gtest/gtest.h>
#include <gmock/gmock.h>
#include <rapidcheck/gtest.h>

#include <crest/geometry/indexed_mesh.hpp>
#include <crest/geometry/refinement.hpp>
#include <crest/basis/lagrange_basis2d.hpp>
#include <crest/basis/homogenized_basis.hpp>

#include <crest/util/algorithms.hpp>

#include <util/eigen_matchers.hpp>
#include <util/test_generators.hpp>

#include <Eigen/Sparse>
#include <Eigen/Dense>

using ::testing::ElementsAreArray;
using ::testing::Pointwise;
using ::testing::Eq;
using ::testing::DoubleEq;
using ::testing::Lt;
using ::testing::Each;

using crest::IndexedMesh;
using crest::LagrangeBasis2d;

typedef IndexedMesh<double, int>::Vertex Vertex;
typedef IndexedMesh<double, int>::Element Element;

TEST(homogenized_basis_test, correctors_are_in_interpolator_kernel_for_threshold_fine_mesh)
{
    // Recall that by definition, I_H c = 0, for any corrector with weights c.
    const std::vector<Vertex> coarse_vertices {
            Vertex(0.0, 0.0),
            Vertex(1.0, 0.0),
            Vertex(1.0, 1.0),
            Vertex(0.0, 1.0)
    };

    const std::vector<Element> coarse_elements {
            Element({3, 0, 1}),
            Element({1, 2, 3})
    };

    const auto initial_mesh = IndexedMesh<double, int>(coarse_vertices, coarse_elements);

    const auto two_scale_meshes = crest::threshold(initial_mesh, 0.5,
                                                   { crest::ReentrantCorner<double, int>(0, 2.35)});

    const auto basis = crest::detail::homogenized_basis_correctors(two_scale_meshes.coarse,
                                                                   two_scale_meshes.fine,
                                                                   4);

    ASSERT_THAT(basis.rows(), Eq(two_scale_meshes.coarse.num_vertices()));
    ASSERT_THAT(basis.cols(), Eq(two_scale_meshes.fine.num_vertices()));

    const auto coarse_interior = two_scale_meshes.coarse.compute_interior_vertices();
    const auto fine_dof = crest::algo::integer_range<int>(0, two_scale_meshes.fine.num_vertices());

    const auto I_H = crest::quasi_interpolator(two_scale_meshes.coarse, two_scale_meshes.fine);
    // We need to test with the coarse interior, because the matrix returned by quasi_interpolator
    // does not impose that I_H x = 0 on the boundary.
    const auto I_H_interior = sparse_submatrix(I_H, coarse_interior, fine_dof);

    // Recall that each row of 'basis' corresponds to weights in the fine space, so by transposing it and
    // left-multiplying by I_H, we effectively compute I_H x_i for each basis function i.
    // Since x_i is in the kernel of I_H, we expect the result to be zero.
    const Eigen::SparseMatrix<double> Z = I_H_interior * basis.transpose();
    const Eigen::MatrixXd Z_dense = Z;

    EXPECT_TRUE(Z_dense.isZero(1e-15));
}

TEST(construct_saddle_point_problem_test, blockwise_correct)
{
    MatrixX<double> A(5, 5);
    A.setZero();
    A(0, 0) = 1.0;
    A(1, 1) = 2.0;
    A(2, 2) = 3.0;
    A(3, 3) = 4.0;
    A(4, 4) = 5.0;

    MatrixX<double> I_H(3, 5);
    I_H.setZero();
    I_H(0, 0) = 6.0;
    I_H(1, 1) = 7.0;
    I_H(2, 2) = 8.0;

    VectorX<double> b(5);
    b.setZero();
    b(0) = 1.0;
    b(1) = 2.0;
    b(2) = 3.0;
    b(3) = 4.0;
    b(4) = 5.0;

    const Eigen::SparseMatrix<double> A_sparse = A.sparseView();
    const Eigen::SparseMatrix<double> I_H_sparse = I_H.sparseView();

    const auto saddle_point = crest::detail::construct_saddle_point_problem(A_sparse, I_H_sparse, b);
    const MatrixX<double> C = saddle_point.first;
    const VectorX<double> c = saddle_point.second;

    ASSERT_THAT(C.rows(), 8);
    ASSERT_THAT(C.cols(), 8);
    ASSERT_THAT(c.rows(), 8);

    const MatrixX<double> bottomRightCorner = C.bottomRightCorner(3, 3);
    const auto bottomRightCornerElements = std::vector<double>(bottomRightCorner.data(), bottomRightCorner.data() + 9);

    // TODO: The current MatrixEq is not entirely reliable, I think. Improve!
    EXPECT_THAT(C.topLeftCorner(5, 5), MatrixEq(A));
    EXPECT_THAT(C.bottomLeftCorner(3, 5), MatrixEq(I_H));
    EXPECT_THAT(C.topRightCorner(5, 3), MatrixEq(I_H.transpose()));
    EXPECT_THAT(bottomRightCornerElements, Each(DoubleEq(0.0)));

    EXPECT_THAT(c.topRows(5), MatrixEq(b));
    EXPECT_THAT(c.bottomRows(3)(0), DoubleEq(0.0));
    EXPECT_THAT(c.bottomRows(3)(1), DoubleEq(0.0));
    EXPECT_THAT(c.bottomRows(3)(2), DoubleEq(0.0));
}

RC_GTEST_PROP(homogenized_basis_test, correctors_are_in_interpolator_kernel, ())
{
    const auto oversampling = *rc::gen::arbitrary<unsigned int>();
    const auto coarse_mesh = *crest::gen::arbitrary_unit_square_mesh(4);

    // Make sure the fine mesh does not coincide with the coarse mesh
    const auto fine_mesh = *rc::gen::suchThat(crest::gen::arbitrary_refinement(coarse_mesh, 1, 4),
                                              [coarse_mesh] (auto mesh)
                                              {
                                                  return mesh.num_elements() != coarse_mesh.num_elements();
                                              });

    const auto basis = crest::detail::homogenized_basis_correctors(coarse_mesh, fine_mesh, oversampling);

    ASSERT_THAT(basis.rows(), Eq(coarse_mesh.num_vertices()));
    ASSERT_THAT(basis.cols(), Eq(fine_mesh.num_vertices()));

    const auto coarse_interior = coarse_mesh.compute_interior_vertices();
    const auto fine_dof = crest::algo::integer_range<int>(0, fine_mesh.num_vertices());

    const auto I_H = crest::quasi_interpolator(coarse_mesh, fine_mesh);

    // We need to test with the coarse interior, because the matrix returned by quasi_interpolator
    // does not impose that I_H x = 0 on the boundary.
    const auto I_H_interior = sparse_submatrix(I_H, coarse_interior, fine_dof);

    // Recall that each row of 'basis' corresponds to weights in the fine space, so by transposing it and
    // left-multiplying by I_H, we effectively compute I_H x_i for each basis function i.
    // Since x_i is in the kernel of I_H, we expect the result to be zero.
    const Eigen::SparseMatrix<double> Z = I_H_interior * basis.transpose();
    const Eigen::MatrixXd Z_dense = Z;

    RC_ASSERT(Z_dense.isZero(1e-14));
}

RC_GTEST_PROP(homogenized_basis_test, correctors_are_zero_with_no_refinement, ())
{
    const auto oversampling = static_cast<unsigned int>(*rc::gen::inRange(0, 6));
    const auto mesh = *crest::gen::arbitrary_unit_square_mesh();
    const auto basis = crest::detail::homogenized_basis_correctors(mesh, mesh, oversampling);

    ASSERT_THAT(basis.rows(), Eq(mesh.num_vertices()));
    ASSERT_THAT(basis.cols(), Eq(mesh.num_vertices()));

    double max_abs = 0.0;
    for (int o = 0; o < basis.outerSize(); ++o)
    {
        for (Eigen::SparseMatrix<double>::InnerIterator it(basis, o); it; ++it)
        {
            max_abs = std::max(max_abs, std::abs(it.value()));
        }
    }

    RC_ASSERT(max_abs < 1e-14);
}

RC_GTEST_PROP(homogenized_basis_test, corrected_basis_is_orthogonal_to_fine_space, ())
{
    // Let W be such that W_ij corrected basis function i is defined by
    //
    // lambda_i = sum_j W_ij phi_j
    //
    // where phi_j denotes the fine-scale basis functions. Then the a-orthogonality of the corrected space V_H and
    // the fine-scale finite element space implies that
    //
    // A W^T = 0
    //
    // where A is the fine-scale stiffness matrix (interior nodes).
    //
    // Our corrected basis is constructed by a localized approximation, so this does not hold exactly.
    // We can, however, choose an oversampling parameter so large that each corrector problem is a global problem.
    // In this case, the above property should still hold.

    // Make sure interior is not empty, because Eigen's SVD fails on empty matrices
    const auto coarse = *rc::gen::suchThat(crest::gen::arbitrary_unit_square_mesh(),
                                           [] (const auto & mesh) {
                                               return !mesh.compute_interior_vertices().empty();
                                           }).as("coarse mesh");
    const auto fine = *crest::gen::arbitrary_refinement(coarse).as("fine mesh");

    const auto fine_interior = fine.compute_interior_vertices();
    const auto coarse_interior = coarse.compute_interior_vertices();
    const auto fine_basis = LagrangeBasis2d<double>(fine);

    const auto oversampling = static_cast<unsigned int>(coarse.num_vertices());
    const auto basis_weights = crest::detail::corrected_basis_coefficients(coarse, fine, oversampling);
    const auto basis_interior_weights = sparse_submatrix(basis_weights, coarse_interior, fine_interior);

    // Construct a basis for W_H, the kernel of I_H.
    const Eigen::MatrixXd I_H = sparse_submatrix(crest::quasi_interpolator(coarse, fine), coarse_interior, fine_interior);
    const Eigen::JacobiSVD<Eigen::MatrixXd> svd(I_H, Eigen::ComputeFullV);
    const auto r = svd.rank();
    const auto n = I_H.cols();
    const Eigen::MatrixXd W_H_kernel_basis = svd.matrixV().rightCols(n - r);

    const Eigen::MatrixXd W = basis_interior_weights;
    const Eigen::MatrixXd A = sparse_submatrix(fine_basis.assemble().stiffness, fine_interior, fine_interior);
    const auto & N = W_H_kernel_basis;

    const Eigen::MatrixXd product = W * A * N;

    RC_LOG() << "W:   " << std::endl << W << std::endl << std::endl;
    RC_LOG() << "A:   " << std::endl << A << std::endl << std::endl;
    RC_LOG() << "N:   " << std::endl << N << std::endl << std::endl;
    RC_LOG() << "WAN: " << std::endl << product << std::endl << std::endl;

    RC_ASSERT(product.isZero(1e-14));
}

TEST(standard_coarse_basis_in_fine_space, basic_mesh)
{
    const std::vector<Vertex> coarse_vertices{
            Vertex(0.0, 0.0),
            Vertex(1.0, 0.0),
            Vertex(1.0, 1.0),
            Vertex(0.0, 1.0)
    };

    const std::vector<Element> coarse_elements {
            Element({3, 0, 1}),
            Element({1, 2, 3})
    };

    const std::vector<Vertex> fine_vertices {
            Vertex(0.0, 0.0),
            Vertex(1.0, 0.0),
            Vertex(1.0, 1.0),
            Vertex(0.0, 1.0),
            Vertex(0.5, 0.5)
    };

    const std::vector<Element> fine_elements {
            Element({0, 4, 3}),
            Element({2, 4, 1}),
            Element({1, 4, 0}),
            Element({2, 4, 3})
    };

    const auto coarse_mesh = IndexedMesh<double, int>(coarse_vertices, coarse_elements);
    const auto fine_mesh = IndexedMesh<double, int>(fine_vertices, fine_elements);

    const Eigen::Matrix<double, 4, 5> coarse_basis_in_fine =
            crest::detail::standard_coarse_basis_in_fine_space(coarse_mesh, fine_mesh);

    Eigen::Matrix<double, 4, 5> expected;
    expected <<
             1.0, 0.0, 0.0, 0.0, 0.0,
            0.0, 1.0, 0.0, 0.0, 0.5,
            0.0, 0.0, 1.0, 0.0, 0.0,
            0.0, 0.0, 0.0, 1.0, 0.5;

    ASSERT_THAT(coarse_basis_in_fine, MatrixEq(expected));
}

RC_GTEST_PROP(standard_coarse_basis_in_fine_space, quasi_interpolation_recovers_coarse_basis, ())
{
    // Let v_H be the weights of a function in the coarse space. Then v_H is also a function in the fine space,
    // and since I_H is a quasi interpolation, we have that I_H v_h = v_H, where v_h are the fine weights of
    // the coarse function.
    // We can verify that this property holds for all basis functions as represented in the fine space.

    const auto coarse = *crest::gen::arbitrary_unit_square_mesh().as("Coarse mesh");
    const auto fine = *crest::gen::arbitrary_refinement(coarse, 0, 5).as("Fine mesh");

    const auto basis = crest::detail::standard_coarse_basis_in_fine_space(coarse, fine);
    const auto I_H = crest::quasi_interpolator(coarse, fine);

    const Eigen::SparseMatrix<double> interpolated = I_H * basis.transpose();
    const Eigen::MatrixXd interpolated_dense = interpolated;

    RC_LOG() << interpolated_dense << std::endl;

    // We could have had a shaper tolerance here, but it appears Eigen does not do an elementwise comparison
    // for isIdentity (which is very annoying, by the way), as it can spuriously fail even when all non-diagonal
    // elements are smaller than 3e-16.
    RC_ASSERT(interpolated_dense.isIdentity(1e-14));
}
<commit_msg>Correct explanation in orthogonality test<commit_after>#include <gtest/gtest.h>
#include <gmock/gmock.h>
#include <rapidcheck/gtest.h>

#include <crest/geometry/indexed_mesh.hpp>
#include <crest/geometry/refinement.hpp>
#include <crest/basis/lagrange_basis2d.hpp>
#include <crest/basis/homogenized_basis.hpp>

#include <crest/util/algorithms.hpp>

#include <util/eigen_matchers.hpp>
#include <util/test_generators.hpp>

#include <Eigen/Sparse>
#include <Eigen/Dense>

using ::testing::ElementsAreArray;
using ::testing::Pointwise;
using ::testing::Eq;
using ::testing::DoubleEq;
using ::testing::Lt;
using ::testing::Each;

using crest::IndexedMesh;
using crest::LagrangeBasis2d;

typedef IndexedMesh<double, int>::Vertex Vertex;
typedef IndexedMesh<double, int>::Element Element;

TEST(homogenized_basis_test, correctors_are_in_interpolator_kernel_for_threshold_fine_mesh)
{
    // Recall that by definition, I_H c = 0, for any corrector with weights c.
    const std::vector<Vertex> coarse_vertices {
            Vertex(0.0, 0.0),
            Vertex(1.0, 0.0),
            Vertex(1.0, 1.0),
            Vertex(0.0, 1.0)
    };

    const std::vector<Element> coarse_elements {
            Element({3, 0, 1}),
            Element({1, 2, 3})
    };

    const auto initial_mesh = IndexedMesh<double, int>(coarse_vertices, coarse_elements);

    const auto two_scale_meshes = crest::threshold(initial_mesh, 0.5,
                                                   { crest::ReentrantCorner<double, int>(0, 2.35)});

    const auto basis = crest::detail::homogenized_basis_correctors(two_scale_meshes.coarse,
                                                                   two_scale_meshes.fine,
                                                                   4);

    ASSERT_THAT(basis.rows(), Eq(two_scale_meshes.coarse.num_vertices()));
    ASSERT_THAT(basis.cols(), Eq(two_scale_meshes.fine.num_vertices()));

    const auto coarse_interior = two_scale_meshes.coarse.compute_interior_vertices();
    const auto fine_dof = crest::algo::integer_range<int>(0, two_scale_meshes.fine.num_vertices());

    const auto I_H = crest::quasi_interpolator(two_scale_meshes.coarse, two_scale_meshes.fine);
    // We need to test with the coarse interior, because the matrix returned by quasi_interpolator
    // does not impose that I_H x = 0 on the boundary.
    const auto I_H_interior = sparse_submatrix(I_H, coarse_interior, fine_dof);

    // Recall that each row of 'basis' corresponds to weights in the fine space, so by transposing it and
    // left-multiplying by I_H, we effectively compute I_H x_i for each basis function i.
    // Since x_i is in the kernel of I_H, we expect the result to be zero.
    const Eigen::SparseMatrix<double> Z = I_H_interior * basis.transpose();
    const Eigen::MatrixXd Z_dense = Z;

    EXPECT_TRUE(Z_dense.isZero(1e-15));
}

TEST(construct_saddle_point_problem_test, blockwise_correct)
{
    MatrixX<double> A(5, 5);
    A.setZero();
    A(0, 0) = 1.0;
    A(1, 1) = 2.0;
    A(2, 2) = 3.0;
    A(3, 3) = 4.0;
    A(4, 4) = 5.0;

    MatrixX<double> I_H(3, 5);
    I_H.setZero();
    I_H(0, 0) = 6.0;
    I_H(1, 1) = 7.0;
    I_H(2, 2) = 8.0;

    VectorX<double> b(5);
    b.setZero();
    b(0) = 1.0;
    b(1) = 2.0;
    b(2) = 3.0;
    b(3) = 4.0;
    b(4) = 5.0;

    const Eigen::SparseMatrix<double> A_sparse = A.sparseView();
    const Eigen::SparseMatrix<double> I_H_sparse = I_H.sparseView();

    const auto saddle_point = crest::detail::construct_saddle_point_problem(A_sparse, I_H_sparse, b);
    const MatrixX<double> C = saddle_point.first;
    const VectorX<double> c = saddle_point.second;

    ASSERT_THAT(C.rows(), 8);
    ASSERT_THAT(C.cols(), 8);
    ASSERT_THAT(c.rows(), 8);

    const MatrixX<double> bottomRightCorner = C.bottomRightCorner(3, 3);
    const auto bottomRightCornerElements = std::vector<double>(bottomRightCorner.data(), bottomRightCorner.data() + 9);

    // TODO: The current MatrixEq is not entirely reliable, I think. Improve!
    EXPECT_THAT(C.topLeftCorner(5, 5), MatrixEq(A));
    EXPECT_THAT(C.bottomLeftCorner(3, 5), MatrixEq(I_H));
    EXPECT_THAT(C.topRightCorner(5, 3), MatrixEq(I_H.transpose()));
    EXPECT_THAT(bottomRightCornerElements, Each(DoubleEq(0.0)));

    EXPECT_THAT(c.topRows(5), MatrixEq(b));
    EXPECT_THAT(c.bottomRows(3)(0), DoubleEq(0.0));
    EXPECT_THAT(c.bottomRows(3)(1), DoubleEq(0.0));
    EXPECT_THAT(c.bottomRows(3)(2), DoubleEq(0.0));
}

RC_GTEST_PROP(homogenized_basis_test, correctors_are_in_interpolator_kernel, ())
{
    const auto oversampling = *rc::gen::arbitrary<unsigned int>();
    const auto coarse_mesh = *crest::gen::arbitrary_unit_square_mesh(4);

    // Make sure the fine mesh does not coincide with the coarse mesh
    const auto fine_mesh = *rc::gen::suchThat(crest::gen::arbitrary_refinement(coarse_mesh, 1, 4),
                                              [coarse_mesh] (auto mesh)
                                              {
                                                  return mesh.num_elements() != coarse_mesh.num_elements();
                                              });

    const auto basis = crest::detail::homogenized_basis_correctors(coarse_mesh, fine_mesh, oversampling);

    ASSERT_THAT(basis.rows(), Eq(coarse_mesh.num_vertices()));
    ASSERT_THAT(basis.cols(), Eq(fine_mesh.num_vertices()));

    const auto coarse_interior = coarse_mesh.compute_interior_vertices();
    const auto fine_dof = crest::algo::integer_range<int>(0, fine_mesh.num_vertices());

    const auto I_H = crest::quasi_interpolator(coarse_mesh, fine_mesh);

    // We need to test with the coarse interior, because the matrix returned by quasi_interpolator
    // does not impose that I_H x = 0 on the boundary.
    const auto I_H_interior = sparse_submatrix(I_H, coarse_interior, fine_dof);

    // Recall that each row of 'basis' corresponds to weights in the fine space, so by transposing it and
    // left-multiplying by I_H, we effectively compute I_H x_i for each basis function i.
    // Since x_i is in the kernel of I_H, we expect the result to be zero.
    const Eigen::SparseMatrix<double> Z = I_H_interior * basis.transpose();
    const Eigen::MatrixXd Z_dense = Z;

    RC_ASSERT(Z_dense.isZero(1e-14));
}

RC_GTEST_PROP(homogenized_basis_test, correctors_are_zero_with_no_refinement, ())
{
    const auto oversampling = static_cast<unsigned int>(*rc::gen::inRange(0, 6));
    const auto mesh = *crest::gen::arbitrary_unit_square_mesh();
    const auto basis = crest::detail::homogenized_basis_correctors(mesh, mesh, oversampling);

    ASSERT_THAT(basis.rows(), Eq(mesh.num_vertices()));
    ASSERT_THAT(basis.cols(), Eq(mesh.num_vertices()));

    double max_abs = 0.0;
    for (int o = 0; o < basis.outerSize(); ++o)
    {
        for (Eigen::SparseMatrix<double>::InnerIterator it(basis, o); it; ++it)
        {
            max_abs = std::max(max_abs, std::abs(it.value()));
        }
    }

    RC_ASSERT(max_abs < 1e-14);
}

RC_GTEST_PROP(homogenized_basis_test, corrected_basis_is_orthogonal_to_fine_space, ())
{
    // Let W be such that W_ij corrected basis function i is defined by
    //
    // lambda_i = sum_j W_ij phi_j
    //
    // where phi_j denotes the fine-scale basis functions. Then the a-orthogonality of the corrected space V_H and
    // the fine-scale finite element space implies that
    //
    // W A N = 0
    //
    // where A is the fine-scale stiffness matrix (interior nodes) and N is a basis for W_H = ker(I_H).
    // In real applications, obtaining a basis for W_H is normally infeasible due to sparsity requirements.
    // However, since we only do small-scale tests here, we have the luxury of just simply building a dense
    // basis through the SVD.
    //
    // Our corrected basis is constructed by a localized approximation, so the above property does not hold exactly.
    // We can, however, choose an oversampling parameter so large that each corrector problem is a global problem.
    // In this case, the above property must still hold.

    // Make sure interior is not empty, because Eigen's SVD fails on empty matrices
    const auto coarse = *rc::gen::suchThat(crest::gen::arbitrary_unit_square_mesh(),
                                           [] (const auto & mesh) {
                                               return !mesh.compute_interior_vertices().empty();
                                           }).as("coarse mesh");
    const auto fine = *crest::gen::arbitrary_refinement(coarse).as("fine mesh");

    const auto fine_interior = fine.compute_interior_vertices();
    const auto coarse_interior = coarse.compute_interior_vertices();
    const auto fine_basis = LagrangeBasis2d<double>(fine);

    const auto oversampling = static_cast<unsigned int>(coarse.num_vertices());
    const auto basis_weights = crest::detail::corrected_basis_coefficients(coarse, fine, oversampling);
    const auto basis_interior_weights = sparse_submatrix(basis_weights, coarse_interior, fine_interior);

    // Construct a basis for W_H, the kernel of I_H.
    const Eigen::MatrixXd I_H = sparse_submatrix(crest::quasi_interpolator(coarse, fine), coarse_interior, fine_interior);
    const Eigen::JacobiSVD<Eigen::MatrixXd> svd(I_H, Eigen::ComputeFullV);
    const auto r = svd.rank();
    const auto n = I_H.cols();
    const Eigen::MatrixXd W_H_kernel_basis = svd.matrixV().rightCols(n - r);

    const Eigen::MatrixXd W = basis_interior_weights;
    const Eigen::MatrixXd A = sparse_submatrix(fine_basis.assemble().stiffness, fine_interior, fine_interior);
    const auto & N = W_H_kernel_basis;

    const Eigen::MatrixXd product = W * A * N;

    RC_LOG() << "W:   " << std::endl << W << std::endl << std::endl;
    RC_LOG() << "A:   " << std::endl << A << std::endl << std::endl;
    RC_LOG() << "N:   " << std::endl << N << std::endl << std::endl;
    RC_LOG() << "WAN: " << std::endl << product << std::endl << std::endl;

    RC_ASSERT(product.isZero(1e-14));
}

TEST(standard_coarse_basis_in_fine_space, basic_mesh)
{
    const std::vector<Vertex> coarse_vertices{
            Vertex(0.0, 0.0),
            Vertex(1.0, 0.0),
            Vertex(1.0, 1.0),
            Vertex(0.0, 1.0)
    };

    const std::vector<Element> coarse_elements {
            Element({3, 0, 1}),
            Element({1, 2, 3})
    };

    const std::vector<Vertex> fine_vertices {
            Vertex(0.0, 0.0),
            Vertex(1.0, 0.0),
            Vertex(1.0, 1.0),
            Vertex(0.0, 1.0),
            Vertex(0.5, 0.5)
    };

    const std::vector<Element> fine_elements {
            Element({0, 4, 3}),
            Element({2, 4, 1}),
            Element({1, 4, 0}),
            Element({2, 4, 3})
    };

    const auto coarse_mesh = IndexedMesh<double, int>(coarse_vertices, coarse_elements);
    const auto fine_mesh = IndexedMesh<double, int>(fine_vertices, fine_elements);

    const Eigen::Matrix<double, 4, 5> coarse_basis_in_fine =
            crest::detail::standard_coarse_basis_in_fine_space(coarse_mesh, fine_mesh);

    Eigen::Matrix<double, 4, 5> expected;
    expected <<
             1.0, 0.0, 0.0, 0.0, 0.0,
            0.0, 1.0, 0.0, 0.0, 0.5,
            0.0, 0.0, 1.0, 0.0, 0.0,
            0.0, 0.0, 0.0, 1.0, 0.5;

    ASSERT_THAT(coarse_basis_in_fine, MatrixEq(expected));
}

RC_GTEST_PROP(standard_coarse_basis_in_fine_space, quasi_interpolation_recovers_coarse_basis, ())
{
    // Let v_H be the weights of a function in the coarse space. Then v_H is also a function in the fine space,
    // and since I_H is a quasi interpolation, we have that I_H v_h = v_H, where v_h are the fine weights of
    // the coarse function.
    // We can verify that this property holds for all basis functions as represented in the fine space.

    const auto coarse = *crest::gen::arbitrary_unit_square_mesh().as("Coarse mesh");
    const auto fine = *crest::gen::arbitrary_refinement(coarse, 0, 5).as("Fine mesh");

    const auto basis = crest::detail::standard_coarse_basis_in_fine_space(coarse, fine);
    const auto I_H = crest::quasi_interpolator(coarse, fine);

    const Eigen::SparseMatrix<double> interpolated = I_H * basis.transpose();
    const Eigen::MatrixXd interpolated_dense = interpolated;

    RC_LOG() << interpolated_dense << std::endl;

    // We could have had a shaper tolerance here, but it appears Eigen does not do an elementwise comparison
    // for isIdentity (which is very annoying, by the way), as it can spuriously fail even when all non-diagonal
    // elements are smaller than 3e-16.
    RC_ASSERT(interpolated_dense.isIdentity(1e-14));
}
<|endoftext|>


<commit_before>/**
 * This file is part of the "FnordMetric" project
 *   Copyright (c) 2014 Paul Asmuth, Google Inc.
 *
 * FnordMetric is free software: you can redistribute it and/or modify it under
 * the terms of the GNU General Public License v3.0. You should have received a
 * copy of the GNU General Public License along with this program. If not, see
 * <http://www.gnu.org/licenses/>.
 */
#include <fnordmetric/environment.h>
#include <fnordmetric/metricdb/backends/disk/binaryformat.h>
#include <fnordmetric/metricdb/backends/disk/metric.h>
#include <fnordmetric/metricdb/backends/disk/tableref.h>
#include <fnordmetric/metricdb/backends/disk/samplewriter.h>
#include <fnordmetric/util/runtimeexception.h>
#include <fnordmetric/util/freeondestroy.h>
#include <fnordmetric/util/wallclock.h>
#include <string.h>

using namespace fnord;
using fnord::util::WallClock;

namespace fnordmetric {
namespace metricdb {
namespace disk_backend {

Metric::Metric(
    const std::string& key,
    io::FileRepository* file_repo) :
    IMetric(key),
    file_repo_(file_repo),
    head_(nullptr),
    max_generation_(0),
    live_table_max_size_(kLiveTableMaxSize),
    live_table_idle_time_micros_(kLiveTableIdleTimeMicros),
    last_insert_(0) {}

Metric::Metric(
    const std::string& key,
    io::FileRepository* file_repo,
    std::vector<std::unique_ptr<TableRef>>&& tables) :
    IMetric(key),
    file_repo_(file_repo),
    live_table_max_size_(kLiveTableMaxSize),
    live_table_idle_time_micros_(kLiveTableIdleTimeMicros),
    last_insert_(0) {
  TableRef* head_table = nullptr;
  std::vector<uint64_t> generations;

  for (auto& table : tables) {
    if (head_table == nullptr ||
        table->generation() > head_table->generation()) {
      head_table = table.get();
    }
  }

  generations = head_table->parents();
  generations.emplace_back(head_table->generation());

  if (env()->verbose()) {
    env()->logger()->printf(
        "DEBUG",
        "Reopening metric: '%s' with %i table(s), generation: %" PRIu64,
        key.c_str(),
        (int) generations.size(),
        head_table->generation());
  }

  std::shared_ptr<MetricSnapshot> snapshot(new MetricSnapshot());
  for (const auto gen : generations) {
    for (auto& table : tables) {
      if (table.get() == nullptr) {
        continue;
      }

      if (table->generation() == gen) {
        table->import(&token_index_, &label_index_);
        snapshot->appendTable(std::move(table));
        table.reset(nullptr);
      }
    }
  }

  head_ = snapshot;
  max_generation_ = head_table->generation();

  for (auto& table : tables) {
    if (table.get() != nullptr) {
      env()->logger()->printf("INFO", "Orphaned table...");
    }
  }
}

std::shared_ptr<MetricSnapshot> Metric::getSnapshot() const {
  std::lock_guard<std::mutex> lock_holder(head_mutex_);
  return head_;
}

// Must hold append_mutex_ to call this!
std::shared_ptr<MetricSnapshot> Metric::getOrCreateSnapshot() {
  {
    std::lock_guard<std::mutex> lock_holder(head_mutex_);

    if (head_.get() != nullptr &&
        head_->isWritable() &&
        head_->tables().back()->isWritable() &&
        head_->tables().back()->bodySize() < live_table_max_size_) {
      return head_;
    }
  }

  std::lock_guard<std::mutex> lock_holder(head_mutex_);
  auto new_snapshot = createSnapshot(true);
  head_ = new_snapshot;
  return head_;
}

void Metric::insertSample(
    double value,
    const std::vector<std::pair<std::string, std::string>>& labels) {
  SampleWriter writer(&token_index_);
  writer.writeValue(value);
  for (const auto& label : labels) {
    writer.writeLabel(label.first, label.second);
    label_index_.addLabel(label.first);
  }

  std::lock_guard<std::mutex> lock_holder(append_mutex_);
  auto snapshot = getOrCreateSnapshot();
  auto& table = snapshot->tables().back();

  uint64_t now = fnord::util::WallClock::unixMicros();
  table->addSample(&writer, now);
  last_insert_ = now;
}

// FIXPAUL misnomer...it creates a new snapshot + appends a new, clean table
std::shared_ptr<MetricSnapshot> Metric::createSnapshot(bool writable) {
  std::shared_ptr<MetricSnapshot> snapshot;
  std::vector<uint64_t> parents;

  if (head_.get() == nullptr) {
    snapshot.reset(new MetricSnapshot());
  } else {
    snapshot = head_->clone();
    for (const auto& tbl : snapshot->tables()) {
      parents.emplace_back(tbl->generation());
    }
  }

  snapshot->setWritable(writable);

  if (writable) {
    // open new file
    auto fileref = file_repo_->createFile();
    auto file = io::File::openFile(
        fileref.absolute_path,
        io::File::O_READ | io::File::O_WRITE | io::File::O_CREATE);

    snapshot->appendTable(TableRef::createTable(
        fileref.absolute_path,
        key_,
        std::move(file),
        ++max_generation_,
        parents));
  }

  return snapshot;
}

void Metric::scanSamples(
    const fnord::util::DateTime& time_begin,
    const fnord::util::DateTime& time_end,
    std::function<bool (Sample* sample)> callback) {
  auto snapshot = getSnapshot();
  if (snapshot.get() == nullptr) {
    return;
  }

  MetricCursor cursor(snapshot, &token_index_);
  while (cursor.valid()) {
    //callback(&cursor);

    if (!cursor.next()) {
      break;
    }
  }
}

void Metric::compact(CompactionPolicy* compaction /* = nullptr */) {
  if (!compaction_mutex_.try_lock()) {
    return;
  }

  std::lock_guard<std::mutex>(compaction_mutex_, std::adopt_lock);

  if (env()->verbose()) {
    env()->logger()->printf(
        "DEBUG",
        "Running compaction for metric: '%s'",
        key_.c_str());
  }

  // create a new snapshot and make all tables immutable
  std::shared_ptr<MetricSnapshot> snapshot;
  {
    std::lock_guard<std::mutex> append_lock_holder(append_mutex_);
    std::lock_guard<std::mutex> head_lock_holder(head_mutex_);
    snapshot = createSnapshot(false);
  }

  auto old_tables = snapshot->tables();
  auto idle_time_micros = WallClock::unixMicros() - last_insert_;

  if (idle_time_micros < live_table_idle_time_micros_) {
    old_tables.pop_back();
  }

  std::vector<std::shared_ptr<TableRef>> new_tables;

  // finalize unfinished sstables
  for (auto& table : old_tables) {
    if (table->isWritable()) {
      if (table->bodySize() == 0) {
        if (env()->verbose()) {
          env()->logger()->printf(
              "DEBUG",
              "SSTable '%s' (%s) is empty, deleting...",
              table->filename().c_str(),
              table->metricKey().c_str());
        }

        continue;
      } else {
        if (env()->verbose()) {
          env()->logger()->printf(
              "DEBUG",
              "Finalizing sstable '%s' (%s)",
              table->filename().c_str(),
              table->metricKey().c_str());
        }

        table->finalize(&token_index_, &label_index_);
        new_tables.emplace_back(new ReadonlyTableRef(*table));
      }
    } else {
      new_tables.emplace_back(table);
    }
  }

  // run the compaction
  if (compaction != nullptr) {
    compaction->compact(&new_tables);
  }

  // create a new snapshot and commit modifications
  {
    std::lock_guard<std::mutex> append_lock_holder(append_mutex_);
    std::lock_guard<std::mutex> head_lock_holder(head_mutex_);
    auto new_snapshot = new MetricSnapshot();
    new_snapshot->setWritable(head_->isWritable());

    for (const auto& table : new_tables) {
      new_snapshot->appendTable(table);
    }

    for (int i = old_tables.size(); i < head_->tables().size(); ++i) {
      new_snapshot->appendTable(head_->tables()[i]);
    }

    head_.reset(new_snapshot);
  }
}

void Metric::setLiveTableMaxSize(size_t max_size) {
  live_table_max_size_ = max_size;
}

void Metric::setLiveTableIdleTimeMicros(uint64_t idle_time_micros) {
  live_table_idle_time_micros_ = idle_time_micros;
}

size_t Metric::numTables() const {
  auto snapshot = getSnapshot();
  return snapshot->tables().size();
}

size_t Metric::totalBytes() const {
  auto snapshot = getSnapshot();

  size_t bytes = 0;
  for (const auto& tbl : snapshot->tables()) {
    bytes += tbl->bodySize();
  }

  return bytes;
}

DateTime Metric::lastInsertTime() const {
  return DateTime(last_insert_);
}

std::set<std::string> Metric::labels() const {
  return label_index_.labels();
}

bool Metric::hasLabel(const std::string& label) const {
  return label_index_.hasLabel(label);
}

}
}
}

<commit_msg>disk_backend::Metric::scanSamples<commit_after>/**
 * This file is part of the "FnordMetric" project
 *   Copyright (c) 2014 Paul Asmuth, Google Inc.
 *
 * FnordMetric is free software: you can redistribute it and/or modify it under
 * the terms of the GNU General Public License v3.0. You should have received a
 * copy of the GNU General Public License along with this program. If not, see
 * <http://www.gnu.org/licenses/>.
 */
#include <fnordmetric/environment.h>
#include <fnordmetric/metricdb/backends/disk/binaryformat.h>
#include <fnordmetric/metricdb/backends/disk/metric.h>
#include <fnordmetric/metricdb/backends/disk/tableref.h>
#include <fnordmetric/metricdb/backends/disk/samplewriter.h>
#include <fnordmetric/util/runtimeexception.h>
#include <fnordmetric/util/freeondestroy.h>
#include <fnordmetric/util/wallclock.h>
#include <string.h>

using namespace fnord;
using fnord::util::WallClock;

namespace fnordmetric {
namespace metricdb {
namespace disk_backend {

Metric::Metric(
    const std::string& key,
    io::FileRepository* file_repo) :
    IMetric(key),
    file_repo_(file_repo),
    head_(nullptr),
    max_generation_(0),
    live_table_max_size_(kLiveTableMaxSize),
    live_table_idle_time_micros_(kLiveTableIdleTimeMicros),
    last_insert_(0) {}

Metric::Metric(
    const std::string& key,
    io::FileRepository* file_repo,
    std::vector<std::unique_ptr<TableRef>>&& tables) :
    IMetric(key),
    file_repo_(file_repo),
    live_table_max_size_(kLiveTableMaxSize),
    live_table_idle_time_micros_(kLiveTableIdleTimeMicros),
    last_insert_(0) {
  TableRef* head_table = nullptr;
  std::vector<uint64_t> generations;

  for (auto& table : tables) {
    if (head_table == nullptr ||
        table->generation() > head_table->generation()) {
      head_table = table.get();
    }
  }

  generations = head_table->parents();
  generations.emplace_back(head_table->generation());

  if (env()->verbose()) {
    env()->logger()->printf(
        "DEBUG",
        "Reopening metric: '%s' with %i table(s), generation: %" PRIu64,
        key.c_str(),
        (int) generations.size(),
        head_table->generation());
  }

  std::shared_ptr<MetricSnapshot> snapshot(new MetricSnapshot());
  for (const auto gen : generations) {
    for (auto& table : tables) {
      if (table.get() == nullptr) {
        continue;
      }

      if (table->generation() == gen) {
        table->import(&token_index_, &label_index_);
        snapshot->appendTable(std::move(table));
        table.reset(nullptr);
      }
    }
  }

  head_ = snapshot;
  max_generation_ = head_table->generation();

  for (auto& table : tables) {
    if (table.get() != nullptr) {
      env()->logger()->printf("INFO", "Orphaned table...");
    }
  }
}

std::shared_ptr<MetricSnapshot> Metric::getSnapshot() const {
  std::lock_guard<std::mutex> lock_holder(head_mutex_);
  return head_;
}

// Must hold append_mutex_ to call this!
std::shared_ptr<MetricSnapshot> Metric::getOrCreateSnapshot() {
  {
    std::lock_guard<std::mutex> lock_holder(head_mutex_);

    if (head_.get() != nullptr &&
        head_->isWritable() &&
        head_->tables().back()->isWritable() &&
        head_->tables().back()->bodySize() < live_table_max_size_) {
      return head_;
    }
  }

  std::lock_guard<std::mutex> lock_holder(head_mutex_);
  auto new_snapshot = createSnapshot(true);
  head_ = new_snapshot;
  return head_;
}

void Metric::insertSample(
    double value,
    const std::vector<std::pair<std::string, std::string>>& labels) {
  SampleWriter writer(&token_index_);
  writer.writeValue(value);
  for (const auto& label : labels) {
    writer.writeLabel(label.first, label.second);
    label_index_.addLabel(label.first);
  }

  std::lock_guard<std::mutex> lock_holder(append_mutex_);
  auto snapshot = getOrCreateSnapshot();
  auto& table = snapshot->tables().back();

  uint64_t now = fnord::util::WallClock::unixMicros();
  table->addSample(&writer, now);
  last_insert_ = now;
}

// FIXPAUL misnomer...it creates a new snapshot + appends a new, clean table
std::shared_ptr<MetricSnapshot> Metric::createSnapshot(bool writable) {
  std::shared_ptr<MetricSnapshot> snapshot;
  std::vector<uint64_t> parents;

  if (head_.get() == nullptr) {
    snapshot.reset(new MetricSnapshot());
  } else {
    snapshot = head_->clone();
    for (const auto& tbl : snapshot->tables()) {
      parents.emplace_back(tbl->generation());
    }
  }

  snapshot->setWritable(writable);

  if (writable) {
    // open new file
    auto fileref = file_repo_->createFile();
    auto file = io::File::openFile(
        fileref.absolute_path,
        io::File::O_READ | io::File::O_WRITE | io::File::O_CREATE);

    snapshot->appendTable(TableRef::createTable(
        fileref.absolute_path,
        key_,
        std::move(file),
        ++max_generation_,
        parents));
  }

  return snapshot;
}

void Metric::scanSamples(
    const fnord::util::DateTime& time_begin,
    const fnord::util::DateTime& time_end,
    std::function<bool (Sample* sample)> callback) {
  auto snapshot = getSnapshot();
  if (snapshot.get() == nullptr) {
    return;
  }

  MetricCursor cursor(snapshot, &token_index_);
  while (cursor.valid()) {
    auto sample = cursor.sample<double>();
    auto time = cursor.time();

    if (time >= static_cast<uint64_t>(time_end)) {
      break;
    }

    if (time >= static_cast<uint64_t>(time_begin)) {
      Sample cb_sample(
          time,
          sample->value(),
          sample->labels());

      callback(&cb_sample);
    }

    if (!cursor.next()) {
      break;
    }
  }
}

void Metric::compact(CompactionPolicy* compaction /* = nullptr */) {
  if (!compaction_mutex_.try_lock()) {
    return;
  }

  std::lock_guard<std::mutex>(compaction_mutex_, std::adopt_lock);

  if (env()->verbose()) {
    env()->logger()->printf(
        "DEBUG",
        "Running compaction for metric: '%s'",
        key_.c_str());
  }

  // create a new snapshot and make all tables immutable
  std::shared_ptr<MetricSnapshot> snapshot;
  {
    std::lock_guard<std::mutex> append_lock_holder(append_mutex_);
    std::lock_guard<std::mutex> head_lock_holder(head_mutex_);
    snapshot = createSnapshot(false);
  }

  auto old_tables = snapshot->tables();
  auto idle_time_micros = WallClock::unixMicros() - last_insert_;

  if (idle_time_micros < live_table_idle_time_micros_) {
    old_tables.pop_back();
  }

  std::vector<std::shared_ptr<TableRef>> new_tables;

  // finalize unfinished sstables
  for (auto& table : old_tables) {
    if (table->isWritable()) {
      if (table->bodySize() == 0) {
        if (env()->verbose()) {
          env()->logger()->printf(
              "DEBUG",
              "SSTable '%s' (%s) is empty, deleting...",
              table->filename().c_str(),
              table->metricKey().c_str());
        }

        continue;
      } else {
        if (env()->verbose()) {
          env()->logger()->printf(
              "DEBUG",
              "Finalizing sstable '%s' (%s)",
              table->filename().c_str(),
              table->metricKey().c_str());
        }

        table->finalize(&token_index_, &label_index_);
        new_tables.emplace_back(new ReadonlyTableRef(*table));
      }
    } else {
      new_tables.emplace_back(table);
    }
  }

  // run the compaction
  if (compaction != nullptr) {
    compaction->compact(&new_tables);
  }

  // create a new snapshot and commit modifications
  {
    std::lock_guard<std::mutex> append_lock_holder(append_mutex_);
    std::lock_guard<std::mutex> head_lock_holder(head_mutex_);
    auto new_snapshot = new MetricSnapshot();
    new_snapshot->setWritable(head_->isWritable());

    for (const auto& table : new_tables) {
      new_snapshot->appendTable(table);
    }

    for (int i = old_tables.size(); i < head_->tables().size(); ++i) {
      new_snapshot->appendTable(head_->tables()[i]);
    }

    head_.reset(new_snapshot);
  }
}

void Metric::setLiveTableMaxSize(size_t max_size) {
  live_table_max_size_ = max_size;
}

void Metric::setLiveTableIdleTimeMicros(uint64_t idle_time_micros) {
  live_table_idle_time_micros_ = idle_time_micros;
}

size_t Metric::numTables() const {
  auto snapshot = getSnapshot();
  return snapshot->tables().size();
}

size_t Metric::totalBytes() const {
  auto snapshot = getSnapshot();

  size_t bytes = 0;
  for (const auto& tbl : snapshot->tables()) {
    bytes += tbl->bodySize();
  }

  return bytes;
}

DateTime Metric::lastInsertTime() const {
  return DateTime(last_insert_);
}

std::set<std::string> Metric::labels() const {
  return label_index_.labels();
}

bool Metric::hasLabel(const std::string& label) const {
  return label_index_.hasLabel(label);
}

}
}
}

<|endoftext|>


<commit_before>/*
 * Copyright (C) 2003-2005 Tommi Maekitalo
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * As a special exception, you may use this file as part of a free
 * software library without restriction. Specifically, if other files
 * instantiate templates or use macros or inline functions from this
 * file, or you compile this file and link it with other files to
 * produce an executable, this file does not by itself cause the
 * resulting executable to be covered by the GNU General Public
 * License. This exception does not however invalidate any other
 * reasons why the executable file might be covered by the GNU Library
 * General Public License.
 *
 * This library 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
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
 */

#include <tnt/httpparser.h>
#include <tnt/httperror.h>
#include <tnt/httpheader.h>
#include <cxxtools/log.h>
#include <sstream>
#include <algorithm>

#define SET_STATE(new_state)  state = &Parser::new_state

namespace tnt
{
  namespace
  {
    std::string chartoprint(char ch)
    {
      const static char hex[] = "0123456789abcdef";
      if (std::isprint(ch))
        return std::string(1, '\'') + ch + '\'';
      else
        return std::string("'\\x") + hex[(ch >> 4) & 0xf] + hex[ch & 0xf] + '\'';
    }

    inline bool istokenchar(char ch)
    {
      static const char s[] = "\"(),/:;<=>?@[\\]{}";
      return std::isalpha(ch) || std::binary_search(s, s + sizeof(s) - 1, ch);
    }

    inline bool isHexDigit(char ch)
    {
      return (ch >= '0' && ch <= '9')
          || (ch >= 'A' && ch <= 'Z')
          || (ch >= 'a' && ch <= 'z');
    }

    inline unsigned valueOfHexDigit(char ch)
    {
      return ch >= '0' && ch <= '9' ? ch - '0'
           : ch >= 'a' && ch <= 'z' ? ch - 'a' + 10
           : ch >= 'A' && ch <= 'Z' ? ch - 'A' + 10
           : 0;
    }
  }

  log_define("tntnet.httpmessage.parser")

  bool RequestSizeMonitor::post(bool ret)
  {
    if (++requestSize > HttpRequest::getMaxRequestSize()
      && HttpRequest::getMaxRequestSize() > 0)
    {
      requestSizeExceeded();
      return true;
    }
    return ret;
  }

  void RequestSizeMonitor::requestSizeExceeded()
  { }

  void HttpRequest::Parser::reset()
  {
    message.clear();
    SET_STATE(state_cmd0);
    httpCode = HTTP_OK;
    failedFlag = false;
    RequestSizeMonitor::reset();
    headerParser.reset();
  }

  bool HttpRequest::Parser::state_cmd0(char ch)
  {
    if (istokenchar(ch))
    {
      message.method[0] = ch;
      message.methodLen = 1;
      SET_STATE(state_cmd);
    }
    else if (ch != ' ' && ch != '\t')
    {
      log_warn("invalid character " << chartoprint(ch) << " in method");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }
    return failedFlag;
  }

  bool HttpRequest::Parser::state_cmd(char ch)
  {
    if (istokenchar(ch))
    {
      if (message.methodLen >= sizeof(message.method) - 1)
      {
        log_debug("invalid method field; method=" << std::string(message.method, message.methodLen) << ", len=" << message.methodLen);
        throw HttpError(HTTP_BAD_REQUEST, "invalid method field");
      }
      message.method[message.methodLen++] = ch;
    }
    else if (ch == ' ')
    {
      message.method[message.methodLen] = '\0';
      log_debug("method=" << message.method);
      SET_STATE(state_url0);
    }
    else
    {
      log_warn("invalid character " << chartoprint(ch) << " in method");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }
    return failedFlag;
  }

  bool HttpRequest::Parser::state_url0(char ch)
  {
    if (ch == ' ' || ch == '\t')
    {
    }
    else if (ch == '/')
    {
      message.url.clear();
      message.url.reserve(32);
      message.url += ch;
      SET_STATE(state_url);
    }
    else if (std::isalpha(ch))
    {
      SET_STATE(state_protocol);
    }
    else
    {
      log_warn("invalid character " << chartoprint(ch) << " in url");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }

    return failedFlag;
  }

  bool HttpRequest::Parser::state_protocol(char ch)
  {
    if (ch == ':')
      SET_STATE(state_protocol_slash1);
    else if (!std::isalpha(ch))
    {
      log_warn("invalid character " << chartoprint(ch) << " in url");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }

    return failedFlag;
  }

  bool HttpRequest::Parser::state_protocol_slash1(char ch)
  {
    if (ch == '/')
      SET_STATE(state_protocol_slash2);
    else
    {
      log_warn("invalid character " << chartoprint(ch) << " in url");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }

    return failedFlag;
  }

  bool HttpRequest::Parser::state_protocol_slash2(char ch)
  {
    if (ch == '/')
      SET_STATE(state_protocol_host);
    else
    {
      log_warn("invalid character " << chartoprint(ch) << " in url");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }

    return failedFlag;
  }

  bool HttpRequest::Parser::state_protocol_host(char ch)
  {
    if (ch == '/')
    {
      message.url.clear();
      message.url.reserve(32);
      message.url += ch;
      SET_STATE(state_url);
    }
    else if (!std::isalpha(ch)
           && !std::isdigit(ch)
           && ch != '['
           && ch != ']'
           && ch != '.'
           && ch != ':')
    {
      log_warn("invalid character " << chartoprint(ch) << " in url");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }

    return failedFlag;
  }

  bool HttpRequest::Parser::state_url(char ch)
  {
    if (ch == '?')
    {
      log_debug("url=" << message.url);
      SET_STATE(state_qparam);
    }
    else if (ch == '\r')
    {
      log_debug("url=" << message.url);
      SET_STATE(state_end0);
    }
    else if (ch == '\n')
    {
      log_debug("url=" << message.url);
      SET_STATE(state_header);
    }
    else if (ch == ' ' || ch == '\t')
    {
      log_debug("url=" << message.url);
      SET_STATE(state_version);
    }
    else if (ch == '+')
      message.url += ' ';
    else if (ch == '%')
    {
      SET_STATE(state_urlesc);
      message.url += ch;
    }
    else if (ch > ' ')
      message.url += ch;
    else
    {
      log_warn("invalid character " << chartoprint(ch) << " in url");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }
    return failedFlag;
  }

  bool HttpRequest::Parser::state_urlesc(char ch)
  {
    if (isHexDigit(ch))
    {
      if (message.url.size() >= 2 && message.url[message.url.size() - 2] == '%')
      {
        unsigned v = (valueOfHexDigit(message.url[message.url.size() - 1]) << 4) | valueOfHexDigit(ch);
        message.url[message.url.size() - 2] = static_cast<char>(v);
        message.url.resize(message.url.size() - 1);
        SET_STATE(state_url);
      }
      else
      {
        message.url += ch;
      }
      return false;
    }
    else
    {
      SET_STATE(state_url);
      return state_url(ch);
    }
  }

  bool HttpRequest::Parser::state_qparam(char ch)
  {
    if (ch == ' ' || ch == '\t')
    {
      log_debug("queryString=" << message.queryString);
      SET_STATE(state_version);
    }
    else
      message.queryString += ch;
    return false;
  }

  bool HttpRequest::Parser::state_version(char ch)
  {
    if (ch == '/')
    {
      message.setVersion(0, 0);
      skipWs(&Parser::state_version_major);
    }
    else if (ch == '\r')
    {
      log_warn("invalid character " << chartoprint(ch) << " in version");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }
    return failedFlag;
  }

  bool HttpRequest::Parser::state_version_major(char ch)
  {
    if (ch == '.')
      SET_STATE(state_version_minor0);
    else if (std::isdigit(ch))
      message.setVersion(message.getMajorVersion() * 10 + (ch - '0'), message.getMinorVersion());
    else if (ch == ' ' || ch == '\t')
      SET_STATE(state_version_major_sp);
    else
    {
      log_warn("invalid character " << chartoprint(ch) << " in version-major");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }
    return failedFlag;
  }

  bool HttpRequest::Parser::state_version_major_sp(char ch)
  {
    if (ch == '.')
      SET_STATE(state_version_minor0);
    else
    {
      log_warn("invalid character " << chartoprint(ch) << " in version-major");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }
    return failedFlag;
  }

  bool HttpRequest::Parser::state_version_minor0(char ch)
  {
    return ch == ' ' || ch == '\t' ? failedFlag
                                   : state_version_minor(ch);
  }

  bool HttpRequest::Parser::state_version_minor(char ch)
  {
    if (ch == '\n')
      SET_STATE(state_header);
    else if (ch == ' ' || ch == '\t' || ch == '\r')
      SET_STATE(state_end0);
    else if (std::isdigit(ch))
      message.setVersion(message.getMajorVersion(), message.getMinorVersion() * 10 + (ch - '0'));
    else
    {
      log_warn("invalid character " << chartoprint(ch) << " in version-minor");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }
    return failedFlag;
  }

  bool HttpRequest::Parser::state_end0(char ch)
  {
    if (ch == '\n')
      SET_STATE(state_header);
    else if (ch != ' ' && ch != '\t')
    {
      log_warn("invalid character " << chartoprint(ch) << " in end");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }
    return failedFlag;
  }

  bool HttpRequest::Parser::state_header(char ch)
  {
    if (headerParser.parse(ch))
    {
      if (headerParser.failed())
      {
        httpCode = HTTP_BAD_REQUEST;
        failedFlag = true;
        return true;
      }

      const char* content_length_header = message.getHeader(httpheader::contentLength);
      if (*content_length_header)
      {
        bodySize = 0;
        for (const char* c = content_length_header; *c; ++c)
        {
          if (*c > '9' || *c < '0')
            throw HttpError(HTTP_BAD_REQUEST, "invalid Content-Length");
          bodySize = bodySize * 10 + *c - '0';
        }

        if (getMaxRequestSize() > 0
          && getCurrentRequestSize() + bodySize > getMaxRequestSize())
        {
          requestSizeExceeded();
          return true;
        }

        message.contentSize = bodySize;
        if (bodySize == 0)
          return true;
        else
        {
          SET_STATE(state_body);
          message.body.reserve(bodySize);
          return false;
        }
      }

      return true;
    }

    return false;
  }

  bool HttpRequest::Parser::state_body(char ch)
  {
    message.body += ch;
    return --bodySize == 0;
  }

  void HttpRequest::Parser::requestSizeExceeded()
  {
    log_warn("max request size " << getMaxRequestSize() << " exceeded");
    httpCode = HTTP_REQUEST_ENTITY_TOO_LARGE;
    failedFlag = true;
  }
}
<commit_msg>do not translate '+' to space when parsing urls<commit_after>/*
 * Copyright (C) 2003-2005 Tommi Maekitalo
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * As a special exception, you may use this file as part of a free
 * software library without restriction. Specifically, if other files
 * instantiate templates or use macros or inline functions from this
 * file, or you compile this file and link it with other files to
 * produce an executable, this file does not by itself cause the
 * resulting executable to be covered by the GNU General Public
 * License. This exception does not however invalidate any other
 * reasons why the executable file might be covered by the GNU Library
 * General Public License.
 *
 * This library 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
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
 */

#include <tnt/httpparser.h>
#include <tnt/httperror.h>
#include <tnt/httpheader.h>
#include <cxxtools/log.h>
#include <sstream>
#include <algorithm>

#define SET_STATE(new_state)  state = &Parser::new_state

namespace tnt
{
  namespace
  {
    std::string chartoprint(char ch)
    {
      const static char hex[] = "0123456789abcdef";
      if (std::isprint(ch))
        return std::string(1, '\'') + ch + '\'';
      else
        return std::string("'\\x") + hex[(ch >> 4) & 0xf] + hex[ch & 0xf] + '\'';
    }

    inline bool istokenchar(char ch)
    {
      static const char s[] = "\"(),/:;<=>?@[\\]{}";
      return std::isalpha(ch) || std::binary_search(s, s + sizeof(s) - 1, ch);
    }

    inline bool isHexDigit(char ch)
    {
      return (ch >= '0' && ch <= '9')
          || (ch >= 'A' && ch <= 'Z')
          || (ch >= 'a' && ch <= 'z');
    }

    inline unsigned valueOfHexDigit(char ch)
    {
      return ch >= '0' && ch <= '9' ? ch - '0'
           : ch >= 'a' && ch <= 'z' ? ch - 'a' + 10
           : ch >= 'A' && ch <= 'Z' ? ch - 'A' + 10
           : 0;
    }
  }

  log_define("tntnet.httpmessage.parser")

  bool RequestSizeMonitor::post(bool ret)
  {
    if (++requestSize > HttpRequest::getMaxRequestSize()
      && HttpRequest::getMaxRequestSize() > 0)
    {
      requestSizeExceeded();
      return true;
    }
    return ret;
  }

  void RequestSizeMonitor::requestSizeExceeded()
  { }

  void HttpRequest::Parser::reset()
  {
    message.clear();
    SET_STATE(state_cmd0);
    httpCode = HTTP_OK;
    failedFlag = false;
    RequestSizeMonitor::reset();
    headerParser.reset();
  }

  bool HttpRequest::Parser::state_cmd0(char ch)
  {
    if (istokenchar(ch))
    {
      message.method[0] = ch;
      message.methodLen = 1;
      SET_STATE(state_cmd);
    }
    else if (ch != ' ' && ch != '\t')
    {
      log_warn("invalid character " << chartoprint(ch) << " in method");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }
    return failedFlag;
  }

  bool HttpRequest::Parser::state_cmd(char ch)
  {
    if (istokenchar(ch))
    {
      if (message.methodLen >= sizeof(message.method) - 1)
      {
        log_debug("invalid method field; method=" << std::string(message.method, message.methodLen) << ", len=" << message.methodLen);
        throw HttpError(HTTP_BAD_REQUEST, "invalid method field");
      }
      message.method[message.methodLen++] = ch;
    }
    else if (ch == ' ')
    {
      message.method[message.methodLen] = '\0';
      log_debug("method=" << message.method);
      SET_STATE(state_url0);
    }
    else
    {
      log_warn("invalid character " << chartoprint(ch) << " in method");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }
    return failedFlag;
  }

  bool HttpRequest::Parser::state_url0(char ch)
  {
    if (ch == ' ' || ch == '\t')
    {
    }
    else if (ch == '/')
    {
      message.url.clear();
      message.url.reserve(32);
      message.url += ch;
      SET_STATE(state_url);
    }
    else if (std::isalpha(ch))
    {
      SET_STATE(state_protocol);
    }
    else
    {
      log_warn("invalid character " << chartoprint(ch) << " in url");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }

    return failedFlag;
  }

  bool HttpRequest::Parser::state_protocol(char ch)
  {
    if (ch == ':')
      SET_STATE(state_protocol_slash1);
    else if (!std::isalpha(ch))
    {
      log_warn("invalid character " << chartoprint(ch) << " in url");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }

    return failedFlag;
  }

  bool HttpRequest::Parser::state_protocol_slash1(char ch)
  {
    if (ch == '/')
      SET_STATE(state_protocol_slash2);
    else
    {
      log_warn("invalid character " << chartoprint(ch) << " in url");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }

    return failedFlag;
  }

  bool HttpRequest::Parser::state_protocol_slash2(char ch)
  {
    if (ch == '/')
      SET_STATE(state_protocol_host);
    else
    {
      log_warn("invalid character " << chartoprint(ch) << " in url");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }

    return failedFlag;
  }

  bool HttpRequest::Parser::state_protocol_host(char ch)
  {
    if (ch == '/')
    {
      message.url.clear();
      message.url.reserve(32);
      message.url += ch;
      SET_STATE(state_url);
    }
    else if (!std::isalpha(ch)
           && !std::isdigit(ch)
           && ch != '['
           && ch != ']'
           && ch != '.'
           && ch != ':')
    {
      log_warn("invalid character " << chartoprint(ch) << " in url");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }

    return failedFlag;
  }

  bool HttpRequest::Parser::state_url(char ch)
  {
    if (ch == '?')
    {
      log_debug("url=" << message.url);
      SET_STATE(state_qparam);
    }
    else if (ch == '\r')
    {
      log_debug("url=" << message.url);
      SET_STATE(state_end0);
    }
    else if (ch == '\n')
    {
      log_debug("url=" << message.url);
      SET_STATE(state_header);
    }
    else if (ch == ' ' || ch == '\t')
    {
      log_debug("url=" << message.url);
      SET_STATE(state_version);
    }
    else if (ch == '%')
    {
      SET_STATE(state_urlesc);
      message.url += ch;
    }
    else if (ch > ' ')
      message.url += ch;
    else
    {
      log_warn("invalid character " << chartoprint(ch) << " in url");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }
    return failedFlag;
  }

  bool HttpRequest::Parser::state_urlesc(char ch)
  {
    if (isHexDigit(ch))
    {
      if (message.url.size() >= 2 && message.url[message.url.size() - 2] == '%')
      {
        unsigned v = (valueOfHexDigit(message.url[message.url.size() - 1]) << 4) | valueOfHexDigit(ch);
        message.url[message.url.size() - 2] = static_cast<char>(v);
        message.url.resize(message.url.size() - 1);
        SET_STATE(state_url);
      }
      else
      {
        message.url += ch;
      }
      return false;
    }
    else
    {
      SET_STATE(state_url);
      return state_url(ch);
    }
  }

  bool HttpRequest::Parser::state_qparam(char ch)
  {
    if (ch == ' ' || ch == '\t')
    {
      log_debug("queryString=" << message.queryString);
      SET_STATE(state_version);
    }
    else
      message.queryString += ch;
    return false;
  }

  bool HttpRequest::Parser::state_version(char ch)
  {
    if (ch == '/')
    {
      message.setVersion(0, 0);
      skipWs(&Parser::state_version_major);
    }
    else if (ch == '\r')
    {
      log_warn("invalid character " << chartoprint(ch) << " in version");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }
    return failedFlag;
  }

  bool HttpRequest::Parser::state_version_major(char ch)
  {
    if (ch == '.')
      SET_STATE(state_version_minor0);
    else if (std::isdigit(ch))
      message.setVersion(message.getMajorVersion() * 10 + (ch - '0'), message.getMinorVersion());
    else if (ch == ' ' || ch == '\t')
      SET_STATE(state_version_major_sp);
    else
    {
      log_warn("invalid character " << chartoprint(ch) << " in version-major");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }
    return failedFlag;
  }

  bool HttpRequest::Parser::state_version_major_sp(char ch)
  {
    if (ch == '.')
      SET_STATE(state_version_minor0);
    else
    {
      log_warn("invalid character " << chartoprint(ch) << " in version-major");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }
    return failedFlag;
  }

  bool HttpRequest::Parser::state_version_minor0(char ch)
  {
    return ch == ' ' || ch == '\t' ? failedFlag
                                   : state_version_minor(ch);
  }

  bool HttpRequest::Parser::state_version_minor(char ch)
  {
    if (ch == '\n')
      SET_STATE(state_header);
    else if (ch == ' ' || ch == '\t' || ch == '\r')
      SET_STATE(state_end0);
    else if (std::isdigit(ch))
      message.setVersion(message.getMajorVersion(), message.getMinorVersion() * 10 + (ch - '0'));
    else
    {
      log_warn("invalid character " << chartoprint(ch) << " in version-minor");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }
    return failedFlag;
  }

  bool HttpRequest::Parser::state_end0(char ch)
  {
    if (ch == '\n')
      SET_STATE(state_header);
    else if (ch != ' ' && ch != '\t')
    {
      log_warn("invalid character " << chartoprint(ch) << " in end");
      httpCode = HTTP_BAD_REQUEST;
      failedFlag = true;
    }
    return failedFlag;
  }

  bool HttpRequest::Parser::state_header(char ch)
  {
    if (headerParser.parse(ch))
    {
      if (headerParser.failed())
      {
        httpCode = HTTP_BAD_REQUEST;
        failedFlag = true;
        return true;
      }

      const char* content_length_header = message.getHeader(httpheader::contentLength);
      if (*content_length_header)
      {
        bodySize = 0;
        for (const char* c = content_length_header; *c; ++c)
        {
          if (*c > '9' || *c < '0')
            throw HttpError(HTTP_BAD_REQUEST, "invalid Content-Length");
          bodySize = bodySize * 10 + *c - '0';
        }

        if (getMaxRequestSize() > 0
          && getCurrentRequestSize() + bodySize > getMaxRequestSize())
        {
          requestSizeExceeded();
          return true;
        }

        message.contentSize = bodySize;
        if (bodySize == 0)
          return true;
        else
        {
          SET_STATE(state_body);
          message.body.reserve(bodySize);
          return false;
        }
      }

      return true;
    }

    return false;
  }

  bool HttpRequest::Parser::state_body(char ch)
  {
    message.body += ch;
    return --bodySize == 0;
  }

  void HttpRequest::Parser::requestSizeExceeded()
  {
    log_warn("max request size " << getMaxRequestSize() << " exceeded");
    httpCode = HTTP_REQUEST_ENTITY_TOO_LARGE;
    failedFlag = true;
  }
}
<|endoftext|>


<commit_before>/*
 * Copyright 2013-present Facebook, Inc.
 *
 * 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 <cstring>

#include <folly/experimental/TestUtil.h>
#include <folly/experimental/symbolizer/StackTrace.h>
#include <folly/experimental/symbolizer/Symbolizer.h>

#include <glog/logging.h>

#include <folly/portability/GTest.h>

using namespace folly;
using namespace folly::symbolizer;

FOLLY_NOINLINE void foo1();
FOLLY_NOINLINE void foo2();

void verifyStackTraces() {
  constexpr size_t kMaxAddresses = 100;
  FrameArray<kMaxAddresses> fa;
  CHECK(getStackTrace(fa));

  FrameArray<kMaxAddresses> faSafe;
  CHECK(getStackTraceSafe(faSafe));

  CHECK_EQ(fa.frameCount, faSafe.frameCount);

  if (VLOG_IS_ON(1)) {
    Symbolizer symbolizer;
    OStreamSymbolizePrinter printer(std::cerr, SymbolizePrinter::COLOR_IF_TTY);

    symbolizer.symbolize(fa);
    VLOG(1) << "getStackTrace\n";
    printer.println(fa);

    symbolizer.symbolize(faSafe);
    VLOG(1) << "getStackTraceSafe\n";
    printer.println(faSafe);
  }

  // Other than the top 2 frames (this one and getStackTrace /
  // getStackTraceSafe), the stack traces should be identical
  for (size_t i = 2; i < fa.frameCount; ++i) {
    LOG(INFO) << "i=" << i << " " << std::hex << "0x" << fa.addresses[i]
              << " 0x" << faSafe.addresses[i];
    EXPECT_EQ(fa.addresses[i], faSafe.addresses[i]);
  }
}

void foo1() {
  foo2();
}

void foo2() {
  verifyStackTraces();
}

volatile bool handled = false;
void handler(int /* num */, siginfo_t* /* info */, void* /* ctx */) {
  // Yes, getStackTrace and VLOG aren't async-signal-safe, but signals
  // raised with raise() aren't "async" signals.
  foo1();
  handled = true;
}

TEST(StackTraceTest, Simple) {
  foo1();
}

TEST(StackTraceTest, Signal) {
  struct sigaction sa;
  memset(&sa, 0, sizeof(sa));
  sa.sa_sigaction = handler;
  sa.sa_flags = SA_RESETHAND | SA_SIGINFO;
  CHECK_ERR(sigaction(SIGUSR1, &sa, nullptr));
  raise(SIGUSR1);
  EXPECT_TRUE(handled);
}

ssize_t read_all(int fd, uint8_t* buffer, size_t size) {
  uint8_t* pos = buffer;
  ssize_t bytes_read;
  do {
    bytes_read = read(fd, pos, size);
    if (bytes_read < 0) {
      if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) {
        continue;
      }
      return bytes_read;
    }

    pos += bytes_read;
    size -= bytes_read;
  } while (bytes_read > 0 && size > 0);

  return pos - buffer;
}

// Returns the position in the file after done reading.
off_t get_stack_trace(int fd, size_t file_pos, uint8_t* buffer, size_t count) {
  off_t rv = lseek(fd, file_pos, SEEK_SET);
  assert(rv == (off_t)file_pos);

  // Subtract 1 from size of buffer to hold nullptr.
  ssize_t bytes_read = read_all(fd, buffer, count - 1);
  assert(bytes_read > 0);
  buffer[bytes_read] = '\0';
  return lseek(fd, 0, SEEK_CUR);
}

template <class StackTracePrinter>
void testStackTracePrinter(StackTracePrinter& printer, int fd) {
  ASSERT_GT(fd, 0);

  printer.printStackTrace(true);
  printer.flush();

  std::array<uint8_t, 4000> first;
  off_t pos = get_stack_trace(fd, 0, first.data(), first.size());
  ASSERT_GT(pos, 0);

  printer.printStackTrace(true);
  printer.flush();

  std::array<uint8_t, 4000> second;
  get_stack_trace(fd, pos, second.data(), second.size());

  // The first two lines refer to this stack frame, which is different in the
  // two cases, so strip those off.  The rest should be equal.
  ASSERT_STREQ(
      strchr(strchr((const char*)first.data(), '\n') + 1, '\n') + 1,
      strchr(strchr((const char*)second.data(), '\n') + 1, '\n') + 1);
}

TEST(StackTraceTest, SafeStackTracePrinter) {
  test::TemporaryFile file;

  SafeStackTracePrinter printer{10, file.fd()};

  testStackTracePrinter<SafeStackTracePrinter>(printer, file.fd());
}

TEST(StackTraceTest, FastStackTracePrinter) {
  test::TemporaryFile file;

  FastStackTracePrinter printer{
      std::make_unique<FDSymbolizePrinter>(file.fd())};

  testStackTracePrinter<FastStackTracePrinter>(printer, file.fd());
}
<commit_msg>Fix broken opt build due to "error: unused variable 'rv'"<commit_after>/*
 * Copyright 2013-present Facebook, Inc.
 *
 * 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 <cstring>

#include <folly/experimental/TestUtil.h>
#include <folly/experimental/symbolizer/StackTrace.h>
#include <folly/experimental/symbolizer/Symbolizer.h>

#include <glog/logging.h>

#include <folly/portability/GTest.h>

using namespace folly;
using namespace folly::symbolizer;

FOLLY_NOINLINE void foo1();
FOLLY_NOINLINE void foo2();

void verifyStackTraces() {
  constexpr size_t kMaxAddresses = 100;
  FrameArray<kMaxAddresses> fa;
  CHECK(getStackTrace(fa));

  FrameArray<kMaxAddresses> faSafe;
  CHECK(getStackTraceSafe(faSafe));

  CHECK_EQ(fa.frameCount, faSafe.frameCount);

  if (VLOG_IS_ON(1)) {
    Symbolizer symbolizer;
    OStreamSymbolizePrinter printer(std::cerr, SymbolizePrinter::COLOR_IF_TTY);

    symbolizer.symbolize(fa);
    VLOG(1) << "getStackTrace\n";
    printer.println(fa);

    symbolizer.symbolize(faSafe);
    VLOG(1) << "getStackTraceSafe\n";
    printer.println(faSafe);
  }

  // Other than the top 2 frames (this one and getStackTrace /
  // getStackTraceSafe), the stack traces should be identical
  for (size_t i = 2; i < fa.frameCount; ++i) {
    LOG(INFO) << "i=" << i << " " << std::hex << "0x" << fa.addresses[i]
              << " 0x" << faSafe.addresses[i];
    EXPECT_EQ(fa.addresses[i], faSafe.addresses[i]);
  }
}

void foo1() {
  foo2();
}

void foo2() {
  verifyStackTraces();
}

volatile bool handled = false;
void handler(int /* num */, siginfo_t* /* info */, void* /* ctx */) {
  // Yes, getStackTrace and VLOG aren't async-signal-safe, but signals
  // raised with raise() aren't "async" signals.
  foo1();
  handled = true;
}

TEST(StackTraceTest, Simple) {
  foo1();
}

TEST(StackTraceTest, Signal) {
  struct sigaction sa;
  memset(&sa, 0, sizeof(sa));
  sa.sa_sigaction = handler;
  sa.sa_flags = SA_RESETHAND | SA_SIGINFO;
  CHECK_ERR(sigaction(SIGUSR1, &sa, nullptr));
  raise(SIGUSR1);
  EXPECT_TRUE(handled);
}

ssize_t read_all(int fd, uint8_t* buffer, size_t size) {
  uint8_t* pos = buffer;
  ssize_t bytes_read;
  do {
    bytes_read = read(fd, pos, size);
    if (bytes_read < 0) {
      if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) {
        continue;
      }
      return bytes_read;
    }

    pos += bytes_read;
    size -= bytes_read;
  } while (bytes_read > 0 && size > 0);

  return pos - buffer;
}

// Returns the position in the file after done reading.
off_t get_stack_trace(int fd, size_t file_pos, uint8_t* buffer, size_t count) {
  off_t rv = lseek(fd, file_pos, SEEK_SET);
  CHECK_EQ(rv, (off_t)file_pos);

  // Subtract 1 from size of buffer to hold nullptr.
  ssize_t bytes_read = read_all(fd, buffer, count - 1);
  CHECK_GT(bytes_read, 0);
  buffer[bytes_read] = '\0';
  return lseek(fd, 0, SEEK_CUR);
}

template <class StackTracePrinter>
void testStackTracePrinter(StackTracePrinter& printer, int fd) {
  ASSERT_GT(fd, 0);

  printer.printStackTrace(true);
  printer.flush();

  std::array<uint8_t, 4000> first;
  off_t pos = get_stack_trace(fd, 0, first.data(), first.size());
  ASSERT_GT(pos, 0);

  printer.printStackTrace(true);
  printer.flush();

  std::array<uint8_t, 4000> second;
  get_stack_trace(fd, pos, second.data(), second.size());

  // The first two lines refer to this stack frame, which is different in the
  // two cases, so strip those off.  The rest should be equal.
  ASSERT_STREQ(
      strchr(strchr((const char*)first.data(), '\n') + 1, '\n') + 1,
      strchr(strchr((const char*)second.data(), '\n') + 1, '\n') + 1);
}

TEST(StackTraceTest, SafeStackTracePrinter) {
  test::TemporaryFile file;

  SafeStackTracePrinter printer{10, file.fd()};

  testStackTracePrinter<SafeStackTracePrinter>(printer, file.fd());
}

TEST(StackTraceTest, FastStackTracePrinter) {
  test::TemporaryFile file;

  FastStackTracePrinter printer{
      std::make_unique<FDSymbolizePrinter>(file.fd())};

  testStackTracePrinter<FastStackTracePrinter>(printer, file.fd());
}
<|endoftext|>