<commit_before>// BushingForce.cpp
// Author: Ajay Seth
/*
 * Copyright (c)  2010, Stanford University. All rights reserved.
* Use of the OpenSim software in source form is permitted provided that the following
* conditions are met:
* 	1. The software is used only for non-commercial research and education. It may not
*     be used in relation to any commercial activity.
* 	2. The software is not distributed or redistributed.  Software distribution is allowed
*     only through https://simtk.org/home/opensim.
* 	3. Use of the OpenSim software or derivatives must be acknowledged in all publications,
*      presentations, or documents describing work in which OpenSim or derivatives are used.
* 	4. Credits to developers may not be removed from executables
*     created from modifications of the source.
* 	5. Modifications of source code must retain the above copyright notice, this list of
*     conditions and the following disclaimer.
*
*  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY
*  EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
*  OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT
*  SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
*  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
*  TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
*  HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
*  OR BUSINESS INTERRUPTION) OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY
*  WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

//=============================================================================
// INCLUDES
//=============================================================================
#include <OpenSim/Simulation/Model/BodySet.h>
#include <OpenSim/Simulation/Model/Model.h>

#include "BushingForce.h"

//=============================================================================
// STATICS
//=============================================================================
using namespace std;
using namespace SimTK;
using namespace OpenSim;

//=============================================================================
// CONSTRUCTOR(S) AND DESTRUCTOR
//=============================================================================
//_____________________________________________________________________________
/**
 * Destructor.
 */
BushingForce::~BushingForce()
{
}
//_____________________________________________________________________________
/**
 * Default constructor.
 */
BushingForce::BushingForce() :
	Force(),
	_body1Name(_body1NameProp.getValueStr()),
	_body2Name(_body2NameProp.getValueStr()),
	_locationInBody1(_locationInBody1Prop.getValueDblVec3()),
	_orientationInBody1(_orientationInBody1Prop.getValueDblVec3()),
	_locationInBody2(_locationInBody2Prop.getValueDblVec3()),
	_orientationInBody2(_orientationInBody2Prop.getValueDblVec3()),
	_rotStiffness(_rotStiffnessProp.getValueDblVec3()),
	_transStiffness(_transStiffnessProp.getValueDblVec3()),
	_rotDamping(_rotDampingProp.getValueDblVec3()),
	_transDamping(_transDampingProp.getValueDblVec3())
{
	setNull();
}

/* Convenience constructor */
BushingForce::BushingForce( std::string body1Name, SimTK::Vec3 point1, SimTK::Vec3 orientation1,
		          std::string body2Name, SimTK::Vec3 point2, SimTK::Vec3 orientation2,
				  SimTK::Vec3 transStiffness, SimTK::Vec3 rotStiffness, SimTK::Vec3 transDamping, SimTK::Vec3 rotDamping ):
	Force(),
	_body1Name(_body1NameProp.getValueStr()),
	_body2Name(_body2NameProp.getValueStr()),
	_locationInBody1(_locationInBody1Prop.getValueDblVec3()),
	_orientationInBody1(_orientationInBody1Prop.getValueDblVec3()),
	_locationInBody2(_locationInBody2Prop.getValueDblVec3()),
	_orientationInBody2(_orientationInBody2Prop.getValueDblVec3()),
	_rotStiffness(_rotStiffnessProp.getValueDblVec3()),
	_transStiffness(_transStiffnessProp.getValueDblVec3()),
	_rotDamping(_rotDampingProp.getValueDblVec3()),
	_transDamping(_transDampingProp.getValueDblVec3())
{
	setNull();
	_body1Name = body1Name;
	_body2Name = body2Name;
	_locationInBody1 = point1;
	_orientationInBody1 = orientation1;
	_locationInBody2 = point2;
	_orientationInBody2 = orientation2;
	_rotStiffness = rotStiffness;
	_transStiffness = transStiffness;
	_rotDamping = rotDamping;
	_transDamping = transDamping;
}

//_____________________________________________________________________________
/**
 * Copy constructor.
 *
 * @param aForce BushingForce to be copied.
 */
BushingForce::BushingForce(const BushingForce &aForce) :
   Force(aForce),
	_body1Name(_body1NameProp.getValueStr()),
	_body2Name(_body2NameProp.getValueStr()),
	_locationInBody1(_locationInBody1Prop.getValueDblVec3()),
	_orientationInBody1(_orientationInBody1Prop.getValueDblVec3()),
	_locationInBody2(_locationInBody2Prop.getValueDblVec3()),
	_orientationInBody2(_orientationInBody2Prop.getValueDblVec3()),
	_rotStiffness(_rotStiffnessProp.getValueDblVec3()),
	_transStiffness(_transStiffnessProp.getValueDblVec3()),
	_rotDamping(_rotDampingProp.getValueDblVec3()),
	_transDamping(_transDampingProp.getValueDblVec3())
{
	setNull();
	copyData(aForce);
}

//=============================================================================
// CONSTRUCTION
//=============================================================================
//_____________________________________________________________________________
/**
 * Copy this body and return a pointer to the copy.
 * The copy constructor for this class is used.
 *
 * @return Pointer to a copy of this OpenSim::Body.
 */
Object* BushingForce::copy() const
{
	BushingForce *Force = new BushingForce(*this);
	return(Force);
}
//_____________________________________________________________________________
/**
 * Copy data members from one BushingForce to another.
 *
 * @param aForce BushingForce to be copied.
 */
void BushingForce::copyData(const BushingForce &aForce)
{
	Force::copyData(aForce);
	_body1Name = aForce._body1Name;
	_body2Name = aForce._body2Name;
	_locationInBody1 = aForce._locationInBody1;
	_orientationInBody1 = aForce._orientationInBody1;
	_locationInBody2 = aForce._locationInBody2;
	_orientationInBody2 = aForce._orientationInBody2;
	_rotStiffness = aForce._rotStiffness;
	_transStiffness = aForce._transStiffness;
	_rotDamping = aForce._rotDamping;
	_transDamping = aForce._transDamping;
}

//_____________________________________________________________________________
/**
 * Set the data members of this BushingForce to their null values.
 */
void BushingForce::setNull()
{
	setType("BushingForce");
	setupProperties();
}

//_____________________________________________________________________________
/**
 * Connect properties to local pointers.
 */
void BushingForce::setupProperties()
{
	// Body 1 name
	_body1NameProp.setName("body_1");
	_propertySet.append(&_body1NameProp);

	// Body 2 name
	_body2NameProp.setName("body_2");
	_propertySet.append(&_body2NameProp);

	//Default location and orientation (rotation sequence)
	SimTK::Vec3 origin(0.0, 0.0, 0.0);

	// Location in Body 1
	_locationInBody1Prop.setName("location_body_1");
	_locationInBody1Prop.setValue(origin);
	_propertySet.append(&_locationInBody1Prop);

	// Orientation in Body 1
	_orientationInBody1Prop.setName("orientation_body_1");
	_orientationInBody1Prop.setValue(origin);
	_propertySet.append(&_orientationInBody1Prop);

	// Location in Body 2
	_locationInBody2Prop.setName("location_body_2");
	_locationInBody2Prop.setValue(origin);
	_propertySet.append(&_locationInBody2Prop);

	// Orientation in Body 2
	_orientationInBody2Prop.setName("orientation_body_2");
	_orientationInBody2Prop.setValue(origin);
	_propertySet.append(&_orientationInBody2Prop);

	_rotStiffnessProp.setName("rotational_stiffness");
	_rotStiffnessProp.setValue(Vec3(0));
	_propertySet.append(&_rotStiffnessProp);

	_transStiffnessProp.setName("translational_stiffness");
	_transStiffnessProp.setValue(Vec3(0));
	_propertySet.append(&_transStiffnessProp);

	_rotDampingProp.setName("rotational_damping");
	_rotDampingProp.setValue(Vec3(0));
	_propertySet.append(&_rotDampingProp);

	_transDampingProp.setName("rotational_damping");
	_transDampingProp.setValue(Vec3(0));
	_propertySet.append(&_transDampingProp);
}

//_____________________________________________________________________________
/**
 * Perform some set up functions that happen after the
 * object has been deserialized or copied.
 *
 * @param aModel OpenSim model containing this BushingForce.
 */
void BushingForce::setup(Model& aModel)
{
	string errorMessage;

	// Base class
	Force::setup(aModel);

	// Look up the two bodies being connected by bushing by name in the
	// model and might as well keep a pointer to them
	if (!aModel.updBodySet().contains(_body1Name)) {
		errorMessage = "Invalid bushing body1 (" + _body1Name + ") specified in Force " + getName();
		throw (Exception(errorMessage.c_str()));
	}
	if (!aModel.updBodySet().contains(_body2Name)) {
		errorMessage = "Invalid bushing body2 (" + _body2Name + ") specified in Force " + getName();
		throw (Exception(errorMessage.c_str()));
	}
}

void BushingForce::createSystem(SimTK::MultibodySystem& system) const
{
	Body &body1 = _model->updBodySet().get(_body1Name);
	Body &body2 = _model->updBodySet().get(_body2Name);

	// Get underlying mobilized bodies
	SimTK::MobilizedBody b1 = _model->updMatterSubsystem().getMobilizedBody(body1.getIndex());
	SimTK::MobilizedBody b2 = _model->updMatterSubsystem().getMobilizedBody(body2.getIndex());
	// Build the transforms
	SimTK::Rotation r1; r1.setRotationToBodyFixedXYZ(_orientationInBody1);
	SimTK::Rotation r2; r2.setRotationToBodyFixedXYZ(_orientationInBody2);
	SimTK::Transform inb1(r1, _locationInBody1);
	SimTK::Transform inb2(r2, _locationInBody2);

	Vec6 stiffness(_rotStiffness[0], _rotStiffness[1], _rotStiffness[2], _transStiffness[0], _transStiffness[1], _transStiffness[2]);
	Vec6 damping(_rotDamping[0], _rotDamping[1], _rotDamping[2], _transDamping[0], _transDamping[1], _transDamping[2]);

    // Now create a Simbody Force::LinearBushing
    SimTK::Force::LinearBushing simtkForce(_model->updUserForceSubsystem(), b1, inb1, b2, inb2, stiffness, damping);

    // Beyond the const Component get the index so we can access the SimTK::Force later
	BushingForce* mutableThis = const_cast<BushingForce *>(this);
	mutableThis->_index = simtkForce.getForceIndex();
}

//=============================================================================
// OPERATORS
//=============================================================================
//_____________________________________________________________________________
/**
 * Assignment operator.
 *
 * @return Reference to this object.
 */
BushingForce& BushingForce::operator=(const BushingForce &aForce)
{
	Force::operator=(aForce);
	copyData(aForce);
	return(*this);
}

//=============================================================================
// SET
//=============================================================================
//_____________________________________________________________________________
/**
 * Following methods set attributes of the weld Force */
void BushingForce::setBody1ByName(std::string aBodyName)
{
	_body1Name = aBodyName;
}

void BushingForce::setBody2ByName(std::string aBodyName)
{
	_body2Name = aBodyName;
}

/** Set the location and orientation (optional) for weld on body 1*/
void BushingForce::setBody1BushingLocation(Vec3 location, Vec3 orientation)
{
	_locationInBody1 = location;
	_orientationInBody1 = orientation;
}

/** Set the location and orientation (optional) for weld on body 2*/
void BushingForce::setBody2BushingLocation(Vec3 location, Vec3 orientation)
{
	_locationInBody2 = location;
	_orientationInBody2 = orientation;
}

//=============================================================================
// Reporting
//=============================================================================
/**
 * Provide names of the quantities (column labels) of the force value(s) reported
 *
 */
OpenSim::Array<std::string> BushingForce::getRecordLabels() const
{
	OpenSim::Array<std::string> labels("");
	labels.append(getName()+"."+_body1Name+".force.X");
	labels.append(getName()+"."+_body1Name+".force.Y");
	labels.append(getName()+"."+_body1Name+".force.Z");
	labels.append(getName()+"."+_body1Name+".torque.X");
	labels.append(getName()+"."+_body1Name+".torque.Y");
	labels.append(getName()+"."+_body1Name+".torque.Z");
	labels.append(getName()+"."+_body2Name+".force.X");
	labels.append(getName()+"."+_body2Name+".force.Y");
	labels.append(getName()+"."+_body2Name+".force.Z");
	labels.append(getName()+"."+_body2Name+".torque.X");
	labels.append(getName()+"."+_body2Name+".torque.Y");
	labels.append(getName()+"."+_body2Name+".torque.Z");

	return labels;
}
/**
 * Provide the value(s) to be reported that correspond to the labels
 */
