Editable Initial State - Attempt 17

README.md

⬅️ Back

In the previous attempt, we got automatic initial state working by way of reading the initial state on time change, but ahead of evaluation. It wasn't all roses though, as it didn't fit some of the common use-cases such as..

1. Edits ahead of simulation
2. Animated initial state
3. Transition to and from active/passive

..and now it does!

Edits Ahead of Simulation

Animated Initial State

Transitioning between Active & Passive

And that's it! We did it!

Test	Status
DG	✔️
Parallel	✔️
Manipulator	✔️
Parented to each other	✔️
Children preserved	✔️
ChannelBox	✔️
Keyframe	✔️
Autokeyframe	✔️
Hierarchy	✔️

Notes

There's one question we will at some point want to revisit.

Which should evaluate first, the solver or rigid?

1. If the rigid evaluates first, we'll output the previously solved value, but deliver the current animation of the rigid, like whether it is active and whatever it's damping value is.
2. If the solver evalates first, we'll advance the simulation using the previously evaluated animation, but deliver the current position to the rigid.

It'll have to be one way or the other. Here's a visual to clarify the current implementation.

Summary

This (hopefully!) marks the end of our journey. Still needs testing in the actual production example. I've managed to reproduce the Bullet integration behavior in an efficient and robust way.

• Efficient in that the feature does not negatively impact simulation performance, all of what we do extra happens during passive states and before simulation.
• Robust in that we rely on a Maya callback triggering ahead of any and all evaluation, but after time change. So long we this callback is consistent - which it appears to be - we should expect a robust result.

It also won't affect offline use, such as via mayapy or during render. And it works identically under DG. And isn't affected by external nodes, like that weird issue in the previous attempt.

Overview

At a glance, here's how the solver works.

def before_time_changed(rigid):
    if rigid.update_initial_state:
        rigid.initial_state = cmds.getAttr(rigid + ".input_world_matrix")


class Rigid:
    next_state = ...
    input_world_matrix = ...

    def __init__(self, transform):
        self.uid = generate_unique_uid()

        # Record initial state from connected transform
        self.initial_state = cmds.getAttr(self + ".input_world_matrix")

    def compute(self, datablock):
        simulation_step = datablock.input_value(Rigid.next_state)

        if simulation_step == 0:
            self.update_start_state(datablock)
        else:
            self.update_current_state(datablock)

    def update_start_state(self, datablock):
        pass

    def update_current_state(self, datablock):
        pass

    # Todo

1. Assume there is a series of boxes in the scene prior to the solver existing
2. On creation of the solver, the current state of the world is read and stored as initial state
3. On time change, one of two things happen..
   • If moving from frame 0 to frame 1..
     1. Record initial state, ahead of evaluating anything
     2. Pull Rigid plugs, for use on next simulaton step
     3. Step simulation
     4. Output updated simulation to connected Maya transform
   • If moving from frame 1 to frame 2..
     1. Pull rigid plugs, for use on next simulation step
     2. Step simulation

Next Steps

Aside from actually integrating this with the production project, there's still room for optimisation, as we are currently recording the initial state on every frame, even when nothing has changed. I expect we can hook into a worldMatrixChanged callback to provide our system with a hint about whether or not we need to read and write any matrices.

Build & Run

From Powershell.

cd c:\
git clone https://gist.github.com/1c3323d5b3f81a1ece8ed5f86b045068.git initialStateNode17
mkdir initialStateNode17/build
cd initialStateNode10/build
$env:DEVKIT_LOCATION="c:\path\to\your\devkit"
cmake .. -G Ninja
ninja

Then from the Script Editor.

import os
from maya import cmds

fname = r"c:\initialState17\build\initialStateNode.mll"
cmds.loadPlugin(fname)

time = "time1"
solver = cmds.createNode("solverNode", parent=cmds.createNode("transform", name="solver"))
cmds.setAttr(solver + ".gravity", 0, -9.81, 0)
cmds.setAttr(solver + ".startTime", cmds.getAttr(time + ".outTime"))
cmds.connectAttr(time + ".outTime", solver + ".currentTime")

parent = None
for i in range(3):
    cube, _ = cmds.polyCube()

    if parent:
        cmds.parent(cube, parent, relative=True)
        cmds.move(3, 0, 0, relative=True)
    else:
        cmds.move(-3 * i, 15, -7)

    rigid = cmds.createNode("rigidNode", parent=cube)

    cmds.connectAttr(time + ".outTime", rigid + ".currentTime")
    cmds.connectAttr(cube + ".parentInverseMatrix", rigid + ".inParentInverseMatrix")
    cmds.connectAttr(cube + ".worldMatrix[0]", rigid + ".inWorldMatrix")
    cmds.connectAttr(rigid + ".outputTranslateX", cube + ".translateX")
    cmds.connectAttr(rigid + ".outputTranslateY", cube + ".translateY")
    cmds.connectAttr(rigid + ".outputTranslateZ", cube + ".translateZ")

    # Connect to scene
    cmds.connectAttr(rigid + ".startState", solver + ".startState[%d]" % i)
    cmds.connectAttr(rigid + ".currentState", solver + ".currentState[%d]" % i)
    cmds.connectAttr(solver + ".nextState[%d]" % i, rigid + ".nextState")

    parent = cube


