Effigy/Rig/WeightPaintLayer.cs

WeightPaintLayer is a small model class that stores per-vertex painted skin weights tied to a mesh topology. It records painted weights as lists of (bone name, weight), preserves topology id, marks staleness when topology no longer matches, can capture from indexed BoneWeight arrays, apply stored paints onto SkinWeights resolving bone names against a Skeleton, rebase topology, rename bones, clone itself and serialize/deserialize rows.

Native Interop
using System;
using System.Collections.Generic;

namespace Effigy;

/// <summary>
/// Painted weights that survive a rebuild — the part of weight painting that is actually hard.
///
/// THE PROBLEM IS THE SAME ONE THE SCULPT STAGE HAD. Effigy never stores a rig by vertex index; a
/// rebuild re-derives every weight from the body-to-bone assignments (`SkinBinder.BindBodies`), and
/// that is exactly what lets the parametric history stay alive under a rig. Painting fights that
/// head on: a painted vertex IS an edit expressed per vertex, and the next rebuild would throw it
/// away. Recomputing from assignments would then look like the brush had never worked.
///
/// THE ANSWER IS THE SCULPT ANSWER, and it is worth seeing that it is the same shape. A sculpt is
/// stored as deltas keyed on the cage's TOPOLOGY (`MultiresSculpt.TopologyId` — counts and face
/// indices, deliberately not positions, because positions are what a parametric edit is expected to
/// change). Topology unchanged means the vertex numbering means the same thing it did, so the edits
/// go straight back on. Topology changed means it does not, so the edits are KEPT and marked stale
/// rather than silently misapplied. That is the rule here too, and this class is deliberately small
/// because that rule is the whole of it.
///
/// TWO THINGS ARE STORED DIFFERENTLY FROM A SCULPT, and both matter.
///
/// **Bones are stored by NAME, never by index.** Bone names are the stable identifier this whole
/// project already committed to — Rig Control keys clips by name, and renaming one is a documented
/// breaking change. Indices are not stable at all: `Skeleton.RemoveBone` re-indexes everything after
/// the hole, so a layer holding index 4 would silently move to a different bone the first time
/// somebody deleted a bone above it. Names cost a dictionary lookup on apply and remove the entire
/// class of bug.
///
/// **The painted RESULT is stored, not a delta from the auto weights.** A delta of a normalised
/// quantity is not a well-defined thing: adding one breaks the sum, and re-normalising afterwards
/// gives an answer that depends on what the auto-binder happened to produce that run. Storing what
/// the vertex should end up as makes a repaint idempotent and makes "what did I paint here" a
/// question with an answer.
/// </summary>
public sealed class WeightPaintLayer
{
	/// <summary>Painted vertices, by index into the mesh this layer was captured on. Valid only
	/// while <see cref="Topology"/> matches — see the class comment.</summary>
	readonly Dictionary<int, List<(string Bone, float Weight)>> _painted = new();

	/// <summary>The topology this layer's vertex numbers mean something on.</summary>
	public long Topology { get; private set; }

	/// <summary>How many vertices carry paint.</summary>
	public int Count => _painted.Count;

	public bool IsEmpty => _painted.Count == 0;

	/// <summary>
	/// Whether the layer has been kept across a topology change it cannot be applied over.
	///
	/// Set rather than thrown, and the paint is NOT discarded, because the change may well be
	/// undone: the user pushes a fillet radius past what the boolean can take, the cage's face count
	/// moves, and pulling it back should not have cost them their paint. Same reasoning as
	/// SculptFeature's stale flag.
	/// </summary>
	public bool IsStale { get; private set; }

	public WeightPaintLayer() { }

	public WeightPaintLayer( PolyMesh mesh )
	{
		Topology = MultiresSculpt.TopologyId( mesh );
	}

	/// <summary>Rebuild a layer from a saved document. Topology is stored, not re-derived, because
	/// the mesh it was painted on is not in the file.</summary>
	public static WeightPaintLayer FromSaved( long topology,
		IEnumerable<(int Vertex, List<(string Bone, float Weight)> Weights)> rows )
	{
		var layer = new WeightPaintLayer { Topology = topology };

		if ( rows is null )
			return layer;

		foreach ( var (vertex, weights) in rows )
		{
			if ( weights is { Count: > 0 } )
				layer._painted[vertex] = new List<(string, float)>( weights );
		}

		return layer;
	}

	/// <summary>Every painted vertex, for a writer or a UI that wants to list them.</summary>
	public IEnumerable<(int Vertex, IReadOnlyList<(string Bone, float Weight)> Weights)> Painted
	{
		get
		{
			foreach ( var (vertex, weights) in _painted )
				yield return (vertex, weights);
		}
	}

	/// <summary>
	/// Record what a vertex now weighs, resolving indices to names against the skeleton it was
	/// painted with.
	///
	/// A vertex naming a bone the skeleton does not have is a caller fault rather than a shape, so
	/// it throws: the alternative is a layer that silently forgets an influence and a painted vertex
	/// whose weights no longer sum to one.
	/// </summary>
	public void Capture( int vertex, BoneWeight[] weights, Skeleton skeleton )
	{
		if ( skeleton is null )
			throw new ArgumentNullException( nameof( skeleton ) );

		if ( weights is null || weights.Length == 0 )
		{
			_painted.Remove( vertex );
			return;
		}

		var named = new List<(string, float)>( weights.Length );

		foreach ( var w in weights )
		{
			if ( w.Bone < 0 || w.Bone >= skeleton.Count )
				throw new ArgumentException( $"vertex {vertex} is weighted to bone {w.Bone}, and the skeleton has {skeleton.Count}" );

			named.Add( (skeleton.Bones[w.Bone].Name, w.Weight) );
		}

		_painted[vertex] = named;
	}

