<commit_before>#include "physics.h"
#include "GLDebugDrawer.h"
#include "../state/state.h"
#include <iostream>

using namespace Physics;

Simulation::Simulation()
{
	m_broadphase = new btDbvtBroadphase();
	m_collisionConfiguration = new btDefaultCollisionConfiguration();
	m_dispatcher = new btCollisionDispatcher(m_collisionConfiguration);
	m_solver = new btSequentialImpulseConstraintSolver;
	m_world = new btDiscreteDynamicsWorld(m_dispatcher, m_broadphase, m_solver, m_collisionConfiguration);

	m_world->setGravity(btVector3(0,-10,0));

	mb.request(Events::EventType::Input);
}

Simulation::~Simulation()
{
	/*
	delete m_world;
	delete m_solver;
	delete m_broadphase;
	delete m_dispatcher;
	delete m_collisionConfiguration;
	*/
}

// Tuning
// Credit to:
// http://bullet.googlecode.com/svn-history/r2704/trunk/Demos/ForkLiftDemo/ForkLiftDemo.cpp
float	gEngineForce = 0.f;

float	defaultBreakingForce = 10.f;
float	gBreakingForce = 100.f;

float	maxEngineForce = 1000.f;//this should be engine/velocity dependent
float	maxBreakingForce = 100.f;

float	gVehicleSteering = 0.f;
float	wheelFriction = 1000;//BT_LARGE_FLOAT;
float	suspensionStiffness = 2.f;
float	suspensionDamping = 1.3f;
float	suspensionCompression = 1.1f;
float	rollInfluence = 1.01f;//1.0f;
btScalar suspensionRestLength(0.2);

btRigidBody *Simulation::addRigidBody(double mass, const btTransform& startTransform, btCollisionShape* shape)
{
	btVector3 localInertia(0, 0, 0);
	if (mass != 0.0) {
		shape->calculateLocalInertia(mass, localInertia);
	}

	btDefaultMotionState* myMotionState = new btDefaultMotionState(startTransform);

	btRigidBody::btRigidBodyConstructionInfo cInfo(mass,myMotionState,shape,localInertia);

	btRigidBody* body = new btRigidBody(cInfo);
	body->setContactProcessingThreshold(1);

	m_world->addRigidBody(body);

	return body;
}

int Simulation::loadWorld()
{
	// Create car
#define CAR_WIDTH (0.15)
#define CAR_LENGTH (0.25)
#define CAR_MASS (400.0)

	btCollisionShape* chassisShape = new btBoxShape(btVector3(CAR_WIDTH, CAR_WIDTH, CAR_LENGTH));
	btCompoundShape* compound = new btCompoundShape();
	m_collisionShapes.push_back(chassisShape);
	m_collisionShapes.push_back(compound);

	btTransform localTrans; // shift gravity to center of car
	localTrans.setIdentity();
	localTrans.setOrigin(btVector3(0,CAR_WIDTH*1, 0));
	compound->addChildShape(localTrans, chassisShape);

	btTransform tr;
	tr.setIdentity();
	tr.setOrigin(btVector3(0,5,0));
	m_carChassis = addRigidBody(CAR_MASS, tr, compound);

	m_vehicleRayCaster = new btDefaultVehicleRaycaster(m_world);
	m_tuning.m_maxSuspensionTravelCm = 204.8f;
	m_tuning.m_suspensionCompression = 4.4f;
	m_tuning.m_suspensionDamping = 2.3f;
	m_tuning.m_frictionSlip = 10000.0f;
	m_tuning.m_suspensionStiffness = 5.0f;
	m_tuning.m_maxSuspensionForce = 1000.0;
	m_vehicle = new btRaycastVehicle(m_tuning,m_carChassis, m_vehicleRayCaster);
	m_carChassis->setActivationState(DISABLE_DEACTIVATION);
	m_world->addVehicle(m_vehicle);

	float connectionHeight = 0.1;
	btVector3 wheelDirectionCS0(0,-1,0);
	btVector3 wheelAxleCS(-1,0,0);

#define CON1 (CAR_WIDTH)
#define CON2 (CAR_LENGTH)
	float	wheelRadius = 0.075f;
	float	wheelWidth = 0.03f;
	bool isFrontWheel=true;
	btVector3 connectionPointCS0(CON1,connectionHeight,CON2);
	m_vehicle->addWheel(connectionPointCS0,wheelDirectionCS0,wheelAxleCS,suspensionRestLength,wheelRadius,m_tuning,isFrontWheel);

	connectionPointCS0 = btVector3(-CON1,connectionHeight,CON2);
	m_vehicle->addWheel(connectionPointCS0,wheelDirectionCS0,wheelAxleCS,suspensionRestLength,wheelRadius,m_tuning,isFrontWheel);

	isFrontWheel = false;

	connectionPointCS0 = btVector3(-CON1,connectionHeight,-CON2);
	m_vehicle->addWheel(connectionPointCS0,wheelDirectionCS0,wheelAxleCS,suspensionRestLength,wheelRadius,m_tuning,isFrontWheel);

	connectionPointCS0 = btVector3(CON1,connectionHeight,-CON2);
	m_vehicle->addWheel(connectionPointCS0,wheelDirectionCS0,wheelAxleCS,suspensionRestLength,wheelRadius,m_tuning,isFrontWheel);

	for (int i=0;i<m_vehicle->getNumWheels();i++)
	{
		btWheelInfo& wheel = m_vehicle->getWheelInfo(i);
		wheel.m_suspensionStiffness = suspensionStiffness;
		wheel.m_wheelsDampingRelaxation = suspensionDamping;
		wheel.m_wheelsDampingCompression = suspensionCompression;
		wheel.m_frictionSlip = wheelFriction;
		wheel.m_rollInfluence = rollInfluence;
	}

	// Add map
	StateData *state = GetMutState();
	btBvhTriangleMeshShape *arenaShape = new btBvhTriangleMeshShape(state->bttmArena, true, true);
	m_arena = addRigidBody(0.0, btTransform(btQuaternion(0,0,0,1),btVector3(0,0,0)), arenaShape);

	m_vehicle->resetSuspension();
	for (int i=0;i<m_vehicle->getNumWheels();i++)
	{
		//synchronize the wheels with the (interpolated) chassis worldtransform
		m_vehicle->updateWheelTransform(i,true);
	}

	return 0;
}

void Simulation::step(double seconds)
{
	// DEBUGOUT("RUNING PHYSICS %lf\n", seconds);

#define STEER_MAX_ANGLE (40.0)
#define ENGINE_MAX_FORCE (5000)
#define BRAKE_MAX_FORCE (3000)

	for ( Events::Event *event : (mb.checkMail()) )
	{
		// Hack: Sorry
		switch ( event->type )
		{
		case Events::EventType::Input:
		{
			Events::InputEvent *input = (Events::InputEvent *)event;
			//double steer = STEER_MAX_ANGLE * input->leftThumbStickRL;
			//if (steer) {
			//	gVehicleSteering = steer;
			//}
			////DEBUGOUT("STEER %lf, ", gVehicleSteering);
			//double force = ENGINE_MAX_FORCE * input->rightTrigger;
			//if(force) {
			//	gEngineForce = force;
			//}
			////DEBUGOUT("FORCE %lf\n", gEngineForce);
			//double breakingForce = BRAKE_MAX_FORCE * input->leftTrigger;
			//if(breakingForce)
			//{
			//	//gBreakingForce = breakingForce;
			//	gBreakingForce = 0.0;
			//	gEngineForce -= breakingForce;
			//}

			gBreakingForce = 0.0;	// Not using this yet, maybe for E-brake or something.
			gVehicleSteering = -STEER_MAX_ANGLE * input->leftThumbStickRL;
			gEngineForce = ENGINE_MAX_FORCE * input->rightTrigger - BRAKE_MAX_FORCE * input->leftTrigger;

			DEBUGOUT("Bforce: %lf, Eforce: %lf, Steer: %f\n", gBreakingForce, gEngineForce, gVehicleSteering);
		}
			break;
		default:
			break;
		}
	}
	mb.emptyMail();

	// Apply steering to front wheels
	m_vehicle->setSteeringValue(gVehicleSteering, 0);
	m_vehicle->setSteeringValue(gVehicleSteering, 1);

	// Apply braking and engine force to rear wheels
	m_vehicle->applyEngineForce(gEngineForce, 2);
	m_vehicle->applyEngineForce(gEngineForce, 3);
	m_vehicle->setBrake(gBreakingForce, 2);
	m_vehicle->setBrake(gBreakingForce, 3);

	m_world->stepSimulation((btScalar)seconds, 10);

	// Update the placement of the car in the world state
	StateData *state = GetMutState();
	btTransform car1 = m_vehicle->getChassisWorldTransform();
	btVector3 pos = car1.getOrigin();
	state->Karts[0].vPos.x = pos.getX();
	state->Karts[0].vPos.y = pos.getY();
	state->Karts[0].vPos.z = pos.getZ();
	btQuaternion rot = car1.getRotation();
	btVector3 axis = rot.getAxis();
	Vector3 v;
	v.x = axis.getX();
	v.y = axis.getY();
	v.z = axis.getZ();
	Real angle = -rot.getAngle();
	state->Karts[0].qOrient.RotateAxisAngle(v, angle);
}

void Simulation::enableDebugView()
{
}
<commit_msg>Fixed steering, ITS IN RADIANS!!!!111!!!111<commit_after>#include "physics.h"
#include "GLDebugDrawer.h"
#include "../state/state.h"
#include <iostream>

using namespace Physics;

Simulation::Simulation()
{
	m_broadphase = new btDbvtBroadphase();
	m_collisionConfiguration = new btDefaultCollisionConfiguration();
	m_dispatcher = new btCollisionDispatcher(m_collisionConfiguration);
	m_solver = new btSequentialImpulseConstraintSolver;
	m_world = new btDiscreteDynamicsWorld(m_dispatcher, m_broadphase, m_solver, m_collisionConfiguration);

	m_world->setGravity(btVector3(0,-10,0));

	mb.request(Events::EventType::Input);
}

Simulation::~Simulation()
{
	/*
	delete m_world;
	delete m_solver;
	delete m_broadphase;
	delete m_dispatcher;
	delete m_collisionConfiguration;
	*/
}

// Tuning
// Credit to:
// http://bullet.googlecode.com/svn-history/r2704/trunk/Demos/ForkLiftDemo/ForkLiftDemo.cpp
float	gEngineForce = 0.f;

float	defaultBreakingForce = 10.f;
float	gBreakingForce = 100.f;

float	maxEngineForce = 1000.f;//this should be engine/velocity dependent
float	maxBreakingForce = 100.f;

float	gVehicleSteering = 0.f;
float	wheelFriction = 1000;//BT_LARGE_FLOAT;
float	suspensionStiffness = 2.f;
float	suspensionDamping = 1.3f;
float	suspensionCompression = 1.1f;
float	rollInfluence = 1.01f;//1.0f;
btScalar suspensionRestLength(0.2);

btRigidBody *Simulation::addRigidBody(double mass, const btTransform& startTransform, btCollisionShape* shape)
{
	btVector3 localInertia(0, 0, 0);
	if (mass != 0.0) {
		shape->calculateLocalInertia(mass, localInertia);
	}

	btDefaultMotionState* myMotionState = new btDefaultMotionState(startTransform);

	btRigidBody::btRigidBodyConstructionInfo cInfo(mass,myMotionState,shape,localInertia);

	btRigidBody* body = new btRigidBody(cInfo);
	body->setContactProcessingThreshold(1);

	m_world->addRigidBody(body);

	return body;
}