OpenSim::Array<double> BushingForce::getRecordValues(const SimTK::State& state) const
{
	OpenSim::Array<double> values(1);

	const SimTK::Force::LinearBushing &simtkSpring = (SimTK::Force::LinearBushing &)(_model->getUserForceSubsystem().getForce(_index));

	SimTK::Vector_<SimTK::SpatialVec> bodyForces(0);
	SimTK::Vector_<SimTK::Vec3> particleForces(0);
	SimTK::Vector mobilityForces(0);

	//get the net force added to the system contributed by the bushing
	simtkSpring.calcForceContribution(state, bodyForces, particleForces, mobilityForces);
	SimTK::Vec3 forces = bodyForces(_model->getBodySet().get(_body1Name).getIndex())[1];
	SimTK::Vec3 torques = bodyForces(_model->getBodySet().get(_body1Name).getIndex())[0];
	values.append(3, &forces[0]);
	values.append(3, &torques[0]);

	forces = bodyForces(_model->getBodySet().get(_body2Name).getIndex())[1];
	torques = bodyForces(_model->getBodySet().get(_body2Name).getIndex())[0];

	values.append(3, &forces[0]);
	values.append(3, &torques[0]);

	return values;
}
<commit_msg>Fixed a copy-paste bug that was overwriting rotational_damping properties with translational_damping property values and failing to have translational_damping.<commit_after>// BushingForce.cpp
// Author: Ajay Seth
/*
 * Copyright (c)  2010, Stanford University. All rights reserved.
* Use of the OpenSim software in source form is permitted provided that the following
* conditions are met:
* 	1. The software is used only for non-commercial research and education. It may not
*     be used in relation to any commercial activity.
* 	2. The software is not distributed or redistributed.  Software distribution is allowed
*     only through https://simtk.org/home/opensim.
* 	3. Use of the OpenSim software or derivatives must be acknowledged in all publications,
*      presentations, or documents describing work in which OpenSim or derivatives are used.
* 	4. Credits to developers may not be removed from executables
*     created from modifications of the source.
* 	5. Modifications of source code must retain the above copyright notice, this list of
*     conditions and the following disclaimer.
*
*  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY
*  EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
*  OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT
*  SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
*  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
*  TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
*  HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
*  OR BUSINESS INTERRUPTION) OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY
*  WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

//=============================================================================
// INCLUDES
//=============================================================================
#include <OpenSim/Simulation/Model/BodySet.h>
#include <OpenSim/Simulation/Model/Model.h>

#include "BushingForce.h"

//=============================================================================
// STATICS
//=============================================================================
using namespace std;
using namespace SimTK;
using namespace OpenSim;

//=============================================================================
// CONSTRUCTOR(S) AND DESTRUCTOR
//=============================================================================
//_____________________________________________________________________________
/**
 * Destructor.
 */
BushingForce::~BushingForce()
{
}
//_____________________________________________________________________________
/**
 * Default constructor.
 */
BushingForce::BushingForce() :
	Force(),
	_body1Name(_body1NameProp.getValueStr()),
	_body2Name(_body2NameProp.getValueStr()),
	_locationInBody1(_locationInBody1Prop.getValueDblVec3()),
	_orientationInBody1(_orientationInBody1Prop.getValueDblVec3()),
	_locationInBody2(_locationInBody2Prop.getValueDblVec3()),
	_orientationInBody2(_orientationInBody2Prop.getValueDblVec3()),
	_rotStiffness(_rotStiffnessProp.getValueDblVec3()),
	_transStiffness(_transStiffnessProp.getValueDblVec3()),
	_rotDamping(_rotDampingProp.getValueDblVec3()),
	_transDamping(_transDampingProp.getValueDblVec3())
{
	setNull();
}

/* Convenience constructor */
BushingForce::BushingForce( std::string body1Name, SimTK::Vec3 point1, SimTK::Vec3 orientation1,
		          std::string body2Name, SimTK::Vec3 point2, SimTK::Vec3 orientation2,
				  SimTK::Vec3 transStiffness, SimTK::Vec3 rotStiffness, SimTK::Vec3 transDamping, SimTK::Vec3 rotDamping ):
	Force(),
	_body1Name(_body1NameProp.getValueStr()),
	_body2Name(_body2NameProp.getValueStr()),
	_locationInBody1(_locationInBody1Prop.getValueDblVec3()),
	_orientationInBody1(_orientationInBody1Prop.getValueDblVec3()),
	_locationInBody2(_locationInBody2Prop.getValueDblVec3()),
	_orientationInBody2(_orientationInBody2Prop.getValueDblVec3()),
	_rotStiffness(_rotStiffnessProp.getValueDblVec3()),
	_transStiffness(_transStiffnessProp.getValueDblVec3()),
	_rotDamping(_rotDampingProp.getValueDblVec3()),
	_transDamping(_transDampingProp.getValueDblVec3())
{
	setNull();
	_body1Name = body1Name;
	_body2Name = body2Name;
	_locationInBody1 = point1;
	_orientationInBody1 = orientation1;
	_locationInBody2 = point2;
	_orientationInBody2 = orientation2;
	_rotStiffness = rotStiffness;
	_transStiffness = transStiffness;
	_rotDamping = rotDamping;
	_transDamping = transDamping;
}

//_____________________________________________________________________________
/**
 * Copy constructor.
 *
 * @param aForce BushingForce to be copied.
 */
BushingForce::BushingForce(const BushingForce &aForce) :
   Force(aForce),
	_body1Name(_body1NameProp.getValueStr()),
	_body2Name(_body2NameProp.getValueStr()),
	_locationInBody1(_locationInBody1Prop.getValueDblVec3()),
	_orientationInBody1(_orientationInBody1Prop.getValueDblVec3()),
	_locationInBody2(_locationInBody2Prop.getValueDblVec3()),
	_orientationInBody2(_orientationInBody2Prop.getValueDblVec3()),
	_rotStiffness(_rotStiffnessProp.getValueDblVec3()),
	_transStiffness(_transStiffnessProp.getValueDblVec3()),
	_rotDamping(_rotDampingProp.getValueDblVec3()),
	_transDamping(_transDampingProp.getValueDblVec3())
{
	setNull();
	copyData(aForce);
}

//=============================================================================
// CONSTRUCTION
//=============================================================================
//_____________________________________________________________________________
/**
 * Copy this body and return a pointer to the copy.
 * The copy constructor for this class is used.
 *
 * @return Pointer to a copy of this OpenSim::Body.
 */
Object* BushingForce::copy() const
{
	BushingForce *Force = new BushingForce(*this);
	return(Force);
}
//_____________________________________________________________________________
/**
 * Copy data members from one BushingForce to another.
 *
 * @param aForce BushingForce to be copied.
 */
void BushingForce::copyData(const BushingForce &aForce)
{
	Force::copyData(aForce);
	_body1Name = aForce._body1Name;
	_body2Name = aForce._body2Name;
	_locationInBody1 = aForce._locationInBody1;
	_orientationInBody1 = aForce._orientationInBody1;
	_locationInBody2 = aForce._locationInBody2;
	_orientationInBody2 = aForce._orientationInBody2;
	_rotStiffness = aForce._rotStiffness;
	_transStiffness = aForce._transStiffness;
	_rotDamping = aForce._rotDamping;
	_transDamping = aForce._transDamping;
}

//_____________________________________________________________________________
/**
 * Set the data members of this BushingForce to their null values.
 */
void BushingForce::setNull()
{
	setType("BushingForce");
	setupProperties();
}

//_____________________________________________________________________________
/**
 * Connect properties to local pointers.
 */
void BushingForce::setupProperties()
{
	// Body 1 name
	_body1NameProp.setName("body_1");
	_propertySet.append(&_body1NameProp);

	// Body 2 name
	_body2NameProp.setName("body_2");
	_propertySet.append(&_body2NameProp);

	//Default location and orientation (rotation sequence)
	SimTK::Vec3 origin(0.0, 0.0, 0.0);

	// Location in Body 1
	_locationInBody1Prop.setName("location_body_1");
	_locationInBody1Prop.setValue(origin);
	_propertySet.append(&_locationInBody1Prop);

	// Orientation in Body 1
	_orientationInBody1Prop.setName("orientation_body_1");
	_orientationInBody1Prop.setValue(origin);
	_propertySet.append(&_orientationInBody1Prop);

	// Location in Body 2
	_locationInBody2Prop.setName("location_body_2");
	_locationInBody2Prop.setValue(origin);
	_propertySet.append(&_locationInBody2Prop);

	// Orientation in Body 2
	_orientationInBody2Prop.setName("orientation_body_2");
	_orientationInBody2Prop.setValue(origin);
	_propertySet.append(&_orientationInBody2Prop);

	_rotStiffnessProp.setName("rotational_stiffness");
	_rotStiffnessProp.setValue(Vec3(0));
	_propertySet.append(&_rotStiffnessProp);

	_transStiffnessProp.setName("translational_stiffness");
	_transStiffnessProp.setValue(Vec3(0));
	_propertySet.append(&_transStiffnessProp);

	_rotDampingProp.setName("rotational_damping");
	_rotDampingProp.setValue(Vec3(0));
	_propertySet.append(&_rotDampingProp);

	_transDampingProp.setName("translational_damping");
	_transDampingProp.setValue(Vec3(0));
	_propertySet.append(&_transDampingProp);
}

//_____________________________________________________________________________
/**
 * Perform some set up functions that happen after the
 * object has been deserialized or copied.
 *
 * @param aModel OpenSim model containing this BushingForce.
 */
void BushingForce::setup(Model& aModel)
{
	string errorMessage;

	// Base class
	Force::setup(aModel);

	// Look up the two bodies being connected by bushing by name in the
	// model and might as well keep a pointer to them
	if (!aModel.updBodySet().contains(_body1Name)) {
		errorMessage = "Invalid bushing body1 (" + _body1Name + ") specified in Force " + getName();
		throw (Exception(errorMessage.c_str()));
	}
	if (!aModel.updBodySet().contains(_body2Name)) {
		errorMessage = "Invalid bushing body2 (" + _body2Name + ") specified in Force " + getName();
		throw (Exception(errorMessage.c_str()));
	}
}

void BushingForce::createSystem(SimTK::MultibodySystem& system) const
{
	Body &body1 = _model->updBodySet().get(_body1Name);
	Body &body2 = _model->updBodySet().get(_body2Name);

	// Get underlying mobilized bodies
	SimTK::MobilizedBody b1 = _model->updMatterSubsystem().getMobilizedBody(body1.getIndex());
	SimTK::MobilizedBody b2 = _model->updMatterSubsystem().getMobilizedBody(body2.getIndex());
	// Build the transforms
	SimTK::Rotation r1; r1.setRotationToBodyFixedXYZ(_orientationInBody1);
	SimTK::Rotation r2; r2.setRotationToBodyFixedXYZ(_orientationInBody2);
	SimTK::Transform inb1(r1, _locationInBody1);
	SimTK::Transform inb2(r2, _locationInBody2);

	Vec6 stiffness(_rotStiffness[0], _rotStiffness[1], _rotStiffness[2], _transStiffness[0], _transStiffness[1], _transStiffness[2]);
	Vec6 damping(_rotDamping[0], _rotDamping[1], _rotDamping[2], _transDamping[0], _transDamping[1], _transDamping[2]);

    // Now create a Simbody Force::LinearBushing
    SimTK::Force::LinearBushing simtkForce(_model->updUserForceSubsystem(), b1, inb1, b2, inb2, stiffness, damping);

    // Beyond the const Component get the index so we can access the SimTK::Force later
	BushingForce* mutableThis = const_cast<BushingForce *>(this);
	mutableThis->_index = simtkForce.getForceIndex();
}

//=============================================================================
// OPERATORS
//=============================================================================
//_____________________________________________________________________________
/**
 * Assignment operator.
 *
 * @return Reference to this object.
 */
BushingForce& BushingForce::operator=(const BushingForce &aForce)
{
	Force::operator=(aForce);
	copyData(aForce);
	return(*this);
}

//=============================================================================
// SET
//=============================================================================
//_____________________________________________________________________________
/**
 * Following methods set attributes of the weld Force */
void BushingForce::setBody1ByName(std::string aBodyName)
{
	_body1Name = aBodyName;
}

void BushingForce::setBody2ByName(std::string aBodyName)
{
	_body2Name = aBodyName;
}

/** Set the location and orientation (optional) for weld on body 1*/
void BushingForce::setBody1BushingLocation(Vec3 location, Vec3 orientation)
{
	_locationInBody1 = location;
	_orientationInBody1 = orientation;
}

/** Set the location and orientation (optional) for weld on body 2*/
void BushingForce::setBody2BushingLocation(Vec3 location, Vec3 orientation)
{
	_locationInBody2 = location;
	_orientationInBody2 = orientation;
}