	/// <summary>Forget one vertex's paint, so the auto weights show through it again.</summary>
	public bool Clear( int vertex ) => _painted.Remove( vertex );

	public void ClearAll()
	{
		_painted.Clear();
		IsStale = false;
	}

	/// <summary>
	/// Whether this layer can be applied to a mesh, and if not, why — in the shape a diagnostic
	/// wants: what stopped it, with both models' numbers, and what would work.
	/// </summary>
	public bool CanApply( PolyMesh mesh, out string why )
	{
		if ( mesh is null )
			throw new ArgumentNullException( nameof( mesh ) );

		if ( _painted.Count == 0 )
		{
			why = null;
			return true;
		}

		if ( MultiresSculpt.TopologyId( mesh ) == Topology )
		{
			why = null;
			return true;
		}

		why = $"The rebuilt mesh has {mesh.VertexCount} vertices and {mesh.FaceCount} faces, and its "
			+ "topology is not the one this weight paint was made on. Painted weights are stored per "
			+ $"vertex, so the {_painted.Count} painted vertices cannot be placed on it. The paint is "
			+ "kept: undo the feature edit that changed the topology and it comes straight back.";
		return false;
	}

	/// <summary>
	/// Put the paint back on top of freshly auto-bound weights, and say how many vertices took it.
	///
	/// **After the auto-bind, never instead of it.** Everything the brush never touched has to come
	/// from `BindBodies` as it always did, or a rebuild that moved geometry would leave the unpainted
	/// nine tenths of the model weighted for where it used to be.
	///
	/// A bone that has since been DELETED is dropped from the vertex and what is left renormalised —
	/// which is the only defensible answer, because the alternatives are a vertex whose weights do
	/// not sum to one and a hard failure over a bone the user removed on purpose. `Missing` names
	/// them so the tool can say so.
	/// </summary>
	public int Apply( PolyMesh mesh, SkinWeights weights, Skeleton skeleton, out List<string> missing )
	{
		if ( mesh is null )
			throw new ArgumentNullException( nameof( mesh ) );

		if ( weights is null )
			throw new ArgumentNullException( nameof( weights ) );

		if ( skeleton is null )
			throw new ArgumentNullException( nameof( skeleton ) );

		missing = new List<string>();

		if ( _painted.Count == 0 )
		{
			IsStale = false;
			return 0;
		}

		if ( !CanApply( mesh, out _ ) )
		{
			IsStale = true;
			return 0;
		}

		if ( weights.Count != mesh.VertexCount )
			throw new ArgumentException( $"weights ({weights.Count}) and mesh ({mesh.VertexCount}) disagree" );

		var index = new Dictionary<string, int>( StringComparer.Ordinal );

		for ( var i = 0; i < skeleton.Count; i++ )
			index[skeleton.Bones[i].Name] = i;

		var applied = 0;

		foreach ( var (vertex, named) in _painted )
		{
			if ( vertex < 0 || vertex >= weights.Count )
				continue;

			var resolved = new List<BoneWeight>( named.Count );

			foreach ( var (name, weight) in named )
			{
				if ( index.TryGetValue( name, out var bone ) )
					resolved.Add( new BoneWeight( bone, weight ) );
				else if ( !missing.Contains( name ) )
					missing.Add( name );
			}

			// Every bone this vertex was painted to is gone. Leaving it as the auto-bind made it is
			// better than leaving it bound to nothing, and it is what the user would get if they had
			// never painted it - which is exactly the state deleting those bones put them back in.
			if ( resolved.Count == 0 )
				continue;

			weights[vertex] = SkinWeights.Blend( new[] { (resolved.ToArray(), 1f) } );
			applied++;
		}

		IsStale = false;
		return applied;
	}

	/// <summary>Re-key the layer onto a mesh whose topology it now matches. Called after a rebuild
	/// that changed the topology and was then undone.</summary>
	public void Rebase( PolyMesh mesh )
	{
		Topology = MultiresSculpt.TopologyId( mesh );
		IsStale = false;
	}

	/// <summary>Follow a bone rename, so paint is not lost to a rename the rest of the project
	/// treats as an ordinary edit.</summary>
	public int RenameBone( string from, string to )
	{
		if ( string.IsNullOrEmpty( from ) || string.IsNullOrEmpty( to ) || from == to )
			return 0;

		var touched = 0;

		foreach ( var (_, named) in _painted )
		{
			for ( var i = 0; i < named.Count; i++ )
			{
				if ( named[i].Bone != from )
					continue;

				named[i] = (to, named[i].Weight);
				touched++;
			}
		}

		return touched;
	}

	public WeightPaintLayer Clone()
	{
		var copy = new WeightPaintLayer { Topology = Topology, IsStale = IsStale };

		foreach ( var (vertex, named) in _painted )
			copy._painted[vertex] = new List<(string, float)>( named );

		return copy;
	}

	public override string ToString() =>
		_painted.Count == 0 ? "no paint" : $"{_painted.Count} painted vertices{(IsStale ? ", stale" : "")}";
}