cmds.file(new=True, force=True)
cmds.unloadPlugin(os.path.basename(fname))

CMakeLists.txt

cmake_minimum_required(VERSION 2.8)

include($ENV{DEVKIT_LOCATION}/cmake/pluginEntry.cmake)

set(PROJECT_NAME initialStateNode)

set(SOURCE_FILES
   main.cpp
)

set(LIBRARIES
    OpenMaya
    OpenMayaRender
    OpenMayaAnim
    OpenMayaUI
    Foundation
)

add_compile_definitions("/D_DEBUG /Zi")

build_plugin()

main.cpp

#include <string.h>
#include <unordered_map>
#include <vector>
#include <cassert>

#include <maya/MPxLocatorNode.h>

#include <maya/MFnNumericAttribute.h>
#include <maya/MFnMatrixAttribute.h>
#include <maya/MFnUnitAttribute.h>
#include <maya/MFnPlugin.h>
#include <maya/MFnDagNode.h>
#include <maya/MFnTransform.h>
#include <maya/MDagPath.h>
#include <maya/MVector.h>
#include <maya/MQuaternion.h>

#include <maya/MTypeId.h> 
#include <maya/MPlug.h>
#include <maya/MString.h>
#include <maya/MVector.h>
#include <maya/MTransformationMatrix.h>
#include <maya/MMatrix.h>
#include <maya/MDataBlock.h>
#include <maya/MDataHandle.h>
#include <maya/MStreamUtils.h>
#include <maya/MNodeMessage.h>
#include <maya/MFnDependencyNode.h>
#include <maya/MFnMatrixData.h>
#include <maya/MCallbackIdArray.h>
#include <maya/MGlobal.h>
#include <maya/MMessage.h>
#include <maya/MDGMessage.h>
#include <maya/MDagMessage.h>
#include <maya/MAnimControl.h>
#include <maya/MEvaluationNode.h>
#include <maya/MEvaluationManager.h>

#define Print                           MStreamUtils::stdErrorStream
#define Translate(mat)                  MTransformationMatrix(mat).translation(MSpace::kTransform)
#define AddAttribute(attr)              assert(addAttribute(attr) == MS::kSuccess);
#define AttributeAffects(attr, affects) assert(attributeAffects(attr, affects) == MS::kSuccess);

static bool debug() {
    static bool debugInfoExists { false };
    const bool debugInfo = MGlobal::optionVarIntValue("debugInfo", &debugInfoExists);
    return debugInfoExists ? debugInfo : false;
}

// Global counter, for a unique UID per rigid
static int gUniqueIdCounter { 1000 };


struct Pose {
    MMatrix matrix;
    MVector velocity;
};

struct RigidBody {
    int       uid { -1 };
    MString   name;

    // Example animated property
    float     damping { 0.5f };

    MVector   velocity { 0, 0, 0 };
    bool      active { true };
    MMatrix   parentInverseMatrix {};
    MMatrix   restMatrix {};
    MMatrix   worldMatrix {};

    // Keep parent matrix from getting cancelled
    // out during standstill
    bool      worldMatrixChanged { true };

    // Monitor changes to initial state
    bool      updateInitialState { false };

    bool      drivenBySimulation { true };

    std::vector<Pose> history;
};

using RigidMap = std::unordered_map<int, RigidBody>;

// For purposes of demonstration, rather
// than passing data between nodes, we'll
// use a global variable and pass an index
// into this map between nodes instead.
static RigidMap gRigids;


// Helper functions

inline MDataHandle InputValue(MDataBlock& datablock, MObject attr) {
    MStatus status;
    MDataHandle handle = datablock.inputValue(attr, &status);
    assert(status == MS::kSuccess);
    return handle;
}

inline MDataHandle OutputValue(MDataBlock& datablock, MObject attr) {
    MStatus status;
    MDataHandle handle = datablock.outputValue(attr, &status);
    assert(status == MS::kSuccess);
    return handle;
}


#include "rigid.hpp"
#include "solver.hpp"


MStatus initializePlugin(MObject obj) { 
    MStatus   status;
    MFnPlugin plugin(obj, PLUGIN_COMPANY, "3.0", "Any");

    status = plugin.registerNode(SolverNode::name,
                                 SolverNode::id,
                                 SolverNode::creator,
                                 SolverNode::initialize,
                                 MPxNode::kLocatorNode,
                                 &SolverNode::drawDbClassification);

    status = plugin.registerNode(RigidNode::name,
                                 RigidNode::id,
                                 RigidNode::creator,
                                 RigidNode::initialize,
                                 MPxNode::kLocatorNode,
                                 &RigidNode::drawDbClassification);

    Print() << "==============\n";
    Print() << "Initialising\n";

    return status;
}