//=============================================================================
// Reporting
//=============================================================================
/**
 * Provide names of the quantities (column labels) of the force value(s) reported
 *
 */
OpenSim::Array<std::string> BushingForce::getRecordLabels() const
{
	OpenSim::Array<std::string> labels("");
	labels.append(getName()+"."+_body1Name+".force.X");
	labels.append(getName()+"."+_body1Name+".force.Y");
	labels.append(getName()+"."+_body1Name+".force.Z");
	labels.append(getName()+"."+_body1Name+".torque.X");
	labels.append(getName()+"."+_body1Name+".torque.Y");
	labels.append(getName()+"."+_body1Name+".torque.Z");
	labels.append(getName()+"."+_body2Name+".force.X");
	labels.append(getName()+"."+_body2Name+".force.Y");
	labels.append(getName()+"."+_body2Name+".force.Z");
	labels.append(getName()+"."+_body2Name+".torque.X");
	labels.append(getName()+"."+_body2Name+".torque.Y");
	labels.append(getName()+"."+_body2Name+".torque.Z");

	return labels;
}
/**
 * Provide the value(s) to be reported that correspond to the labels
 */
OpenSim::Array<double> BushingForce::getRecordValues(const SimTK::State& state) const
{
	OpenSim::Array<double> values(1);

	const SimTK::Force::LinearBushing &simtkSpring = (SimTK::Force::LinearBushing &)(_model->getUserForceSubsystem().getForce(_index));

	SimTK::Vector_<SimTK::SpatialVec> bodyForces(0);
	SimTK::Vector_<SimTK::Vec3> particleForces(0);
	SimTK::Vector mobilityForces(0);

	//get the net force added to the system contributed by the bushing
	simtkSpring.calcForceContribution(state, bodyForces, particleForces, mobilityForces);
	SimTK::Vec3 forces = bodyForces(_model->getBodySet().get(_body1Name).getIndex())[1];
	SimTK::Vec3 torques = bodyForces(_model->getBodySet().get(_body1Name).getIndex())[0];
	values.append(3, &forces[0]);
	values.append(3, &torques[0]);

	forces = bodyForces(_model->getBodySet().get(_body2Name).getIndex())[1];
	torques = bodyForces(_model->getBodySet().get(_body2Name).getIndex())[0];

	values.append(3, &forces[0]);
	values.append(3, &torques[0]);

	return values;
}
<|endoftext|>


<commit_before>/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h"

#include "absl/strings/ascii.h"
#include "absl/strings/escaping.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/tf2tensorrt/convert/convert_graph.h"
#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
#include "tensorflow/core/grappler/clusters/cluster.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h"
#include "tensorflow/core/lib/strings/numbers.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/stacktrace.h"

#if GOOGLE_CUDA && GOOGLE_TENSORRT
namespace tensorflow {
namespace tensorrt {
namespace convert {
// TODO(sami): Remove VLOG messages once the code matures
using absl::AsciiStrToUpper;
using absl::StrAppend;
using absl::StrCat;

Status TRTOptimizationPass::Init(
    const RewriterConfig_CustomGraphOptimizer* config) {
  VLOG(1) << "Called INIT for " << name_ << " with config = " << config;
  if (config == nullptr) {
    return Status::OK();
  }
  VLOG(1) << "config = " << config->DebugString();
  const auto params = config->parameter_map();
  if (params.count("minimum_segment_size")) {
    minimum_segment_size_ = params.at("minimum_segment_size").i();
  }
  if (params.count("max_batch_size")) {
    maximum_batch_size_ = params.at("max_batch_size").i();
  }
  if (params.count("is_dynamic_op")) {
    is_dynamic_op_ = params.at("is_dynamic_op").b();
  }
  if (params.count("maximum_cached_engines")) {
    max_cached_batches_ = params.at("maximum_cached_engines").i();
  }
  if (params.count("max_workspace_size_bytes")) {
    max_workspace_size_bytes_ = params.at("max_workspace_size_bytes").i();
  }
  if (params.count("precision_mode")) {
    TF_RETURN_IF_ERROR(TrtPrecisionModeFromName(
        AsciiStrToUpper(params.at("precision_mode").s()), &precision_mode_));
  }
  if (params.count("use_calibration")) {
    use_calibration_ = params.at("use_calibration").b();
  }
  if (params.count("trt_logger")) {
    trt_logger_name_ = params.at("trt_logger").s();
  }
  if (params.count("allow_build_at_runtime")) {
    allow_build_at_runtime_ = params.at("allow_build_at_runtime").b();
  }
  if (params.count("use_implicit_batch")) {
    use_implicit_batch_ = params.at("use_implicit_batch").b();
  }
  return Status::OK();
}

void TRTOptimizationPass::PrintDebugInfo(grappler::Cluster* cluster,
                                         const grappler::GrapplerItem& item) {
  LOG(INFO) << "Cluster = " << cluster;
  string offset("  ");
  string offset2 = StrCat(offset, offset);
  string offset3 = StrCat(offset2, offset);
  string offset4 = StrCat(offset2, offset2);

  if (cluster) {
    LOG(INFO) << offset << "type             = " << cluster->type();
    LOG(INFO) << offset << "num warmup steps = " << cluster->NumWarmupSteps();
    const auto dev_names = cluster->GetDeviceNames();
    if (!dev_names.empty()) {
      LOG(INFO) << offset << " Device names:";
      for (const auto& s : dev_names) {
        LOG(INFO) << offset2 << s;
      }
    }
    std::unordered_map<string, uint64> peak_mem;
    auto status = cluster->GetPeakMemoryUsage(&peak_mem);
    if (status == Status::OK()) {
      LOG(INFO) << offset << "Peak Memory Usage :";
      for (const auto& s : peak_mem) {
        LOG(INFO) << offset2 << s.first << " = " << s.second;
      }
    }

    const auto dev_props = cluster->GetDevices();
    if (!dev_props.empty()) {
      LOG(INFO) << offset << "Device properties:";
      for (const auto& k : dev_props) {
        LOG(INFO) << offset2 << k.first;
        const auto& dt = k.second;
        LOG(INFO) << offset3 << "type          = " << dt.type();
        LOG(INFO) << offset3 << "vendor        = " << dt.vendor();
        LOG(INFO) << offset3 << "model         = " << dt.model();
        LOG(INFO) << offset3 << "frequency     = " << dt.frequency();
        LOG(INFO) << offset3 << "num cores     = " << dt.num_cores();
        LOG(INFO) << offset3 << "num registers = " << dt.num_registers();
        LOG(INFO) << offset3 << "L1 cache size = " << dt.l1_cache_size();
        LOG(INFO) << offset3 << "L2 cache size = " << dt.l2_cache_size();
        LOG(INFO) << offset3 << "L3 cache size = " << dt.l3_cache_size();
        LOG(INFO) << offset3 << "SHMem per SMP = "
                  << dt.shared_memory_size_per_multiprocessor();
        LOG(INFO) << offset3 << "memory size   = " << dt.memory_size();
        LOG(INFO) << offset3 << "bandwidth     = " << dt.bandwidth();
        if (dt.environment_size()) {
          LOG(INFO) << offset3 << "environment   :";
          for (const auto& e : dt.environment()) {
            LOG(INFO) << offset4 << e.first << " = " << e.second;
          }
        }
      }
    }

    if (cluster->GetDeviceSet()) {
      for (const auto dev : cluster->GetDeviceSet()->devices()) {
        LOG(INFO) << "Device name= " << dev->name() << "Pased name= "
                  << DeviceNameUtils::ParsedNameToString(dev->parsed_name());
      }
    }
  }

  LOG(INFO) << "item: " << item.id;
  if (!item.feed.empty()) {
    LOG(INFO) << offset << "Feeds  :";
    for (const auto& f : item.feed) {
      const auto& shape = f.second.shape();
      LOG(INFO) << offset2 << f.first << " = shaped " << shape.DebugString();
    }
  } else {
    LOG(INFO) << offset << "No Feeds";
  }
  if (!item.fetch.empty()) {
    LOG(INFO) << offset << "Fetches  :";
    for (const auto& f : item.fetch) {
      LOG(INFO) << offset2 << f;
    }
  } else {
    LOG(INFO) << offset << "No Fetches";
  }

  if (!item.init_ops.empty()) {
    LOG(INFO) << offset << "init ops  :";
    for (const auto& f : item.init_ops) {
      LOG(INFO) << offset2 << f;
    }
  } else {
    LOG(INFO) << offset << "No init ops";
  }
  LOG(INFO) << "Save Op = " << item.save_op;
  LOG(INFO) << "Restore Op = " << item.restore_op;
  LOG(INFO) << "save_restore_loc_tensor = " << item.save_restore_loc_tensor;
  if (!item.keep_ops.empty()) {
    LOG(INFO) << offset << "keep ops  :";
    for (const auto& f : item.keep_ops) {
      LOG(INFO) << offset2 << f;
    }
  } else {
    LOG(INFO) << offset << "No keep ops";
  }
}

Status TRTOptimizationPass::Optimize(grappler::Cluster* cluster,
                                     const grappler::GrapplerItem& item,
                                     GraphDef* optimized_graph) {
  VLOG(1) << "Called TRTOptimization Pass " << name_;
  // This is a hack to workaround optimizer issue. MetaOptimizer calls
  // optimization passes on function objects as well, we should not modify
  // generated funcdefs! This is fragile but we don't have any other option
  // until framework fixes it.
  if (item.id != "tf_graph") {
    VLOG(1) << "Called TRTOptimization Pass " << name_
            << " on a grappler item with id=" << item.id
            << ", which is probably a function object (funcdef). "
            << "Skipping optimization because TensorRTOptimizer "
            << "should not be called on function objects.";
    *optimized_graph = item.graph;
    return Status::OK();
  }
  if (VLOG_IS_ON(3)) {
    LOG(INFO) << CurrentStackTrace();
    PrintDebugInfo(cluster, item);
  }
  if (!is_dynamic_op_) {
    int max_batch_dim = -1;
    if (!item.feed.empty()) {
      for (const auto& f : item.feed) {
        const auto& shape = f.second.shape();
        if (shape.dims() > 0) {
          if (shape.dim_size(0) > max_batch_dim)
            max_batch_dim = shape.dim_size(0);
          VLOG(2) << "Setting max_batch_dim to " << max_batch_dim
                  << " using batch dimension of " << f.first << " with shape "
                  << shape;
        }
      }
    }
    if (max_batch_dim > maximum_batch_size_) {
      return errors::InvalidArgument(
          "Specified max_batch_size=", maximum_batch_size_,
          " is less than maximum batch dimension of inputs (", max_batch_dim,
          "). ", "To continue, set max_batch_size to >= ", max_batch_dim);
    } else if (max_batch_dim < maximum_batch_size_) {
      LOG(INFO) << "Specified max_batch_size=" << maximum_batch_size_
                << " is larger than maximum batch dimension of inputs ("
                << max_batch_dim << "). "
                << "This can result in poor performance.";
    }
  }

  if (use_calibration_ && precision_mode_ != TrtPrecisionMode::INT8) {
    VLOG(1) << "Calibration with FP32 or FP16 is not implemented. "
            << "Falling back to use_calibration = False."
            << "Note that the default value of use_calibration is True.";
    use_calibration_ = false;
  }

  std::vector<string> nodes_to_preserve;
  for (const auto& n : item.NodesToPreserve()) {
    auto tokens = str_util::Split(n, ":");
    string s = tokens.at(0);
    for (int i = 1; i < tokens.size() - 1; ++i) {
      StrAppend(&s, ":", tokens.at(i));
    }
    int dumm_port = -1;
    // If the last token is not an integer, it must be part of the name.
    // Otherwise it is port number.
    if (tokens.size() > 1 &&
        !strings::safe_strto32(tokens.back(), &dumm_port)) {  // non-absl ok
      StrAppend(&s, ":", tokens.back());
    }
    nodes_to_preserve.push_back(s);
  }

  ConversionParams cp;
  cp.grappler_item = &item;
  cp.output_names = &nodes_to_preserve;
  cp.trt_logger_name = trt_logger_name_;
  cp.max_batch_size = maximum_batch_size_;
  cp.max_workspace_size_bytes = max_workspace_size_bytes_;
  cp.output_graph_def = optimized_graph;
  cp.precision_mode = precision_mode_;
  cp.minimum_segment_size = minimum_segment_size_;
  cp.cluster = cluster;
  cp.is_dyn_op = is_dynamic_op_;
  cp.max_cached_engines = max_cached_batches_;
  cp.use_calibration = use_calibration_;
  cp.use_implicit_batch = use_implicit_batch_;
  cp.allow_build_at_runtime = allow_build_at_runtime_;
  auto status = ConvertAfterShapes(cp);
  VLOG(1) << "Returning from " << name_;
  return status;
}

void TRTOptimizationPass::Feedback(grappler::Cluster* cluster,
                                   const grappler::GrapplerItem& item,
                                   const GraphDef& optimized_graph,
                                   double result) {}

class VerboseCustomGraphOptimizerRegistrar
    : public grappler::CustomGraphOptimizerRegistrar {
 public:
  VerboseCustomGraphOptimizerRegistrar(
      const grappler::CustomGraphOptimizerRegistry::Creator& cr,
      const string& name)
      : grappler::CustomGraphOptimizerRegistrar(cr, name) {
    VLOG(1) << "Constructing a CustomOptimizationPass registration object for "
            << name;
  }
};

static VerboseCustomGraphOptimizerRegistrar TRTOptimizationPass_Registrar(
    []() {
      VLOG(1)
          << "Instantiating CustomOptimizationPass object TensorRTOptimizer";
      return new TRTOptimizationPass("TensorRTOptimizer");
    },
    ("TensorRTOptimizer"));

}  // namespace convert
}  // namespace tensorrt
}  // namespace tensorflow