int Simulation::loadWorld()
{
	// Create car
#define CAR_WIDTH (0.15)
#define CAR_LENGTH (0.25)
#define CAR_MASS (400.0)

	btCollisionShape* chassisShape = new btBoxShape(btVector3(CAR_WIDTH, CAR_WIDTH, CAR_LENGTH));
	btCompoundShape* compound = new btCompoundShape();
	m_collisionShapes.push_back(chassisShape);
	m_collisionShapes.push_back(compound);

	btTransform localTrans; // shift gravity to center of car
	localTrans.setIdentity();
	localTrans.setOrigin(btVector3(0,CAR_WIDTH*1, 0));
	compound->addChildShape(localTrans, chassisShape);

	btTransform tr;
	tr.setIdentity();
	tr.setOrigin(btVector3(0,5,0));
	m_carChassis = addRigidBody(CAR_MASS, tr, compound);

	m_vehicleRayCaster = new btDefaultVehicleRaycaster(m_world);
	m_tuning.m_maxSuspensionTravelCm = 204.8f;
	m_tuning.m_suspensionCompression = 4.4f;
	m_tuning.m_suspensionDamping = 2.3f;
	m_tuning.m_frictionSlip = 10000.0f;
	m_tuning.m_suspensionStiffness = 5.0f;
	m_tuning.m_maxSuspensionForce = 1000.0;
	m_vehicle = new btRaycastVehicle(m_tuning,m_carChassis, m_vehicleRayCaster);
	m_carChassis->setActivationState(DISABLE_DEACTIVATION);
	m_world->addVehicle(m_vehicle);

	float connectionHeight = 0.1;
	btVector3 wheelDirectionCS0(0,-1,0);
	btVector3 wheelAxleCS(-1,0,0);

#define CON1 (CAR_WIDTH)
#define CON2 (CAR_LENGTH)
	float	wheelRadius = 0.075f;
	float	wheelWidth = 0.03f;
	bool isFrontWheel=true;
	btVector3 connectionPointCS0(CON1,connectionHeight,CON2);
	m_vehicle->addWheel(connectionPointCS0,wheelDirectionCS0,wheelAxleCS,suspensionRestLength,wheelRadius,m_tuning,isFrontWheel);

	connectionPointCS0 = btVector3(-CON1,connectionHeight,CON2);
	m_vehicle->addWheel(connectionPointCS0,wheelDirectionCS0,wheelAxleCS,suspensionRestLength,wheelRadius,m_tuning,isFrontWheel);

	isFrontWheel = false;

	connectionPointCS0 = btVector3(-CON1,connectionHeight,-CON2);
	m_vehicle->addWheel(connectionPointCS0,wheelDirectionCS0,wheelAxleCS,suspensionRestLength,wheelRadius,m_tuning,isFrontWheel);

	connectionPointCS0 = btVector3(CON1,connectionHeight,-CON2);
	m_vehicle->addWheel(connectionPointCS0,wheelDirectionCS0,wheelAxleCS,suspensionRestLength,wheelRadius,m_tuning,isFrontWheel);

	for (int i=0;i<m_vehicle->getNumWheels();i++)
	{
		btWheelInfo& wheel = m_vehicle->getWheelInfo(i);
		wheel.m_suspensionStiffness = suspensionStiffness;
		wheel.m_wheelsDampingRelaxation = suspensionDamping;
		wheel.m_wheelsDampingCompression = suspensionCompression;
		wheel.m_frictionSlip = wheelFriction;
		wheel.m_rollInfluence = rollInfluence;
	}

	// Add map
	StateData *state = GetMutState();
	btBvhTriangleMeshShape *arenaShape = new btBvhTriangleMeshShape(state->bttmArena, true, true);
	m_arena = addRigidBody(0.0, btTransform(btQuaternion(0,0,0,1),btVector3(0,0,0)), arenaShape);

	m_vehicle->resetSuspension();
	for (int i=0;i<m_vehicle->getNumWheels();i++)
	{
		//synchronize the wheels with the (interpolated) chassis worldtransform
		m_vehicle->updateWheelTransform(i,true);
	}

	return 0;
}

void Simulation::step(double seconds)
{
	// DEBUGOUT("RUNING PHYSICS %lf\n", seconds);

#define STEER_MAX_ANGLE (40)
#define ENGINE_MAX_FORCE (3000)
#define BRAKE_MAX_FORCE (3000)

	for ( Events::Event *event : (mb.checkMail()) )
	{
		// Hack: Sorry
		switch ( event->type )
		{
		case Events::EventType::Input:
		{
			Events::InputEvent *input = (Events::InputEvent *)event;
			//double steer = STEER_MAX_ANGLE * input->leftThumbStickRL;
			//if (steer) {
			//	gVehicleSteering = steer;
			//}
			////DEBUGOUT("STEER %lf, ", gVehicleSteering);
			//double force = ENGINE_MAX_FORCE * input->rightTrigger;
			//if(force) {
			//	gEngineForce = force;
			//}
			////DEBUGOUT("FORCE %lf\n", gEngineForce);
			//double breakingForce = BRAKE_MAX_FORCE * input->leftTrigger;
			//if(breakingForce)
			//{
			//	//gBreakingForce = breakingForce;
			//	gBreakingForce = 0.0;
			//	gEngineForce -= breakingForce;
			//}

			gBreakingForce = 0.0;	// Not using this yet, maybe for E-brake or something.
			gVehicleSteering = DEGTORAD(STEER_MAX_ANGLE) * input->leftThumbStickRL;
			gEngineForce = ENGINE_MAX_FORCE * input->rightTrigger - BRAKE_MAX_FORCE * input->leftTrigger;

			DEBUGOUT("Bforce: %lf, Eforce: %lf, Steer: %f\n", gBreakingForce, gEngineForce, gVehicleSteering);

			if( GetState().key_map['r'] )
			{
				btTransform trans;
				trans.setOrigin( btVector3( 0, 5, 0 ) );
				m_vehicle->getRigidBody()->setWorldTransform( trans );
			}
		}
			break;
		default:
			break;
		}
	}
	mb.emptyMail();

	// Apply steering to front wheels
	m_vehicle->setSteeringValue(gVehicleSteering, 0);
	m_vehicle->setSteeringValue(gVehicleSteering, 1);

	// Apply braking and engine force to rear wheels
	m_vehicle->applyEngineForce(gEngineForce, 2);
	m_vehicle->applyEngineForce(gEngineForce, 3);
	m_vehicle->setBrake(gBreakingForce, 2);
	m_vehicle->setBrake(gBreakingForce, 3);

	m_world->stepSimulation((btScalar)seconds, 10);

	// Update the placement of the car in the world state
	StateData *state = GetMutState();
	btTransform car1 = m_vehicle->getChassisWorldTransform();

	btVector3 pos = car1.getOrigin();
	state->Karts[0].vPos.x = (Real)pos.getX();
	state->Karts[0].vPos.y = (Real)pos.getY();
	state->Karts[0].vPos.z = (Real)pos.getZ();

	btQuaternion rot = car1.getRotation();
	state->Karts[0].qOrient.x = (Real)rot.getX();
	state->Karts[0].qOrient.y = (Real)-rot.getY();
	state->Karts[0].qOrient.z = (Real)rot.getZ();
	state->Karts[0].qOrient.w = (Real)rot.getW();
}

void Simulation::enableDebugView()
{
}
<|endoftext|>


<commit_before>//===- FileCheck.cpp - Check that File's Contents match what is expected --===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// FileCheck does a line-by line check of a file that validates whether it
// contains the expected content.  This is useful for regression tests etc.
//
// This program exits with an error status of 2 on error, exit status of 0 if
// the file matched the expected contents, and exit status of 1 if it did not
// contain the expected contents.
//
//===----------------------------------------------------------------------===//

#include "llvm/Support/CommandLine.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/PrettyStackTrace.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/System/Signals.h"
using namespace llvm;

static cl::opt<std::string>
CheckFilename(cl::Positional, cl::desc("<check-file>"), cl::Required);

static cl::opt<std::string>
InputFilename("input-file", cl::desc("File to check (defaults to stdin)"),
              cl::init("-"), cl::value_desc("filename"));

static cl::opt<std::string>
CheckPrefix("check-prefix", cl::init("CHECK"),
            cl::desc("Prefix to use from check file (defaults to 'CHECK')"));

static cl::opt<bool>
NoCanonicalizeWhiteSpace("strict-whitespace",
              cl::desc("Do not treat all horizontal whitespace as equivalent"));

/// CheckString - This is a check that we found in the input file.
struct CheckString {
  /// Str - The string to match.
  std::string Str;

  /// Loc - The location in the match file that the check string was specified.
  SMLoc Loc;

  /// IsCheckNext - This is true if this is a CHECK-NEXT: directive (as opposed
  /// to a CHECK: directive.
  bool IsCheckNext;

  CheckString(const std::string &S, SMLoc L, bool isCheckNext)
    : Str(S), Loc(L), IsCheckNext(isCheckNext) {}
};

/// FindFixedStringInBuffer - This works like strstr, except for two things:
/// 1) it handles 'nul' characters in memory buffers.  2) it returns the end of
/// the memory buffer on match failure instead of null.
static const char *FindFixedStringInBuffer(StringRef Str, const char *CurPtr,
                                           const MemoryBuffer &MB) {
  assert(!Str.empty() && "Can't find an empty string");
  const char *BufEnd = MB.getBufferEnd();

  while (1) {
    // Scan for the first character in the match string.
    CurPtr = (char*)memchr(CurPtr, Str[0], BufEnd-CurPtr);

    // If we didn't find the first character of the string, then we failed to
    // match.
    if (CurPtr == 0) return BufEnd;

    // If the match string is one character, then we win.
    if (Str.size() == 1) return CurPtr;

    // Otherwise, verify that the rest of the string matches.
    if (Str.size() <= unsigned(BufEnd-CurPtr) &&
        memcmp(CurPtr+1, Str.data()+1, Str.size()-1) == 0)
      return CurPtr;

    // If not, advance past this character and try again.
    ++CurPtr;
  }
}

/// ReadCheckFile - Read the check file, which specifies the sequence of
/// expected strings.  The strings are added to the CheckStrings vector.
static bool ReadCheckFile(SourceMgr &SM,
                          std::vector<CheckString> &CheckStrings) {
  // Open the check file, and tell SourceMgr about it.
  std::string ErrorStr;
  MemoryBuffer *F =
    MemoryBuffer::getFileOrSTDIN(CheckFilename.c_str(), &ErrorStr);
  if (F == 0) {
    errs() << "Could not open check file '" << CheckFilename << "': "
           << ErrorStr << '\n';
    return true;
  }
  SM.AddNewSourceBuffer(F, SMLoc());

  // Find all instances of CheckPrefix followed by : in the file.
  const char *CurPtr = F->getBufferStart(), *BufferEnd = F->getBufferEnd();

  while (1) {
    // See if Prefix occurs in the memory buffer.
    const char *Ptr = FindFixedStringInBuffer(CheckPrefix, CurPtr, *F);

    // If we didn't find a match, we're done.
    if (Ptr == BufferEnd)
      break;

    const char *CheckPrefixStart = Ptr;

    // When we find a check prefix, keep track of whether we find CHECK: or
    // CHECK-NEXT:
    bool IsCheckNext;

    // Verify that the : is present after the prefix.
    if (Ptr[CheckPrefix.size()] == ':') {
      Ptr += CheckPrefix.size()+1;
      IsCheckNext = false;
    } else if (BufferEnd-Ptr > 6 &&
               memcmp(Ptr+CheckPrefix.size(), "-NEXT:", 6) == 0) {
      Ptr += CheckPrefix.size()+7;
      IsCheckNext = true;
    } else {
      CurPtr = Ptr+1;
      continue;
    }

    // Okay, we found the prefix, yay.  Remember the rest of the line, but
    // ignore leading and trailing whitespace.
    while (*Ptr == ' ' || *Ptr == '\t')
      ++Ptr;

    // Scan ahead to the end of line.
    CurPtr = Ptr;
    while (CurPtr != BufferEnd && *CurPtr != '\n' && *CurPtr != '\r')
      ++CurPtr;

    // Ignore trailing whitespace.
    while (CurPtr[-1] == ' ' || CurPtr[-1] == '\t')
      --CurPtr;

    // Check that there is something on the line.
    if (Ptr >= CurPtr) {
      SM.PrintMessage(SMLoc::getFromPointer(CurPtr),
                      "found empty check string with prefix '"+CheckPrefix+":'",
                      "error");
      return true;
    }

    // Verify that CHECK-NEXT lines have at least one CHECK line before them.
    if (IsCheckNext && CheckStrings.empty()) {
      SM.PrintMessage(SMLoc::getFromPointer(CheckPrefixStart),
                      "found '"+CheckPrefix+"-NEXT:' without previous '"+
                      CheckPrefix+ ": line", "error");
      return true;
    }

    // Okay, add the string we captured to the output vector and move on.
    CheckStrings.push_back(CheckString(std::string(Ptr, CurPtr),
                                       SMLoc::getFromPointer(Ptr),
                                       IsCheckNext));
  }

  if (CheckStrings.empty()) {
    errs() << "error: no check strings found with prefix '" << CheckPrefix
           << ":'\n";
    return true;
  }

  return false;
}

// CanonicalizeCheckStrings - Replace all sequences of horizontal whitespace in
// the check strings with a single space.
static void CanonicalizeCheckStrings(std::vector<CheckString> &CheckStrings) {
  for (unsigned i = 0, e = CheckStrings.size(); i != e; ++i) {
    std::string &Str = CheckStrings[i].Str;

    for (unsigned C = 0; C != Str.size(); ++C) {
      // If C is not a horizontal whitespace, skip it.
      if (Str[C] != ' ' && Str[C] != '\t')
        continue;

      // Replace the character with space, then remove any other space
      // characters after it.
      Str[C] = ' ';

      while (C+1 != Str.size() &&
             (Str[C+1] == ' ' || Str[C+1] == '\t'))
        Str.erase(Str.begin()+C+1);
    }
  }
}

/// CanonicalizeInputFile - Remove duplicate horizontal space from the specified
/// memory buffer, free it, and return a new one.
static MemoryBuffer *CanonicalizeInputFile(MemoryBuffer *MB) {
  SmallVector<char, 16> NewFile;
  NewFile.reserve(MB->getBufferSize());

  for (const char *Ptr = MB->getBufferStart(), *End = MB->getBufferEnd();
       Ptr != End; ++Ptr) {
    // If C is not a horizontal whitespace, skip it.
    if (*Ptr != ' ' && *Ptr != '\t') {
      NewFile.push_back(*Ptr);
      continue;
    }

    // Otherwise, add one space and advance over neighboring space.
    NewFile.push_back(' ');
    while (Ptr+1 != End &&
           (Ptr[1] == ' ' || Ptr[1] == '\t'))
      ++Ptr;
  }

  // Free the old buffer and return a new one.
  MemoryBuffer *MB2 =
    MemoryBuffer::getMemBufferCopy(NewFile.data(),
                                   NewFile.data() + NewFile.size(),
                                   MB->getBufferIdentifier());

  delete MB;
  return MB2;
}