MStatus uninitializePlugin(MObject obj) {
    MStatus   status;
    MFnPlugin plugin(obj);

    status = plugin.deregisterNode(RigidNode::id);
    status = plugin.deregisterNode(SolverNode::id);

    Print() << "Deinitialising.\n";
    Print() << "==============\n";

    return status;
}

rigid.hpp

/**
 * Individual rigid body
 *
 */
class RigidNode : public MPxLocatorNode {
public:
    RigidNode() {};
    ~RigidNode() override;

    MStatus compute(const MPlug&, MDataBlock&) override;
    MStatus computeStartState(const MPlug&, MDataBlock&);
    MStatus computeCurrentState(const MPlug&, MDataBlock&);
    MStatus computeOutputTranslate(const MPlug&, MDataBlock&);
    MStatus computeOutputWorldMatrix(const MPlug&, MDataBlock&);

    MStatus updateStartState(const MPlug&, MDataBlock&);
    MStatus updateCurrentState(const MPlug&, MDataBlock&);

    static  void* creator() { return new RigidNode(); }
    static  MStatus initialize();
    void    postConstructor();
    void    reset(MDataBlock&);
    int     uid() { return _uid; }
    bool    getInternalValue(const MPlug&, MDataHandle&) override;
    bool    isPassiveOutput(const MPlug&) const override;

    // Does nothing
    MStatus preEvaluation(const MDGContext&, const MEvaluationNode&) override;
    MStatus setDependentsDirty(const MPlug&, MPlugArray&) override;

public:
    static  MString       name;
    static  MTypeId       id;
    static const MString  drawDbClassification;
    static const MString  drawRegistrantId;

    // Attributes
    static  MObject   aUid;
    static  MObject   aDamping;
    static  MObject   aActive;
    static  MObject   aCurrentTime;
    static  MObject   aCachedWorldMatrix;
    static  MObject   aInputWorldMatrix;
    static  MObject   aInputParentInverseMatrix;
    static  MObject   aOutputWorldMatrix;
    static  MObject   aOutputTranslate;
    static  MObject   aOutputTranslateX;
    static  MObject   aOutputTranslateY;
    static  MObject   aOutputTranslateZ;

    // <-- Pull on this to fetch latest goods from the solver
    static MObject aNextState;

    // --> Deliver initial state of this rigid to the solver
    static MObject aStartState;

    // --> Deliver animated values, like damping, to the solver
    static MObject aCurrentState;

private:
    int               _uid { -1 };
    MCallbackIdArray  _callbackIds;
};


MString       RigidNode::name("rigidNode");
MTypeId       RigidNode::id(0x80012);
const MString RigidNode::drawDbClassification("drawdb/geometry/RigidNode");
const MString RigidNode::drawRegistrantId("RigidNodePlugin");
MObject       RigidNode::aUid;
MObject       RigidNode::aDamping;
MObject       RigidNode::aActive;
MObject       RigidNode::aCachedWorldMatrix;
MObject       RigidNode::aInputWorldMatrix;
MObject       RigidNode::aInputParentInverseMatrix;
MObject       RigidNode::aOutputWorldMatrix;
MObject       RigidNode::aOutputTranslate;
MObject       RigidNode::aOutputTranslateX;
MObject       RigidNode::aOutputTranslateY;
MObject       RigidNode::aOutputTranslateZ;
MObject       RigidNode::aCurrentTime;
MObject       RigidNode::aNextState;
MObject       RigidNode::aStartState;
MObject       RigidNode::aCurrentState;



RigidNode::~RigidNode() {
    // Any node created will have an entry in this map
    gRigids.erase(_uid);

    for (unsigned int ii=0; ii<_callbackIds.length(); ii++) 
        MMessage::removeCallback(_callbackIds[ii]);

    _callbackIds.clear();
}