#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
<commit_msg>[TF:TRT] Removes maximum batch size violation detection for static engines.<commit_after>/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h"

#include "absl/strings/ascii.h"
#include "absl/strings/escaping.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/tf2tensorrt/convert/convert_graph.h"
#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
#include "tensorflow/core/grappler/clusters/cluster.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h"
#include "tensorflow/core/lib/strings/numbers.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/stacktrace.h"

#if GOOGLE_CUDA && GOOGLE_TENSORRT
namespace tensorflow {
namespace tensorrt {
namespace convert {
// TODO(sami): Remove VLOG messages once the code matures
using absl::AsciiStrToUpper;
using absl::StrAppend;
using absl::StrCat;

Status TRTOptimizationPass::Init(
    const RewriterConfig_CustomGraphOptimizer* config) {
  VLOG(1) << "Called INIT for " << name_ << " with config = " << config;
  if (config == nullptr) {
    return Status::OK();
  }
  VLOG(1) << "config = " << config->DebugString();
  const auto params = config->parameter_map();
  if (params.count("minimum_segment_size")) {
    minimum_segment_size_ = params.at("minimum_segment_size").i();
  }
  if (params.count("max_batch_size")) {
    maximum_batch_size_ = params.at("max_batch_size").i();
  }
  if (params.count("is_dynamic_op")) {
    is_dynamic_op_ = params.at("is_dynamic_op").b();
  }
  if (params.count("maximum_cached_engines")) {
    max_cached_batches_ = params.at("maximum_cached_engines").i();
  }
  if (params.count("max_workspace_size_bytes")) {
    max_workspace_size_bytes_ = params.at("max_workspace_size_bytes").i();
  }
  if (params.count("precision_mode")) {
    TF_RETURN_IF_ERROR(TrtPrecisionModeFromName(
        AsciiStrToUpper(params.at("precision_mode").s()), &precision_mode_));
  }
  if (params.count("use_calibration")) {
    use_calibration_ = params.at("use_calibration").b();
  }
  if (params.count("trt_logger")) {
    trt_logger_name_ = params.at("trt_logger").s();
  }
  if (params.count("allow_build_at_runtime")) {
    allow_build_at_runtime_ = params.at("allow_build_at_runtime").b();
  }
  if (params.count("use_implicit_batch")) {
    use_implicit_batch_ = params.at("use_implicit_batch").b();
  }
  return Status::OK();
}

void TRTOptimizationPass::PrintDebugInfo(grappler::Cluster* cluster,
                                         const grappler::GrapplerItem& item) {
  LOG(INFO) << "Cluster = " << cluster;
  string offset("  ");
  string offset2 = StrCat(offset, offset);
  string offset3 = StrCat(offset2, offset);
  string offset4 = StrCat(offset2, offset2);

  if (cluster) {
    LOG(INFO) << offset << "type             = " << cluster->type();
    LOG(INFO) << offset << "num warmup steps = " << cluster->NumWarmupSteps();
    const auto dev_names = cluster->GetDeviceNames();
    if (!dev_names.empty()) {
      LOG(INFO) << offset << " Device names:";
      for (const auto& s : dev_names) {
        LOG(INFO) << offset2 << s;
      }
    }
    std::unordered_map<string, uint64> peak_mem;
    auto status = cluster->GetPeakMemoryUsage(&peak_mem);
    if (status == Status::OK()) {
      LOG(INFO) << offset << "Peak Memory Usage :";
      for (const auto& s : peak_mem) {
        LOG(INFO) << offset2 << s.first << " = " << s.second;
      }
    }

    const auto dev_props = cluster->GetDevices();
    if (!dev_props.empty()) {
      LOG(INFO) << offset << "Device properties:";
      for (const auto& k : dev_props) {
        LOG(INFO) << offset2 << k.first;
        const auto& dt = k.second;
        LOG(INFO) << offset3 << "type          = " << dt.type();
        LOG(INFO) << offset3 << "vendor        = " << dt.vendor();
        LOG(INFO) << offset3 << "model         = " << dt.model();
        LOG(INFO) << offset3 << "frequency     = " << dt.frequency();
        LOG(INFO) << offset3 << "num cores     = " << dt.num_cores();
        LOG(INFO) << offset3 << "num registers = " << dt.num_registers();
        LOG(INFO) << offset3 << "L1 cache size = " << dt.l1_cache_size();
        LOG(INFO) << offset3 << "L2 cache size = " << dt.l2_cache_size();
        LOG(INFO) << offset3 << "L3 cache size = " << dt.l3_cache_size();
        LOG(INFO) << offset3 << "SHMem per SMP = "
                  << dt.shared_memory_size_per_multiprocessor();
        LOG(INFO) << offset3 << "memory size   = " << dt.memory_size();
        LOG(INFO) << offset3 << "bandwidth     = " << dt.bandwidth();
        if (dt.environment_size()) {
          LOG(INFO) << offset3 << "environment   :";
          for (const auto& e : dt.environment()) {
            LOG(INFO) << offset4 << e.first << " = " << e.second;
          }
        }
      }
    }

    if (cluster->GetDeviceSet()) {
      for (const auto dev : cluster->GetDeviceSet()->devices()) {
        LOG(INFO) << "Device name= " << dev->name() << "Pased name= "
                  << DeviceNameUtils::ParsedNameToString(dev->parsed_name());
      }
    }
  }

  LOG(INFO) << "item: " << item.id;
  if (!item.feed.empty()) {
    LOG(INFO) << offset << "Feeds  :";
    for (const auto& f : item.feed) {
      const auto& shape = f.second.shape();
      LOG(INFO) << offset2 << f.first << " = shaped " << shape.DebugString();
    }
  } else {
    LOG(INFO) << offset << "No Feeds";
  }
  if (!item.fetch.empty()) {
    LOG(INFO) << offset << "Fetches  :";
    for (const auto& f : item.fetch) {
      LOG(INFO) << offset2 << f;
    }
  } else {
    LOG(INFO) << offset << "No Fetches";
  }

  if (!item.init_ops.empty()) {
    LOG(INFO) << offset << "init ops  :";
    for (const auto& f : item.init_ops) {
      LOG(INFO) << offset2 << f;
    }
  } else {
    LOG(INFO) << offset << "No init ops";
  }
  LOG(INFO) << "Save Op = " << item.save_op;
  LOG(INFO) << "Restore Op = " << item.restore_op;
  LOG(INFO) << "save_restore_loc_tensor = " << item.save_restore_loc_tensor;
  if (!item.keep_ops.empty()) {
    LOG(INFO) << offset << "keep ops  :";
    for (const auto& f : item.keep_ops) {
      LOG(INFO) << offset2 << f;
    }
  } else {
    LOG(INFO) << offset << "No keep ops";
  }
}

Status TRTOptimizationPass::Optimize(grappler::Cluster* cluster,
                                     const grappler::GrapplerItem& item,
                                     GraphDef* optimized_graph) {
  VLOG(1) << "Called TRTOptimization Pass " << name_;
  // This is a hack to workaround optimizer issue. MetaOptimizer calls
  // optimization passes on function objects as well, we should not modify
  // generated funcdefs! This is fragile but we don't have any other option
  // until framework fixes it.
  if (item.id != "tf_graph") {
    VLOG(1) << "Called TRTOptimization Pass " << name_
            << " on a grappler item with id=" << item.id
            << ", which is probably a function object (funcdef). "
            << "Skipping optimization because TensorRTOptimizer "
            << "should not be called on function objects.";
    *optimized_graph = item.graph;
    return Status::OK();
  }
  if (VLOG_IS_ON(3)) {
    LOG(INFO) << CurrentStackTrace();
    PrintDebugInfo(cluster, item);
  }

  if (use_calibration_ && precision_mode_ != TrtPrecisionMode::INT8) {
    VLOG(1) << "Calibration with FP32 or FP16 is not implemented. "
            << "Falling back to use_calibration = False."
            << "Note that the default value of use_calibration is True.";
    use_calibration_ = false;
  }

  std::vector<string> nodes_to_preserve;
  for (const auto& n : item.NodesToPreserve()) {
    auto tokens = str_util::Split(n, ":");
    string s = tokens.at(0);
    for (int i = 1; i < tokens.size() - 1; ++i) {
      StrAppend(&s, ":", tokens.at(i));
    }
    int dumm_port = -1;
    // If the last token is not an integer, it must be part of the name.
    // Otherwise it is port number.
    if (tokens.size() > 1 &&
        !strings::safe_strto32(tokens.back(), &dumm_port)) {  // non-absl ok
      StrAppend(&s, ":", tokens.back());
    }
    nodes_to_preserve.push_back(s);
  }

  ConversionParams cp;
  cp.grappler_item = &item;
  cp.output_names = &nodes_to_preserve;
  cp.trt_logger_name = trt_logger_name_;
  cp.max_batch_size = maximum_batch_size_;
  cp.max_workspace_size_bytes = max_workspace_size_bytes_;
  cp.output_graph_def = optimized_graph;
  cp.precision_mode = precision_mode_;
  cp.minimum_segment_size = minimum_segment_size_;
  cp.cluster = cluster;
  cp.is_dyn_op = is_dynamic_op_;
  cp.max_cached_engines = max_cached_batches_;
  cp.use_calibration = use_calibration_;
  cp.use_implicit_batch = use_implicit_batch_;
  cp.allow_build_at_runtime = allow_build_at_runtime_;
  auto status = ConvertAfterShapes(cp);
  VLOG(1) << "Returning from " << name_;
  return status;
}

void TRTOptimizationPass::Feedback(grappler::Cluster* cluster,
                                   const grappler::GrapplerItem& item,
                                   const GraphDef& optimized_graph,
                                   double result) {}

class VerboseCustomGraphOptimizerRegistrar
    : public grappler::CustomGraphOptimizerRegistrar {
 public:
  VerboseCustomGraphOptimizerRegistrar(
      const grappler::CustomGraphOptimizerRegistry::Creator& cr,
      const string& name)
      : grappler::CustomGraphOptimizerRegistrar(cr, name) {
    VLOG(1) << "Constructing a CustomOptimizationPass registration object for "
            << name;
  }
};

static VerboseCustomGraphOptimizerRegistrar TRTOptimizationPass_Registrar(
    []() {
      VLOG(1)
          << "Instantiating CustomOptimizationPass object TensorRTOptimizer";
      return new TRTOptimizationPass("TensorRTOptimizer");
    },
    ("TensorRTOptimizer"));

}  // namespace convert
}  // namespace tensorrt
}  // namespace tensorflow

#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
<|endoftext|>


<commit_before>/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/lite/delegates/gpu/common/transformations/add_bias.h"

#include "absl/memory/memory.h"
#include "absl/strings/str_cat.h"
#include "absl/types/any.h"
#include "tensorflow/lite/delegates/gpu/common/data_type.h"
#include "tensorflow/lite/delegates/gpu/common/model.h"
#include "tensorflow/lite/delegates/gpu/common/operations.h"
#include "tensorflow/lite/delegates/gpu/common/status.h"