static void PrintCheckFailed(const SourceMgr &SM, const CheckString &CheckStr,
                             const char *CurPtr, const char *BufferEnd) {
  // Otherwise, we have an error, emit an error message.
  SM.PrintMessage(CheckStr.Loc, "expected string not found in input",
                  "error");

  // Print the "scanning from here" line.  If the current position is at the
  // end of a line, advance to the start of the next line.
  const char *Scan = CurPtr;
  while (Scan != BufferEnd &&
         (*Scan == ' ' || *Scan == '\t'))
    ++Scan;
  if (*Scan == '\n' || *Scan == '\r')
    CurPtr = Scan+1;

  SM.PrintMessage(SMLoc::getFromPointer(CurPtr), "scanning from here",
                  "note");
}

static unsigned CountNumNewlinesBetween(const char *Start, const char *End) {
  unsigned NumNewLines = 0;
  for (; Start != End; ++Start) {
    // Scan for newline.
    if (Start[0] != '\n' && Start[0] != '\r')
      continue;

    ++NumNewLines;

    // Handle \n\r and \r\n as a single newline.
    if (Start+1 != End &&
        (Start[0] == '\n' || Start[0] == '\r') &&
        (Start[0] != Start[1]))
      ++Start;
  }

  return NumNewLines;
}

int main(int argc, char **argv) {
  sys::PrintStackTraceOnErrorSignal();
  PrettyStackTraceProgram X(argc, argv);
  cl::ParseCommandLineOptions(argc, argv);

  SourceMgr SM;

  // Read the expected strings from the check file.
  std::vector<CheckString> CheckStrings;
  if (ReadCheckFile(SM, CheckStrings))
    return 2;

  // Remove duplicate spaces in the check strings if requested.
  if (!NoCanonicalizeWhiteSpace)
    CanonicalizeCheckStrings(CheckStrings);

  // Open the file to check and add it to SourceMgr.
  std::string ErrorStr;
  MemoryBuffer *F =
    MemoryBuffer::getFileOrSTDIN(InputFilename.c_str(), &ErrorStr);
  if (F == 0) {
    errs() << "Could not open input file '" << InputFilename << "': "
           << ErrorStr << '\n';
    return true;
  }

  // Remove duplicate spaces in the input file if requested.
  if (!NoCanonicalizeWhiteSpace)
    F = CanonicalizeInputFile(F);

  SM.AddNewSourceBuffer(F, SMLoc());

  // Check that we have all of the expected strings, in order, in the input
  // file.
  const char *CurPtr = F->getBufferStart(), *BufferEnd = F->getBufferEnd();

  const char *LastMatch = 0;
  for (unsigned StrNo = 0, e = CheckStrings.size(); StrNo != e; ++StrNo) {
    const CheckString &CheckStr = CheckStrings[StrNo];

    // Find StrNo in the file.
    const char *Ptr = FindFixedStringInBuffer(CheckStr.Str, CurPtr, *F);

    // If we didn't find a match, reject the input.
    if (Ptr == BufferEnd) {
      PrintCheckFailed(SM, CheckStr, CurPtr, BufferEnd);
      return 1;
    }

    // If this check is a "CHECK-NEXT", verify that the previous match was on
    // the previous line (i.e. that there is one newline between them).
    if (CheckStr.IsCheckNext) {
      // Count the number of newlines between the previous match and this one.
      assert(LastMatch && "CHECK-NEXT can't be the first check in a file");

      unsigned NumNewLines = CountNumNewlinesBetween(LastMatch, Ptr);
      if (NumNewLines == 0) {
        SM.PrintMessage(CheckStr.Loc,
                    CheckPrefix+"-NEXT: is on the same line as previous match",
                        "error");
        SM.PrintMessage(SMLoc::getFromPointer(Ptr),
                        "'next' match was here", "note");
        SM.PrintMessage(SMLoc::getFromPointer(LastMatch),
                        "previous match was here", "note");
        return 1;
      }

      if (NumNewLines != 1) {
        SM.PrintMessage(CheckStr.Loc,
                        CheckPrefix+
                        "-NEXT: is not on the line after the previous match",
                        "error");
        SM.PrintMessage(SMLoc::getFromPointer(Ptr),
                        "'next' match was here", "note");
        SM.PrintMessage(SMLoc::getFromPointer(LastMatch),
                        "previous match was here", "note");
        return 1;
      }
    }

    // Otherwise, everything is good.  Remember this as the last match and move
    // on to the next one.
    LastMatch = Ptr;
    CurPtr = Ptr + CheckStr.Str.size();
  }

  return 0;
}
<commit_msg>rewrite FileCheck in terms of StringRef instead of manual pointer pairs.<commit_after>//===- FileCheck.cpp - Check that File's Contents match what is expected --===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// FileCheck does a line-by line check of a file that validates whether it
// contains the expected content.  This is useful for regression tests etc.
//
// This program exits with an error status of 2 on error, exit status of 0 if
// the file matched the expected contents, and exit status of 1 if it did not
// contain the expected contents.
//
//===----------------------------------------------------------------------===//

#include "llvm/Support/CommandLine.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/PrettyStackTrace.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/System/Signals.h"
using namespace llvm;

static cl::opt<std::string>
CheckFilename(cl::Positional, cl::desc("<check-file>"), cl::Required);

static cl::opt<std::string>
InputFilename("input-file", cl::desc("File to check (defaults to stdin)"),
              cl::init("-"), cl::value_desc("filename"));

static cl::opt<std::string>
CheckPrefix("check-prefix", cl::init("CHECK"),
            cl::desc("Prefix to use from check file (defaults to 'CHECK')"));

static cl::opt<bool>
NoCanonicalizeWhiteSpace("strict-whitespace",
              cl::desc("Do not treat all horizontal whitespace as equivalent"));

/// CheckString - This is a check that we found in the input file.
struct CheckString {
  /// Str - The string to match.
  std::string Str;

  /// Loc - The location in the match file that the check string was specified.
  SMLoc Loc;

  /// IsCheckNext - This is true if this is a CHECK-NEXT: directive (as opposed
  /// to a CHECK: directive.
  bool IsCheckNext;

  CheckString(const std::string &S, SMLoc L, bool isCheckNext)
    : Str(S), Loc(L), IsCheckNext(isCheckNext) {}
};

/// ReadCheckFile - Read the check file, which specifies the sequence of
/// expected strings.  The strings are added to the CheckStrings vector.
static bool ReadCheckFile(SourceMgr &SM,
                          std::vector<CheckString> &CheckStrings) {
  // Open the check file, and tell SourceMgr about it.
  std::string ErrorStr;
  MemoryBuffer *F =
    MemoryBuffer::getFileOrSTDIN(CheckFilename.c_str(), &ErrorStr);
  if (F == 0) {
    errs() << "Could not open check file '" << CheckFilename << "': "
           << ErrorStr << '\n';
    return true;
  }
  SM.AddNewSourceBuffer(F, SMLoc());

  // Find all instances of CheckPrefix followed by : in the file.
  StringRef Buffer = F->getBuffer();

  while (1) {
    // See if Prefix occurs in the memory buffer.
    Buffer = Buffer.substr(Buffer.find(CheckPrefix));

    // If we didn't find a match, we're done.
    if (Buffer.empty())
      break;

    const char *CheckPrefixStart = Buffer.data();

    // When we find a check prefix, keep track of whether we find CHECK: or
    // CHECK-NEXT:
    bool IsCheckNext;

    // Verify that the : is present after the prefix.
    if (Buffer[CheckPrefix.size()] == ':') {
      Buffer = Buffer.substr(CheckPrefix.size()+1);
      IsCheckNext = false;
    } else if (Buffer.size() > CheckPrefix.size()+6 &&
               memcmp(Buffer.data()+CheckPrefix.size(), "-NEXT:", 6) == 0) {
      Buffer = Buffer.substr(CheckPrefix.size()+7);
      IsCheckNext = true;
    } else {
      Buffer = Buffer.substr(1);
      continue;
    }

    // Okay, we found the prefix, yay.  Remember the rest of the line, but
    // ignore leading and trailing whitespace.
    while (!Buffer.empty() && (Buffer[0] == ' ' || Buffer[0] == '\t'))
      Buffer = Buffer.substr(1);

    // Scan ahead to the end of line.
    size_t EOL = Buffer.find_first_of("\n\r");
    if (EOL == StringRef::npos) EOL = Buffer.size();

    // Ignore trailing whitespace.
    while (EOL && (Buffer[EOL-1] == ' ' || Buffer[EOL-1] == '\t'))
      --EOL;

    // Check that there is something on the line.
    if (EOL == 0) {
      SM.PrintMessage(SMLoc::getFromPointer(Buffer.data()),
                      "found empty check string with prefix '"+CheckPrefix+":'",
                      "error");
      return true;
    }

    // Verify that CHECK-NEXT lines have at least one CHECK line before them.
    if (IsCheckNext && CheckStrings.empty()) {
      SM.PrintMessage(SMLoc::getFromPointer(CheckPrefixStart),
                      "found '"+CheckPrefix+"-NEXT:' without previous '"+
                      CheckPrefix+ ": line", "error");
      return true;
    }

    // Okay, add the string we captured to the output vector and move on.
    CheckStrings.push_back(CheckString(std::string(Buffer.data(),
                                                   Buffer.data()+EOL),
                                       SMLoc::getFromPointer(Buffer.data()),
                                       IsCheckNext));

    Buffer = Buffer.substr(EOL);
  }

  if (CheckStrings.empty()) {
    errs() << "error: no check strings found with prefix '" << CheckPrefix
           << ":'\n";
    return true;
  }

  return false;
}

// CanonicalizeCheckStrings - Replace all sequences of horizontal whitespace in
// the check strings with a single space.
static void CanonicalizeCheckStrings(std::vector<CheckString> &CheckStrings) {
  for (unsigned i = 0, e = CheckStrings.size(); i != e; ++i) {
    std::string &Str = CheckStrings[i].Str;

    for (unsigned C = 0; C != Str.size(); ++C) {
      // If C is not a horizontal whitespace, skip it.
      if (Str[C] != ' ' && Str[C] != '\t')
        continue;

      // Replace the character with space, then remove any other space
      // characters after it.
      Str[C] = ' ';

      while (C+1 != Str.size() &&
             (Str[C+1] == ' ' || Str[C+1] == '\t'))
        Str.erase(Str.begin()+C+1);
    }
  }
}

/// CanonicalizeInputFile - Remove duplicate horizontal space from the specified
/// memory buffer, free it, and return a new one.
static MemoryBuffer *CanonicalizeInputFile(MemoryBuffer *MB) {
  SmallVector<char, 16> NewFile;
  NewFile.reserve(MB->getBufferSize());

  for (const char *Ptr = MB->getBufferStart(), *End = MB->getBufferEnd();
       Ptr != End; ++Ptr) {
    // If C is not a horizontal whitespace, skip it.
    if (*Ptr != ' ' && *Ptr != '\t') {
      NewFile.push_back(*Ptr);
      continue;
    }

    // Otherwise, add one space and advance over neighboring space.
    NewFile.push_back(' ');
    while (Ptr+1 != End &&
           (Ptr[1] == ' ' || Ptr[1] == '\t'))
      ++Ptr;
  }

  // Free the old buffer and return a new one.
  MemoryBuffer *MB2 =
    MemoryBuffer::getMemBufferCopy(NewFile.data(),
                                   NewFile.data() + NewFile.size(),
                                   MB->getBufferIdentifier());

  delete MB;
  return MB2;
}

static void PrintCheckFailed(const SourceMgr &SM, const CheckString &CheckStr,
                             StringRef Buffer) {
  // Otherwise, we have an error, emit an error message.
  SM.PrintMessage(CheckStr.Loc, "expected string not found in input",
                  "error");

  // Print the "scanning from here" line.  If the current position is at the
  // end of a line, advance to the start of the next line.
  Buffer = Buffer.substr(Buffer.find_first_not_of(" \t\n\r"));

  SM.PrintMessage(SMLoc::getFromPointer(Buffer.data()), "scanning from here",
                  "note");
}

static unsigned CountNumNewlinesBetween(const char *Start, const char *End) {
  unsigned NumNewLines = 0;
  for (; Start != End; ++Start) {
    // Scan for newline.
    if (Start[0] != '\n' && Start[0] != '\r')
      continue;

    ++NumNewLines;

    // Handle \n\r and \r\n as a single newline.
    if (Start+1 != End &&
        (Start[0] == '\n' || Start[0] == '\r') &&
        (Start[0] != Start[1]))
      ++Start;
  }

  return NumNewLines;
}