MStatus RigidNode::initialize() {
    MStatus status;
    MFnNumericAttribute numFn;
    MFnMatrixAttribute matFn;
    MFnUnitAttribute uniFn;

    aUid = numFn.create("uid", "uid", MFnNumericData::kInt, -1, &status);
    numFn.setWritable(false);
    numFn.setStorable(false);
    numFn.setInternal(true);  // We'll manage this ourselves

    aNextState = numFn.create("nextState", "nest", MFnNumericData::kInt, -1, &status);
    numFn.setStorable(false);
    numFn.setReadable(false);

    aStartState = numFn.create("startState", "stst", MFnNumericData::kInt, -1, &status);
    numFn.setStorable(false);

    aCurrentState = numFn.create("currentState", "cust", MFnNumericData::kInt, -1, &status);
    numFn.setStorable(false);

    aDamping = numFn.create("damping", "damp", MFnNumericData::kFloat, 0.5f, &status);
    numFn.setKeyable(true);

    aActive = numFn.create("active", "acti", MFnNumericData::kBoolean, true, &status);
    numFn.setKeyable(true);

    aCurrentTime = uniFn.create("currentTime", "cuti", MFnUnitAttribute::kTime, 0.0, &status);
    uniFn.setReadable(false);
    uniFn.setStorable(false);
    uniFn.setHidden(true);

    aInputParentInverseMatrix = matFn.create("inParentInverseMatrix", "ipim");
    matFn.setStorable(false);
    matFn.setReadable(false);

    aInputWorldMatrix = matFn.create("inWorldMatrix", "inma");
    matFn.setStorable(false);
    matFn.setReadable(false);

    aCachedWorldMatrix = matFn.create("inCachedMatrix", "incm");
    matFn.setStorable(true);
    matFn.setReadable(false);
    matFn.setConnectable(false);
    matFn.setHidden(true);

    aOutputWorldMatrix = matFn.create("outWorldMatrix", "ouma");
    matFn.setWritable(false);

    aOutputTranslateX = uniFn.create("outputTranslateX", "outx", MFnUnitAttribute::kDistance, 0.0, &status);
    uniFn.setWritable(false);
    uniFn.setStorable(false);
    uniFn.setChannelBox(true);
    aOutputTranslateY = uniFn.create("outputTranslateY", "outy", MFnUnitAttribute::kDistance, 0.0, &status);
    uniFn.setWritable(false);
    uniFn.setStorable(false);
    uniFn.setChannelBox(true);
    aOutputTranslateZ = uniFn.create("outputTranslateZ", "outz", MFnUnitAttribute::kDistance, 0.0, &status);
    uniFn.setWritable(false);
    uniFn.setStorable(false);
    uniFn.setChannelBox(true);
    aOutputTranslate = numFn.create("outputTranslate", "outr", aOutputTranslateX, aOutputTranslateY, aOutputTranslateZ);

    AddAttribute(aUid);
    AddAttribute(aActive);
    AddAttribute(aDamping);
    AddAttribute(aNextState);
    AddAttribute(aStartState);
    AddAttribute(aCurrentState);
    AddAttribute(aCurrentTime);
    AddAttribute(aCachedWorldMatrix);
    AddAttribute(aInputWorldMatrix);
    AddAttribute(aInputParentInverseMatrix);
    AddAttribute(aOutputWorldMatrix);
    AddAttribute(aOutputTranslate);

    AttributeAffects(aInputParentInverseMatrix, aOutputWorldMatrix);
    AttributeAffects(aNextState,                aOutputWorldMatrix);
    AttributeAffects(aDamping,                  aOutputWorldMatrix);
    AttributeAffects(aActive,                  aOutputWorldMatrix);
    AttributeAffects(aCurrentTime,              aOutputWorldMatrix);

    AttributeAffects(aCurrentTime,              aOutputTranslate);
    AttributeAffects(aOutputWorldMatrix,        aOutputTranslate);

    return status;
}


bool RigidNode::isPassiveOutput(const MPlug& plug) const {
    /**
     * Let anything connected to these still be editable
     *
     */
    return plug == aOutputTranslateX ||
           plug == aOutputTranslateY ||
           plug == aOutputTranslateZ;
}


MStatus RigidNode::preEvaluation(const MDGContext& context, const MEvaluationNode& evaluationNode) {
    return MS::kSuccess;
}


MStatus RigidNode::setDependentsDirty(const MPlug& plugBeingDirtied, MPlugArray& affectedPlugs) {
    return MPxLocatorNode::setDependentsDirty(plugBeingDirtied, affectedPlugs);
}


/**
 * Respond to changes in time
 *
 * When time has changed, but before any evaluation takes place,
 * read and record the world matrix of this rigid as the new initial state.
 *
 */
static void onTimeChanged(MTime&, void* clientData) {
    MStatus status { MS::kSuccess };

    RigidNode* node = static_cast<RigidNode*>(clientData);
    RigidBody& rigid = gRigids.at(node->uid());

    if (rigid.updateInitialState) {
        MPlug worldMatrixPlug { node->thisMObject(), RigidNode::aInputWorldMatrix };
        MObject obj = worldMatrixPlug.asMObject(&status);
        assert(status == MS::kSuccess);

        MMatrix mat = MFnMatrixData(obj).matrix(&status);
        assert(status == MS::kSuccess);

        rigid.restMatrix = mat;
        rigid.worldMatrix = mat;
        rigid.worldMatrixChanged = true;

        // Cache value for when time rewinds and file save
        MPlug restMatrixPlug { node->thisMObject(), RigidNode::aCachedWorldMatrix };
        obj = MFnMatrixData().create(mat);
        restMatrixPlug.setMObject(obj);

        Print() << "(+) Successfully updated initial state\n";
        rigid.updateInitialState = false;
    }
}


/**
 * Respond to newly connected inWorldMatrix
 *
 * In order to preserve the position of the connected Maya transform,
 * we'll read and store whatever the world matrix was at the time of connecting.
 *
 * If we don't do this, then the `outputTranslate` would be zero once
 * evaluation starts.
 *
 */