namespace tflite {
namespace gpu {
namespace {

TransformResult FillBias(
    int output_channels,
    tflite::gpu::Tensor<Linear, DataType::FLOAT32>* biases) {
  if (biases->data.empty()) {
    *biases =
        MakeZeroTensor<Linear, DataType::FLOAT32>(Linear(output_channels));
    return {TransformStatus::APPLIED, "Added bias"};
  }
  if (biases->shape.v != output_channels) {
    float last_value = biases->data.back();
    biases->shape.v = output_channels;
    biases->data.resize(output_channels, last_value);
    return {TransformStatus::APPLIED, "Bias extended"};
  }
  return {TransformStatus::SKIPPED, ""};
}

class AddBias : public NodeTransformation {
 public:
  TransformResult ApplyToNode(Node* node, GraphFloat32* graph) final {
    if (node->operation.type == ToString(OperationType::CONVOLUTION_2D)) {
      auto& attr =
          absl::any_cast<Convolution2DAttributes&>(node->operation.attributes);
      return FillBias(attr.weights.shape.o, &attr.bias);
    }
    if (node->operation.type ==
        ToString(OperationType::CONVOLUTION_TRANSPOSED)) {
      auto& attr = absl::any_cast<ConvolutionTransposedAttributes&>(
          node->operation.attributes);
      return FillBias(attr.weights.shape.o, &attr.bias);
    }
    if (node->operation.type ==
        ToString(OperationType::DEPTHWISE_CONVOLUTION)) {
      auto& attr = absl::any_cast<DepthwiseConvolution2DAttributes&>(
          node->operation.attributes);
      return FillBias(attr.weights.shape.o * attr.weights.shape.i, &attr.bias);
    }
    if (node->operation.type == ToString(OperationType::FULLY_CONNECTED)) {
      auto& attr =
          absl::any_cast<FullyConnectedAttributes&>(node->operation.attributes);
      return FillBias(attr.weights.shape.o, &attr.bias);
    }
    return {TransformStatus::SKIPPED, ""};
  }
};

}  // namespace

std::unique_ptr<NodeTransformation> NewAddBias() {
  return absl::make_unique<AddBias>();
}

}  // namespace gpu
}  // namespace tflite
<commit_msg>Fixed add bias transformation. Added check for convolution with dynamic weights.<commit_after>/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/lite/delegates/gpu/common/transformations/add_bias.h"

#include "absl/memory/memory.h"
#include "absl/strings/str_cat.h"
#include "absl/types/any.h"
#include "tensorflow/lite/delegates/gpu/common/data_type.h"
#include "tensorflow/lite/delegates/gpu/common/model.h"
#include "tensorflow/lite/delegates/gpu/common/operations.h"
#include "tensorflow/lite/delegates/gpu/common/status.h"

namespace tflite {
namespace gpu {
namespace {

TransformResult FillBias(
    int output_channels,
    tflite::gpu::Tensor<Linear, DataType::FLOAT32>* biases) {
  if (biases->data.empty()) {
    *biases =
        MakeZeroTensor<Linear, DataType::FLOAT32>(Linear(output_channels));
    return {TransformStatus::APPLIED, "Added bias"};
  }
  if (biases->shape.v != output_channels) {
    float last_value = biases->data.back();
    biases->shape.v = output_channels;
    biases->data.resize(output_channels, last_value);
    return {TransformStatus::APPLIED, "Bias extended"};
  }
  return {TransformStatus::SKIPPED, ""};
}

class AddBias : public NodeTransformation {
 public:
  TransformResult ApplyToNode(Node* node, GraphFloat32* graph) final {
    if (node->operation.type == ToString(OperationType::CONVOLUTION_2D)) {
      if (graph->FindInputs(node->id).size() != 1) {
        return {TransformStatus::DECLINED,
                "This transformation is only applicable to conv with one "
                "runtime input."};
      }
      auto& attr =
          absl::any_cast<Convolution2DAttributes&>(node->operation.attributes);
      return FillBias(attr.weights.shape.o, &attr.bias);
    }
    if (node->operation.type ==
        ToString(OperationType::CONVOLUTION_TRANSPOSED)) {
      auto& attr = absl::any_cast<ConvolutionTransposedAttributes&>(
          node->operation.attributes);
      return FillBias(attr.weights.shape.o, &attr.bias);
    }
    if (node->operation.type ==
        ToString(OperationType::DEPTHWISE_CONVOLUTION)) {
      auto& attr = absl::any_cast<DepthwiseConvolution2DAttributes&>(
          node->operation.attributes);
      return FillBias(attr.weights.shape.o * attr.weights.shape.i, &attr.bias);
    }
    if (node->operation.type == ToString(OperationType::FULLY_CONNECTED)) {
      auto& attr =
          absl::any_cast<FullyConnectedAttributes&>(node->operation.attributes);
      return FillBias(attr.weights.shape.o, &attr.bias);
    }
    return {TransformStatus::SKIPPED, ""};
  }
};

}  // namespace

std::unique_ptr<NodeTransformation> NewAddBias() {
  return absl::make_unique<AddBias>();
}

}  // namespace gpu
}  // namespace tflite
<|endoftext|>


<commit_before>AliAnalysisTaskCorrForFlowFMD* AddCorrForFlowTaskFMD(TString name = "name", TString efficiencyFile = "",TString calibrationFile = "",const char* suffix = "")
{
    AliAnalysisManager *mgr = AliAnalysisManager::GetAnalysisManager();
    if (!mgr) {
        return 0x0;
    }
    if (!mgr->GetInputEventHandler()) {
        return 0x0;
    }
    TString fileName = AliAnalysisManager::GetCommonFileName();
    fileName += ":CorrForFlow";
    TString taskName = Form("%s_%s",name.Data(),suffix);

    Bool_t useEfficiency = kFALSE;
    if(!efficiencyFile.IsNull()) useEfficiency = kTRUE;

    Bool_t useCalibration = kFALSE;
    if(!calibrationFile.IsNull()) useCalibration = kTRUE;

    AliAnalysisTaskCorrForFlowFMD* task = new AliAnalysisTaskCorrForFlowFMD(taskName.Data(), useEfficiency, useCalibration);
    if(!task) return 0x0;
    mgr->AddTask(task);
    mgr->ConnectInput(task,0,mgr->GetCommonInputContainer());
    mgr->ConnectOutput(task,1,mgr->CreateContainer(Form("Charged_%s",taskName.Data()), TList::Class(), AliAnalysisManager::kOutputContainer, fileName.Data()));

    if(useEfficiency) {
      TObjArray* taskContainers = mgr->GetContainers();
      if(!taskContainers) { printf("Task containers does not exists!\n"); return NULL; }

      AliAnalysisDataContainer* efficiency = (AliAnalysisDataContainer*) taskContainers->FindObject("inputEfficiency");
      if(!efficiency) {
        if(efficiencyFile.Contains("alien://")) { gGrid->Connect("alien://"); }

        printf("input file name: %s \n", efficiencyFile.Data());

        TFile* efficiency_file = TFile::Open(efficiencyFile.Data(),"READ");
        if(!efficiency_file) { printf("Input file with efficiency not found!\n"); return NULL; }

        TList* efficiency_list = (TList*) efficiency_file->Get("EffAndFD");
        if(!efficiency_list) { printf("E-AddTask: Input list with efficiency not found!\n"); efficiency_file->ls(); return NULL; }

        AliAnalysisDataContainer* cInputEfficiency = mgr->CreateContainer("inputEfficiency",TList::Class(), AliAnalysisManager::kInputContainer);
        cInputEfficiency->SetData(efficiency_list);
        mgr->ConnectInput(task,1,cInputEfficiency);
      }
      else {
        mgr->ConnectInput(task,1,efficiency);
      }
    }

    if(useCalibration) {
      TObjArray* taskContainers = mgr->GetContainers();
      if(!taskContainers) { printf("Task containers does not exists!\n"); return NULL; }

      AliAnalysisDataContainer* calib = (AliAnalysisDataContainer*) taskContainers->FindObject("calib");
      if(!calib) {
        if(calibrationFile.Contains("alien://")) { gGrid->Connect("alien://"); }

        printf("input file name: %s \n", calibrationFile.Data());

        TFile* calib_file = TFile::Open(calibrationFile.Data(),"READ");
        if(!calib_file) { printf("Input file with AMPT centrality calibration not found!\n"); return NULL; }

        TH1D* calib_histo = (TH1D*)calib_file->Get("hcent");
        if(!calib_histo) { printf("E-AddTask: Input list with AMPT centrality calibration not found!\n"); calib_file->ls(); return NULL; }

        AliAnalysisDataContainer* cInputCalib = mgr->CreateContainer("calib",TH1D::Class(), AliAnalysisManager::kInputContainer);
        cInputCalib->SetData(calib_histo);
        if(useEfficiency) mgr->ConnectInput(task,2,cInputCalib);
        else mgr->ConnectInput(task,1,cInputCalib);
      }
      else {
        if(useEfficiency) mgr->ConnectInput(task,2,calib);
        else mgr->ConnectInput(task,1,calib);
      }
    }

    return task;
}
<commit_msg>Reading the efficiency file<commit_after>AliAnalysisTaskCorrForFlowFMD* AddCorrForFlowTaskFMD(TString name = "name", TString efficiencyFile = "",TString calibrationFile = "",const char* suffix = "")
{
    AliAnalysisManager *mgr = AliAnalysisManager::GetAnalysisManager();
    if (!mgr) {
        return 0x0;
    }
    if (!mgr->GetInputEventHandler()) {
        return 0x0;
    }
    TString fileName = AliAnalysisManager::GetCommonFileName();
    fileName += ":CorrForFlow";
    TString taskName = Form("%s_%s",name.Data(),suffix);

    Bool_t useEfficiency = kFALSE;
    if(!efficiencyFile.IsNull()) useEfficiency = kTRUE;

    Bool_t useCalibration = kFALSE;
    if(!calibrationFile.IsNull()) useCalibration = kTRUE;

    AliAnalysisTaskCorrForFlowFMD* task = new AliAnalysisTaskCorrForFlowFMD(taskName.Data(), useEfficiency, useCalibration);
    if(!task) return 0x0;
    mgr->AddTask(task);
    mgr->ConnectInput(task,0,mgr->GetCommonInputContainer());
    mgr->ConnectOutput(task,1,mgr->CreateContainer(Form("Charged_%s",taskName.Data()), TList::Class(), AliAnalysisManager::kOutputContainer, fileName.Data()));

    if(useEfficiency) {
      TObjArray* taskContainers = mgr->GetContainers();
      if(!taskContainers) { printf("Task containers does not exists!\n"); return NULL; }

      AliAnalysisDataContainer* efficiency = (AliAnalysisDataContainer*) taskContainers->FindObject("inputEfficiency");
      if(!efficiency) {
        if(efficiencyFile.Contains("alien://")) { gGrid->Connect("alien://"); }

        printf("input file name: %s \n", efficiencyFile.Data());

        TFile* efficiency_file = TFile::Open(efficiencyFile.Data(),"READ");
        if(!efficiency_file) { printf("Input file with efficiency not found!\n"); return NULL; }

        TList* efficiency_list = (TList*) efficiency_file->Get("Efficiency2D_wFD");
        if(!efficiency_list) { printf("E-AddTask: Input list with efficiency not found!\n"); efficiency_file->ls(); return NULL; }

        AliAnalysisDataContainer* cInputEfficiency = mgr->CreateContainer("inputEfficiency",TList::Class(), AliAnalysisManager::kInputContainer);
        cInputEfficiency->SetData(efficiency_list);
        mgr->ConnectInput(task,1,cInputEfficiency);
      }
      else {
        mgr->ConnectInput(task,1,efficiency);
      }
    }

    if(useCalibration) {
      TObjArray* taskContainers = mgr->GetContainers();
      if(!taskContainers) { printf("Task containers does not exists!\n"); return NULL; }

      AliAnalysisDataContainer* calib = (AliAnalysisDataContainer*) taskContainers->FindObject("calib");
      if(!calib) {
        if(calibrationFile.Contains("alien://")) { gGrid->Connect("alien://"); }

        printf("input file name: %s \n", calibrationFile.Data());

        TFile* calib_file = TFile::Open(calibrationFile.Data(),"READ");
        if(!calib_file) { printf("Input file with AMPT centrality calibration not found!\n"); return NULL; }

        TH1D* calib_histo = (TH1D*)calib_file->Get("hcent");
        if(!calib_histo) { printf("E-AddTask: Input list with AMPT centrality calibration not found!\n"); calib_file->ls(); return NULL; }

        AliAnalysisDataContainer* cInputCalib = mgr->CreateContainer("calib",TH1D::Class(), AliAnalysisManager::kInputContainer);
        cInputCalib->SetData(calib_histo);
        if(useEfficiency) mgr->ConnectInput(task,2,cInputCalib);
        else mgr->ConnectInput(task,1,cInputCalib);
      }
      else {
        if(useEfficiency) mgr->ConnectInput(task,2,calib);
        else mgr->ConnectInput(task,1,calib);
      }
    }

    return task;
}
<|endoftext|>