int main(int argc, char **argv) {
  sys::PrintStackTraceOnErrorSignal();
  PrettyStackTraceProgram X(argc, argv);
  cl::ParseCommandLineOptions(argc, argv);

  SourceMgr SM;

  // Read the expected strings from the check file.
  std::vector<CheckString> CheckStrings;
  if (ReadCheckFile(SM, CheckStrings))
    return 2;

  // Remove duplicate spaces in the check strings if requested.
  if (!NoCanonicalizeWhiteSpace)
    CanonicalizeCheckStrings(CheckStrings);

  // Open the file to check and add it to SourceMgr.
  std::string ErrorStr;
  MemoryBuffer *F =
    MemoryBuffer::getFileOrSTDIN(InputFilename.c_str(), &ErrorStr);
  if (F == 0) {
    errs() << "Could not open input file '" << InputFilename << "': "
           << ErrorStr << '\n';
    return true;
  }

  // Remove duplicate spaces in the input file if requested.
  if (!NoCanonicalizeWhiteSpace)
    F = CanonicalizeInputFile(F);

  SM.AddNewSourceBuffer(F, SMLoc());

  // Check that we have all of the expected strings, in order, in the input
  // file.
  StringRef Buffer = F->getBuffer();

  const char *LastMatch = 0;
  for (unsigned StrNo = 0, e = CheckStrings.size(); StrNo != e; ++StrNo) {
    const CheckString &CheckStr = CheckStrings[StrNo];

    StringRef SearchFrom = Buffer;

    // Find StrNo in the file.
    Buffer = Buffer.substr(Buffer.find(CheckStr.Str));

    // If we didn't find a match, reject the input.
    if (Buffer.empty()) {
      PrintCheckFailed(SM, CheckStr, SearchFrom);
      return 1;
    }

    // If this check is a "CHECK-NEXT", verify that the previous match was on
    // the previous line (i.e. that there is one newline between them).
    if (CheckStr.IsCheckNext) {
      // Count the number of newlines between the previous match and this one.
      assert(LastMatch && "CHECK-NEXT can't be the first check in a file");

      unsigned NumNewLines = CountNumNewlinesBetween(LastMatch, Buffer.data());
      if (NumNewLines == 0) {
        SM.PrintMessage(CheckStr.Loc,
                    CheckPrefix+"-NEXT: is on the same line as previous match",
                        "error");
        SM.PrintMessage(SMLoc::getFromPointer(Buffer.data()),
                        "'next' match was here", "note");
        SM.PrintMessage(SMLoc::getFromPointer(LastMatch),
                        "previous match was here", "note");
        return 1;
      }

      if (NumNewLines != 1) {
        SM.PrintMessage(CheckStr.Loc,
                        CheckPrefix+
                        "-NEXT: is not on the line after the previous match",
                        "error");
        SM.PrintMessage(SMLoc::getFromPointer(Buffer.data()),
                        "'next' match was here", "note");
        SM.PrintMessage(SMLoc::getFromPointer(LastMatch),
                        "previous match was here", "note");
        return 1;
      }
    }

    // Otherwise, everything is good.  Remember this as the last match and move
    // on to the next one.
    LastMatch = Buffer.data();
    Buffer = Buffer.substr(CheckStr.Str.size());
  }

  return 0;
}
<|endoftext|>


<commit_before>/*=========================================================================

  Library   : Image Registration Toolkit (IRTK)
  Module    : $Id$
  Copyright : Imperial College, Department of Computing
              Visual Information Processing (VIP), 2009 onwards
  Date      : $Date$
  Version   : $Revision$
  Changes   : $Author$

=========================================================================*/

#include <irtkImage.h>
#include <irtkImageFunction.h>

#ifdef HAS_VTK

#include <vtkPointData.h>
#include <vtkPolyData.h>
#include <vtkPolyDataReader.h>
#include <vtkPolyDataWriter.h>
#include <vtkDoubleArray.h>

char *input_name = NULL, *output_name = NULL, *image_name = NULL;

int main(int argc, char **argv)
{
  int i, j, n;
  double *profile;
  double x, y, z, ds, point[3], normal[3];

  if (argc != 4) {
    cerr << "Usage: msample [input] [image] [output] -n [number of steps (default 10)] -ds [step size (default 1)]" << endl;
    exit(1);
  }

  // Parse filenames
  input_name = argv[1];
  argv++;
  argc--;
  image_name = argv[1];
  argv++;
  argc--;
  output_name = argv[1];
  argv++;
  argc--;

  n  = 10;
  ds = 1;

  // Allocate memory for intensity profile
  profile = new double[2*n+1];

  // Read model
  vtkPolyDataReader *reader = vtkPolyDataReader::New();
  reader->SetFileName(input_name);
  reader->Modified();
  reader->Update();
  vtkPolyData *model = vtkPolyData::New();
  model = reader->GetOutput();
  model->Update();

  // Extract normals
  if (model->GetPointData()->GetNormals() == NULL) {
    cerr << "Model has no normals" << endl;
    exit(1);
  }

  // Read image
  irtkGreyImage image;
  image.Read(image_name);
  irtkLinearInterpolateImageFunction interpolator;
  interpolator.SetInput(&image);
  interpolator.Initialize();

  // Allocate memory
  vtkDoubleArray *array = vtkDoubleArray::New();
  array->SetNumberOfTuples(model->GetNumberOfPoints());
  array->SetNumberOfComponents(2*n+1);
  array->SetName("IntensityProfile");

  for (i = 0; i < model->GetNumberOfPoints(); i++) {
    model->GetPoints()->GetPoint (i, point);
    model->GetPointData()->GetNormals()->GetTuple(i, normal);
    image.WorldToImage(point[0], point[1], point[2]);
    for (j = 0; j < 2*n+1; j++) {
      x = point[0] + (j - n) * ds * normal[0];
      y = point[1] + (j - n) * ds * normal[1];
      z = point[2] + (j - n) * ds * normal[2];
      profile[j] = interpolator.Evaluate(x, y, z);
    }
    array->InsertTupleValue(i, profile);
  }
  model->GetPointData()->SetScalars(array);

  vtkPolyDataWriter *writer = vtkPolyDataWriter::New();
  writer->SetInput(model);
  writer->SetFileName(output_name);
  writer->Write();
}

#else

#include <irtkImage.h>

int main( int argc, char *argv[] )
{
  cerr << argv[0] << " this program needs to be compiled with vtk enabled." << endl;
  return 0;
}

#endif
<commit_msg>fixed a bug. world to image wrong place...all mess up previously<commit_after>/*=========================================================================

  Library   : Image Registration Toolkit (IRTK)
  Module    : $Id$
  Copyright : Imperial College, Department of Computing
              Visual Information Processing (VIP), 2009 onwards
  Date      : $Date$
  Version   : $Revision$
  Changes   : $Author$

=========================================================================*/

#include <irtkImage.h>
#include <irtkImageFunction.h>
#include <irtkUtil.h>

#ifdef HAS_VTK

#include <vtkPointData.h>
#include <vtkPolyData.h>
#include <vtkPolyDataReader.h>
#include <vtkPolyDataWriter.h>
#include <vtkDoubleArray.h>
#include <vtkPointLocator.h>
#include <vtkIdList.h>

char *input_name = NULL, *output_name = NULL, *image_name = NULL, *count_name = NULL;

void usage(){
	cerr << "Usage: msample [input] [image] [output] " << endl;
	cerr << "-n [number of steps (default 10)]" << endl;
	cerr << "-ds [step size (default 1)]" << endl;
	cerr << "-scalar use one scalar instead of array on the norm"<<endl;
	cerr << "-abs take the abs value of the scalar"<<endl;
	cerr << "-count [another mesh] sample along between two mesh, evaluate percentage > 0"<<endl;
	cerr << "-blur blur the result or not"<<endl;
	exit(1);
}

int main(int argc, char **argv)
{
  int i, j, n, son, abson, blur, ok;
  double *profile, count;
  double x, y, z, ds, point[3], normal[3];

  if (argc < 4) {
    usage();
  }

  // Parse filenames
  input_name = argv[1];
  argv++;
  argc--;
  image_name = argv[1];
  argv++;
  argc--;
  output_name = argv[1];
  argv++;
  argc--;

  n  = 10;
  ds = 1;
  son = 0;
  abson = 0;
  count = 0;
  blur = 0;

  while (argc > 1) {
    ok = false;
	if ((ok == false) && (strcmp(argv[1], "-n") == 0)) {
      argc--;
      argv++;
      n = atoi(argv[1]);
	  argc--;
      argv++;
      ok = true;
    }
	if ((ok == false) && (strcmp(argv[1], "-count") == 0)) {
      argc--;
      argv++;
	  count_name = argv[1];
	  argc--;
      argv++;
      ok = true;
    }
	if ((ok == false) && (strcmp(argv[1], "-ds") == 0)) {
      argc--;
      argv++;
      ds = atof(argv[1]);
	  argc--;
      argv++;
      ok = true;
    }
	if ((ok == false) && (strcmp(argv[1], "-scalar") == 0)) {
      argc--;
      argv++;
      son = 1;
      ok = true;
    }
	if ((ok == false) && (strcmp(argv[1], "-abs") == 0)) {
      argc--;
      argv++;
      abson = 1;
      ok = true;
    }
	if ((ok == false) && (strcmp(argv[1], "-blur") == 0)) {
      argc--;
      argv++;
      blur = 1;
      ok = true;
    }
    if (ok == false) {
      cerr << "Can't parse argument " << argv[1] << endl;
      usage();
    }
  }

  // Allocate memory for intensity profile
  if(son)
  profile = new double[1];
  else
  profile = new double[2*n+1];

  // Read model
  vtkPolyDataReader *reader = vtkPolyDataReader::New();
  reader->SetFileName(input_name);
  reader->Modified();
  reader->Update();
  vtkPolyData *model = vtkPolyData::New();
  model = reader->GetOutput();
  model->Update();

  // Extract normals
  if (model->GetPointData()->GetNormals() == NULL) {
    cerr << "Model has no normals" << endl;
    exit(1);
  }

  // Read image
  irtkGreyImage image;
  image.Read(image_name);
  irtkLinearInterpolateImageFunction interpolator;
  interpolator.SetInput(&image);
  interpolator.Initialize();

  // Allocate memory
  if(!son){
	  if(count_name){
		  vtkPolyDataReader *reader_count = vtkPolyDataReader::New();
		  reader_count->SetFileName(count_name);
		  reader_count->Modified();
		  reader_count->Update();
		  vtkPolyData *model_count = vtkPolyData::New();
		  model_count = reader_count->GetOutput();
		  model_count->Update();
		  vtkDoubleArray *array = vtkDoubleArray::New();
		  array->SetNumberOfTuples(model->GetNumberOfPoints());
		  array->SetNumberOfComponents(1);
		  array->SetName("IntensityCountProfile");
		  vtkDoubleArray *narray = vtkDoubleArray::New();
		  narray->SetNumberOfTuples(model->GetNumberOfPoints());
		  narray->SetNumberOfComponents(1);
		  narray->SetName("IntensityCountProfile");
		  int count2;

		  // Build a locator
		  vtkPointLocator *pointLocator = vtkPointLocator::New();
		  pointLocator->SetDataSet(reader_count->GetOutput());
		  pointLocator->BuildLocator();

		  for (i = 0; i < model->GetNumberOfPoints(); i++) {
			  double distance,point2[3];
			  model->GetPoints()->GetPoint (i, point);
			  vtkIdType ptId;
			  ptId = pointLocator->FindClosestPoint(point);
			  model_count->GetPoints()->GetPoint(ptId,point2);
			  normal[0] = point2[0] - point[0];
			  normal[1] = point2[1] - point[1];
			  normal[2] = point2[2] - point[2];
			  distance = sqrt(pow(normal[0],2)+pow(normal[1],2)+pow(normal[2],2));
			  if(distance > 0){
				  for(j = 0; j < 3; j++){
					  normal[j] = normal[j] / distance;
				  }
				  ds = distance / n;
				  count = 0; count2 = 0;
				  for (j = 0; j < n; j++) {
					  x = point[0] + (j + 1) * ds * normal[0];
					  y = point[1] + (j + 1) * ds * normal[1];
					  z = point[2] + (j + 1) * ds * normal[2];
					  image.WorldToImage(x,y,z);
					  if(interpolator.Evaluate(x, y, z) >= 0){
						  count += interpolator.Evaluate(x, y, z);
						  count2 ++;
					  }
				  }
			  }
			  if(count2 > 0){
				  count = count / count2;
			  }
			  array->InsertTupleValue(i, &count);
		  }
		  // blur
		  if(blur == 1){
			  for (i = 0; i < model->GetNumberOfPoints(); i++) {
				  vtkIdList *list = vtkIdList::New();
				  //Find neighbor
				  GetConnectedVertices(model,i,list);
				  //Get number of neighbor
				  n = list->GetNumberOfIds();
				  double value = 0, *tmp;
				  for (j = 0; j < n; j++){
					  tmp = array->GetTuple(list->GetId(j));
					  value += *tmp;
				  }
				  //stretch measure relative change of the vertices spacing
				  count = value / n;
				  narray->InsertTupleValue(i, &count);
				  list->Delete();
			  }
		  }
		  model->GetPointData()->SetScalars(array);
		  array->Delete();
		  narray->Delete();
		  reader_count->Delete();
	  }else{
		  vtkDoubleArray *array = vtkDoubleArray::New();
		  array->SetNumberOfTuples(model->GetNumberOfPoints());
		  array->SetNumberOfComponents(2*n+1);
		  array->SetName("IntensityProfile");

		  for (i = 0; i < model->GetNumberOfPoints(); i++) {
			  model->GetPoints()->GetPoint (i, point);
			  model->GetPointData()->GetNormals()->GetTuple(i, normal);
			  for (j = 0; j < 2*n+1; j++) {
				  x = point[0] + (j - n) * ds * normal[0];
				  y = point[1] + (j - n) * ds * normal[1];
				  z = point[2] + (j - n) * ds * normal[2];
				  image.WorldToImage(x,y,z);
				  if(abson)
					  profile[j] = abs(interpolator.Evaluate(x, y, z));
				  else
					  profile[j] = interpolator.Evaluate(x, y, z);
			  }
			  array->InsertTupleValue(i, profile);
		  }
		  model->GetPointData()->SetScalars(array);
		  array->Delete();
	  }
  }else{
	  vtkDoubleArray *array = vtkDoubleArray::New();
	  array->SetNumberOfTuples(model->GetNumberOfPoints());
	  array->SetNumberOfComponents(1);
	  array->SetName("IntensityValue");

	  for (i = 0; i < model->GetNumberOfPoints(); i++) {
		  model->GetPoints()->GetPoint (i, point);
		  image.WorldToImage(point[0], point[1], point[2]);
		  if(abson)
			*profile = abs(interpolator.Evaluate(point[0], point[1], point[2]));
		  else
			*profile = interpolator.Evaluate(point[0], point[1], point[2]);
		  array->InsertTupleValue(i, profile);
	  }
	  model->GetPointData()->SetScalars(array);
	  array->Delete();
  }

  vtkPolyDataWriter *writer = vtkPolyDataWriter::New();
  writer->SetInput(model);
  writer->SetFileName(output_name);
  writer->Write();
  reader->Delete();

  delete []profile;
}