void onAttributeChanged(MNodeMessage::AttributeMessage msg, MPlug& plug, MPlug& otherPlug, void*) {
    MStatus status { MS::kSuccess };

    if (msg & MNodeMessage::kConnectionMade) {
        if (plug == RigidNode::aInputWorldMatrix) {
            MObject obj = plug.asMObject(&status);
            assert(status == MS::kSuccess);

            MMatrix mat = MFnMatrixData(obj).matrix(&status);
            assert(status == MS::kSuccess);

            MPlug restMatrixPlug { plug.node(), RigidNode::aCachedWorldMatrix };
            obj = MFnMatrixData().create(mat);
            restMatrixPlug.setMObject(obj);

            Print() << "(+) Successfully cached newly connected inWorldMatrix\n";
        }
    }
}


void RigidNode::postConstructor() {
    MStatus status { MS::kSuccess };

    // Find ourselves amongst rigids hosted by the solver
    _uid = ++gUniqueIdCounter;

    // Initialise our data, in a place accessible to our solver
    gRigids[_uid] = { _uid };

    MObject tempobj { thisMObject() };

    _callbackIds.append(MDGMessage::addForceUpdateCallback(onTimeChanged, this, &status));
    assert(status == MS::kSuccess);

    _callbackIds.append(MNodeMessage::addAttributeChangedCallback(tempobj, onAttributeChanged, this, &status));
    assert(status == MS::kSuccess);
}


bool RigidNode::getInternalValue(const MPlug& plug, MDataHandle& datahandle) {
    if (plug == aUid) {
        datahandle.setInt(_uid);
        return true;
    }
    return false;
}


/**
 * Restore initial state
 *
 */
MStatus RigidNode::computeStartState(const MPlug& plug, MDataBlock& datablock) {
    Print() << "RigidNode:: Rigid[" << _uid << "].computeStartState() ===============\n";
    MStatus status { MS::kSuccess };

    OutputValue(datablock, aStartState).set(_uid);

    return status;
}


/**
 * Read values that change over time
 *
 * We don't really have any in this example, but for example linearDamping
 *
 */
MStatus RigidNode::computeCurrentState(const MPlug& plug, MDataBlock& datablock) {
    Print() << "RigidNode:: Rigid[" << _uid << "].computeCurrentState() ===============\n";
    MStatus status { MS::kSuccess };

    OutputValue(datablock, aCurrentState).set(_uid);

    return status;
}


MStatus RigidNode::updateStartState(const MPlug& plug, MDataBlock& datablock) {
    Print() << "RigidNode:: Rigid[" << _uid << "].updateStartState() ===============\n";
    MStatus status { MS::kSuccess };

    RigidBody& rigid = gRigids.at(_uid);

    // Restore initial state, do not evaluate anything
    MMatrix worldMatrix = InputValue(datablock, aCachedWorldMatrix).asMatrix();
    rigid.restMatrix = worldMatrix;
    rigid.worldMatrix = rigid.restMatrix;
    rigid.active = InputValue(datablock, aActive).asBool();
    rigid.damping = InputValue(datablock, aDamping).asFloat();

    // Tell the next frame we'd like for *this* to be
    // our new initial state, whatever it may be.
    rigid.updateInitialState = true;

    MFnDagNode fn { thisMObject() };
    rigid.name = fn.name();

    return status;
}


MStatus RigidNode::updateCurrentState(const MPlug& plug, MDataBlock& datablock) {
    Print() << "RigidNode:: Rigid[" << _uid << "].updateCurrentState() ===============\n";
    MStatus status { MS::kSuccess };

    RigidBody& rigid = gRigids.at(_uid);
    rigid.active = InputValue(datablock, aActive).asBool();
    rigid.damping = InputValue(datablock, aDamping).asFloat();

    if (!rigid.active) {
        rigid.updateInitialState = true;
    }

    return status;
}


MStatus RigidNode::computeOutputWorldMatrix(const MPlug& plug, MDataBlock& datablock) {
    Print() << "RigidNode::computeOutputWorldMatrix()\n";

    MStatus status { MS::kSuccess };
    RigidBody& rigid = gRigids.at(_uid);

    const int simulatedFrame = InputValue(datablock, aNextState).asInt();

    if (simulatedFrame <= 0) {
        updateStartState(plug, datablock);
    }
    else {
        updateCurrentState(plug, datablock);
    }

    // Preserve hierarchy when manipulating an object during standstill
    // Try removing this and rotate the parent of a cube, you will find that
    // children remain stationary as the parent moves around. That's bad.
    if (rigid.worldMatrixChanged) {
        rigid.parentInverseMatrix = InputValue(datablock, aInputParentInverseMatrix).asMatrix();
        rigid.worldMatrixChanged = false;
    }

    MDataHandle outputHandle = OutputValue(datablock, aOutputWorldMatrix);
    outputHandle.setMMatrix(rigid.worldMatrix);

    datablock.setClean(plug);

    return status;
}