<commit_before>#include "tcp_socket.h"
#include "thread_manager.h"

#include <IPAddress.h>
#include <debug.h>
#include <Arduino.h>
#include <lwip/dns.h>
#include <lwip/err.h>

#define ONCE 0

void onError(void* arg, err_t err)
{
    TCPSocket* socket = reinterpret_cast<TCPSocket*>(arg);
    if(socket) {
        socket->OnError(err);
    }
}

void dnsFoundCallback(const char *name, const ip_addr_t *ipaddr, void *callback_arg)
{
    uint32_t threadId = (uint32_t)(*reinterpret_cast<IPAddress*>(callback_arg));
    if(ipaddr) {
        (*reinterpret_cast<IPAddress*>(callback_arg)) = ipaddr->addr;
    }
    ThreadManager::GetInstance().ResumeThread(threadId);
}

err_t onTcpConnect(void* arg, tcp_pcb* tpcb,  err_t err)
{
    TCPSocket* socket = reinterpret_cast<TCPSocket*>(arg);
    if(socket) {
        return socket->OnConnect(tpcb, err);
    }

    return ERR_VAL;
}

err_t onTCPRecv(void *arg, struct tcp_pcb *tpcb, struct pbuf *pb, err_t err)
{
    TCPSocket* socket = reinterpret_cast<TCPSocket*>(arg);
    if(socket) {
        return socket->OnTCPRecv(tpcb, pb, err);
    }

    return ERR_VAL;
}

err_t onTCPSent(void *arg, struct tcp_pcb *tpcb, uint16_t len)
{
    TCPSocket* socket = reinterpret_cast<TCPSocket*>(arg);
    if(socket) {
        return socket->OnTCPSent(tpcb, len);
    }

    return ERR_VAL;
}

void onSendError(void *arg)
{
    TCPSocket* socket = reinterpret_cast<TCPSocket*>(arg);
    if(socket) {
        return socket->OnError(ERR_ABRT);
    }
}

TCPSocket::TCPSocket()
: m_buf_offset(0)
, m_buf(0)
, m_sentSize(0)
, m_suspendedThread(-1)
, m_suspendedRecvThread(-1)
, m_suspendedSendThread(-1)
{
	Initialize();
}

TCPSocket::~TCPSocket()
{
	Close();
}

bool TCPSocket::Bind(const String& host, int port)
{
	if(m_socket == INVALID_SOCKET)
	{
		return false;
	}

    ip_addr_t ipaddr;
	if(host.length() == 0)
	{
        ipaddr = ip_addr_any;
	}
	else
	{
       IPAddress ipaddress(ThreadManager::GetInstance().GetCurrentThreadId());
       if(!ipaddress.fromString(host)) {
          err_t err = dns_gethostbyname(host.c_str(), &ipaddr, &dnsFoundCallback, &ipaddress);
          if(err == ERR_INPROGRESS) {
               ThreadManager::GetInstance().SuspendThread(ThreadManager::GetInstance().GetCurrentThreadId());
               // will return here when dns_found_callback fires
          } else if(err != ERR_OK) {
              return false;
          }
       }

       ipaddr.addr = ipaddress;
	}

	return tcp_bind(m_socket, &ipaddr, port) == ERR_OK;
}

bool TCPSocket::Connect(const String& host, int port)
{
	if(m_socket == INVALID_SOCKET || host.length() == 0)
	{
		return false;
	}

    ip_addr_t ipaddr;
    if(host.length() == 0)
    {
        ipaddr = ip_addr_any;
    }
    else
    {
        IPAddress ipaddress;
        if(!ipaddress.fromString(host)) {
            err_t err = dns_gethostbyname(host.c_str(), &ipaddr, &dnsFoundCallback, &ipaddress);
            if(err == ERR_INPROGRESS) {
                m_suspendedThread = ThreadManager::GetInstance().GetCurrentThreadId();
                ThreadManager::GetInstance().SuspendThread(m_suspendedThread);
            } else if(err != ERR_OK) {
                return false;
            }
        }

        ipaddr.addr = ipaddress;
    }

    err_t err = tcp_connect(m_socket, &ipaddr, port, &onTcpConnect);
    if(err == ERR_OK || err == ERR_INPROGRESS) {
        m_suspendedThread = ThreadManager::GetInstance().GetCurrentThreadId();
        ThreadManager::GetInstance().SuspendThread(m_suspendedThread);
    }
    return m_lastError == ERR_OK;
}

bool TCPSocket::Send(char* data, int len)
{
    //wait until other thread will send data
    while(m_suspendedSendThread > -1) {
        ThreadManager::GetInstance().DelayThread(ThreadManager::GetInstance().GetCurrentThreadId(), 100);
    }

    while(len) {
        if(m_socket == INVALID_SOCKET)
        {
            return false;
        }

        size_t room = tcp_sndbuf(m_socket);
        size_t will_send = (room < len) ? room : len;

        err_t err = tcp_write(m_socket, data, will_send, 0);
        if(err != ERR_OK) {
            return false;
        }

        data += will_send;
        len -= will_send;
        m_sentSize = 0;
        tcp_output(m_socket);

        // waiting to send tcp data
        while(m_sentSize != will_send && m_socket != INVALID_SOCKET) {
            os_timer_disarm(&m_sendErrorTimer);
            os_timer_setfn(&m_sendErrorTimer, &onSendError, this);
            os_timer_arm(&m_sendErrorTimer, 1500, ONCE);
            m_suspendedSendThread = ThreadManager::GetInstance().GetCurrentThreadId();
            ThreadManager::GetInstance().SuspendThread(m_suspendedSendThread);
        }
    }

	return true;
}

bool TCPSocket::Recv(char* data, int len)
{
    while(len) {
        if(m_socket == INVALID_SOCKET)
        {
            return false;
        }

        if(!m_buf) {
            m_suspendedRecvThread = ThreadManager::GetInstance().GetCurrentThreadId();
            ThreadManager::GetInstance().SuspendThread(m_suspendedRecvThread);
        }

        if(m_buf) {
            size_t max_size = m_buf->tot_len - m_buf_offset;
            size_t size = (len < max_size) ? len : max_size;
            size_t buf_size = m_buf->len - m_buf_offset;
            size_t copy_size = (size < buf_size) ? size : buf_size;
            os_memcpy(data, reinterpret_cast<char*>(m_buf->payload) + m_buf_offset, copy_size);
            data += copy_size;
            len -= copy_size;
            Consume(copy_size);
        }
    }

	return true;
}
void TCPSocket::Consume(size_t size)
{
    ptrdiff_t left = m_buf->len - m_buf_offset - size;
    if(left > 0) {
        m_buf_offset += size;
    } else if(!m_buf->next) {
        if(m_socket) tcp_recved(m_socket, m_buf->len);
        pbuf_free(m_buf);
        m_buf = 0;
        m_buf_offset = 0;
    } else {
        auto head = m_buf;
        m_buf = m_buf->next;
        m_buf_offset = 0;
        pbuf_ref(m_buf);
        if(m_socket) tcp_recved(m_socket, head->len);
        pbuf_free(head);
    }
}

bool TCPSocket::Close()
{
	if(m_socket == INVALID_SOCKET)
	{
		return false;
	}

    tcp_arg(m_socket, NULL);
    tcp_sent(m_socket, NULL);
    tcp_recv(m_socket, NULL);
    tcp_err(m_socket, NULL);
    tcp_close(m_socket);
    m_socket = INVALID_SOCKET;
    return true;
}

void TCPSocket::GetIPPort(String& ip, int& port)
{
    IPAddress address(m_socket->remote_ip.addr);
    ip = address.toString();
    port = m_socket->remote_port;
}

int TCPSocket::GetSocket()
{
	return *(int*)m_socket;
}

void TCPSocket::Initialize()
{
	m_socket = tcp_new();
    if(m_socket) {
        tcp_arg(m_socket, this);
        tcp_err(m_socket, &onError);
    }
}

int TCPSocket::GetMaxBufferSize()
{
	if(m_socket == INVALID_SOCKET)
	{
		return 0;
	}

	return m_socket->snd_buf;
}

void TCPSocket::OnError(err_t err)
{
    m_lastError = err;
    tcp_arg(m_socket, NULL);
    tcp_sent(m_socket, NULL);
    tcp_recv(m_socket, NULL);
    tcp_err(m_socket, NULL);
    m_socket = INVALID_SOCKET;
    if(m_suspendedThread != -1) {
        ThreadManager::GetInstance().ResumeThread(m_suspendedThread);
        m_suspendedThread = -1;
    } else if(m_suspendedRecvThread != -1) {
        ThreadManager::GetInstance().ResumeThread(m_suspendedRecvThread);
        m_suspendedRecvThread = -1;
    } else if(m_suspendedSendThread != -1) {
        ThreadManager::GetInstance().ResumeThread(m_suspendedSendThread);
        m_suspendedSendThread = -1;
    }
}

err_t TCPSocket::OnConnect(tcp_pcb* tpcb, err_t err)
{
    m_lastError = ERR_OK;
    m_socket = tpcb;
    tcp_setprio(m_socket, TCP_PRIO_MIN);
    tcp_arg(m_socket, this);
    tcp_recv(m_socket, (tcp_recv_fn)&onTCPRecv);
    tcp_sent(m_socket, &onTCPSent);
    tcp_err(m_socket, &onError);
    if(m_suspendedThread != -1) {
        ThreadManager::GetInstance().ResumeThread(m_suspendedThread);
        m_suspendedThread = -1;
    }
    return ERR_OK;
}

err_t TCPSocket::OnTCPRecv(tcp_pcb* tpcb, pbuf* pb, err_t err)
{
    err_t ret;
    if(!pb) {
        tcp_arg(m_socket, NULL);
        tcp_sent(m_socket, NULL);
        tcp_recv(m_socket, NULL);
        tcp_err(m_socket, NULL);
        tcp_abort(m_socket);
        m_socket = INVALID_SOCKET;
        ret = ERR_ABRT;
    } else {
        m_lastError = err;
        if(m_buf) {
            pbuf_cat(m_buf, pb);
        } else {
            m_buf = pb;
        }
        ret = ERR_OK;
    }
    if(m_suspendedRecvThread != -1) {
        ThreadManager::GetInstance().ResumeThread(m_suspendedRecvThread);
        m_suspendedRecvThread = -1;
    }
    return ret;
}

err_t TCPSocket::OnTCPSent(tcp_pcb* tpcb, uint16_t len)
{
    m_lastError = ERR_OK;
    m_sentSize = len;
    os_timer_disarm(&m_sendErrorTimer);
    if(m_suspendedSendThread != -1) {
        ThreadManager::GetInstance().ResumeThread(m_suspendedSendThread);
        m_suspendedSendThread = -1;
    }
    return ERR_OK;
}
<commit_msg>fix freeze<commit_after>#include "tcp_socket.h"
#include "thread_manager.h"

#include <IPAddress.h>
#include <debug.h>
#include <Arduino.h>
#include <lwip/dns.h>
#include <lwip/err.h>

#define ONCE 0

void onError(void* arg, err_t err)
{
    TCPSocket* socket = reinterpret_cast<TCPSocket*>(arg);
    if(socket) {
        socket->OnError(err);
    }
}

void dnsFoundCallback(const char *name, const ip_addr_t *ipaddr, void *callback_arg)
{
    uint32_t threadId = (uint32_t)(*reinterpret_cast<IPAddress*>(callback_arg));
    if(ipaddr) {
        (*reinterpret_cast<IPAddress*>(callback_arg)) = ipaddr->addr;
    }
    ThreadManager::GetInstance().ResumeThread(threadId);
}

err_t onTcpConnect(void* arg, tcp_pcb* tpcb,  err_t err)
{
    TCPSocket* socket = reinterpret_cast<TCPSocket*>(arg);
    if(socket) {
        return socket->OnConnect(tpcb, err);
    }

    return ERR_VAL;
}

err_t onTCPRecv(void *arg, struct tcp_pcb *tpcb, struct pbuf *pb, err_t err)
{
    TCPSocket* socket = reinterpret_cast<TCPSocket*>(arg);
    if(socket) {
        return socket->OnTCPRecv(tpcb, pb, err);
    }

    return ERR_VAL;
}

err_t onTCPSent(void *arg, struct tcp_pcb *tpcb, uint16_t len)
{
    TCPSocket* socket = reinterpret_cast<TCPSocket*>(arg);
    if(socket) {
        return socket->OnTCPSent(tpcb, len);
    }

    return ERR_VAL;
}

void onSendError(void *arg)
{
    TCPSocket* socket = reinterpret_cast<TCPSocket*>(arg);
    if(socket) {
        return socket->OnError(ERR_ABRT);
    }
}