#else

#include <irtkImage.h>

int main( int argc, char *argv[] )
{
  cerr << argv[0] << " this program needs to be compiled with vtk enabled." << endl;
  return 0;
}

#endif
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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 "arrow/gpu/cuda_memory.h"

#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <memory>
#include <mutex>

#include <cuda.h>

#include "arrow/buffer.h"
#include "arrow/io/memory.h"
#include "arrow/status.h"
#include "arrow/util/logging.h"

#include "arrow/gpu/cuda_common.h"
#include "arrow/gpu/cuda_context.h"

namespace arrow {
namespace cuda {

// ----------------------------------------------------------------------
// CUDA IPC memory handle

struct CudaIpcMemHandle::CudaIpcMemHandleImpl {
  explicit CudaIpcMemHandleImpl(const uint8_t* handle) {
    memcpy(&memory_size, handle, sizeof(memory_size));
    if (memory_size != 0)
      memcpy(&ipc_handle, handle + sizeof(memory_size), sizeof(CUipcMemHandle));
  }

  explicit CudaIpcMemHandleImpl(int64_t memory_size, const void* cu_handle)
      : memory_size(memory_size) {
    if (memory_size != 0) memcpy(&ipc_handle, cu_handle, sizeof(CUipcMemHandle));
  }

  CUipcMemHandle ipc_handle;  /// initialized only when memory_size != 0
  int64_t memory_size;        /// size of the memory that ipc_handle refers to
};

CudaIpcMemHandle::CudaIpcMemHandle(const void* handle) {
  impl_.reset(new CudaIpcMemHandleImpl(reinterpret_cast<const uint8_t*>(handle)));
}

CudaIpcMemHandle::CudaIpcMemHandle(int64_t memory_size, const void* cu_handle) {
  impl_.reset(new CudaIpcMemHandleImpl(memory_size, cu_handle));
}

CudaIpcMemHandle::~CudaIpcMemHandle() {}

Status CudaIpcMemHandle::FromBuffer(const void* opaque_handle,
                                    std::shared_ptr<CudaIpcMemHandle>* handle) {
  *handle = std::shared_ptr<CudaIpcMemHandle>(new CudaIpcMemHandle(opaque_handle));
  return Status::OK();
}

Status CudaIpcMemHandle::Serialize(MemoryPool* pool, std::shared_ptr<Buffer>* out) const {
  std::shared_ptr<Buffer> buffer;
  int64_t size = impl_->memory_size;
  size_t kHandleSize =
      (size > 0 ? sizeof(int64_t) + sizeof(CUipcMemHandle) : sizeof(int64_t));
  RETURN_NOT_OK(AllocateBuffer(pool, static_cast<int64_t>(kHandleSize), &buffer));
  memcpy(buffer->mutable_data(), &impl_->memory_size, sizeof(impl_->memory_size));
  if (size > 0)
    memcpy(buffer->mutable_data() + sizeof(impl_->memory_size), &impl_->ipc_handle,
           sizeof(impl_->ipc_handle));
  *out = buffer;
  return Status::OK();
}

const void* CudaIpcMemHandle::handle() const { return &impl_->ipc_handle; }

int64_t CudaIpcMemHandle::memory_size() const { return impl_->memory_size; }

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

CudaBuffer::CudaBuffer(uint8_t* data, int64_t size,
                       const std::shared_ptr<CudaContext>& context, bool own_data,
                       bool is_ipc)
    : Buffer(data, size), context_(context), own_data_(own_data), is_ipc_(is_ipc) {
  is_mutable_ = true;
  mutable_data_ = data;
}

CudaBuffer::~CudaBuffer() { DCHECK(Close().ok()); }

Status CudaBuffer::Close() {
  if (own_data_) {
    if (is_ipc_) {
      return context_->CloseIpcBuffer(this);
    } else {
      return context_->Free(mutable_data_, size_);
    }
  }
  return Status::OK();
}

CudaBuffer::CudaBuffer(const std::shared_ptr<CudaBuffer>& parent, const int64_t offset,
                       const int64_t size)
    : Buffer(parent, offset, size),
      context_(parent->context()),
      own_data_(false),
      is_ipc_(false) {
  if (parent->is_mutable()) {
    is_mutable_ = true;
    mutable_data_ = const_cast<uint8_t*>(data_);
  }
}

Status CudaBuffer::FromBuffer(std::shared_ptr<Buffer> buffer,
                              std::shared_ptr<CudaBuffer>* out) {
  int64_t offset = 0, size = buffer->size();
  bool is_mutable = buffer->is_mutable();
  // The original CudaBuffer may have been wrapped in another Buffer
  // (for example through slicing).
  while (!(*out = std::dynamic_pointer_cast<CudaBuffer>(buffer))) {
    const std::shared_ptr<Buffer> parent = buffer->parent();
    if (!parent) {
      return Status::TypeError("buffer is not backed by a CudaBuffer");
    }
    offset += buffer->data() - parent->data();
    buffer = parent;
  }
  // Re-slice to represent the same memory area
  if (offset != 0 || (*out)->size() != size || !is_mutable) {
    *out = std::make_shared<CudaBuffer>(*out, offset, size);
    (*out)->is_mutable_ = is_mutable;
  }
  return Status::OK();
}

Status CudaBuffer::CopyToHost(const int64_t position, const int64_t nbytes,
                              void* out) const {
  return context_->CopyDeviceToHost(out, data_ + position, nbytes);
}

Status CudaBuffer::CopyFromHost(const int64_t position, const void* data,
                                int64_t nbytes) {
  DCHECK_LE(nbytes, size_ - position) << "Copy would overflow buffer";
  return context_->CopyHostToDevice(mutable_data_ + position, data, nbytes);
}

Status CudaBuffer::CopyFromDevice(const int64_t position, const void* data,
                                  int64_t nbytes) {
  DCHECK_LE(nbytes, size_ - position) << "Copy would overflow buffer";
  return context_->CopyDeviceToDevice(mutable_data_ + position, data, nbytes);
}

Status CudaBuffer::ExportForIpc(std::shared_ptr<CudaIpcMemHandle>* handle) {
  if (is_ipc_) {
    return Status::Invalid("Buffer has already been exported for IPC");
  }
  RETURN_NOT_OK(context_->ExportIpcBuffer(mutable_data_, size_, handle));
  own_data_ = false;
  return Status::OK();
}

CudaHostBuffer::~CudaHostBuffer() {
  CudaDeviceManager* manager = nullptr;
  DCHECK(CudaDeviceManager::GetInstance(&manager).ok());
  DCHECK(manager->FreeHost(mutable_data_, size_).ok());
}

// ----------------------------------------------------------------------
// CudaBufferReader

CudaBufferReader::CudaBufferReader(const std::shared_ptr<Buffer>& buffer)
    : io::BufferReader(buffer) {
  if (!CudaBuffer::FromBuffer(buffer, &cuda_buffer_).ok()) {
    throw std::bad_cast();
  }
  context_ = cuda_buffer_->context();
}

CudaBufferReader::~CudaBufferReader() {}

Status CudaBufferReader::Read(int64_t nbytes, int64_t* bytes_read, void* buffer) {
  nbytes = std::min(nbytes, size_ - position_);
  *bytes_read = nbytes;
  RETURN_NOT_OK(context_->CopyDeviceToHost(buffer, data_ + position_, nbytes));
  position_ += nbytes;
  return Status::OK();
}

Status CudaBufferReader::Read(int64_t nbytes, std::shared_ptr<Buffer>* out) {
  int64_t size = std::min(nbytes, size_ - position_);
  *out = std::make_shared<CudaBuffer>(cuda_buffer_, position_, size);
  position_ += size;
  return Status::OK();
}

// ----------------------------------------------------------------------
// CudaBufferWriter

class CudaBufferWriter::CudaBufferWriterImpl {
 public:
  explicit CudaBufferWriterImpl(const std::shared_ptr<CudaBuffer>& buffer)
      : context_(buffer->context()),
        buffer_(buffer),
        buffer_size_(0),
        buffer_position_(0) {
    buffer_ = buffer;
    DCHECK(buffer->is_mutable()) << "Must pass mutable buffer";
    mutable_data_ = buffer->mutable_data();
    size_ = buffer->size();
    position_ = 0;
  }

  Status Seek(int64_t position) {
    if (position < 0 || position >= size_) {
      return Status::IOError("position out of bounds");
    }
    position_ = position;
    return Status::OK();
  }

  Status Close() {
    if (!closed_) {
      closed_ = true;
      RETURN_NOT_OK(Flush());
    }
    return Status::OK();
  }

  Status Flush() {
    if (buffer_size_ > 0 && buffer_position_ > 0) {
      // Only need to flush when the write has been buffered
      RETURN_NOT_OK(
          context_->CopyHostToDevice(mutable_data_ + position_ - buffer_position_,
                                     host_buffer_data_, buffer_position_));
      buffer_position_ = 0;
    }
    return Status::OK();
  }

  bool closed() const { return closed_; }

  Status Tell(int64_t* position) const {
    *position = position_;
    return Status::OK();
  }

  Status Write(const void* data, int64_t nbytes) {
    if (nbytes == 0) {
      return Status::OK();
    }

    if (buffer_size_ > 0) {
      if (nbytes + buffer_position_ >= buffer_size_) {
        // Reach end of buffer, write everything
        RETURN_NOT_OK(Flush());
        RETURN_NOT_OK(
            context_->CopyHostToDevice(mutable_data_ + position_, data, nbytes));
      } else {
        // Write bytes to buffer
        std::memcpy(host_buffer_data_ + buffer_position_, data, nbytes);
        buffer_position_ += nbytes;
      }
    } else {
      // Unbuffered write
      RETURN_NOT_OK(context_->CopyHostToDevice(mutable_data_ + position_, data, nbytes));
    }
    position_ += nbytes;
    return Status::OK();
  }

  Status WriteAt(int64_t position, const void* data, int64_t nbytes) {
    std::lock_guard<std::mutex> guard(lock_);
    RETURN_NOT_OK(Seek(position));
    return Write(data, nbytes);
  }

  Status SetBufferSize(const int64_t buffer_size) {
    if (buffer_position_ > 0) {
      // Flush any buffered data
      RETURN_NOT_OK(Flush());
    }
    RETURN_NOT_OK(AllocateCudaHostBuffer(context_.get()->device_number(), buffer_size,
                                         &host_buffer_));
    host_buffer_data_ = host_buffer_->mutable_data();
    buffer_size_ = buffer_size;
    return Status::OK();
  }

  int64_t buffer_size() const { return buffer_size_; }

  int64_t buffer_position() const { return buffer_position_; }

 private:
  std::shared_ptr<CudaContext> context_;
  std::shared_ptr<CudaBuffer> buffer_;
  std::mutex lock_;
  uint8_t* mutable_data_;
  int64_t size_;
  int64_t position_;
  bool closed_;

  // Pinned host buffer for buffering writes on CPU before calling cudaMalloc
  int64_t buffer_size_;
  int64_t buffer_position_;
  std::shared_ptr<CudaHostBuffer> host_buffer_;
  uint8_t* host_buffer_data_;
};

CudaBufferWriter::CudaBufferWriter(const std::shared_ptr<CudaBuffer>& buffer) {
  impl_.reset(new CudaBufferWriterImpl(buffer));
}

CudaBufferWriter::~CudaBufferWriter() {}

Status CudaBufferWriter::Close() { return impl_->Close(); }

bool CudaBufferWriter::closed() const { return impl_->closed(); }

Status CudaBufferWriter::Flush() { return impl_->Flush(); }

Status CudaBufferWriter::Seek(int64_t position) {
  if (impl_->buffer_position() > 0) {
    RETURN_NOT_OK(Flush());
  }
  return impl_->Seek(position);
}

Status CudaBufferWriter::Tell(int64_t* position) const { return impl_->Tell(position); }

Status CudaBufferWriter::Write(const void* data, int64_t nbytes) {
  return impl_->Write(data, nbytes);
}

Status CudaBufferWriter::WriteAt(int64_t position, const void* data, int64_t nbytes) {
  return impl_->WriteAt(position, data, nbytes);
}

Status CudaBufferWriter::SetBufferSize(const int64_t buffer_size) {
  return impl_->SetBufferSize(buffer_size);
}

int64_t CudaBufferWriter::buffer_size() const { return impl_->buffer_size(); }

int64_t CudaBufferWriter::num_bytes_buffered() const { return impl_->buffer_position(); }

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

Status AllocateCudaHostBuffer(int device_number, const int64_t size,
                              std::shared_ptr<CudaHostBuffer>* out) {
  CudaDeviceManager* manager = nullptr;
  RETURN_NOT_OK(CudaDeviceManager::GetInstance(&manager));
  return manager->AllocateHost(device_number, size, out);
}

}  // namespace cuda
}  // namespace arrow
<commit_msg>ARROW-3879: [C++] Fix uninitialized member in CudaBufferWriter<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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 "arrow/gpu/cuda_memory.h"