MStatus RigidNode::computeOutputTranslate(const MPlug& plug, MDataBlock& datablock) {
    Print() << "RigidNode::computeOutputTranslate() -------------------------------\n";

    MStatus status { MS::kSuccess };

    RigidBody& rigid = gRigids.at(_uid);

    MMatrix worldMatrix = InputValue(datablock, aOutputWorldMatrix).asMatrix();
    MMatrix localMatrix = worldMatrix * rigid.parentInverseMatrix;

    MTransformationMatrix tm { localMatrix };
    MVector translate = tm.translation(MSpace::kTransform);
    MDataHandle tx = OutputValue(datablock, aOutputTranslateX),
                ty = OutputValue(datablock, aOutputTranslateY),
                tz = OutputValue(datablock, aOutputTranslateZ);

    tx.set(translate.x);
    ty.set(translate.y);
    tz.set(translate.z);

    datablock.setClean(plug);

    return status;
}



MStatus RigidNode::compute(const MPlug& plug, MDataBlock& datablock) {
    MStatus status { MS::kUnknownParameter };

    // Print() << "RigidNode::compute() - " << plug.name() << "\n";

    // For some reason, we need to pull on time, even though we don't need it
    // in order for its dependents to continue to be dirtied when it changes.
    // The alternative is having it drawn in the channel box or AE.
    InputValue(datablock, aCurrentTime).asTime();

    if      (plug == aStartState) {
        status = computeStartState(plug, datablock);
    }

    else if (plug == aCurrentState) {
        status = computeCurrentState(plug, datablock);
    }

    else if (plug == aOutputWorldMatrix) {
        status = computeOutputWorldMatrix(plug, datablock);
    }

    else if (plug == aOutputTranslate ||
             plug == aOutputTranslateX ||
             plug == aOutputTranslateY ||
             plug == aOutputTranslateZ) {
        status = computeOutputTranslate(plug, datablock);
    }

    return status;
}

solver.hpp

/**
 * Pull on all rigid at once, and solve them in one big batch.
 * 
 */

class SolverNode : public MPxLocatorNode {
public:
    SolverNode() {};
    ~SolverNode() override;

    static  void* creator() { return new SolverNode(); }
    static  MStatus initialize();
    MStatus compute(const MPlug& plug, MDataBlock& data) override;
    MStatus computeOutputRigids(const MPlug& plug, MDataBlock& data);
    MStatus addRigids(const MPlug& plug, MDataBlock& datablock);
    MStatus updateRigids(const MPlug& plug, MDataBlock& datablock);
    MStatus setDependentsDirty(const MPlug&, MPlugArray&) override;
    void    postConstructor() override;

    MStatus stepSimulation(MDataBlock& datablock);
    MStatus restoreSimulation(MDataBlock& datablock);

public:
    static MString        name;
    static MTypeId        id;
    static const MString  drawDbClassification;
    static const MString  drawRegistrantId;

    // Attributes
    static MObject   aStartTime;
    static MObject   aCurrentTime;

    static MObject   aGravity;
    static MObject   aGravityX;
    static MObject   aGravityY;
    static MObject   aGravityZ;

    // Initial state of each rigid, e.g. position and orientation
    static MObject   aStartState;

    // Current values of rigid plugs, e.g. damping
    static MObject   aCurrentState;

    // Resulting transforms after simulation
    static MObject   aNextState;

private:
    bool    _isStartTime { false };
    bool    _beforeStartTime { false };
    bool    _isFirstTime { false };
    MTime   _lastTime { };
    MTime   _currentTime { };
    double  _deltaTime { 0.0 };

    // 0-based time
    int     _startFrame { };
    int     _currentFrame { };
    int     _simulationFrame { 0 };
    MCallbackIdArray  _callbackIds;

};


MString       SolverNode::name("solverNode");
MTypeId       SolverNode::id(0x80013);
const MString SolverNode::drawDbClassification("drawdb/geometry/SolverNode");
const MString SolverNode::drawRegistrantId("SolverNodePlugin");
MObject       SolverNode::aCurrentState;
MObject       SolverNode::aStartState;
MObject       SolverNode::aNextState;
MObject       SolverNode::aStartTime;
MObject       SolverNode::aCurrentTime;
MObject       SolverNode::aGravity;
MObject       SolverNode::aGravityX;
MObject       SolverNode::aGravityY;
MObject       SolverNode::aGravityZ;


/**
 * To help separate print statements in the Output Window
 *
 */
static void onForceUpdate(MTime& time, void* clientData) {
    Print() << "\n";
    Print() << "-------------------------------------------------------------------\n";
    Print() << "(!) onForceUpdate() " << time.value() << "\n";
    Print() << "-------------------------------------------------------------------\n";
    Print() << "\n";
}


SolverNode::~SolverNode() {
    for (unsigned int ii=0; ii<_callbackIds.length(); ii++) 
        MMessage::removeCallback(_callbackIds[ii]);

    _callbackIds.clear();
}