TCPSocket::TCPSocket()
: m_buf_offset(0)
, m_buf(0)
, m_sentSize(0)
, m_suspendedThread(-1)
, m_suspendedRecvThread(-1)
, m_suspendedSendThread(-1)
{
	Initialize();
}

TCPSocket::~TCPSocket()
{
	Close();
}

bool TCPSocket::Bind(const String& host, int port)
{
	if(m_socket == INVALID_SOCKET)
	{
		return false;
	}

    ip_addr_t ipaddr;
	if(host.length() == 0)
	{
        ipaddr = ip_addr_any;
	}
	else
	{
       IPAddress ipaddress(ThreadManager::GetInstance().GetCurrentThreadId());
       if(!ipaddress.fromString(host)) {
          err_t err = dns_gethostbyname(host.c_str(), &ipaddr, &dnsFoundCallback, &ipaddress);
          if(err == ERR_INPROGRESS) {
               ThreadManager::GetInstance().SuspendThread(ThreadManager::GetInstance().GetCurrentThreadId());
               // will return here when dns_found_callback fires
          } else if(err != ERR_OK) {
              return false;
          }
       }

       ipaddr.addr = ipaddress;
	}

	return tcp_bind(m_socket, &ipaddr, port) == ERR_OK;
}

bool TCPSocket::Connect(const String& host, int port)
{
	if(m_socket == INVALID_SOCKET || host.length() == 0)
	{
		return false;
	}

    ip_addr_t ipaddr;
    if(host.length() == 0)
    {
        ipaddr = ip_addr_any;
    }
    else
    {
        IPAddress ipaddress;
        if(!ipaddress.fromString(host)) {
            err_t err = dns_gethostbyname(host.c_str(), &ipaddr, &dnsFoundCallback, &ipaddress);
            if(err == ERR_INPROGRESS) {
                m_suspendedThread = ThreadManager::GetInstance().GetCurrentThreadId();
                ThreadManager::GetInstance().SuspendThread(m_suspendedThread);
            } else if(err != ERR_OK) {
                return false;
            }
        }

        ipaddr.addr = ipaddress;
    }

    err_t err = tcp_connect(m_socket, &ipaddr, port, &onTcpConnect);
    if(err == ERR_OK || err == ERR_INPROGRESS) {
        m_suspendedThread = ThreadManager::GetInstance().GetCurrentThreadId();
        ThreadManager::GetInstance().SuspendThread(m_suspendedThread);
    }
    return m_lastError == ERR_OK;
}

bool TCPSocket::Send(char* data, int len)
{
    //wait until other thread will send data
    while(m_suspendedSendThread > -1 && m_socket != INVALID_SOCKET) {
        ThreadManager::GetInstance().DelayThread(ThreadManager::GetInstance().GetCurrentThreadId(), 100);
    }

    while(len) {
        if(m_socket == INVALID_SOCKET)
        {
            return false;
        }

        size_t room = tcp_sndbuf(m_socket);
        size_t will_send = (room < len) ? room : len;

        err_t err = tcp_write(m_socket, data, will_send, 0);
        if(err != ERR_OK) {
            return false;
        }

        data += will_send;
        len -= will_send;
        m_sentSize = 0;
        tcp_output(m_socket);

        os_timer_disarm(&m_sendErrorTimer);
        os_timer_setfn(&m_sendErrorTimer, &onSendError, this);
        os_timer_arm(&m_sendErrorTimer, 1000, ONCE);

        // waiting to send tcp data
        while(m_sentSize != will_send && m_socket != INVALID_SOCKET) {
            m_suspendedSendThread = ThreadManager::GetInstance().GetCurrentThreadId();
            ThreadManager::GetInstance().SuspendThread(m_suspendedSendThread);
        }
    }

	return true;
}

bool TCPSocket::Recv(char* data, int len)
{
    while(len) {
        if(m_socket == INVALID_SOCKET)
        {
            return false;
        }

        if(!m_buf) {
            m_suspendedRecvThread = ThreadManager::GetInstance().GetCurrentThreadId();
            ThreadManager::GetInstance().SuspendThread(m_suspendedRecvThread);
        }

        if(m_buf) {
            size_t max_size = m_buf->tot_len - m_buf_offset;
            size_t size = (len < max_size) ? len : max_size;
            size_t buf_size = m_buf->len - m_buf_offset;
            size_t copy_size = (size < buf_size) ? size : buf_size;
            os_memcpy(data, reinterpret_cast<char*>(m_buf->payload) + m_buf_offset, copy_size);
            data += copy_size;
            len -= copy_size;
            Consume(copy_size);
        }
    }

	return true;
}
void TCPSocket::Consume(size_t size)
{
    ptrdiff_t left = m_buf->len - m_buf_offset - size;
    if(left > 0) {
        m_buf_offset += size;
    } else if(!m_buf->next) {
        if(m_socket) tcp_recved(m_socket, m_buf->len);
        pbuf_free(m_buf);
        m_buf = 0;
        m_buf_offset = 0;
    } else {
        auto head = m_buf;
        m_buf = m_buf->next;
        m_buf_offset = 0;
        pbuf_ref(m_buf);
        if(m_socket) tcp_recved(m_socket, head->len);
        pbuf_free(head);
    }
}

bool TCPSocket::Close()
{
	if(m_socket == INVALID_SOCKET)
	{
		return false;
	}

    tcp_arg(m_socket, NULL);
    tcp_sent(m_socket, NULL);
    tcp_recv(m_socket, NULL);
    tcp_err(m_socket, NULL);
    tcp_close(m_socket);
    m_socket = INVALID_SOCKET;
    return true;
}

void TCPSocket::GetIPPort(String& ip, int& port)
{
    IPAddress address(m_socket->remote_ip.addr);
    ip = address.toString();
    port = m_socket->remote_port;
}

int TCPSocket::GetSocket()
{
	return *(int*)m_socket;
}

void TCPSocket::Initialize()
{
	m_socket = tcp_new();
    if(m_socket) {
        tcp_arg(m_socket, this);
        tcp_err(m_socket, &onError);
    }
}

int TCPSocket::GetMaxBufferSize()
{
	if(m_socket == INVALID_SOCKET)
	{
		return 0;
	}

	return m_socket->snd_buf;
}

void TCPSocket::OnError(err_t err)
{
    m_lastError = err;
    tcp_arg(m_socket, NULL);
    tcp_sent(m_socket, NULL);
    tcp_recv(m_socket, NULL);
    tcp_err(m_socket, NULL);
    m_socket = INVALID_SOCKET;
    int thread_id;
    if(m_suspendedThread != -1) {
        thread_id = m_suspendedThread;
        m_suspendedThread = -1;
        ThreadManager::GetInstance().ResumeThread(thread_id);
    } else if(m_suspendedRecvThread != -1) {
        thread_id = m_suspendedRecvThread;
        m_suspendedRecvThread = -1;
        ThreadManager::GetInstance().ResumeThread(thread_id);
    } else if(m_suspendedSendThread != -1) {
        thread_id = m_suspendedSendThread;
        m_suspendedSendThread = -1;
        ThreadManager::GetInstance().ResumeThread(thread_id);
    }
}

err_t TCPSocket::OnConnect(tcp_pcb* tpcb, err_t err)
{
    m_lastError = ERR_OK;
    m_socket = tpcb;
    tcp_setprio(m_socket, TCP_PRIO_MIN);
    tcp_arg(m_socket, this);
    tcp_recv(m_socket, (tcp_recv_fn)&onTCPRecv);
    tcp_sent(m_socket, &onTCPSent);
    tcp_err(m_socket, &onError);
    if(m_suspendedThread != -1) {
        int thread_id = m_suspendedThread;
        m_suspendedThread = -1;
        ThreadManager::GetInstance().ResumeThread(thread_id);
    }
    return ERR_OK;
}

err_t TCPSocket::OnTCPRecv(tcp_pcb* tpcb, pbuf* pb, err_t err)
{
    err_t ret;
    if(!pb) {
        tcp_arg(m_socket, NULL);
        tcp_sent(m_socket, NULL);
        tcp_recv(m_socket, NULL);
        tcp_err(m_socket, NULL);
        tcp_abort(m_socket);
        m_socket = INVALID_SOCKET;
        ret = ERR_ABRT;
    } else {
        m_lastError = err;
        if(m_buf) {
            pbuf_cat(m_buf, pb);
        } else {
            m_buf = pb;
        }
        ret = ERR_OK;
    }
    if(m_suspendedRecvThread != -1) {
        int thread_id = m_suspendedRecvThread;
        m_suspendedRecvThread = -1;
        ThreadManager::GetInstance().ResumeThread(thread_id);
    }
    return ret;
}

err_t TCPSocket::OnTCPSent(tcp_pcb* tpcb, uint16_t len)
{
    m_lastError = ERR_OK;
    m_sentSize = len;
    os_timer_disarm(&m_sendErrorTimer);
    if(m_suspendedSendThread != -1) {
        int thread_id = m_suspendedSendThread;
        m_suspendedSendThread = -1;
        ThreadManager::GetInstance().ResumeThread(thread_id);
    }
    return ERR_OK;
}
<|endoftext|>


<commit_before>/*
 *  Copyright (c) 2012 The WebM project authors. All Rights Reserved.
 *
 *  Use of this source code is governed by a BSD-style license
 *  that can be found in the LICENSE file in the root of the source
 *  tree. An additional intellectual property rights grant can be found
 *  in the file PATENTS.  All contributing project authors may
 *  be found in the AUTHORS file in the root of the source tree.
 */

#include <string>

#include "third_party/googletest/src/include/gtest/gtest.h"

#include "./vpx_config.h"
#include "./y4menc.h"
#include "test/md5_helper.h"
#include "test/util.h"
#include "test/y4m_video_source.h"