#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <memory>
#include <mutex>

#include <cuda.h>

#include "arrow/buffer.h"
#include "arrow/io/memory.h"
#include "arrow/status.h"
#include "arrow/util/logging.h"

#include "arrow/gpu/cuda_common.h"
#include "arrow/gpu/cuda_context.h"

namespace arrow {
namespace cuda {

// ----------------------------------------------------------------------
// CUDA IPC memory handle

struct CudaIpcMemHandle::CudaIpcMemHandleImpl {
  explicit CudaIpcMemHandleImpl(const uint8_t* handle) {
    memcpy(&memory_size, handle, sizeof(memory_size));
    if (memory_size != 0)
      memcpy(&ipc_handle, handle + sizeof(memory_size), sizeof(CUipcMemHandle));
  }

  explicit CudaIpcMemHandleImpl(int64_t memory_size, const void* cu_handle)
      : memory_size(memory_size) {
    if (memory_size != 0) memcpy(&ipc_handle, cu_handle, sizeof(CUipcMemHandle));
  }

  CUipcMemHandle ipc_handle;  /// initialized only when memory_size != 0
  int64_t memory_size;        /// size of the memory that ipc_handle refers to
};

CudaIpcMemHandle::CudaIpcMemHandle(const void* handle) {
  impl_.reset(new CudaIpcMemHandleImpl(reinterpret_cast<const uint8_t*>(handle)));
}

CudaIpcMemHandle::CudaIpcMemHandle(int64_t memory_size, const void* cu_handle) {
  impl_.reset(new CudaIpcMemHandleImpl(memory_size, cu_handle));
}

CudaIpcMemHandle::~CudaIpcMemHandle() {}

Status CudaIpcMemHandle::FromBuffer(const void* opaque_handle,
                                    std::shared_ptr<CudaIpcMemHandle>* handle) {
  *handle = std::shared_ptr<CudaIpcMemHandle>(new CudaIpcMemHandle(opaque_handle));
  return Status::OK();
}

Status CudaIpcMemHandle::Serialize(MemoryPool* pool, std::shared_ptr<Buffer>* out) const {
  std::shared_ptr<Buffer> buffer;
  int64_t size = impl_->memory_size;
  size_t kHandleSize =
      (size > 0 ? sizeof(int64_t) + sizeof(CUipcMemHandle) : sizeof(int64_t));
  RETURN_NOT_OK(AllocateBuffer(pool, static_cast<int64_t>(kHandleSize), &buffer));
  memcpy(buffer->mutable_data(), &impl_->memory_size, sizeof(impl_->memory_size));
  if (size > 0)
    memcpy(buffer->mutable_data() + sizeof(impl_->memory_size), &impl_->ipc_handle,
           sizeof(impl_->ipc_handle));
  *out = buffer;
  return Status::OK();
}

const void* CudaIpcMemHandle::handle() const { return &impl_->ipc_handle; }

int64_t CudaIpcMemHandle::memory_size() const { return impl_->memory_size; }

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

CudaBuffer::CudaBuffer(uint8_t* data, int64_t size,
                       const std::shared_ptr<CudaContext>& context, bool own_data,
                       bool is_ipc)
    : Buffer(data, size), context_(context), own_data_(own_data), is_ipc_(is_ipc) {
  is_mutable_ = true;
  mutable_data_ = data;
}

CudaBuffer::~CudaBuffer() { DCHECK(Close().ok()); }

Status CudaBuffer::Close() {
  if (own_data_) {
    if (is_ipc_) {
      return context_->CloseIpcBuffer(this);
    } else {
      return context_->Free(mutable_data_, size_);
    }
  }
  return Status::OK();
}

CudaBuffer::CudaBuffer(const std::shared_ptr<CudaBuffer>& parent, const int64_t offset,
                       const int64_t size)
    : Buffer(parent, offset, size),
      context_(parent->context()),
      own_data_(false),
      is_ipc_(false) {
  if (parent->is_mutable()) {
    is_mutable_ = true;
    mutable_data_ = const_cast<uint8_t*>(data_);
  }
}

Status CudaBuffer::FromBuffer(std::shared_ptr<Buffer> buffer,
                              std::shared_ptr<CudaBuffer>* out) {
  int64_t offset = 0, size = buffer->size();
  bool is_mutable = buffer->is_mutable();
  // The original CudaBuffer may have been wrapped in another Buffer
  // (for example through slicing).
  while (!(*out = std::dynamic_pointer_cast<CudaBuffer>(buffer))) {
    const std::shared_ptr<Buffer> parent = buffer->parent();
    if (!parent) {
      return Status::TypeError("buffer is not backed by a CudaBuffer");
    }
    offset += buffer->data() - parent->data();
    buffer = parent;
  }
  // Re-slice to represent the same memory area
  if (offset != 0 || (*out)->size() != size || !is_mutable) {
    *out = std::make_shared<CudaBuffer>(*out, offset, size);
    (*out)->is_mutable_ = is_mutable;
  }
  return Status::OK();
}

Status CudaBuffer::CopyToHost(const int64_t position, const int64_t nbytes,
                              void* out) const {
  return context_->CopyDeviceToHost(out, data_ + position, nbytes);
}

Status CudaBuffer::CopyFromHost(const int64_t position, const void* data,
                                int64_t nbytes) {
  DCHECK_LE(nbytes, size_ - position) << "Copy would overflow buffer";
  return context_->CopyHostToDevice(mutable_data_ + position, data, nbytes);
}

Status CudaBuffer::CopyFromDevice(const int64_t position, const void* data,
                                  int64_t nbytes) {
  DCHECK_LE(nbytes, size_ - position) << "Copy would overflow buffer";
  return context_->CopyDeviceToDevice(mutable_data_ + position, data, nbytes);
}

Status CudaBuffer::ExportForIpc(std::shared_ptr<CudaIpcMemHandle>* handle) {
  if (is_ipc_) {
    return Status::Invalid("Buffer has already been exported for IPC");
  }
  RETURN_NOT_OK(context_->ExportIpcBuffer(mutable_data_, size_, handle));
  own_data_ = false;
  return Status::OK();
}

CudaHostBuffer::~CudaHostBuffer() {
  CudaDeviceManager* manager = nullptr;
  DCHECK(CudaDeviceManager::GetInstance(&manager).ok());
  DCHECK(manager->FreeHost(mutable_data_, size_).ok());
}

// ----------------------------------------------------------------------
// CudaBufferReader

CudaBufferReader::CudaBufferReader(const std::shared_ptr<Buffer>& buffer)
    : io::BufferReader(buffer) {
  if (!CudaBuffer::FromBuffer(buffer, &cuda_buffer_).ok()) {
    throw std::bad_cast();
  }
  context_ = cuda_buffer_->context();
}

CudaBufferReader::~CudaBufferReader() {}

Status CudaBufferReader::Read(int64_t nbytes, int64_t* bytes_read, void* buffer) {
  nbytes = std::min(nbytes, size_ - position_);
  *bytes_read = nbytes;
  RETURN_NOT_OK(context_->CopyDeviceToHost(buffer, data_ + position_, nbytes));
  position_ += nbytes;
  return Status::OK();
}

Status CudaBufferReader::Read(int64_t nbytes, std::shared_ptr<Buffer>* out) {
  int64_t size = std::min(nbytes, size_ - position_);
  *out = std::make_shared<CudaBuffer>(cuda_buffer_, position_, size);
  position_ += size;
  return Status::OK();
}

// ----------------------------------------------------------------------
// CudaBufferWriter

class CudaBufferWriter::CudaBufferWriterImpl {
 public:
  explicit CudaBufferWriterImpl(const std::shared_ptr<CudaBuffer>& buffer)
      : context_(buffer->context()),
        buffer_(buffer),
        buffer_size_(0),
        buffer_position_(0) {
    buffer_ = buffer;
    DCHECK(buffer->is_mutable()) << "Must pass mutable buffer";
    mutable_data_ = buffer->mutable_data();
    size_ = buffer->size();
    position_ = 0;
    closed_ = false;
  }

#define CHECK_CLOSED()                                              \
  if (closed_) {                                                    \
    return Status::Invalid("Operation on closed CudaBufferWriter"); \
  }

  Status Seek(int64_t position) {
    CHECK_CLOSED();
    if (position < 0 || position >= size_) {
      return Status::IOError("position out of bounds");
    }
    position_ = position;
    return Status::OK();
  }

  Status Close() {
    if (!closed_) {
      closed_ = true;
      RETURN_NOT_OK(FlushInternal());
    }
    return Status::OK();
  }

  Status Flush() {
    CHECK_CLOSED();
    return FlushInternal();
  }

  Status FlushInternal() {
    if (buffer_size_ > 0 && buffer_position_ > 0) {
      // Only need to flush when the write has been buffered
      RETURN_NOT_OK(
          context_->CopyHostToDevice(mutable_data_ + position_ - buffer_position_,
                                     host_buffer_data_, buffer_position_));
      buffer_position_ = 0;
    }
    return Status::OK();
  }

  bool closed() const { return closed_; }

  Status Tell(int64_t* position) const {
    CHECK_CLOSED();
    *position = position_;
    return Status::OK();
  }

  Status Write(const void* data, int64_t nbytes) {
    CHECK_CLOSED();
    if (nbytes == 0) {
      return Status::OK();
    }

    if (buffer_size_ > 0) {
      if (nbytes + buffer_position_ >= buffer_size_) {
        // Reach end of buffer, write everything
        RETURN_NOT_OK(Flush());
        RETURN_NOT_OK(
            context_->CopyHostToDevice(mutable_data_ + position_, data, nbytes));
      } else {
        // Write bytes to buffer
        std::memcpy(host_buffer_data_ + buffer_position_, data, nbytes);
        buffer_position_ += nbytes;
      }
    } else {
      // Unbuffered write
      RETURN_NOT_OK(context_->CopyHostToDevice(mutable_data_ + position_, data, nbytes));
    }
    position_ += nbytes;
    return Status::OK();
  }

  Status WriteAt(int64_t position, const void* data, int64_t nbytes) {
    std::lock_guard<std::mutex> guard(lock_);
    CHECK_CLOSED();
    RETURN_NOT_OK(Seek(position));
    return Write(data, nbytes);
  }

  Status SetBufferSize(const int64_t buffer_size) {
    CHECK_CLOSED();
    if (buffer_position_ > 0) {
      // Flush any buffered data
      RETURN_NOT_OK(Flush());
    }
    RETURN_NOT_OK(AllocateCudaHostBuffer(context_.get()->device_number(), buffer_size,
                                         &host_buffer_));
    host_buffer_data_ = host_buffer_->mutable_data();
    buffer_size_ = buffer_size;
    return Status::OK();
  }

  int64_t buffer_size() const { return buffer_size_; }

  int64_t buffer_position() const { return buffer_position_; }

#undef CHECK_CLOSED

 private:
  std::shared_ptr<CudaContext> context_;
  std::shared_ptr<CudaBuffer> buffer_;
  std::mutex lock_;
  uint8_t* mutable_data_;
  int64_t size_;
  int64_t position_;
  bool closed_;

  // Pinned host buffer for buffering writes on CPU before calling cudaMalloc
  int64_t buffer_size_;
  int64_t buffer_position_;
  std::shared_ptr<CudaHostBuffer> host_buffer_;
  uint8_t* host_buffer_data_;
};

CudaBufferWriter::CudaBufferWriter(const std::shared_ptr<CudaBuffer>& buffer) {
  impl_.reset(new CudaBufferWriterImpl(buffer));
}

CudaBufferWriter::~CudaBufferWriter() {}

Status CudaBufferWriter::Close() { return impl_->Close(); }

bool CudaBufferWriter::closed() const { return impl_->closed(); }

Status CudaBufferWriter::Flush() { return impl_->Flush(); }

Status CudaBufferWriter::Seek(int64_t position) {
  if (impl_->buffer_position() > 0) {
    RETURN_NOT_OK(Flush());
  }
  return impl_->Seek(position);
}

Status CudaBufferWriter::Tell(int64_t* position) const { return impl_->Tell(position); }

Status CudaBufferWriter::Write(const void* data, int64_t nbytes) {
  return impl_->Write(data, nbytes);
}

Status CudaBufferWriter::WriteAt(int64_t position, const void* data, int64_t nbytes) {
  return impl_->WriteAt(position, data, nbytes);
}

Status CudaBufferWriter::SetBufferSize(const int64_t buffer_size) {
  return impl_->SetBufferSize(buffer_size);
}

int64_t CudaBufferWriter::buffer_size() const { return impl_->buffer_size(); }

int64_t CudaBufferWriter::num_bytes_buffered() const { return impl_->buffer_position(); }

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

Status AllocateCudaHostBuffer(int device_number, const int64_t size,
                              std::shared_ptr<CudaHostBuffer>* out) {
  CudaDeviceManager* manager = nullptr;
  RETURN_NOT_OK(CudaDeviceManager::GetInstance(&manager));
  return manager->AllocateHost(device_number, size, out);
}

}  // namespace cuda
}  // namespace arrow
<|endoftext|>