MStatus SolverNode::initialize() {
    MStatus status;
    MFnNumericAttribute numFn;
    MFnUnitAttribute uniFn;

    aStartTime = uniFn.create("startTime", "stti", MFnUnitAttribute::kTime, 0.0, &status);
    uniFn.setKeyable(true);

    aCurrentTime = uniFn.create("currentTime", "cuti", MFnUnitAttribute::kTime, 0.0, &status);
    uniFn.setKeyable(true);

    aStartState = numFn.create("startState", "stst", MFnNumericData::kInt, -1, &status);
    numFn.setStorable(false);
    numFn.setReadable(true);
    numFn.setWritable(true);
    numFn.setArray(true);
    numFn.setDisconnectBehavior(MFnUnitAttribute::kDelete);

    aCurrentState = numFn.create("currentState", "cust", MFnNumericData::kInt, -1, &status);
    numFn.setStorable(false);
    numFn.setReadable(true);
    numFn.setWritable(true);
    numFn.setArray(true);
    numFn.setDisconnectBehavior(MFnUnitAttribute::kDelete);

    aNextState = numFn.create("nextState", "nest", MFnNumericData::kInt, -1, &status);
    numFn.setStorable(false);
    numFn.setWritable(false);
    numFn.setArray(true);
    numFn.setDisconnectBehavior(MFnUnitAttribute::kDelete);

    aGravityX = uniFn.create("gravityX", "grvx", MFnUnitAttribute::kDistance, 0.0, &status);
    aGravityY = uniFn.create("gravityY", "grvy", MFnUnitAttribute::kDistance, -981.0, &status);
    aGravityZ = uniFn.create("gravityZ", "grvz", MFnUnitAttribute::kDistance, 0.0, &status);
    aGravity = numFn.create("gravity", "grav", aGravityX, aGravityY, aGravityZ);
    numFn.setKeyable(true);

    AddAttribute(aStartTime);
    AddAttribute(aCurrentTime);
    AddAttribute(aGravity);
    AddAttribute(aCurrentState);
    AddAttribute(aStartState);
    AddAttribute(aNextState);

    AttributeAffects(aStartTime, aNextState);
    AttributeAffects(aCurrentTime, aNextState);
    AttributeAffects(aGravity, aNextState);
    AttributeAffects(aStartState, aNextState);
    AttributeAffects(aCurrentState, aNextState);

    return status;
}

void SolverNode::postConstructor() {
    MStatus status;

    /**
     * We aren't actually using this, but it's interesting to see
     * in what order this is called relative nextState etc.
     *
     */
    _callbackIds.append(
        MDGMessage::addForceUpdateCallback(onForceUpdate, this, &status)
    );
    assert(status == MS::kSuccess);
}


MStatus SolverNode::setDependentsDirty(const MPlug& plugBeingDirtied, MPlugArray& affectedPlugs) {
    return MPxLocatorNode::setDependentsDirty(plugBeingDirtied, affectedPlugs);
}


MStatus SolverNode::compute(const MPlug& plug, MDataBlock& datablock) {
    MStatus status { MS::kUnknownParameter };

    if (plug == aNextState && plug.isElement()){
        const MTime currentTime = InputValue(datablock, aCurrentTime).asTime();
        const MTime startTime = InputValue(datablock, aStartTime).asTime();

        _currentTime = currentTime;
        _startFrame = static_cast<int>(startTime.value());
        _currentFrame = static_cast<int>(currentTime.value()) - _startFrame;
        _deltaTime = (_currentTime - _lastTime).as(MTime::kSeconds);
        _isFirstTime = _isStartTime && _deltaTime > 0.0;
        _isStartTime = currentTime == startTime;
        _beforeStartTime = currentTime < startTime;

        status = computeOutputRigids(plug, datablock);

        _lastTime = currentTime;
    }

    datablock.setClean(plug);

    return status;
}


MStatus SolverNode::computeOutputRigids(const MPlug& plug, MDataBlock& datablock) {
    if (debug()) Print() << "SolverNode::computeOutputRigids()\n";
    MStatus status { MS::kSuccess };

    auto noChange = [&]() -> MStatus {
        MStatus status;

        const int index = plug.logicalIndex(&status);
        assert(status == MS::kSuccess);

        MArrayDataHandle nextStateHandle = datablock.outputArrayValue(aNextState, &status);
        assert(status == MS::kSuccess);

        nextStateHandle.jumpToElement(index);
        MDataHandle outputHandle = nextStateHandle.outputValue(&status);
        assert(status == MS::kSuccess);

        outputHandle.setInt(_currentFrame);

        return status;
    };

    if (_lastTime > _currentTime && _currentFrame <= 0) {
        for (auto& pair : gRigids) {
            RigidBody& rigid = pair.second;
            rigid.worldMatrixChanged = true;
            rigid.velocity = {0, 0, 0};
        }
        Print() << "(!) Solver rewound..\n";
    }

         if (_beforeStartTime) status = noChange();
    else if (_isStartTime)     status = addRigids(plug, datablock);
    else                       status = updateRigids(plug, datablock);

    // These should all have succeeded
    assert(status == MS::kSuccess);

    return status;
}