namespace {

using std::string;

static const unsigned int kWidth  = 160;
static const unsigned int kHeight = 90;
static const unsigned int kFrames = 10;

struct Y4mTestParam {
  const char *filename;
  unsigned int bit_depth;
  vpx_img_fmt format;
  const char *md5raw;
};

const Y4mTestParam kY4mTestVectors[] = {
  {"park_joy_90p_8_420.y4m", 8, VPX_IMG_FMT_I420,
    "e5406275b9fc6bb3436c31d4a05c1cab"},
  {"park_joy_90p_8_422.y4m", 8, VPX_IMG_FMT_I422,
    "284a47a47133b12884ec3a14e959a0b6"},
  {"park_joy_90p_8_444.y4m", 8, VPX_IMG_FMT_I444,
    "90517ff33843d85de712fd4fe60dbed0"},
  {"park_joy_90p_10_420.y4m", 10, VPX_IMG_FMT_I42016,
    "63f21f9f717d8b8631bd2288ee87137b"},
  {"park_joy_90p_10_422.y4m", 10, VPX_IMG_FMT_I42216,
    "48ab51fb540aed07f7ff5af130c9b605"},
  {"park_joy_90p_10_444.y4m", 10, VPX_IMG_FMT_I44416,
    "067bfd75aa85ff9bae91fa3e0edd1e3e"},
  {"park_joy_90p_12_420.y4m", 12, VPX_IMG_FMT_I42016,
    "9e6d8f6508c6e55625f6b697bc461cef"},
  {"park_joy_90p_12_422.y4m", 12, VPX_IMG_FMT_I42216,
    "b239c6b301c0b835485be349ca83a7e3"},
  {"park_joy_90p_12_444.y4m", 12, VPX_IMG_FMT_I44416,
    "5a6481a550821dab6d0192f5c63845e9"},
};

static void write_image_file(const vpx_image_t *img, FILE *file) {
  int plane, y;
  for (plane = 0; plane < 3; ++plane) {
    const unsigned char *buf = img->planes[plane];
    const int stride = img->stride[plane];
    const int bytes_per_sample = (img->fmt & VPX_IMG_FMT_HIGHBITDEPTH) ? 2 : 1;
    const int h = (plane ? (img->d_h + img->y_chroma_shift) >>
                   img->y_chroma_shift : img->d_h);
    const int w = (plane ? (img->d_w + img->x_chroma_shift) >>
                   img->x_chroma_shift : img->d_w);
    for (y = 0; y < h; ++y) {
      fwrite(buf, bytes_per_sample, w, file);
      buf += stride;
    }
  }
}

class Y4mVideoSourceTest
    : public ::testing::TestWithParam<Y4mTestParam>,
      public ::libvpx_test::Y4mVideoSource {
 protected:
  Y4mVideoSourceTest() : Y4mVideoSource("", 0, 0) {}

  virtual ~Y4mVideoSourceTest() {
    CloseSource();
  }

  virtual void Init(const std::string &file_name, int limit) {
    file_name_ = file_name;
    start_ = 0;
    limit_ = limit;
    frame_ = 0;
    Begin();
  }

  // Checks y4m header information
  void HeaderChecks(unsigned int bit_depth, vpx_img_fmt_t fmt) {
    ASSERT_TRUE(input_file_ != NULL);
    ASSERT_EQ(y4m_.pic_w, (int)kWidth);
    ASSERT_EQ(y4m_.pic_h, (int)kHeight);
    ASSERT_EQ(img()->d_w, kWidth);
    ASSERT_EQ(img()->d_h, kHeight);
    ASSERT_EQ(y4m_.bit_depth, bit_depth);
    ASSERT_EQ(y4m_.vpx_fmt, fmt);
    if (fmt == VPX_IMG_FMT_I420 || fmt == VPX_IMG_FMT_I42016) {
      ASSERT_EQ(y4m_.bps, (int)y4m_.bit_depth * 3 / 2);
      ASSERT_EQ(img()->x_chroma_shift, 1U);
      ASSERT_EQ(img()->y_chroma_shift, 1U);
    }
    if (fmt == VPX_IMG_FMT_I422 || fmt == VPX_IMG_FMT_I42216) {
      ASSERT_EQ(y4m_.bps, (int)y4m_.bit_depth * 2);
      ASSERT_EQ(img()->x_chroma_shift, 1U);
      ASSERT_EQ(img()->y_chroma_shift, 0U);
    }
    if (fmt == VPX_IMG_FMT_I444 || fmt == VPX_IMG_FMT_I44416) {
      ASSERT_EQ(y4m_.bps, (int)y4m_.bit_depth * 3);
      ASSERT_EQ(img()->x_chroma_shift, 0U);
      ASSERT_EQ(img()->y_chroma_shift, 0U);
    }
  }

  // Checks MD5 of the raw frame data
  void Md5Check(const string &expected_md5) {
    ASSERT_TRUE(input_file_ != NULL);
    libvpx_test::MD5 md5;
    for (unsigned int i = start_; i < limit_; i++) {
      md5.Add(img());
      Next();
    }
    ASSERT_EQ(string(md5.Get()), expected_md5);
  }
};

TEST_P(Y4mVideoSourceTest, SourceTest) {
  const Y4mTestParam t = GetParam();
  Init(t.filename, kFrames);
  HeaderChecks(t.bit_depth, t.format);
  Md5Check(t.md5raw);
}

INSTANTIATE_TEST_CASE_P(C, Y4mVideoSourceTest,
                        ::testing::ValuesIn(kY4mTestVectors));

class Y4mVideoWriteTest
    : public Y4mVideoSourceTest {
 protected:
  Y4mVideoWriteTest() {}

  virtual ~Y4mVideoWriteTest() {
    delete tmpfile_;
    input_file_ = NULL;
  }

  void ReplaceInputFile(FILE *input_file) {
    CloseSource();
    frame_ = 0;
    input_file_ = input_file;
    rewind(input_file_);
    ReadSourceToStart();
  }

  // Writes out a y4m file and then reads it back
  void WriteY4mAndReadBack() {
    ASSERT_TRUE(input_file_ != NULL);
    char buf[Y4M_BUFFER_SIZE] = {0};
    const struct VpxRational framerate = {y4m_.fps_n, y4m_.fps_d};
    tmpfile_ = new libvpx_test::TempOutFile;
    ASSERT_TRUE(tmpfile_->file() != NULL);
    y4m_write_file_header(buf, sizeof(buf),
                          kWidth, kHeight,
                          &framerate, y4m_.vpx_fmt,
                          y4m_.bit_depth);
    fputs(buf, tmpfile_->file());
    for (unsigned int i = start_; i < limit_; i++) {
      y4m_write_frame_header(buf, sizeof(buf));
      fputs(buf, tmpfile_->file());
      write_image_file(img(), tmpfile_->file());
      Next();
    }
    ReplaceInputFile(tmpfile_->file());
  }

  virtual void Init(const std::string &file_name, int limit) {
    Y4mVideoSourceTest::Init(file_name, limit);
    WriteY4mAndReadBack();
  }
  libvpx_test::TempOutFile *tmpfile_;
};

TEST_P(Y4mVideoWriteTest, WriteTest) {
  const Y4mTestParam t = GetParam();
  Init(t.filename, kFrames);
  HeaderChecks(t.bit_depth, t.format);
  Md5Check(t.md5raw);
}

INSTANTIATE_TEST_CASE_P(C, Y4mVideoWriteTest,
                        ::testing::ValuesIn(kY4mTestVectors));
}  // namespace
<commit_msg>y4m_test: fix segfault if test files are missing<commit_after>/*
 *  Copyright (c) 2012 The WebM project authors. All Rights Reserved.
 *
 *  Use of this source code is governed by a BSD-style license
 *  that can be found in the LICENSE file in the root of the source
 *  tree. An additional intellectual property rights grant can be found
 *  in the file PATENTS.  All contributing project authors may
 *  be found in the AUTHORS file in the root of the source tree.
 */

#include <string>

#include "third_party/googletest/src/include/gtest/gtest.h"

#include "./vpx_config.h"
#include "./y4menc.h"
#include "test/md5_helper.h"
#include "test/util.h"
#include "test/y4m_video_source.h"

namespace {

using std::string;

static const unsigned int kWidth  = 160;
static const unsigned int kHeight = 90;
static const unsigned int kFrames = 10;

struct Y4mTestParam {
  const char *filename;
  unsigned int bit_depth;
  vpx_img_fmt format;
  const char *md5raw;
};

const Y4mTestParam kY4mTestVectors[] = {
  {"park_joy_90p_8_420.y4m", 8, VPX_IMG_FMT_I420,
    "e5406275b9fc6bb3436c31d4a05c1cab"},
  {"park_joy_90p_8_422.y4m", 8, VPX_IMG_FMT_I422,
    "284a47a47133b12884ec3a14e959a0b6"},
  {"park_joy_90p_8_444.y4m", 8, VPX_IMG_FMT_I444,
    "90517ff33843d85de712fd4fe60dbed0"},
  {"park_joy_90p_10_420.y4m", 10, VPX_IMG_FMT_I42016,
    "63f21f9f717d8b8631bd2288ee87137b"},
  {"park_joy_90p_10_422.y4m", 10, VPX_IMG_FMT_I42216,
    "48ab51fb540aed07f7ff5af130c9b605"},
  {"park_joy_90p_10_444.y4m", 10, VPX_IMG_FMT_I44416,
    "067bfd75aa85ff9bae91fa3e0edd1e3e"},
  {"park_joy_90p_12_420.y4m", 12, VPX_IMG_FMT_I42016,
    "9e6d8f6508c6e55625f6b697bc461cef"},
  {"park_joy_90p_12_422.y4m", 12, VPX_IMG_FMT_I42216,
    "b239c6b301c0b835485be349ca83a7e3"},
  {"park_joy_90p_12_444.y4m", 12, VPX_IMG_FMT_I44416,
    "5a6481a550821dab6d0192f5c63845e9"},
};

static void write_image_file(const vpx_image_t *img, FILE *file) {
  int plane, y;
  for (plane = 0; plane < 3; ++plane) {
    const unsigned char *buf = img->planes[plane];
    const int stride = img->stride[plane];
    const int bytes_per_sample = (img->fmt & VPX_IMG_FMT_HIGHBITDEPTH) ? 2 : 1;
    const int h = (plane ? (img->d_h + img->y_chroma_shift) >>
                   img->y_chroma_shift : img->d_h);
    const int w = (plane ? (img->d_w + img->x_chroma_shift) >>
                   img->x_chroma_shift : img->d_w);
    for (y = 0; y < h; ++y) {
      fwrite(buf, bytes_per_sample, w, file);
      buf += stride;
    }
  }
}

class Y4mVideoSourceTest
    : public ::testing::TestWithParam<Y4mTestParam>,
      public ::libvpx_test::Y4mVideoSource {
 protected:
  Y4mVideoSourceTest() : Y4mVideoSource("", 0, 0) {}

  virtual ~Y4mVideoSourceTest() {
    CloseSource();
  }

  virtual void Init(const std::string &file_name, int limit) {
    file_name_ = file_name;
    start_ = 0;
    limit_ = limit;
    frame_ = 0;
    Begin();
  }

  // Checks y4m header information
  void HeaderChecks(unsigned int bit_depth, vpx_img_fmt_t fmt) {
    ASSERT_TRUE(input_file_ != NULL);
    ASSERT_EQ(y4m_.pic_w, (int)kWidth);
    ASSERT_EQ(y4m_.pic_h, (int)kHeight);
    ASSERT_EQ(img()->d_w, kWidth);
    ASSERT_EQ(img()->d_h, kHeight);
    ASSERT_EQ(y4m_.bit_depth, bit_depth);
    ASSERT_EQ(y4m_.vpx_fmt, fmt);
    if (fmt == VPX_IMG_FMT_I420 || fmt == VPX_IMG_FMT_I42016) {
      ASSERT_EQ(y4m_.bps, (int)y4m_.bit_depth * 3 / 2);
      ASSERT_EQ(img()->x_chroma_shift, 1U);
      ASSERT_EQ(img()->y_chroma_shift, 1U);
    }
    if (fmt == VPX_IMG_FMT_I422 || fmt == VPX_IMG_FMT_I42216) {
      ASSERT_EQ(y4m_.bps, (int)y4m_.bit_depth * 2);
      ASSERT_EQ(img()->x_chroma_shift, 1U);
      ASSERT_EQ(img()->y_chroma_shift, 0U);
    }
    if (fmt == VPX_IMG_FMT_I444 || fmt == VPX_IMG_FMT_I44416) {
      ASSERT_EQ(y4m_.bps, (int)y4m_.bit_depth * 3);
      ASSERT_EQ(img()->x_chroma_shift, 0U);
      ASSERT_EQ(img()->y_chroma_shift, 0U);
    }
  }

  // Checks MD5 of the raw frame data
  void Md5Check(const string &expected_md5) {
    ASSERT_TRUE(input_file_ != NULL);
    libvpx_test::MD5 md5;
    for (unsigned int i = start_; i < limit_; i++) {
      md5.Add(img());
      Next();
    }
    ASSERT_EQ(string(md5.Get()), expected_md5);
  }
};

TEST_P(Y4mVideoSourceTest, SourceTest) {
  const Y4mTestParam t = GetParam();
  Init(t.filename, kFrames);
  HeaderChecks(t.bit_depth, t.format);
  Md5Check(t.md5raw);
}

INSTANTIATE_TEST_CASE_P(C, Y4mVideoSourceTest,
                        ::testing::ValuesIn(kY4mTestVectors));

class Y4mVideoWriteTest
    : public Y4mVideoSourceTest {
 protected:
  Y4mVideoWriteTest() : tmpfile_(NULL) {}

  virtual ~Y4mVideoWriteTest() {
    delete tmpfile_;
    input_file_ = NULL;
  }

  void ReplaceInputFile(FILE *input_file) {
    CloseSource();
    frame_ = 0;
    input_file_ = input_file;
    rewind(input_file_);
    ReadSourceToStart();
  }

  // Writes out a y4m file and then reads it back
  void WriteY4mAndReadBack() {
    ASSERT_TRUE(input_file_ != NULL);
    char buf[Y4M_BUFFER_SIZE] = {0};
    const struct VpxRational framerate = {y4m_.fps_n, y4m_.fps_d};
    tmpfile_ = new libvpx_test::TempOutFile;
    ASSERT_TRUE(tmpfile_->file() != NULL);
    y4m_write_file_header(buf, sizeof(buf),
                          kWidth, kHeight,
                          &framerate, y4m_.vpx_fmt,
                          y4m_.bit_depth);
    fputs(buf, tmpfile_->file());
    for (unsigned int i = start_; i < limit_; i++) {
      y4m_write_frame_header(buf, sizeof(buf));
      fputs(buf, tmpfile_->file());
      write_image_file(img(), tmpfile_->file());
      Next();
    }
    ReplaceInputFile(tmpfile_->file());
  }

  virtual void Init(const std::string &file_name, int limit) {
    Y4mVideoSourceTest::Init(file_name, limit);
    WriteY4mAndReadBack();
  }
  libvpx_test::TempOutFile *tmpfile_;
};

TEST_P(Y4mVideoWriteTest, WriteTest) {
  const Y4mTestParam t = GetParam();
  Init(t.filename, kFrames);
  HeaderChecks(t.bit_depth, t.format);
  Md5Check(t.md5raw);
}

INSTANTIATE_TEST_CASE_P(C, Y4mVideoWriteTest,
                        ::testing::ValuesIn(kY4mTestVectors));
}  // namespace
<|endoftext|>