<commit_before>/* Copyright 2016-present Facebook, Inc.
 * Licensed under the Apache License, Version 2.0 */

#include "WatchmanConnection.h"

#include <folly/ExceptionWrapper.h>
#include <folly/SocketAddress.h>
#include <folly/Subprocess.h>
#include <folly/experimental/bser/Bser.h>
#include <folly/futures/InlineExecutor.h>

namespace watchman {

using namespace folly::bser;
using namespace folly;

// Ordered with the most likely kind first
static const std::vector<dynamic> kUnilateralLabels{"subscription", "log"};

static const dynamic kError("error");
static const dynamic kCapabilities("capabilities");

// We'll just dispatch bser decodes and callbacks inline unless they
// give us an alternative environment
static InlineExecutor inlineExecutor;

WatchmanConnection::WatchmanConnection(
    EventBase* eventBase,
    Optional<std::string>&& sockPath,
    Optional<WatchmanConnection::Callback>&& callback,
    Executor* cpuExecutor)
    : eventBase_(eventBase),
      sockPath_(std::move(sockPath)),
      callback_(std::move(callback)),
      cpuExecutor_(cpuExecutor ? cpuExecutor : &inlineExecutor),
      versionCmd_(nullptr),
      bufQ_(IOBufQueue::cacheChainLength()) {
  CHECK_NOTNULL(eventBase);
}

WatchmanConnection::~WatchmanConnection() {
  close();
}

folly::Future<std::string> WatchmanConnection::getSockPath() {
  // Take explicit configuration first
  if (sockPath_.hasValue()) {
    return makeFuture(sockPath_.value());
  }

  // Else use the environmental variable used by watchman to report
  // the active socket path
  auto var = getenv("WATCHMAN_SOCK");
  if (var && *var) {
    return makeFuture(std::string(var));
  }

  return via(cpuExecutor_, [] {
    // Else discover it from the CLI
    folly::Subprocess proc(
        {"watchman", "--output-encoding=bser", "get-sockname"},
        folly::Subprocess::Options().pipeStdout().pipeStderr().usePath());
    SCOPE_FAIL {
      // Always clean up to avoid Subprocess asserting on destruction
      proc.kill();
      proc.wait();
    };
    auto out_pair = proc.communicate();
    auto result = parseBser(out_pair.first);
    proc.waitChecked();
    return result["sockname"].asString();
  });
}

Future<dynamic> WatchmanConnection::connect(folly::dynamic versionArgs) {
  if (!versionArgs.isObject()) {
    throw WatchmanError("versionArgs must be object");
  }
  versionCmd_ = folly::dynamic::array("version", versionArgs);

  auto res = getSockPath().then(
    [shared_this=shared_from_this()] (std::string&& path) {
      shared_this->eventBase_->runInEventBaseThread([=] {
        folly::SocketAddress addr;
        addr.setFromPath(path);

        shared_this->sock_ =
          folly::AsyncSocket::newSocket(shared_this->eventBase_);
        shared_this->sock_->connect(shared_this.get(), addr);
      }
    );

    return shared_this->connectPromise_.getFuture();
  });
  return res;
}

void WatchmanConnection::close() {
  if (closing_) {
    return;
  }
  closing_ = true;
  if (sock_) {
    eventBase_->runImmediatelyOrRunInEventBaseThreadAndWait([this] {
      sock_->close();
      sock_.reset();
    });
  }
  failQueuedCommands(
      make_exception_wrapper<WatchmanError>(
          "WatchmanConnection::close() was called"));
}

// The convention for Watchman responses is that they represent
// an error if they contain the "error" key.  We want to report
// those as exceptions, but it is easier to do that via a Try
Try<dynamic> WatchmanConnection::watchmanResponseToTry(dynamic&& value) {
  auto error = value.get_ptr(kError);
  if (error) {
    return Try<dynamic>(make_exception_wrapper<WatchmanResponseError>(value));
  }
  return Try<dynamic>(std::move(value));
}

void WatchmanConnection::connectSuccess() noexcept {
  try {
    sock_->setReadCB(this);
    sock_->setCloseOnExec();

    run(versionCmd_).then(
      [shared_this=shared_from_this()] (dynamic&& result) {
        // If there is no "capabilities" key then the version of
        // watchman is too old; treat this as an error
        if (!result.get_ptr(kCapabilities)) {
          result["error"] =
              "This watchman server has no support for capabilities, "
              "please upgrade to the current stable version of watchman";
          shared_this->connectPromise_.setTry(
            shared_this->watchmanResponseToTry(std::move(result)));
          return;
        }
        shared_this->connectPromise_.setValue(std::move(result));
      }
    ).onError(
      [shared_this=shared_from_this()]
      (const folly::exception_wrapper& e) {
        shared_this->connectPromise_.setException(e);
      }
    );
  } catch(const std::exception& e) {
    connectPromise_.setException(
      folly::exception_wrapper(std::current_exception(), e));
  } catch(...) {
    connectPromise_.setException(
      folly::exception_wrapper(std::current_exception()));
  }
}

void WatchmanConnection::connectErr(
    const folly::AsyncSocketException& ex) noexcept {
  connectPromise_.setException(ex);
}

WatchmanConnection::QueuedCommand::QueuedCommand(const dynamic& command)
    : cmd(command) {}

Future<dynamic> WatchmanConnection::run(const dynamic& command) noexcept {
  auto cmd = std::make_shared<QueuedCommand>(command);
  if (broken_) {
    cmd->promise.setException(WatchmanError("The connection was broken"));
    return cmd->promise.getFuture();
  }
  if (!sock_) {
    cmd->promise.setException(WatchmanError(
        "No socket (did you call connect() and check result for exceptions?)"));
    return cmd->promise.getFuture();
  }

  bool shouldWrite;
  {
    std::lock_guard<std::mutex> g(mutex_);
    // We only need to call sendCommand if we don't have a command in
    // progress; the completion handler will trigger it once we receive
    // the response
    shouldWrite = commandQ_.empty();
    commandQ_.push_back(cmd);
  }

  if (shouldWrite) {
    eventBase_->runInEventBaseThread(
      [shared_this=shared_from_this()] {
        shared_this->sendCommand();
      }
    );
  }

  return cmd->promise.getFuture();
}

// Generate a failure for all queued commands
void WatchmanConnection::failQueuedCommands(
    const folly::exception_wrapper& ex) {
  std::lock_guard<std::mutex> g(mutex_);
  auto q = commandQ_;
  commandQ_.clear();

  broken_ = true;
  for (auto& cmd : q) {
    if (!cmd->promise.isFulfilled()) {
      cmd->promise.setException(ex);
    }
  }

  // If the user has explicitly closed the connection no need for callback
  if (callback_ && !closing_) {
    cpuExecutor_->add([shared_this=shared_from_this(), ex] {
      (*(shared_this->callback_))(folly::Try<folly::dynamic>(ex));
    });
  }
}

// Sends the next eligible command to the Watchman service
void WatchmanConnection::sendCommand(bool pop) {
  std::shared_ptr<QueuedCommand> cmd;

  {
    std::lock_guard<std::mutex> g(mutex_);

    if (pop) {
      // We finished processing this one, discard it and focus
      // on the next item, if any.
      commandQ_.pop_front();
    }
    if (commandQ_.empty()) {
      return;
    }
    cmd = commandQ_.front();
  }

  sock_->writeChain(this, toBserIOBuf(cmd->cmd, serialization_opts()));
}

void WatchmanConnection::popAndSendCommand() {
  sendCommand(/* pop = */ true);
}

// Called when AsyncSocket::writeChain completes
void WatchmanConnection::writeSuccess() noexcept {
  // Don't care particularly
}

// Called when AsyncSocket::writeChain fails
void WatchmanConnection::writeErr(
    size_t,
    const folly::AsyncSocketException& ex) noexcept {
  failQueuedCommands(ex);
}

// Called when AsyncSocket wants to give us data
void WatchmanConnection::getReadBuffer(void** bufReturn, size_t* lenReturn) {
  std::lock_guard<std::mutex> g(mutex_);
  const auto ret = bufQ_.preallocate(2048, 2048);
  *bufReturn = ret.first;
  *lenReturn = ret.second;
}

// Called when AsyncSocket gave us data
void WatchmanConnection::readDataAvailable(size_t len) noexcept {
  {
    std::lock_guard<std::mutex> g(mutex_);
    bufQ_.postallocate(len);
  }
  cpuExecutor_->add([shared_this=shared_from_this()] {
    shared_this->decodeNextResponse();
  });
}

std::unique_ptr<folly::IOBuf> WatchmanConnection::splitNextPdu() {
  std::lock_guard<std::mutex> g(mutex_);
  if (!bufQ_.front()) {
    return nullptr;
  }

  // Do we have enough data to decode the next item?
  size_t pdu_len = 0;
  try {
    pdu_len = decodePduLength(bufQ_.front());
  } catch (const std::out_of_range&) {
    // Don't have enough data yet
    return nullptr;
  }

  if (pdu_len > bufQ_.chainLength()) {
    // Don't have enough data yet
    return nullptr;
  }

  // Remove the PDU blob from the front of the chain
  return bufQ_.split(pdu_len);
}

// Try to peel off one or more PDU's from our buffer queue.
// Decode each complete PDU from BSER -> dynamic and dispatch
// either the associated QueuedCommand or to the callback_ for
// unilateral responses.
// This is executed via the cpuExecutor.  We only allow one
// thread to carry out the decoding at a time so that the callbacks
// are triggered in the order that they are received.  It is possible
// for us to receive a large PDU followed by a small one and for the
// small one to finish decoding before the large one, so we must
// serialize the dispatching.
void WatchmanConnection::decodeNextResponse() {
  {
    std::lock_guard<std::mutex> g(mutex_);
    if (decoding_) {
      return;
    }
    decoding_ = true;
  }

  SCOPE_EXIT {
    std::lock_guard<std::mutex> g(mutex_);
    decoding_ = false;
  };

  while (true) {
    auto pdu = splitNextPdu();
    if (!pdu) {
      return;
    }

    try {
      auto decoded = parseBser(pdu.get());

      bool is_unilateral = false;
      // Check for a unilateral response
      for (const auto& k : kUnilateralLabels) {
        if (decoded.get_ptr(k)) {
          // This is a unilateral response
          if (callback_.hasValue()) {
            callback_.value()(watchmanResponseToTry(std::move(decoded)));
            is_unilateral = true;
            break;
          }
          // No callback; usage error :-/
          failQueuedCommands(
              std::runtime_error("No unilateral callback has been installed"));
          return;
        }
      }
      if (is_unilateral) {
        continue;
      }

      // It's actually a command response; get the cmd so that we
      // can fulfil its promise
      std::shared_ptr<QueuedCommand> cmd;
      {
        std::lock_guard<std::mutex> g(mutex_);
        if (commandQ_.empty()) {
          failQueuedCommands(
              std::runtime_error("No commands have been queued"));
          return;
        }
        cmd = commandQ_.front();
      }

      // Dispatch outside of the lock in case it tries to send another
      // command
      cmd->promise.setTry(watchmanResponseToTry(std::move(decoded)));

      // Now we're in a position to send the next queued command.
      // We remove it after dispatching the try above in case that
      // queued up more commands; we want to be the one thing that
      // is responsible for sending the next queued command here
      popAndSendCommand();
    } catch (const std::exception& ex) {
      failQueuedCommands(ex);
      return;
    }
  }
}

// Called when AsyncSocket hits EOF
void WatchmanConnection::readEOF() noexcept {
  failQueuedCommands(
      std::system_error(ENOTCONN, std::system_category(), "connection closed"));
}

// Called when AsyncSocket has a read error
void WatchmanConnection::readErr(
    const folly::AsyncSocketException& ex) noexcept {
  failQueuedCommands(ex);
}
} // namespace watchman
<commit_msg>move InlineExecutor, ManualExecutor, and GlobalThreadPoolList to<commit_after>/* Copyright 2016-present Facebook, Inc.
 * Licensed under the Apache License, Version 2.0 */

#include "WatchmanConnection.h"

#include <folly/ExceptionWrapper.h>
#include <folly/SocketAddress.h>
#include <folly/Subprocess.h>
#include <folly/executors/InlineExecutor.h>
#include <folly/experimental/bser/Bser.h>