MStatus SolverNode::addRigids(const MPlug& plug, MDataBlock& datablock) {
    MStatus status { MS::kUnknownParameter };

    // In DG mode, Maya appears to cache previously-visited frames
    // resulting in compute() *sometimes* not being called.
    if (!MEvaluationManager::evaluationManagerActive(MDGContext::current())) {
        Print() << "ogs -reset\n";
        MGlobal::executeCommand("ogs -reset");
    }

    _simulationFrame = 0;

    const int index = plug.logicalIndex(&status);
    assert(status == MS::kSuccess);

    // Only consider rigids that are still connected
    MPlug startStatePlug { thisMObject(), aStartState };
    MPlug elementPlug = startStatePlug.elementByLogicalIndex(index, &status);
    assert(status == MS::kSuccess);

    if (!elementPlug.isConnected()) {
        Print() << "SolverNode: Disconnected index: " << index << "\n";
        return MS::kFailure; 
    }

    // Pull on *all* startStates
    MArrayDataHandle startStateHandle = datablock.inputArrayValue(aStartState, &status);
    assert(status == MS::kSuccess);

    MArrayDataHandle nextStateHandle = datablock.outputArrayValue(aNextState, &status);
    assert(status == MS::kSuccess);

    status = startStateHandle.jumpToElement(index);
    assert(status == MS::kSuccess);
    status = nextStateHandle.jumpToElement(index);
    assert(status == MS::kSuccess);

    // Pull on rigid.outputStartState
    int uid = startStateHandle.inputValue(&status).asInt();
    assert(status == MS::kSuccess && uid >= 0);

    // Tell rigid about which frame we just solved, i.e. the start frame.
    MDataHandle outputHandle = nextStateHandle.outputValue(&status);
    assert(status == MS::kSuccess);
    outputHandle.setInt(_currentFrame);

    Print() << "SolverNode: Solver pulled on solver.index[" << index << "] --> rigid.next[" << uid << "]\n";

    return status;
}


MStatus SolverNode::updateRigids(const MPlug& plug, MDataBlock& datablock) {
    // Print() << "SolverNode::updateRigids()\n";
    MStatus status { MS::kSuccess };

    const int index = plug.logicalIndex(&status);
    assert(status == MS::kSuccess);

    // Pull on *all* currentStates
    MArrayDataHandle currentStateHandle = datablock.inputArrayValue(aCurrentState, &status);
    assert(status == MS::kSuccess);

    MArrayDataHandle nextStateHandle = datablock.outputArrayValue(aNextState, &status);
    assert(status == MS::kSuccess);

    status = currentStateHandle.jumpToElement(index);
    assert(status == MS::kSuccess);
    status = nextStateHandle.jumpToElement(index);
    assert(status == MS::kSuccess);

    MDataHandle inputHandle = currentStateHandle.inputValue(&status);
    assert(status == MS::kSuccess);
    inputHandle.asInt();

    if (_currentFrame > _simulationFrame) {
        stepSimulation(datablock);
    }

    MDataHandle outputHandle = nextStateHandle.outputValue(&status);
    assert(status == MS::kSuccess);
    outputHandle.setInt(_currentFrame);

    return status;
}


MStatus SolverNode::restoreSimulation(MDataBlock& datablock) {
    MStatus status { MS::kSuccess };

    for (auto& pair : gRigids) {
        RigidBody& rigid = pair.second;
        Pose& pose = rigid.history.at(_currentFrame - 1);
        rigid.worldMatrix = pose.matrix;
        rigid.velocity = pose.velocity;
        rigid.worldMatrixChanged = true;
    }

    _simulationFrame = _currentFrame;

    return status;
}


/**
 * Apply gravity and update position
 *
 */
MStatus SolverNode::stepSimulation(MDataBlock& datablock) {
    MStatus status { MS::kSuccess };
    MVector gravity = InputValue(datablock, aGravity).asDouble3();

    for (auto& pair : gRigids) {
        RigidBody& rigid = pair.second;

        if (rigid.active) {
            MTransformationMatrix tm { rigid.worldMatrix };

            rigid.velocity += gravity * _deltaTime;
            tm.addTranslation(rigid.velocity - (rigid.velocity * rigid.damping), MSpace::kTransform);

            rigid.worldMatrix = tm.asMatrix();
            rigid.worldMatrixChanged = true;
        }

        else {
            // Here's a good opportunity to compute kinematic
            // velocity for when we transition *into* simulation
            // and want to preserve the animated velocity.
            rigid.velocity = { 0, 0, 0 };
        }

        rigid.history.resize(_simulationFrame + 1);
        rigid.history[_simulationFrame] = { rigid.worldMatrix, rigid.velocity };
    }

    _simulationFrame = _currentFrame;

    return status;
}