namespace watchman {

using namespace folly::bser;
using namespace folly;

// Ordered with the most likely kind first
static const std::vector<dynamic> kUnilateralLabels{"subscription", "log"};

static const dynamic kError("error");
static const dynamic kCapabilities("capabilities");

// We'll just dispatch bser decodes and callbacks inline unless they
// give us an alternative environment
static InlineExecutor inlineExecutor;

WatchmanConnection::WatchmanConnection(
    EventBase* eventBase,
    Optional<std::string>&& sockPath,
    Optional<WatchmanConnection::Callback>&& callback,
    Executor* cpuExecutor)
    : eventBase_(eventBase),
      sockPath_(std::move(sockPath)),
      callback_(std::move(callback)),
      cpuExecutor_(cpuExecutor ? cpuExecutor : &inlineExecutor),
      versionCmd_(nullptr),
      bufQ_(IOBufQueue::cacheChainLength()) {
  CHECK_NOTNULL(eventBase);
}

WatchmanConnection::~WatchmanConnection() {
  close();
}

folly::Future<std::string> WatchmanConnection::getSockPath() {
  // Take explicit configuration first
  if (sockPath_.hasValue()) {
    return makeFuture(sockPath_.value());
  }

  // Else use the environmental variable used by watchman to report
  // the active socket path
  auto var = getenv("WATCHMAN_SOCK");
  if (var && *var) {
    return makeFuture(std::string(var));
  }

  return via(cpuExecutor_, [] {
    // Else discover it from the CLI
    folly::Subprocess proc(
        {"watchman", "--output-encoding=bser", "get-sockname"},
        folly::Subprocess::Options().pipeStdout().pipeStderr().usePath());
    SCOPE_FAIL {
      // Always clean up to avoid Subprocess asserting on destruction
      proc.kill();
      proc.wait();
    };
    auto out_pair = proc.communicate();
    auto result = parseBser(out_pair.first);
    proc.waitChecked();
    return result["sockname"].asString();
  });
}

Future<dynamic> WatchmanConnection::connect(folly::dynamic versionArgs) {
  if (!versionArgs.isObject()) {
    throw WatchmanError("versionArgs must be object");
  }
  versionCmd_ = folly::dynamic::array("version", versionArgs);

  auto res = getSockPath().then(
    [shared_this=shared_from_this()] (std::string&& path) {
      shared_this->eventBase_->runInEventBaseThread([=] {
        folly::SocketAddress addr;
        addr.setFromPath(path);

        shared_this->sock_ =
          folly::AsyncSocket::newSocket(shared_this->eventBase_);
        shared_this->sock_->connect(shared_this.get(), addr);
      }
    );

    return shared_this->connectPromise_.getFuture();
  });
  return res;
}

void WatchmanConnection::close() {
  if (closing_) {
    return;
  }
  closing_ = true;
  if (sock_) {
    eventBase_->runImmediatelyOrRunInEventBaseThreadAndWait([this] {
      sock_->close();
      sock_.reset();
    });
  }
  failQueuedCommands(
      make_exception_wrapper<WatchmanError>(
          "WatchmanConnection::close() was called"));
}

// The convention for Watchman responses is that they represent
// an error if they contain the "error" key.  We want to report
// those as exceptions, but it is easier to do that via a Try
Try<dynamic> WatchmanConnection::watchmanResponseToTry(dynamic&& value) {
  auto error = value.get_ptr(kError);
  if (error) {
    return Try<dynamic>(make_exception_wrapper<WatchmanResponseError>(value));
  }
  return Try<dynamic>(std::move(value));
}

void WatchmanConnection::connectSuccess() noexcept {
  try {
    sock_->setReadCB(this);
    sock_->setCloseOnExec();

    run(versionCmd_).then(
      [shared_this=shared_from_this()] (dynamic&& result) {
        // If there is no "capabilities" key then the version of
        // watchman is too old; treat this as an error
        if (!result.get_ptr(kCapabilities)) {
          result["error"] =
              "This watchman server has no support for capabilities, "
              "please upgrade to the current stable version of watchman";
          shared_this->connectPromise_.setTry(
            shared_this->watchmanResponseToTry(std::move(result)));
          return;
        }
        shared_this->connectPromise_.setValue(std::move(result));
      }
    ).onError(
      [shared_this=shared_from_this()]
      (const folly::exception_wrapper& e) {
        shared_this->connectPromise_.setException(e);
      }
    );
  } catch(const std::exception& e) {
    connectPromise_.setException(
      folly::exception_wrapper(std::current_exception(), e));
  } catch(...) {
    connectPromise_.setException(
      folly::exception_wrapper(std::current_exception()));
  }
}

void WatchmanConnection::connectErr(
    const folly::AsyncSocketException& ex) noexcept {
  connectPromise_.setException(ex);
}

WatchmanConnection::QueuedCommand::QueuedCommand(const dynamic& command)
    : cmd(command) {}

Future<dynamic> WatchmanConnection::run(const dynamic& command) noexcept {
  auto cmd = std::make_shared<QueuedCommand>(command);
  if (broken_) {
    cmd->promise.setException(WatchmanError("The connection was broken"));
    return cmd->promise.getFuture();
  }
  if (!sock_) {
    cmd->promise.setException(WatchmanError(
        "No socket (did you call connect() and check result for exceptions?)"));
    return cmd->promise.getFuture();
  }

  bool shouldWrite;
  {
    std::lock_guard<std::mutex> g(mutex_);
    // We only need to call sendCommand if we don't have a command in
    // progress; the completion handler will trigger it once we receive
    // the response
    shouldWrite = commandQ_.empty();
    commandQ_.push_back(cmd);
  }

  if (shouldWrite) {
    eventBase_->runInEventBaseThread(
      [shared_this=shared_from_this()] {
        shared_this->sendCommand();
      }
    );
  }

  return cmd->promise.getFuture();
}

// Generate a failure for all queued commands
void WatchmanConnection::failQueuedCommands(
    const folly::exception_wrapper& ex) {
  std::lock_guard<std::mutex> g(mutex_);
  auto q = commandQ_;
  commandQ_.clear();

  broken_ = true;
  for (auto& cmd : q) {
    if (!cmd->promise.isFulfilled()) {
      cmd->promise.setException(ex);
    }
  }

  // If the user has explicitly closed the connection no need for callback
  if (callback_ && !closing_) {
    cpuExecutor_->add([shared_this=shared_from_this(), ex] {
      (*(shared_this->callback_))(folly::Try<folly::dynamic>(ex));
    });
  }
}

// Sends the next eligible command to the Watchman service
void WatchmanConnection::sendCommand(bool pop) {
  std::shared_ptr<QueuedCommand> cmd;

  {
    std::lock_guard<std::mutex> g(mutex_);

    if (pop) {
      // We finished processing this one, discard it and focus
      // on the next item, if any.
      commandQ_.pop_front();
    }
    if (commandQ_.empty()) {
      return;
    }
    cmd = commandQ_.front();
  }

  sock_->writeChain(this, toBserIOBuf(cmd->cmd, serialization_opts()));
}

void WatchmanConnection::popAndSendCommand() {
  sendCommand(/* pop = */ true);
}

// Called when AsyncSocket::writeChain completes
void WatchmanConnection::writeSuccess() noexcept {
  // Don't care particularly
}

// Called when AsyncSocket::writeChain fails
void WatchmanConnection::writeErr(
    size_t,
    const folly::AsyncSocketException& ex) noexcept {
  failQueuedCommands(ex);
}

// Called when AsyncSocket wants to give us data
void WatchmanConnection::getReadBuffer(void** bufReturn, size_t* lenReturn) {
  std::lock_guard<std::mutex> g(mutex_);
  const auto ret = bufQ_.preallocate(2048, 2048);
  *bufReturn = ret.first;
  *lenReturn = ret.second;
}

// Called when AsyncSocket gave us data
void WatchmanConnection::readDataAvailable(size_t len) noexcept {
  {
    std::lock_guard<std::mutex> g(mutex_);
    bufQ_.postallocate(len);
  }
  cpuExecutor_->add([shared_this=shared_from_this()] {
    shared_this->decodeNextResponse();
  });
}

std::unique_ptr<folly::IOBuf> WatchmanConnection::splitNextPdu() {
  std::lock_guard<std::mutex> g(mutex_);
  if (!bufQ_.front()) {
    return nullptr;
  }

  // Do we have enough data to decode the next item?
  size_t pdu_len = 0;
  try {
    pdu_len = decodePduLength(bufQ_.front());
  } catch (const std::out_of_range&) {
    // Don't have enough data yet
    return nullptr;
  }

  if (pdu_len > bufQ_.chainLength()) {
    // Don't have enough data yet
    return nullptr;
  }

  // Remove the PDU blob from the front of the chain
  return bufQ_.split(pdu_len);
}

// Try to peel off one or more PDU's from our buffer queue.
// Decode each complete PDU from BSER -> dynamic and dispatch
// either the associated QueuedCommand or to the callback_ for
// unilateral responses.
// This is executed via the cpuExecutor.  We only allow one
// thread to carry out the decoding at a time so that the callbacks
// are triggered in the order that they are received.  It is possible
// for us to receive a large PDU followed by a small one and for the
// small one to finish decoding before the large one, so we must
// serialize the dispatching.
void WatchmanConnection::decodeNextResponse() {
  {
    std::lock_guard<std::mutex> g(mutex_);
    if (decoding_) {
      return;
    }
    decoding_ = true;
  }

  SCOPE_EXIT {
    std::lock_guard<std::mutex> g(mutex_);
    decoding_ = false;
  };

  while (true) {
    auto pdu = splitNextPdu();
    if (!pdu) {
      return;
    }

    try {
      auto decoded = parseBser(pdu.get());

      bool is_unilateral = false;
      // Check for a unilateral response
      for (const auto& k : kUnilateralLabels) {
        if (decoded.get_ptr(k)) {
          // This is a unilateral response
          if (callback_.hasValue()) {
            callback_.value()(watchmanResponseToTry(std::move(decoded)));
            is_unilateral = true;
            break;
          }
          // No callback; usage error :-/
          failQueuedCommands(
              std::runtime_error("No unilateral callback has been installed"));
          return;
        }
      }
      if (is_unilateral) {
        continue;
      }

      // It's actually a command response; get the cmd so that we
      // can fulfil its promise
      std::shared_ptr<QueuedCommand> cmd;
      {
        std::lock_guard<std::mutex> g(mutex_);
        if (commandQ_.empty()) {
          failQueuedCommands(
              std::runtime_error("No commands have been queued"));
          return;
        }
        cmd = commandQ_.front();
      }

      // Dispatch outside of the lock in case it tries to send another
      // command
      cmd->promise.setTry(watchmanResponseToTry(std::move(decoded)));

      // Now we're in a position to send the next queued command.
      // We remove it after dispatching the try above in case that
      // queued up more commands; we want to be the one thing that
      // is responsible for sending the next queued command here
      popAndSendCommand();
    } catch (const std::exception& ex) {
      failQueuedCommands(ex);
      return;
    }
  }
}

// Called when AsyncSocket hits EOF
void WatchmanConnection::readEOF() noexcept {
  failQueuedCommands(
      std::system_error(ENOTCONN, std::system_category(), "connection closed"));
}

// Called when AsyncSocket has a read error
void WatchmanConnection::readErr(
    const folly::AsyncSocketException& ex) noexcept {
  failQueuedCommands(ex);
}
} // namespace watchman
<|endoftext|>


<commit_before>/**
 *  Multiplicative hashing
 *  ======================
 *
 *  Based on the golden ratio, so A.K.A. Fibonacci hashing.  Simple, fast, but
 *  frequent to collide.
 */

#include <iostream>

uint32_t
mulhash32(uint32_t h, int shift)
{
	return (h * 2654435769u) >> shift;
}

uint64_t
mulhash64(uint64_t h, int shift)
{
	return (h * 11400714819323198485llu) >> shift;
}

int
main()
{
	int shift = 64 - 3;

	// Naively check the distribution.
	std::cout << mulhash32(0, shift) << std::endl;
	std::cout << mulhash32(1, shift) << std::endl;
	std::cout << mulhash32(2, shift) << std::endl;
	std::cout << mulhash32(3, shift) << std::endl;
	std::cout << mulhash32(4, shift) << std::endl;
	std::cout << mulhash32(5, shift) << std::endl;
	std::cout << mulhash32(6, shift) << std::endl;
	std::cout << mulhash32(7, shift) << std::endl;
	std::cout << "---" << std::endl;
	std::cout << mulhash64(0, shift) << std::endl;
	std::cout << mulhash64(1, shift) << std::endl;
	std::cout << mulhash64(2, shift) << std::endl;
	std::cout << mulhash64(3, shift) << std::endl;
	std::cout << mulhash64(4, shift) << std::endl;
	std::cout << mulhash64(5, shift) << std::endl;
	std::cout << mulhash64(6, shift) << std::endl;
	std::cout << mulhash64(7, shift) << std::endl;

	return 0;
}
<commit_msg>updated mulhash: test cases<commit_after>/**
 *  Multiplicative hashing
 *  ======================
 *
 *  Based on the golden ratio, hence A.K.A. Fibonacci hashing.  Simple, fast,
 *  but not well evenly-distributed.
 */

#include <iostream>
#include <array>

uint32_t
mulhash32(uint32_t h, int shift)
{
	return (h * 2654435769u) >> shift;
}

uint64_t
mulhash64(uint64_t h, int shift)
{
	return (h * 11400714819323198485llu) >> shift;
}

#define LEN 8
#define N 64

template <typename T>
void
debug(std::array<T, LEN> arr)
{
	if constexpr (std::is_same_v<T, uint32_t>) {
		std::cout << "mulhash32=";
	} else {
		std::cout << "mulhash64=";
	}

	std::cout << "[";
	for (const auto& a : arr) {
		std::cout << a << ",";
	}
	std::cout << "\b]" << std::endl;
}

int
main()
{
	int shift = 64 - 3;
	std::array<uint32_t, LEN> arr32{};
	std::array<uint64_t, LEN> arr64{};

	// Naively check the distribution.

	for (int i = 0; i < N; ++i) ++arr32[mulhash32(i, shift)];
	debug(arr32);
	for (int i = 0; i < N; ++i) ++arr64[mulhash64(i, shift)];
	debug(arr64);

	return 0;
}
<|endoftext|>